ke+mm: add TLB shootdown and memory reclamation.

src/kern/amd64/impl.zig

@@ -5,6 +5,8 @@ const amd64 = r.arch;
 
 const std = @import("std");
 
+pub const tlb_max_pages = 33;
+
 pub const init = @import("init.zig");
 pub const early_init = init.early_init;
 
@@ -114,3 +116,19 @@ pub inline fn restore_interrupts(state: bool) void {
 pub inline fn send_resched_ipi(target: u32) void {
     pl.impl.send_resched_ipi(target);
 }
+
+pub inline fn send_tlb_ipi(target: u32) void {
+    pl.impl.send_tlb_ipi(target);
+}
+
+pub inline fn flush_full_tlb() void {
+    const prev_cr3 = amd64.read_cr(3);
+    // Reload CR3 to flush the entire TLB.
+    // This works because we assume global pages are never unmapped.
+    // If they are in the future, this should toggle CR4.PGE instead.
+    amd64.write_cr(3, prev_cr3);
+}
+
+pub inline fn flush_tlb(va: r.VAddr) void {
+    amd64.invlpg(va);
+}

src/kern/amd64/int.zig

@@ -4,6 +4,7 @@ const r = @import("root");
 const ke = r.ke;
 const ki = ke.private;
 const pl = r.pl;
+const mi = r.mm.private;
 
 extern const __interrupt_vectors: [256]usize;
 
@@ -84,6 +85,10 @@ export fn isr_handler_main(frame: *const amd64.IrqFrame) callconv(.{ .x86_64_sys
         ki.timer.clock();
     }
 
+    if (frame.intno == 241) {
+        ki.shootdown.ipi_handler();
+    }
+
     if (frame.intno >= 32) {
         pl.impl.apic.eoi();
     }

src/kern/dpc.zig

@@ -121,6 +121,8 @@ fn dispatch_queue(cpu: u32) void {
     sched_cpu.start_timer = false;
     sched_cpu.preemption_reason = .None;
 
+    ki.shootdown.process_shootdowns();
+
     if (sched_cpu.next_thread != null) {
         ki.sched.handle_preemption(sched_cpu);
     }

src/kern/init.zig

@@ -15,6 +15,7 @@ pub fn init(boot_info: *pl.BootInfo) linksection(r.init) void {
     pl.early_init(boot_info);
     ki.cpu.init_cpu(0);
     ki.turnstile.init_turnstiles();
+    ki.shootdown.init();
 
     thread0.init(@intFromPtr(&stack), r.kib(16), ke.Thread.Priority.idle, ki.sched.idle, null);
     thread0.priority = 0;

src/kern/private.zig

@@ -28,3 +28,4 @@ pub const event = @import("event.zig");
 pub const turnstile = @import("turnstile.zig");
 pub const mutex = @import("mutex.zig");
 pub const queue = @import("queue.zig");
+pub const shootdown = @import("shootdown.zig");

src/kern/root.zig

@@ -17,9 +17,11 @@ pub const Console = p.log.Console;
 pub const CpuLocal = p.cpu.CpuLocal;
 pub const ExportedCpuLocal = p.cpu.ExportedCpuLocal;
 pub const CpuMask = p.cpu.CpuMask;
+pub const AtomicCpuMask = p.cpu.AtomicCpuMask;
 pub const Event = p.event.Event;
 pub const Queue = p.queue.Queue;
 pub const Mutex = p.mutex.Mutex;
+pub const ShootdownState = p.shootdown.ShootdownState;
 
 // === Exported Interfaces ===
 pub const panic = p.panic.panic;
@@ -72,5 +74,10 @@ pub const cpu = struct {
     pub const current = p.cpu.current;
 };
 
+pub const shootdown = struct {
+    pub const submit = p.shootdown.submit;
+    pub const shootdowns = &p.shootdown.shootdowns;
+};
+
 /// Number of CPUs on the system
 pub var ncpus: usize = 0;

src/kern/shootdown.zig

@@ -0,0 +1,199 @@
+//! Kernel-side mechanism for asynchronous TLB shootdown and deferred memory reclamation.
+
+const std = @import("std");
+const config = @import("config");
+const rtl = @import("rtl");
+const r = @import("root");
+const ke = r.ke;
+const ki = ke.private;
+const mm = r.mm;
+
+comptime {
+    if (!@hasDecl(ki.impl, "send_tlb_ipi")) @compileError("impl must provide send_tlb_ipi");
+    if (!@hasDecl(ki.impl, "flush_full_tlb")) @compileError("impl must provide flush_full_tlb");
+    if (!@hasDecl(ki.impl, "flush_tlb")) @compileError("impl must provide flush_tlb");
+}
+
+/// Structure representing a single shootdown request.
+pub const ShootdownState = struct {
+    /// State of the request.
+    /// This is used as a counter when the request is outstanding,
+    /// otherwise it is set to the `slot_*` values below.
+    state: std.atomic.Value(u16),
+    link: rtl.List.Entry,
+    /// Base virtual address of the shootdown.
+    base: r.VAddr,
+    /// Number of pages to flush.
+    npages: u32,
+    /// Additional metadata.
+    payload: [2]usize,
+
+    pub fn release(self: *ShootdownState) void {
+        self.state.store(slot_free, .release);
+    }
+};
+
+const PerCpu = struct {
+    /// States attached to this CPU.
+    states: [64]ShootdownState,
+    /// States relevant to each CPU.
+    /// If CPU0 sent a shootdown to CPU1, then CPU0 will have `valid_states[CPU1]`
+    /// pointing to that shootdown, which will sit in `states`.
+    valid_states: [config.ncpus]std.atomic.Value(u64),
+    /// CPUs that have sent this CPU a shootdown.
+    senders: ke.AtomicCpuMask,
+};
+
+const slot_free: u16 = std.math.maxInt(u16);
+const slot_reserved: u16 = std.math.maxInt(u16) - 1;
+const percpu = ke.CpuLocal(PerCpu, undefined);
+
+pub var shootdowns: ke.Queue = undefined;
+
+var sync_counter: std.atomic.Value(usize) = .init(0);
+var sync_addr: r.VAddr = 0;
+var sync_count: usize = 0;
+var shootdown_lock: ke.SpinLock = .init();
+
+export const tlb_percpu_init linksection(r.percpu_init) = &pcpu_init;
+
+fn pcpu_init() linksection(r.init) callconv(.c) void {
+    const local = percpu.local();
+
+    local.valid_states = @splat(.init(0));
+    local.senders = .init(false);
+    local.states = @splat(.{ .base = 0, .npages = 0, .state = .init(slot_free), .payload = undefined, .link = undefined });
+}
+
+/// Allocate an invalidation state slot on this CPU.
+/// O(N) algorithm but N is tiny.
+fn allocate_slot(cpu: *PerCpu) ?usize {
+    for (0.., &cpu.states) |slot, *state| {
+        if (state.state.load(.monotonic) != slot_free) continue;
+
+        if (state.state.cmpxchgStrong(slot_free, slot_reserved, .acquire, .monotonic) == null) {
+            return slot;
+        }
+    }
+
+    return null;
+}
+
+fn flush_range(va: r.VAddr, npages: usize) void {
+    if (npages > ki.impl.tlb_max_pages) {
+        // Just flush the entire thing.
+        ki.impl.flush_full_tlb();
+        return;
+    }
+
+    for (0..npages) |i| {
+        ki.impl.flush_tlb(va + i * mm.page_size);
+    }
+}
+
+/// Synchronous shootdown path.
+fn do_synchronous_shootdown(state: ShootdownState) void {
+    shootdown_lock.acquire_no_ipl();
+
+    sync_addr = state.base;
+    sync_count = state.npages;
+    sync_counter.store(ke.ncpus - 1, .release);
+
+    const cur = ke.cpu.current();
+
+    for (0..ke.ncpus) |i| {
+        if (i == cur) continue;
+        ki.impl.send_tlb_ipi(@truncate(i));
+    }
+
+    while (sync_counter.load(.monotonic) != 0) {
+        std.atomic.spinLoopHint();
+    }
+
+    shootdown_lock.release_no_ipl();
+}
+
+/// Called by the TLB interrupt handler.
+pub fn ipi_handler() void {
+    flush_range(sync_addr, sync_count);
+    _ = sync_counter.fetchSub(1, .release);
+}
+
+/// Called in a quiescent state to process the pending shootdowns.
+/// This must be called at IPL dispatch.
+pub fn process_shootdowns() void {
+    const cpu = percpu.local();
+    const curcpu = ke.cpu.current();
+
+    if (cpu.senders.is_all(false, .monotonic)) {
+        // Nothing to do.
+        return;
+    }
+
+    var iter = cpu.senders.iter(.monotonic);
+
+    // Go through every sender and get their states.
+    while (iter.next()) |i| {
+        if (i == curcpu) continue;
+
+        cpu.senders.clear(i, .monotonic);
+
+        const sender = percpu.remote(@truncate(i));
+
+        while (true) {
+            var states = sender.valid_states[curcpu].swap(0, .acquire);
+            if (states == 0) break;
+
+            while (states != 0) {
+                // Go through each set bit and invalidate the slot.
+                const bit = @ctz(states);
+                states &= ~(@as(u64, 1) << @intCast(bit));
+
+                const state = &sender.states[bit];
+                flush_range(state.base, state.npages);
+
+                if (state.state.fetchSub(1, .release) == 1) {
+                    shootdowns.insert(&state.link, .Tail);
+                }
+            }
+        }
+    }
+}
+
+/// Submit a shootdown to occur asynchronously on `target_mask`.
+/// Returns an error if a synchronous shootdown was done instead.
+pub fn submit(state: ShootdownState, target_mask: ke.CpuMask) !void {
+    const ipl = ke.ipl.raise(.Dispatch);
+    defer ke.ipl.lower(ipl);
+
+    // Flush on our local TLB.
+    flush_range(state.base, state.npages);
+
+    // Get a slot and fill it.
+    const curcpu = ke.cpu.current();
+    const cpu = percpu.local();
+    const slot = allocate_slot(cpu) orelse {
+        do_synchronous_shootdown(state);
+        return error.NoSlot;
+    };
+
+    cpu.states[slot] = state;
+    cpu.states[slot].state.store(@truncate(target_mask.count()), .release);
+
+    var iter = target_mask.iter();
+
+    while (iter.next()) |bit| {
+        if (bit == curcpu) continue;
+
+        _ = cpu.valid_states[bit].fetchOr(@as(u64, 1) << @truncate(slot), .release);
+
+        ki.ipl.set_softint_pending(@truncate(bit), .Dispatch);
+
+        // Tell him we have something for him.
+        percpu.remote(@truncate(bit)).senders.set(curcpu, .release);
+    }
+}
+
+pub fn init() void {
+    shootdowns.init(ke.ncpus);
+}