ke: add Queue

src/kern/private.zig

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

src/kern/queue.zig

@@ -0,0 +1,107 @@
+//! Scheduler-aware queue.
+//! This file implements a queue onto which threads can block.
+//! The queue is aware of its threads blocking and waking up, and supports
+//! a maximum concurrency cap, controlling how many threads can process items
+//! from the queue concurrently.
+//! This is a bit similar to the KQUEUE data structure in the NT kernel.
+
+const r = @import("root");
+const rtl = @import("rtl");
+const ke = r.ke;
+const ki = ke.private;
+const std = @import("std");
+
+pub const Queue = struct {
+    hdr: ki.wait.DispatchHeader,
+    items: rtl.List,
+    active: usize,
+    max_active: usize,
+
+    /// Initialize a queue.
+    pub fn init(self: *Queue, max_active: usize) void {
+        self.* = .{
+            .items = undefined,
+            .active = 0,
+            .max_active = max_active,
+            .hdr = undefined,
+        };
+
+        self.hdr.init(.Queue);
+        self.items.init();
+    }
+
+    pub const Position = enum { Head, Tail };
+
+    /// Insert `item` into the queue at the desired position.
+    pub fn insert(self: *Queue, item: *rtl.List.Entry, pos: Position) void {
+        const ipl = self.hdr.lock.acquire();
+        defer self.hdr.lock.release(ipl);
+
+        self.hdr.signaled += 1;
+
+        switch (pos) {
+            .Head => {
+                self.items.insert_head(item);
+            },
+
+            .Tail => {
+                self.items.insert_tail(item);
+            },
+        }
+
+        ki.wait.satisfy_wait(&self.hdr);
+    }
+
+    /// Remove the item at the head.
+    /// This blocks until an item is actually popped.
+    pub fn remove(self: *Queue) *rtl.List.Entry {
+        const ipl = self.hdr.lock.acquire();
+        const td = ki.sched.percpu.local().current_thread.?;
+
+        if (td.queue) |q| {
+            std.debug.assert(q == self);
+            std.debug.assert(self.active > 0);
+
+            self.active -= 1;
+            if (self.hdr.signaled > 0 and self.active < self.max_active) {
+                ki.wait.satisfy_wait(&self.hdr);
+            }
+        } else {
+            td.queue = self;
+        }
+
+        td.queue_item = null;
+
+        self.hdr.lock.release(ipl);
+
+        // Wait until the queue has something for us.
+        _ = ke.wait.wait_one(&self.hdr, null) catch unreachable;
+
+        // Grab the queue item and set it to null.
+        const ret = td.queue_item orelse unreachable;
+        td.queue_item = null;
+        return ret;
+    }
+};
+
+/// Called when one of the threads on the queue has blocked on something other
+/// than the queue.
+pub fn signal_wait(queue: *Queue) void {
+    const ipl = queue.hdr.lock.acquire();
+    std.debug.assert(queue.active > 0);
+
+    queue.active -= 1;
+
+    if (queue.hdr.signaled > 0 and queue.active < queue.max_active) {
+        ki.wait.satisfy_wait(&queue.hdr);
+    }
+
+    queue.hdr.lock.release(ipl);
+}
+
+/// Called when one of the threads on the queue has woken back up.
+pub fn signal_wake(queue: *Queue) void {
+    const ipl = queue.hdr.lock.acquire();
+    queue.active += 1;
+    queue.hdr.lock.release(ipl);
+}

src/kern/root.zig

@@ -18,6 +18,7 @@ pub const CpuLocal = p.cpu.CpuLocal;
 pub const ExportedCpuLocal = p.cpu.ExportedCpuLocal;
 pub const CpuMask = p.cpu.CpuMask;
 pub const Event = p.event.Event;
+pub const Queue = p.queue.Queue;
 
 // === Exported Interfaces ===
 pub const panic = p.panic.panic;

src/kern/sched.zig

@@ -256,10 +256,8 @@ pub fn block_locked(curtd: *ke.Thread) void {
     yield(ke.cpu.current());
 }
 
-/// Unblock a thread.
-pub fn unblock(td: *ke.Thread) void {
-    const ipl = td.lock.acquire();
-
+pub fn unblock_locked(td: *ke.Thread) void {
+    std.debug.assert(td.lock.is_locked());
     std.debug.assert(td.state == .Blocked);
 
     td.sleep_time = (ke.time.read_time() - td.sleep_start) / std.time.ns_per_ms;
@@ -273,6 +271,13 @@ pub fn unblock(td: *ke.Thread) void {
     // Enqueue the thread
     const cpu = pick_cpu(td);
     enqueue_on_cpu(cpu, td);
+}
+
+/// Unblock a thread.
+pub fn unblock(td: *ke.Thread) void {
+    const ipl = td.lock.acquire();
+
+    unblock_locked(td);
 
     td.lock.release(ipl);
 }

src/kern/thread.zig

@@ -97,16 +97,17 @@ pub const Thread = struct {
     /// Current wait status.
     wait_status: std.atomic.Value(ki.wait.Status),
     waitblocks: [4]ki.wait.WaitBlock,
-
     /// Timer used for timeouts.
     timer: ke.Timer,
-
     /// Turnstile.
     turnstile: *ki.turnstile.Turnstile,
     turnstile_waiter: ?*ki.turnstile.Waiter,
     turnstiles_owned: rtl.List,
     /// Object this thread is currently blocked on, or null.
     waiting_on: ?*anyopaque,
+    /// Queue this thread is associated with.
+    queue: ?*ke.Queue,
+    queue_item: ?*rtl.List.Entry,
 
     /// Initialize a thread.
     /// - `stack`: Address of the **base** of the stack on which the initial context for the thread is built
@@ -139,6 +140,8 @@ pub const Thread = struct {
             .turnstile_waiter = null,
             .turnstiles_owned = undefined,
             .waiting_on = null,
+            .queue = null,
+            .queue_item = null,
         };
 
         thread.turnstiles_owned.init();

src/kern/wait.zig

@@ -15,6 +15,8 @@ pub const DispatchHeader = struct {
         Notification,
         /// When signaled, `signaled` is decreased until 0, and a single waiter is woken up.
         Synchronization,
+        /// Special type for queue objects.
+        Queue,
     };
 
     /// List of WaitBlocks.
@@ -34,14 +36,36 @@ pub const DispatchHeader = struct {
         obj.waitblocks.init();
     }
 
-    fn consume(self: *DispatchHeader) void {
+    fn consume(self: *DispatchHeader, td: *ke.Thread) void {
         switch (self.type) {
             .Synchronization => self.signaled -= 1,
             .Notification => {
                 // Object remains signaled.
             },
+            .Queue => {
+                const q: *ke.Queue = @fieldParentPtr("hdr", self);
+                const item = q.items.first();
+
+                item.remove();
+                q.hdr.signaled -= 1;
+                q.active += 1;
+
+                td.queue_item = item;
+            },
         }
     }
+
+    fn can_satisfy(self: *DispatchHeader) bool {
+        if (self.signaled == 0) return false;
+
+        return switch (self.type) {
+            .Notification, .Synchronization => true,
+            .Queue => blk: {
+                const q: *ke.Queue = @fieldParentPtr("hdr", self);
+                break :blk q.active < q.max_active;
+            },
+        };
+    }
 };
 
 pub const WaitBlock = struct {
@@ -98,7 +122,8 @@ pub fn wait_any(objects: []*DispatchHeader, timeout: ?r.Nanoseconds, waitblocks:
     const total_count = obj_count + @intFromBool(has_timeout);
     const timer_i = obj_count;
     const timer = &curtd.timer;
-    const timer_block = &curtd.waitblocks[3];
+    const timer_block = &curtd.waitblocks[obj_count];
+    var queue: ?*ke.Queue = null;
 
     var satisfier: ?usize = null;
     curtd.wait_status.store(.InProgress, .release);
@@ -111,16 +136,20 @@ pub fn wait_any(objects: []*DispatchHeader, timeout: ?r.Nanoseconds, waitblocks:
         obj.lock.acquire_no_ipl();
         defer obj.lock.release_no_ipl();
 
-        if (obj.signaled > 0) {
+        if (obj.can_satisfy()) {
             // Object was already signaled. Try consuming it.
             if (curtd.wait_status.cmpxchgStrong(.InProgress, .Satisfied, .acq_rel, .monotonic) == null) {
-                obj.consume();
+                obj.consume(curtd);
+                satisfier = i;
             }
             // We have already been satisfied in the meantime, abort.
-            satisfier = i;
             break;
         }
 
+        if (obj.type == .Queue) {
+            queue = @fieldParentPtr("hdr", obj);
+        }
+
         wb.object = obj;
         wb.thread = curtd;
         wb.status = .Active;
@@ -130,7 +159,9 @@ pub fn wait_any(objects: []*DispatchHeader, timeout: ?r.Nanoseconds, waitblocks:
 
     // Wait was already satisfied, back out.
     if (satisfier != null or (has_timeout and timeout.? == 0)) {
-        std.debug.assert(curtd.wait_status.load(.acquire) == .Satisfied);
+        if (satisfier != null) {
+            std.debug.assert(curtd.wait_status.load(.acquire) == .Satisfied);
+        }
 
         const cleanup_count = satisfier orelse total_count;
 
@@ -161,10 +192,17 @@ pub fn wait_any(objects: []*DispatchHeader, timeout: ?r.Nanoseconds, waitblocks:
     // While we're trying to commit the wait, the object locks have been released, and the state could therefore change.
     // We need to re-check the state of the wait before actually blocking.
     if (curtd.wait_status.cmpxchgStrong(.InProgress, .Committed, .acq_rel, .monotonic) == null) {
+        if (queue == null) {
+            if (curtd.queue) |q| {
+                ki.queue.signal_wait(q);
+            }
+        }
+
         // We're good, now actually block.
-        //std.log.debug("Blocking {*} on {*}", .{ curtd, objects[0] });
         ki.sched.block_locked(curtd);
     } else {
+        queue = null;
+
         // Could not commit.
         curtd.lock.release_no_ipl();
     }
@@ -206,10 +244,11 @@ pub fn satisfy_wait(obj: *DispatchHeader) void {
     const all = obj.type == .Notification;
 
     // Go through (potentially) all wait blocks and try to satisfy them.
-    while (!obj.waitblocks.is_empty() and obj.signaled > 0) {
+    while (!obj.waitblocks.is_empty() and obj.can_satisfy()) {
         // Get the first waitblock from the list.
         var wb: *WaitBlock = @fieldParentPtr("link", obj.waitblocks.first());
         var td = wb.thread;
+
         std.debug.assert(wb.status == .Active);
         std.debug.assert(wb.object == obj);
 
@@ -225,7 +264,7 @@ pub fn satisfy_wait(obj: *DispatchHeader) void {
         if (td.wait_status.cmpxchgStrong(.InProgress, .Satisfied, .acq_rel, .monotonic) == null) {
             // We interrupted the wait while it was being prepared.
             wb.status = .Signaled;
-            obj.consume();
+            obj.consume(td);
         }
 
         // 2.
@@ -233,8 +272,18 @@ pub fn satisfy_wait(obj: *DispatchHeader) void {
             // We interrupted the wait while it was committed.
             // Wake the thread.
             wb.status = .Signaled;
-            obj.consume();
-            ke.sched.unblock(td);
+            obj.consume(td);
+
+            const ipl = td.lock.acquire();
+
+            if (obj.type != .Queue) {
+                if (td.queue) |q| {
+                    ki.queue.signal_wake(q);
+                }
+            }
+
+            ki.sched.unblock_locked(td);
+            td.lock.release(ipl);
         }
 
         // 3.