ke+mm: add TLB shootdown and memory reclamation.

src/mm/heap.zig

@@ -16,8 +16,14 @@ pub fn alloc(size: usize) mm.Error!*anyopaque {
     mi.kernel_space.pmap.map_range_allocating(addr, size, .{
         .read = true,
         .write = true,
-        .global = true,
+        .global = false,
     });
 
     return @ptrFromInt(addr);
 }
+
+pub fn free(va: r.VAddr, size: usize) void {
+    const ipl = mi.kernel_space.lock.acquire();
+    mi.tlb.reclaim_range(&mi.kernel_space, va, size);
+    mi.kernel_space.lock.release(ipl);
+}

src/mm/init.zig

@@ -103,4 +103,5 @@ pub fn init(boot_info: *pl.BootInfo) linksection(r.init) void {
 
 pub fn late_init() linksection(r.init) void {
     mi.zone.late_init();
+    mi.tlb.init();
 }

src/mm/phys.zig

@@ -16,13 +16,11 @@ const ke = r.ke;
 
 const log = std.log.scoped(.@"mm/phys");
 
-// === Early state ===
 var bootstrapped = false;
 var early_alloc_entry_idx: usize = 0;
 var memory_map: *pl.BootInfo.MemMap = undefined;
 var early_allocs: usize = 0;
 
-// === Memory allocator state ===
 /// Head for the global treiber stack.
 const GlobalHead = packed struct(u64) {
     pfn: mm.Pfn,
@@ -41,13 +39,11 @@ const PerCpu = struct {
 
 const percpu = ke.CpuLocal(PerCpu, undefined);
 
-const null_pfn: mm.Pfn = std.math.maxInt(mm.Pfn);
-
 // TODO: make this configurable (or dynamic).
 const batch_size: u8 = 64;
 
 var batch_pool: std.atomic.Value(GlobalHead) = .init(.{
-    .pfn = null_pfn,
+    .pfn = mm.null_pfn,
     .tag = 0,
 });
 
@@ -83,7 +79,7 @@ fn pop_from_pool() ?*mm.Page {
     while (true) {
         const batch_pfn = old_head.pfn;
 
-        if (batch_pfn == null_pfn) {
+        if (batch_pfn == mm.null_pfn) {
             return null;
         }
 
@@ -177,15 +173,15 @@ pub fn free(addr: r.PAddr) void {
 
     const page = mm.pfn_to_struct_page(mm.page_to_pfn(addr));
 
-    page.free.next_pfn = null_pfn;
-    page.free.batch_next = null_pfn;
+    page.free.next_pfn = mm.null_pfn;
+    page.free.batch_next = mm.null_pfn;
     page.free.batch_count = 0;
 
     if (cpu.active_batch) |curr| {
         // Add the page to the current batch.
         page.free.next_pfn = mm.struct_page_to_pfn(curr);
     } else {
-        page.free.next_pfn = null_pfn;
+        page.free.next_pfn = mm.null_pfn;
     }
 
     cpu.active_batch = page;

src/mm/pmap.zig

@@ -103,6 +103,17 @@ pub const PMap = struct {
         self.map_from(va, size, src);
     }
 
+    /// 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 {
+        return self.impl.unmap(va, size);
+    }
+
+    /// Flush the TLB entry of a single page.
+    pub fn flush(self: *Self, va: r.VAddr) void {
+        self.impl.flush(va);
+    }
+
     /// Map a single page. Convenience wrapper around map_contiguous_range.
     pub fn map_page(self: *Self, va: r.VAddr, pa: r.PAddr, flags: mm.MapFlags) void {
         self.map_contiguous_range(va, pa, mm.page_size, flags);

src/mm/private.zig

@@ -7,6 +7,7 @@ pub const vmem = @import("vmem.zig");
 pub const heap = @import("heap.zig");
 pub const zone = @import("zone.zig");
 pub const PMap = pmap.PMap;
+pub const tlb = @import("tlb.zig");
 
 const ke = @import("root").ke;
 
@@ -22,4 +23,8 @@ pub const impl = switch (config.arch) {
     else => @compileError("unsupported architecture"),
 };
 
-pub var kernel_space: Space = undefined;
+pub var kernel_space: Space = .{
+    .pmap = undefined,
+    .arena = undefined,
+    .lock = .init(),
+};

src/mm/radix_pmap.zig

@@ -109,6 +109,46 @@ 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 {
+            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;
+
+            for (0..npages) |_| {
+                var level: usize = num_levels - 1;
+                blk: while (true) {
+                    const idx = index_for_level(c.va, level);
+                    const pte = c.tables[level][idx];
+
+                    if (!pte.present) break :blk;
+
+                    if (level == 0) {
+                        // Found a present entry, add it to the list.
+                        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;
+                        head_pfn = pfn;
+                        c.tables[level][idx] = std.mem.zeroes(Impl.Pte);
+                        break;
+                    }
+
+                    if (!c.walk_down_resolve(level - 1)) return null;
+                    level -= 1;
+                }
+
+                c.advance(mm.page_size);
+            }
+
+            return if (head_pfn == mm.null_pfn) null else head_pfn;
+        }
+
         /// Ensure all intermediate page tables down to `target_level` exist for
         /// the given virtual address range, allocating them as needed. Useful for
         /// pre-populating a shared region (e.g. the kernel half of a user pmap)

src/mm/root.zig

@@ -15,6 +15,7 @@ pub const Space = private.Space;
 // 16 TB of maximum physical memory.
 // This should be fine on consumer hardware for at least a decade :^)
 pub const Pfn = u32;
+pub const null_pfn: Pfn = std.math.maxInt(Pfn);
 
 pub const Page = extern union {
     free: extern struct {

src/mm/tlb.zig

@@ -0,0 +1,101 @@
+//! MM-level TLB shootdown support
+
+const r = @import("root");
+const std = @import("std");
+const rtl = @import("rtl");
+const ke = r.ke;
+const mm = r.mm;
+const mi = mm.private;
+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 base = state.base;
+    const npages = state.npages;
+
+    // Reclaim the virtual address space.
+    const ipl = space.lock.acquire();
+
+    space.arena.free(base, @as(usize, npages) * mm.page_size) catch unreachable;
+    space.lock.release(ipl);
+
+    _ = async_shootdowns.fetchAdd(1, .monotonic);
+
+    // Free the physical pages
+    free_pages(pfn);
+
+    // Release the slot after all copied state has been consumed.
+    state.release();
+}
+
+fn worker_thread(_: ?*anyopaque) void {
+    while (true) {
+        const entry = ke.shootdown.shootdowns.remove();
+        const state: *ke.ShootdownState = @fieldParentPtr("link", entry);
+
+        _ = async_shootdowns.fetchAdd(1, .monotonic);
+        reclaim_state(state);
+    }
+}
+
+/// Unmap and flush a virtual range on the given space.
+/// Space lock is held on entry and IPL is raised to Dispatch.
+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 {
+        // If no physical pages, free the VA directly.
+        space.arena.free(va, size) catch unreachable;
+        return;
+    };
+
+    const state: ke.ShootdownState = .{
+        .base = va,
+        .npages = @truncate(size / mm.page_size),
+        .link = undefined,
+        .state = .init(0),
+        .payload = .{
+            @intFromPtr(space),
+            @intCast(head_pfn),
+        },
+    };
+
+    var mask: ke.CpuMask = .init(false);
+
+    for (0..ke.ncpus) |i| {
+        if (i != ke.cpu.current()) {
+            mask.set(i);
+        }
+    }
+
+    ke.shootdown.submit(state, mask) catch {
+        _ = sync_shootdowns.fetchAdd(1, .monotonic);
+        space.arena.free(va, size) catch unreachable;
+        free_pages(head_pfn);
+        return;
+    };
+}
+
+pub fn init() void {
+    const td = ps.thread.create_kernel(ke.Thread.Priority.default, worker_thread, null) catch unreachable;
+    ke.sched.enqueue(&td.kern);
+}

src/mm/zone.zig

@@ -20,8 +20,10 @@ const Slab = struct {
     /// Reference count indicating how many objects are in use.
     refcount: u16,
     capacity: u16,
-    /// Base address of the slab.
+    /// Base address of the slab + its color.
     base: usize,
+    /// Base address of the slab allocation.
+    buf: usize,
 
     pub fn bitmap(self: *Slab) [*]u64 {
         return @ptrFromInt(@intFromPtr(self) + @sizeOf(Slab));
@@ -623,6 +625,7 @@ pub const Zone = struct {
         if (rem > 0) bm[full_words] = (@as(u64, 1) << @intCast(rem)) - 1;
 
         ret.base = @intFromPtr(buf) + self.color;
+        ret.buf = @intFromPtr(buf);
         ret.capacity = @intCast(capacity);
         ret.refcount = 0;
 
@@ -667,7 +670,7 @@ pub const Zone = struct {
             const alloc_size = @sizeOf(Slab) + Slab.bitmap_bytes(slab.capacity);
             gpa.free(@as([*]u8, @ptrCast(slab))[0..alloc_size]);
 
-            // FIXME: call free_pages on slab.base here!
+            mm.heap.free(slab.buf, self.slab_size);
         } else {
             free_page(@ptrFromInt(std.mem.alignBackward(usize, @intFromPtr(slab), mm.page_size)));
         }
@@ -842,9 +845,8 @@ fn gpa_free(
     _: usize,
 ) void {
     if (buf.len > 2048) {
-        //const pages = std.mem.alignForward(usize, buf.len, mm.page_size);
-        // TODO
-        //mi.heap.free_pages(@ptrCast(buf.ptr), pages);
+        const pages = std.mem.alignForward(usize, buf.len, mm.page_size);
+        mm.heap.free(@intFromPtr(buf.ptr), pages);
         return;
     }