mm: rework page frame allocator

src/mm/phys.zig

@@ -1,14 +1,6 @@
-//! Physical memory allocator:
-//! - Uses a per-CPU magazined allocator with a global Treiber stack as the fallback.
-//! Each CPU has 4 batches of 64 pages each for fast allocation and deallocation without contention.
-//!  When a CPU's active batch is exhausted, it tries to get a new batch from its local depot.
-//!  If the depot is empty, it pops a batch from the global stack.
-//! When freeing pages, they are added to the active batch until it reaches capacity,
-//!  at which point the full batch is moved to the local depot.
-//!  If the depot is full when trying to add a new batch, one batch is pushed back to the global stack to make room.
 const r = @import("root");
 const std = @import("std");
-
+const rtl = @import("rtl");
 const pl = r.pl;
 const mm = r.mm;
 const mi = mm.private;
@@ -20,81 +12,10 @@ var bootstrapped = false;
 var early_alloc_entry_idx: usize = 0;
 var memory_map: *pl.BootInfo.MemMap = undefined;
 var early_allocs: usize = 0;
+var list_lock: ke.SpinLock = undefined;
+var free_list: rtl.List = undefined;
 
-/// Head for the global treiber stack.
+pub var usable_memory: std.atomic.Value(usize) = .init(0);
-const GlobalHead = packed struct(u64) {
-    pfn: mm.Pfn,
-    // A 32-bit ABA tag is more than enough for this,
-    // on 32-bit machines which do not support cmpxchg8b, we could use 12 bits instead.
-    tag: u32,
-};
-
-/// Per-CPU state.
-const PerCpu = struct {
-    depot: [4]*mm.Page,
-    depot_count: u8,
-    active_batch: ?*mm.Page,
-    active_count: u8,
-};
-
-const percpu = ke.CpuLocal(PerCpu, undefined);
-
-// TODO: make this configurable (or dynamic).
-const batch_size: u8 = 64;
-
-var batch_pool: std.atomic.Value(GlobalHead) = .init(.{
-    .pfn = mm.null_pfn,
-    .tag = 0,
-});
-
-export const phys_percpu_init linksection(r.percpu_init) = &init_cpu;
-
-fn init_cpu() callconv(.c) void {
-    const cpu = percpu.local();
-    cpu.active_batch = null;
-    cpu.active_count = 0;
-    cpu.depot_count = 0;
-}
-
-/// Push a batch of free pages to the global pool.
-fn push_to_pool(batch: *mm.Page) void {
-    var old_head = batch_pool.load(.acquire);
-
-    while (true) {
-        batch.free.batch_next = old_head.pfn;
-
-        const new_head = GlobalHead{
-            .pfn = mm.struct_page_to_pfn(batch),
-            .tag = old_head.tag +% 1,
-        };
-
-        old_head = batch_pool.cmpxchgWeak(old_head, new_head, .acq_rel, .acquire) orelse break;
-    }
-}
-
-/// Pop a batch of free pages from the global pool. Returns null if the pool is empty.
-fn pop_from_pool() ?*mm.Page {
-    var old_head = batch_pool.load(.acquire);
-
-    while (true) {
-        const batch_pfn = old_head.pfn;
-
-        if (batch_pfn == mm.null_pfn) {
-            return null;
-        }
-
-        const batch = mm.pfn_to_struct_page(batch_pfn);
-
-        const new_head = GlobalHead{
-            .pfn = batch.free.batch_next,
-            .tag = old_head.tag +% 1,
-        };
-
-        old_head = batch_pool.cmpxchgWeak(old_head, new_head, .acq_rel, .acquire) orelse return batch;
-    }
-
-    return null;
-}
 
 fn early_alloc() usize {
     while (early_alloc_entry_idx < memory_map.entry_count) {
@@ -120,97 +41,57 @@ fn early_alloc() usize {
 /// Allocate a page of physical memory.
 pub fn alloc() r.PAddr {
     if (!bootstrapped) {
+        @branchHint(.unlikely);
         return early_alloc();
     }
 
-    const ipl = ke.ipl.raise(.High);
+    const ipl = list_lock.acquire();
-    const cpu = percpu.local();
 
-    defer ke.ipl.lower(ipl);
+    if (free_list.is_empty()) {
-
+        list_lock.release(ipl);
-    if (cpu.active_count == 0) {
+        // FIXME: wait for free memory
-        // Our batch is empty, get another one.
+        @panic("mm/phys: Out of memory for physical allocation");
-        if (cpu.depot_count > 0) {
-            // Get a new batch from the local depot.
-            cpu.depot_count -= 1;
-            cpu.active_batch = cpu.depot[cpu.depot_count];
-            cpu.active_count = batch_size;
-        } else {
-            // Get one from the global pool.
-            const new_batch = pop_from_pool();
-
-            if (new_batch) |new| {
-                cpu.active_batch = new;
-                cpu.active_count = new.free.batch_count;
-            } else {
-                // TODO: call an IPI here for the other CPUs to fill the global pool before giving up.
-                @panic("OOM");
-            }
-        }
     }
 
-    const batch = cpu.active_batch.?;
+    const head = free_list.first();
+    head.remove();
 
-    if (cpu.active_count > 1) {
+    _ = usable_memory.fetchSub(mm.page_size, .monotonic);
-        // Pop the next page from the batch and set it as the new head.
-        const next_page = mm.pfn_to_struct_page(batch.free.next_pfn);
-        cpu.active_batch = next_page;
-    } else {
-        cpu.active_batch = null;
-    }
 
-    cpu.active_count -= 1;
+    const elem: *mm.PageFree = @fieldParentPtr("link", head);
+    const page: *mm.Page = @ptrCast(elem);
+    const phys_addr = mm.pfn_to_page(mm.struct_page_to_pfn(page));
 
-    return mm.pfn_to_page(mm.struct_page_to_pfn(batch));
+    list_lock.release(ipl);
+
+    return phys_addr;
 }
 
 /// Free a page of physical memory.
 pub fn free(addr: r.PAddr) void {
-    const ipl = ke.ipl.raise(.High);
+    const page: *mm.Page = mm.pfn_to_struct_page(mm.page_to_pfn(addr));
-    defer ke.ipl.lower(ipl);
 
-    const cpu = percpu.local();
+    const ipl = list_lock.acquire();
+    _ = usable_memory.fetchAdd(mm.page_size, .monotonic);
+    free_list.insert_head(&page.free.link);
+    list_lock.release(ipl);
+}
 
-    const page = mm.pfn_to_struct_page(mm.page_to_pfn(addr));
+/// Free a list of pages of physical memory.
+pub fn free_batch(head: *mm.Page, tail: *mm.Page, count: usize) void {
+    const ipl = list_lock.acquire();
 
-    page.free.next_pfn = mm.null_pfn;
+    tail.free.link.next = free_list.first();
-    page.free.batch_next = mm.null_pfn;
+    free_list.insert_head(&head.free.link);
-    page.free.batch_count = 0;
 
-    if (cpu.active_batch) |curr| {
+    _ = usable_memory.fetchAdd(count * mm.page_size, .monotonic);
-        // Add the page to the current batch.
-        page.free.next_pfn = mm.struct_page_to_pfn(curr);
-    } else {
-        page.free.next_pfn = mm.null_pfn;
-    }
 
-    cpu.active_batch = page;
+    list_lock.release(ipl);
-    cpu.active_count += 1;
-
-    if (cpu.active_count == batch_size) {
-        // The batch is full, move it to the depot.
-        page.free.batch_count = @intCast(cpu.active_count);
-
-        if (cpu.depot_count == cpu.depot.len) {
-            // The depot is full, push one batch back to the global pool to make room.
-            push_to_pool(cpu.depot[0]);
-
-            cpu.depot[0] = cpu.depot[1];
-            cpu.depot[1] = cpu.depot[2];
-            cpu.depot[2] = cpu.depot[3];
-            cpu.depot_count -= 1;
-        }
-
-        cpu.depot[cpu.depot_count] = cpu.active_batch.?;
-        cpu.depot_count += 1;
-
-        cpu.active_batch = null;
-        cpu.active_count = 0;
-    }
 }
 
 pub fn init(boot_info: *pl.BootInfo) linksection(r.init) void {
     memory_map = &boot_info.memory_map;
+    free_list.init();
 
     var total_usable_memory: usize = 0;
     log.info("physical memory map:", .{});

src/mm/pmap.zig

@@ -105,7 +105,7 @@ pub const PMap = struct {
 
     /// Unmap a contiguous range of virtual pages.
     /// Returns a list threaded through PFNs of the unmapped physical pages if any.
-    pub fn unmap(self: *Self, va: r.VAddr, size: usize) ?mm.Pfn {
+    pub fn unmap(self: *Self, va: r.VAddr, size: usize) ?mi.PfnList {
         return self.impl.unmap(va, size);
     }
 

src/mm/private.zig

@@ -17,6 +17,11 @@ pub const Space = struct {
     lock: ke.QSpinLock,
 };
 
+pub const PfnList = extern struct {
+    head: u32,
+    tail: u32,
+};
+
 pub const impl = switch (config.arch) {
     .amd64 => @import("amd64/impl.zig"),
     .um => @import("um/impl.zig"),

src/mm/radix_pmap.zig

@@ -25,9 +25,10 @@
 //!     the CPU (e.g. write CR3 or satp).
 
 const r = @import("root");
+const rtl = @import("rtl");
+const std = @import("std");
 const mm = r.mm;
 const mi = mm.private;
-const std = @import("std");
 
 pub const Level = struct {
     shift: u6,
@@ -112,13 +113,14 @@ pub fn RadixPmap(comptime Impl: type) type {
         /// Unmap a contiguous range of virtual pages.
         /// Only small pages are supported.
         /// Return a pfn that points to the physical pages that were unmapped.
-        pub fn unmap(self: *Self, va: r.VAddr, size: usize) ?mm.Pfn {
+        pub fn unmap(self: *Self, va: r.VAddr, size: usize) ?mi.PfnList {
             std.debug.assert(std.mem.isAligned(va, mm.page_size));
 
             var c = self.cursor(va);
 
             const npages: usize = size / mm.page_size;
-            var head_pfn: mm.Pfn = mm.null_pfn;
+            var list: rtl.List = undefined;
+            list.init();
 
             for (0..npages) |_| {
                 var level: usize = num_levels - 1;
@@ -133,8 +135,8 @@ pub fn RadixPmap(comptime Impl: type) type {
                         const addr = pte.address();
                         const pfn = mm.page_to_pfn(addr);
                         const page = mm.pfn_to_struct_page(pfn);
-                        page.free.next_pfn = head_pfn;
+                        list.insert_tail(&page.free.link);
-                        head_pfn = pfn;
+
                         c.tables[level][idx] = std.mem.zeroes(Impl.Pte);
                         break;
                     }
@@ -146,7 +148,12 @@ pub fn RadixPmap(comptime Impl: type) type {
                 c.advance(mm.page_size);
             }
 
-            return if (head_pfn == mm.null_pfn) null else head_pfn;
+            const head_free: *mm.PageFree = @fieldParentPtr("link", list.first());
+            const tail_free: *mm.PageFree = @fieldParentPtr("link", list.last());
+            const head: *mm.Page = @ptrCast(head_free);
+            const tail: *mm.Page = @ptrCast(tail_free);
+
+            return if (list.is_empty()) null else .{ .head = mm.struct_page_to_pfn(head), .tail = mm.struct_page_to_pfn(tail) };
         }
 
         /// Ensure all intermediate page tables down to `target_level` exist for

src/mm/root.zig

@@ -1,6 +1,7 @@
 pub const private = @import("private.zig");
 
 const std = @import("std");
+const rtl = @import("rtl");
 
 pub const page_size = 4096;
 
@@ -17,13 +18,12 @@ pub const Space = private.Space;
 pub const Pfn = u32;
 pub const null_pfn: Pfn = std.math.maxInt(Pfn);
 
+pub const PageFree = extern struct {
+    link: rtl.List.Entry,
+};
+
 pub const Page = extern union {
-    free: extern struct {
+    free: PageFree,
-        // Fields used when the page is free.
-        next_pfn: Pfn,
-        batch_next: Pfn,
-        batch_count: u8,
-    },
 
     alloced: extern struct {
         slab_data: zone.Page,

src/mm/tlb.zig

@@ -11,23 +11,11 @@ const ps = r.ps;
 pub var sync_shootdowns: std.atomic.Value(usize) = .init(0);
 pub var async_shootdowns: std.atomic.Value(usize) = .init(0);
 
-fn free_pages(head_pfn: mm.Pfn) void {
-    var pfn = head_pfn;
-
-    // Free all the physical pages.
-    while (pfn != mm.null_pfn) {
-        const page = mm.pfn_to_struct_page(pfn);
-        const next = page.free.next_pfn;
-        mm.phys.free(mm.pfn_to_page(pfn));
-        pfn = next;
-    }
-}
-
 /// Reclaim the memory associated with a state.
 /// This can only be done when all CPUs have flushed their TLB.
 fn reclaim_state(state: *ke.ShootdownState) void {
     const space: *mm.Space = @ptrFromInt(state.payload[0]);
-    const pfn: u32 = @truncate(state.payload[1]);
+    const pfn_list: mi.PfnList = @bitCast(state.payload[1]);
     const base = state.base;
     const npages = state.npages;
 
@@ -40,7 +28,7 @@ fn reclaim_state(state: *ke.ShootdownState) void {
     _ = async_shootdowns.fetchAdd(1, .monotonic);
 
     // Free the physical pages
-    free_pages(pfn);
+    mm.phys.free_batch(mm.pfn_to_struct_page(pfn_list.head), mm.pfn_to_struct_page(pfn_list.tail), npages);
 
     // Release the slot after all copied state has been consumed.
     state.release();
@@ -62,7 +50,7 @@ pub fn reclaim_range(space: *mm.Space, va: r.VAddr, size: usize) void {
     std.debug.assert(std.mem.isAligned(va, mm.page_size));
 
     // Unmap the virtual addresses and get the backing physical pages.
-    const head_pfn = space.pmap.unmap(va, size) orelse {
+    const list = space.pmap.unmap(va, size) orelse {
         // If no physical pages, free the VA directly.
         space.arena.free(va, size) catch unreachable;
         return;
@@ -75,7 +63,7 @@ pub fn reclaim_range(space: *mm.Space, va: r.VAddr, size: usize) void {
         .state = .init(0),
         .payload = .{
             @intFromPtr(space),
-            @intCast(head_pfn),
+            @bitCast(list),
         },
     };
 
@@ -90,7 +78,9 @@ pub fn reclaim_range(space: *mm.Space, va: r.VAddr, size: usize) void {
     ke.shootdown.submit(state, mask) catch {
         _ = sync_shootdowns.fetchAdd(1, .monotonic);
         space.arena.free(va, size) catch unreachable;
-        free_pages(head_pfn);
+
+        // Free the physical pages
+        mm.phys.free_batch(mm.pfn_to_struct_page(list.head), mm.pfn_to_struct_page(list.tail), size / mm.page_size);
         return;
     };
 }