From 1176e97e90c0068f7d56636a858d4304f44c7a9d Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Lo=C3=AFc=20Gr=C3=A9goire?= <loic@gregoire.dev>
Date: Mon, 7 Sep 2026 20:58:25 -0400
Subject: [PATCH] mm: more aggressive reclamation

---
 src/mm/balance.zig | 81 +++++++++++++++++++++++++++++-----------------
 src/mm/phys.zig    | 11 +++++++
 2 files changed, 63 insertions(+), 29 deletions(-)

diff --git a/src/mm/balance.zig b/src/mm/balance.zig
index 4e89b73..0806c6e 100644
--- a/src/mm/balance.zig
+++ b/src/mm/balance.zig
@@ -30,41 +30,64 @@ pub fn init() void {
     ke.sched.enqueue(&td.kern);
 }
 
+fn reap_stacks() void {
+    if (ke.thread.stack_cache_update()) |lst| {
+        // Reap all the excess stacks.
+        var entry: ?*ke.Stack = lst;
+
+        while (entry != null) {
+            const next = entry.?.next;
+            mm.heap.free(
+                (@intFromPtr(entry) + @sizeOf(ke.Stack)) -
+                    ps.thread.kernel_thread_stack_size,
+                ps.thread.kernel_thread_stack_size,
+            );
+            entry = next;
+        }
+    }
+}
+
 fn balance_manager(_: ?*anyopaque) void {
+    const low_memory = 0;
+    const timeout = 1;
+
+    var timer: ke.Timer = undefined;
+    timer.init();
+
+    var which: usize = timeout;
+
     while (true) {
-        var timer: ke.Timer = undefined;
-        timer.init();
-        ke.timer.set(
-            &timer,
-            .from(r.Milliseconds.init(balance_interval.load())),
-            .{},
-        );
-
-        _ = ke.wait.wait_one(&timer.hdr, "balmgr", .{}) catch unreachable;
-
-        stack_reap_time -= 1;
-        zone_update_time -= 1;
-
-        if (zone_update_time == 0) {
-            zone_update_time = zone_update_interval;
-            mm.zone.update();
+        if (which == timeout) {
+            ke.timer.set(
+                &timer,
+                .from(r.Milliseconds.init(balance_interval.load())),
+                .{},
+            );
         }
 
-        if (stack_reap_time == 0) {
-            stack_reap_time = stack_reap_interval;
+        which = ke.wait.wait_any(&.{
+            &mmp.phys.low_memory_event.hdr,
+            &timer.hdr,
+        }, "balmgr", .{}) catch
+            unreachable;
 
-            if (ke.thread.stack_cache_update()) |lst| {
-                // Reap all the excess stacks.
-                var entry: ?*ke.Stack = lst;
+        if (which == low_memory) {
+            reap_stacks();
+            mm.zone.drain();
+            mmp.phys.low_memory_event.reset();
+        } else if (which == timeout) {
+            stack_reap_time -= 1;
+            zone_update_time -= 1;
 
-                while (entry != null) {
-                    const next = entry.?.next;
-                    mm.heap.free(
-                        (@intFromPtr(entry) + @sizeOf(ke.Stack)) - ps.thread.kernel_thread_stack_size,
-                        ps.thread.kernel_thread_stack_size,
-                    );
-                    entry = next;
-                }
+            if (zone_update_time == 0) {
+                zone_update_time = zone_update_interval;
+                mm.zone.update();
+            }
+
+            if (stack_reap_time == 0) {
+                stack_reap_time = stack_reap_interval;
+
+                reap_stacks();
             }
         }
     }
diff --git a/src/mm/phys.zig b/src/mm/phys.zig
index f3d28b0..73fd0ee 100644
--- a/src/mm/phys.zig
+++ b/src/mm/phys.zig
@@ -17,6 +17,8 @@ pub var min_memory_threshold: usize = 0;
 pub var low_memory_threshold: usize = 0;
 /// Threshold at which reclamation stops.
 pub var high_memory_threshold: usize = 0;
+/// Event triggered when memory is low.
+pub var low_memory_event: ke.Event = undefined;
 
 var memory_map: *r.BootInfo.MemMap = undefined;
 var early_allocs: usize = 0;
@@ -71,6 +73,10 @@ pub fn alloc() r.PAddr {
     _ = usable_memory.fetchSub(mm.page_size, .monotonic);
     free_pages -= 1;
 
+    if (free_pages < min_memory_threshold) {
+        low_memory_event.signal();
+    }
+
     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));
@@ -97,6 +103,10 @@ pub fn alloc_opts(opts: struct { policy: mm.WaitPolicy }) ?r.PAddr {
     _ = usable_memory.fetchSub(mm.page_size, .monotonic);
     free_pages -= 1;
 
+    if (free_pages < min_memory_threshold) {
+        low_memory_event.signal();
+    }
+
     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));
@@ -144,6 +154,7 @@ pub fn init() linksection(r.init) void {
     memory_map = &r.boot_info.memory_map;
     free_list.init();
     free_page_event.init(.Notification);
+    low_memory_event.init(.Notification);
 
     var total_usable_memory: usize = memory_map.loader_memory_used;
     log.info("physical memory map:", .{});