ajhahn.de
← FlashOS
Flash 180 lines
// wait_queue: blocking-syscall wait queue.
// Single source of truth; future signal / do_wait migration reuses it.
//
// Discipline:
//   * Wait-side links the running task at head, flips state to
//     INTERRUPTIBLE, drops preempt, then schedules. The wake-side
//     reverses both (pop + flip to RUNNING + clear wq_next).
//   * `wq_next` lives in TaskStruct (src/task_layout.flash) — a task can
//     only be on one queue at a time, mirroring Linux's wq_node.
//   * Single-core: `preempt_disable` around the head mutation is enough
//     because no two callers can race on the same queue. Future work
//     migrates wake/wait to spinlocks (which disable IRQs on acquire);
//     the API surface stays identical.
//   * IRQ callers (mini-UART RX) are fine on single-core because the
//     entry path masks IRQs and there is no concurrent mutator from
//     EL1h.

// Named module — see build.zig (`task_layout_mod`). Required because
// wait_queue.flash is itself a named module; if it pulled task_layout
// in via relative import while a sibling named module (pipe.flash) did
// the same, Zig 0.16 would diagnose "file exists in two modules".
const layout = #import("task_layout")
const TaskStruct = layout.TaskStruct
const TASK_RUNNING = layout.TASK_RUNNING
const TASK_INTERRUPTIBLE = layout.TASK_INTERRUPTIBLE

extern var current ?*mut TaskStruct
extern fn preempt_disable() void
extern fn preempt_enable() void
extern fn schedule() void

pub const WaitQueue = extern struct {
    head ?*mut TaskStruct = null,

    pub fn prepare_to_wait(self *mut WaitQueue) void {
        preempt_disable()
        if current |c| {
            // Link task if not already on the queue
            if c.wq_next == null && self.head != c {
                c.wq_next = self.head
                self.head = c
            }
            c.state = TASK_INTERRUPTIBLE
        }
        preempt_enable()
    }

    pub fn finish_wait(self *mut WaitQueue) void {
        preempt_disable()
        if current |c| {
            c.state = TASK_RUNNING
            // Unlink if still on queue
            if self.head == c {
                self.head = c.wq_next
                c.wq_next = null
            } else {
                var prev = self.head
                while prev |p| {
                    if p.wq_next == c {
                        p.wq_next = c.wq_next
                        c.wq_next = null
                        break
                    }
                    prev = p.wq_next
                }
            }
        }
        preempt_enable()
    }

    pub fn wait(self *mut WaitQueue) void {
        preempt_disable()
        if current |c| {
            if c.wq_next == null && self.head != c {
                c.wq_next = self.head
                self.head = c
            }
            c.state = TASK_INTERRUPTIBLE
        }
        preempt_enable()
        schedule()
        self.finish_wait()
    }

    pub fn wake_one(self *mut WaitQueue) void {
        preempt_disable()
        if self.head |t| {
            self.head = t.wq_next
            t.wq_next = null
            t.state = TASK_RUNNING
        }
        preempt_enable()
    }

    pub fn wake_all(self *mut WaitQueue) void {
        preempt_disable()
        var node = self.head
        self.head = null
        while node |t| {
            const nxt = t.wq_next
            t.wq_next = null
            t.state = TASK_RUNNING
            node = nxt
        }
        preempt_enable()
    }
}

// ---- Host tests ----
//
// schedule is a no-op stub on the host (see tests/host_stubs.zig).
// These exercise the wake-side directly; the wait-side blocking path
// is covered by the in-kernel pipe scenario.

const std = #import("std")

test "wake_one pops in LIFO order (head-insert)" {
    var t1 TaskStruct = .{}
    var t2 TaskStruct = .{}
    var t3 TaskStruct = .{}
    var q WaitQueue = .{}

    // Manual head-insert mirrors WaitQueue.wait without the schedule round-trip.
    t1.wq_next = null
    q.head = &t1
    t2.wq_next = q.head
    q.head = &t2
    t3.wq_next = q.head
    q.head = &t3

    t1.state = TASK_INTERRUPTIBLE
    t2.state = TASK_INTERRUPTIBLE
    t3.state = TASK_INTERRUPTIBLE

    q.wake_one()
    try std.testing.expectEqual(&t2, q.head.?)
    try std.testing.expectEqual(#as(?*mut TaskStruct, null), t3.wq_next)
    try std.testing.expectEqual(TASK_RUNNING, t3.state)
    try std.testing.expectEqual(TASK_INTERRUPTIBLE, t2.state)
}

test "wake_one on empty queue is a noop" {
    var q WaitQueue = .{}
    q.wake_one()
    try std.testing.expectEqual(#as(?*mut TaskStruct, null), q.head)
}

test "wake_all drains every entry and resets state + wq_next" {
    var t1 TaskStruct = .{}
    var t2 TaskStruct = .{}
    var t3 TaskStruct = .{}
    var q WaitQueue = .{}

    t1.wq_next = null
    q.head = &t1
    t2.wq_next = q.head
    q.head = &t2
    t3.wq_next = q.head
    q.head = &t3

    t1.state = TASK_INTERRUPTIBLE
    t2.state = TASK_INTERRUPTIBLE
    t3.state = TASK_INTERRUPTIBLE

    q.wake_all()
    try std.testing.expectEqual(#as(?*mut TaskStruct, null), q.head)
    try std.testing.expectEqual(#as(?*mut TaskStruct, null), t1.wq_next)
    try std.testing.expectEqual(#as(?*mut TaskStruct, null), t2.wq_next)
    try std.testing.expectEqual(#as(?*mut TaskStruct, null), t3.wq_next)
    try std.testing.expectEqual(TASK_RUNNING, t1.state)
    try std.testing.expectEqual(TASK_RUNNING, t2.state)
    try std.testing.expectEqual(TASK_RUNNING, t3.state)
}

test "wake_all on empty queue is a noop" {
    var q WaitQueue = .{}
    q.wake_all()
    try std.testing.expectEqual(#as(?*mut TaskStruct, null), q.head)
}