danos/library/kernel/thread.zig

//! `runtime.Thread` — threads for danos, shaped like Zig's `std.Thread` but built on
//! danos's private thread ABI (docs/threading.md). Several tasks share one address
//! space; `spawn` starts one, the kernel delivers the closure pointer in the new
//! thread's rdi, a plain Zig trampoline runs the user function and calls `thread_exit`,
//! and `join` blocks on the thread's exit notification. See docs/threading.md for why
//! this mirrors `std.Thread`'s API rather than being the literal type.
//!
//! The closure (the function's captured args) lives at the **top of the thread's own
//! stack**, not the heap — each thread's stack is private, so there is no shared-heap
//! concurrency in the spawn/join machinery (the runtime heap is not yet thread-safe).
//! A binary must be built multi-threaded (`addThreadedUserBinary`) before it may spawn.

const std = @import("std");
const builtin = @import("builtin");
const abi = @import("abi");
const sc = @import("system-call");

// A thread allocates its own stack straight from the mmap syscall (not through the
// `memory` module) so `memory`'s heap can depend on this module's Mutex without a cycle.
inline fn mmapStack(len: usize) usize {
    return sc.systemCall2(.mmap, len, abi.prot_read | abi.prot_write);
}
inline fn mmapFailed(ret: usize) bool {
    return ret > ~@as(usize, 0) - 4095;
}
inline fn munmapStack(base: usize, len: usize) void {
    _ = sc.systemCall2(.munmap, base, len);
}

/// True in a real danos binary; false when this module is compiled for host unit tests.
/// The `Futex` seam and the test blocks below branch on it so the lock/condvar state
/// machines can be exercised on the host against `std.Thread.Futex` (docs/threading-plan.md
/// M11), while the danos build uses the futex syscalls.
const on_danos = builtin.os.tag == .freestanding;

/// A thread stack, if the caller does not override it. 64 KiB of mmap'd, zeroed pages.
pub const default_stack_size: usize = 64 * 1024;

/// Bytes reserved at the top of each thread's stack for its per-thread TLS block (the
/// self-pointer plus scratch slots reachable via `%fs`). docs/threading-plan.md M10.
const tls_block_size: usize = 64;

pub const Thread = struct {
    /// The kernel task id of the spawned thread — what `join` waits on.
    tid: u32,
    /// The mmap'd stack, reclaimed by `join` (or at process exit after `detach`).
    stack_base: usize,
    stack_size: usize,

    pub const Id = u32;

    pub const SpawnConfig = struct {
        /// Bytes of stack, rounded up to whole pages by the kernel's mmap.
        stack_size: usize = default_stack_size,
    };

    pub const SpawnError = error{
        /// The kernel refused the thread, the stack mmap failed, or no endpoint was free.
        SystemResources,
    };

    /// Start `function(args...)` on a new thread sharing this address space. Mirrors
    /// `std.Thread.spawn`. The thread's return value is discarded (as in `std.Thread`);
    /// return data through shared state.
    pub fn spawn(config: SpawnConfig, comptime function: anytype, args: anytype) SpawnError!Thread {
        const Args = @TypeOf(args);
        const Closure = struct {
            tls_base: usize,
            args: Args,
            /// Entered directly by the kernel with `self` in rdi (C ABI). Establishes this
            /// thread's TLS pointer, runs the user function, then ends the thread.
            fn entry(self_addr: usize) callconv(.c) noreturn {
                const self: *@This() = @ptrFromInt(self_addr);
                setThreadPointer(self.tls_base); // per-thread thread pointer before any user code
                @call(.auto, function, self.args);
                exitThread();
            }
        };

        const base = mmapStack(config.stack_size);
        if (mmapFailed(base)) return error.SystemResources;

        // Top of the thread's own stack, downward: the closure, then a small per-thread TLS
        // block (the thread pointer points here; slot 0 is the variant-II self-pointer, the rest is
        // scratch for user TLS), then the stack proper (rsp starts below the TLS block, so
        // the growing stack never overwrites either).
        var closure_addr = (base + config.stack_size) - @sizeOf(Closure);
        closure_addr &= ~@as(usize, @alignOf(Closure) - 1); // align the closure down

        const tls_base = (closure_addr - tls_block_size) & ~@as(usize, 15);
        const tls: [*]usize = @ptrFromInt(tls_base);
        tls[0] = tls_base; // self-pointer (fs:0), as the x86_64 TLS ABI expects

        const closure: *Closure = @ptrFromInt(closure_addr);
        closure.* = .{ .tls_base = tls_base, .args = args };

        var stack_top = tls_base & ~@as(usize, 15); // 16-align below the TLS block
        stack_top -= 8; // ...then rsp % 16 == 8 at the C entry

        const tid = threadSpawn(@intFromPtr(&Closure.entry), stack_top, closure_addr);
        if (threadSpawnFailed(tid)) {
            munmapStack(base, config.stack_size);
            return error.SystemResources;
        }
        return .{ .tid = @intCast(tid), .stack_base = base, .stack_size = config.stack_size };
    }

    /// Block until this thread finishes, then reclaim its stack. Mirrors
    /// `std.Thread.join`. The exit endpoint is private to this thread, so the first
    /// child-exit notification on it is this thread's.
    pub fn join(self: Thread) void {
        _ = sc.systemCall1(.thread_join, self.tid); // block until the thread has exited
        munmapStack(self.stack_base, self.stack_size); // reclaim its (now-vacated) stack
    }

    /// Relinquish the right to join: never wait for or reclaim this thread. Its stack is
    /// reclaimed at process exit (docs/threading-plan.md M3 — kernel-reaper stack reclaim
    /// for detached threads is a later refinement). Mirrors `std.Thread.detach`.
    pub fn detach(self: Thread) void {
        _ = self;
    }

    /// The calling thread's id (its kernel task id). Mirrors `std.Thread.getCurrentId`.
    pub fn getCurrentId() Id {
        return @intCast(sc.systemCall0(.thread_self));
    }

    /// The dense 0-based index of the core the calling thread is running on. A danos
    /// extension beyond `std.Thread`, used to observe genuine cross-core parallelism.
    pub fn currentCore() Id {
        return @intCast(sc.systemCall0(.current_core));
    }

    /// Ask the kernel to end the calling thread. A WORKER never returns from this;
    /// the process's MAIN thread gets the kernel's refusal (-EPERM — the group ends
    /// only through exit, a fault, or process_kill, docs/shared-fate-plan.md) and
    /// the call returns. Exists for exactly that refusal path; workers end through
    /// the spawn trampoline, and a process ends through `system.exit`.
    pub fn tryExitCurrent() void {
        _ = sc.systemCall0(.thread_exit);
    }

    /// `std.Thread.Futex`-shaped block/wake on a `u32` atomic — the primitive the
    /// blocking `Mutex`/`Condition`/`Semaphore` are built on. Waiters park in the
    /// kernel (no busy-wait), so an idle core still halts (docs/halting.md).
    pub const Futex = struct {
        /// Block while `ptr.* == expect`. Returns when woken by `wake`, or promptly if
        /// the value already differs (safe against spurious returns, as in std): the
        /// caller re-checks its condition in a loop.
        pub fn wait(ptr: *const std.atomic.Value(u32), expect: u32) void {
            if (comptime on_danos) {
                _ = futexWait(@intFromPtr(ptr), expect, 0);
            } else {
                // Host unit-test mock: spin+yield until the value changes (`wake` is a
                // no-op — the callers re-check their condition in a loop anyway). Correct,
                // if busy; fine for the state-machine tests.
                while (ptr.load(.acquire) == expect) std.Thread.yield() catch {};
            }
        }

        /// As `wait`, but returns `error.Timeout` if `timeout_ns` elapses first.
        pub fn timedWait(ptr: *const std.atomic.Value(u32), expect: u32, timeout_ns: u64) error{Timeout}!void {
            if (comptime on_danos) {
                if (futexWait(@intFromPtr(ptr), expect, timeout_ns) == abi.futex_timed_out) return error.Timeout;
            } else {
                var spins: u64 = 0;
                const limit = timeout_ns / 1000 + 1;
                while (ptr.load(.acquire) == expect) : (spins += 1) {
                    if (spins >= limit) return error.Timeout;
                    std.Thread.yield() catch {};
                }
            }
        }

        /// Wake up to `max_waiters` threads blocked on `ptr`.
        pub fn wake(ptr: *const std.atomic.Value(u32), max_waiters: u32) void {
            if (comptime on_danos) {
                _ = futexWake(@intFromPtr(ptr), max_waiters);
            } else {
                // host mock: spin-waiters re-check their condition, so no wake is needed.
            }
        }
    };

    /// A mutual-exclusion lock, `std.Thread.Mutex`-shaped. The classic three-state
    /// futex mutex (unlocked / locked / contended): the fast path is a single CAS, and
    /// only a contended lock ever enters the kernel.
    pub const Mutex = struct {
        state: std.atomic.Value(u32) = std.atomic.Value(u32).init(unlocked),

        const unlocked: u32 = 0;
        const locked: u32 = 1;
        const contended: u32 = 2;

        /// Try to take the lock without blocking; returns whether it was acquired.
        pub fn tryLock(m: *Mutex) bool {
            return m.state.cmpxchgStrong(unlocked, locked, .acquire, .monotonic) == null;
        }

        /// Acquire the lock, blocking in the kernel while it is contended.
        pub fn lock(m: *Mutex) void {
            if (m.state.cmpxchgStrong(unlocked, locked, .acquire, .monotonic) != null) m.lockSlow();
        }

        fn lockSlow(m: *Mutex) void {
            @branchHint(.cold);
            // Mark the lock contended and take it as soon as it falls unlocked; park on
            // the futex while it stays contended. Marking contended may cause a spurious
            // wake on unlock (harmless), never a missed one.
            while (m.state.swap(contended, .acquire) != unlocked) {
                Futex.wait(&m.state, contended);
            }
        }

        /// Release the lock; wake one waiter if the lock was contended.
        pub fn unlock(m: *Mutex) void {
            if (m.state.swap(unlocked, .release) == contended) Futex.wake(&m.state, 1);
        }
    };

    /// A condition variable, `std.Thread.Condition`-shaped. Spurious wakeups are
    /// allowed — always wait in a predicate loop with the mutex held. Built on a futex
    /// sequence counter: a waiter samples the seq, drops the mutex, and parks until the
    /// seq changes (a signal that races the unlock bumps the seq, so it is not missed).
    pub const Condition = struct {
        seq: std.atomic.Value(u32) = std.atomic.Value(u32).init(0),

        /// Atomically release `mutex` and block until signalled, then re-acquire it.
        pub fn wait(c: *Condition, mutex: *Mutex) void {
            const seq = c.seq.load(.acquire);
            mutex.unlock();
            Futex.wait(&c.seq, seq);
            mutex.lock();
        }

        /// As `wait`, but returns `error.Timeout` if `timeout_ns` elapses first. The
        /// mutex is re-acquired either way.
        pub fn timedWait(c: *Condition, mutex: *Mutex, timeout_ns: u64) error{Timeout}!void {
            const seq = c.seq.load(.acquire);
            mutex.unlock();
            const timed_out = if (Futex.timedWait(&c.seq, seq, timeout_ns)) |_| false else |_| true;
            mutex.lock();
            if (timed_out) return error.Timeout;
        }

        /// Wake one waiter.
        pub fn signal(c: *Condition) void {
            _ = c.seq.fetchAdd(1, .release);
            Futex.wake(&c.seq, 1);
        }

        /// Wake all waiters.
        pub fn broadcast(c: *Condition) void {
            _ = c.seq.fetchAdd(1, .release);
            Futex.wake(&c.seq, std.math.maxInt(u32));
        }
    };

    /// A counting semaphore, `std.Thread.Semaphore`-shaped: a permit count guarded by a
    /// `Mutex` + `Condition`.
    pub const Semaphore = struct {
        mutex: Mutex = .{},
        cond: Condition = .{},
        permits: usize = 0,

        /// Take a permit, blocking until one is available.
        pub fn wait(s: *Semaphore) void {
            s.mutex.lock();
            defer s.mutex.unlock();
            while (s.permits == 0) s.cond.wait(&s.mutex);
            s.permits -= 1;
        }

        /// Return a permit and wake a waiter.
        pub fn post(s: *Semaphore) void {
            s.mutex.lock();
            defer s.mutex.unlock();
            s.permits += 1;
            s.cond.signal();
        }
    };

    /// A reader/writer lock, `std.Thread.RwLock`-shaped: many concurrent readers OR one
    /// exclusive writer. Reader-preferring (a steady stream of readers can delay a writer),
    /// built on `Mutex` + `Condition` over a signed state: `>0` = that many readers hold
    /// it, `-1` = a writer holds it, `0` = free.
    pub const RwLock = struct {
        mutex: Mutex = .{},
        cond: Condition = .{},
        state: i64 = 0,

        /// Acquire shared (read) access, blocking while a writer holds the lock.
        pub fn lockShared(rw: *RwLock) void {
            rw.mutex.lock();
            defer rw.mutex.unlock();
            while (rw.state < 0) rw.cond.wait(&rw.mutex);
            rw.state += 1;
        }

        /// Try to acquire shared access without blocking.
        pub fn tryLockShared(rw: *RwLock) bool {
            rw.mutex.lock();
            defer rw.mutex.unlock();
            if (rw.state < 0) return false;
            rw.state += 1;
            return true;
        }

        /// Release shared access; wake a waiting writer once the last reader leaves.
        pub fn unlockShared(rw: *RwLock) void {
            rw.mutex.lock();
            defer rw.mutex.unlock();
            rw.state -= 1;
            if (rw.state == 0) rw.cond.broadcast();
        }

        /// Acquire exclusive (write) access, blocking until no readers or writer remain.
        pub fn lock(rw: *RwLock) void {
            rw.mutex.lock();
            defer rw.mutex.unlock();
            while (rw.state != 0) rw.cond.wait(&rw.mutex);
            rw.state = -1;
        }

        /// Try to acquire exclusive access without blocking.
        pub fn tryLock(rw: *RwLock) bool {
            rw.mutex.lock();
            defer rw.mutex.unlock();
            if (rw.state != 0) return false;
            rw.state = -1;
            return true;
        }

        /// Release exclusive access; wake all waiters (they re-check their condition).
        pub fn unlock(rw: *RwLock) void {
            rw.mutex.lock();
            defer rw.mutex.unlock();
            rw.state = 0;
            rw.cond.broadcast();
        }
    };

    /// A `std.Thread.WaitGroup`-shaped counter: `start` before spawning work, `finish` as
    /// each unit completes, `wait` blocks until the count returns to zero.
    pub const WaitGroup = struct {
        mutex: Mutex = .{},
        cond: Condition = .{},
        counter: usize = 0,

        /// Register one pending unit of work.
        pub fn start(wg: *WaitGroup) void {
            wg.mutex.lock();
            defer wg.mutex.unlock();
            wg.counter += 1;
        }

        /// Mark one unit done; wake waiters if that was the last.
        pub fn finish(wg: *WaitGroup) void {
            wg.mutex.lock();
            defer wg.mutex.unlock();
            wg.counter -= 1;
            if (wg.counter == 0) wg.cond.broadcast();
        }

        /// Block until every started unit has finished.
        pub fn wait(wg: *WaitGroup) void {
            wg.mutex.lock();
            defer wg.mutex.unlock();
            while (wg.counter != 0) wg.cond.wait(&wg.mutex);
        }
    };
};

/// thread_spawn(entry, stack_top, arg, exit_endpoint) -> tid, or a wrapped error.
fn threadSpawn(entry: usize, stack_top: usize, arg: usize) usize {
    const exit_endpoint: usize = @intCast(abi.no_cap); // join uses thread_join, not an endpoint
    return sc.systemCall4(.thread_spawn, entry, stack_top, arg, exit_endpoint);
}

/// The kernel returns a real (small) task id on success and a wrapped `-1` on failure;
/// no valid task id ever exceeds a u32.
inline fn threadSpawnFailed(ret: usize) bool {
    return ret > std.math.maxInt(u32);
}

/// End the calling thread. Never returns.
fn exitThread() noreturn {
    _ = sc.systemCall0(.thread_exit);
    unreachable;
}

/// Set the calling thread's FS base (its user TLS thread pointer).
fn setThreadPointer(addr: usize) void {
    _ = sc.systemCall1(.set_thread_pointer, addr);
}

/// futex_wait(addr, expect, timeout_ns) -> status (abi.futex_*).
fn futexWait(addr: usize, expect: u32, timeout_ns: u64) usize {
    return sc.systemCall3(.futex_wait, addr, expect, timeout_ns);
}

/// futex_wake(addr, count) -> number woken.
fn futexWake(addr: usize, count: u32) usize {
    return sc.systemCall2(.futex_wake, addr, count);
}

// --- host unit tests (docs/threading-plan.md M11) ---------------------------
//
// These run under `zig build test` on the host: the `Futex` seam above uses
// `std.Thread.Futex` off-danos, so the lock/condvar state machines can be exercised by
// real host threads. They are never compiled into a danos binary (test blocks only build
// under test), so their `std.Thread` use is fine even though `std.Thread` is unavailable
// on the freestanding target.

test "Mutex serialises concurrent increments across host threads" {
    var m: Thread.Mutex = .{};
    var counter: u64 = 0;
    const workers = 8;
    const per = 20_000;
    const Ctx = struct {
        m: *Thread.Mutex,
        c: *u64,
        fn run(ctx: @This()) void {
            var i: usize = 0;
            while (i < per) : (i += 1) {
                ctx.m.lock();
                ctx.c.* += 1;
                ctx.m.unlock();
            }
        }
    };
    var handles: [workers]std.Thread = undefined;
    for (&handles) |*h| h.* = try std.Thread.spawn(.{}, Ctx.run, .{Ctx{ .m = &m, .c = &counter }});
    for (handles) |h| h.join();
    try std.testing.expectEqual(@as(u64, workers * per), counter);
}

test "RwLock never lets a reader observe a half-written pair" {
    var rw: Thread.RwLock = .{};
    var a: u64 = 0;
    var b: u64 = 0; // invariant while a lock is held: a == b
    var stop = std.atomic.Value(bool).init(false);
    var ok = std.atomic.Value(bool).init(true);

    const Writer = struct {
        rw: *Thread.RwLock,
        a: *u64,
        b: *u64,
        stop: *std.atomic.Value(bool),
        fn run(w: @This()) void {
            var v: u64 = 1;
            while (!w.stop.load(.acquire)) : (v +%= 1) {
                w.rw.lock();
                w.a.* = v; // update both halves under the exclusive lock...
                w.b.* = v;
                w.rw.unlock();
            }
        }
    };
    const Reader = struct {
        rw: *Thread.RwLock,
        a: *u64,
        b: *u64,
        ok: *std.atomic.Value(bool),
        fn run(r: @This()) void {
            var i: usize = 0;
            while (i < 200_000) : (i += 1) {
                r.rw.lockShared();
                if (r.a.* != r.b.*) r.ok.store(false, .release); // ...so a reader must never see them differ
                r.rw.unlockShared();
            }
        }
    };

    var writers: [2]std.Thread = undefined;
    for (&writers) |*w| w.* = try std.Thread.spawn(.{}, Writer.run, .{Writer{ .rw = &rw, .a = &a, .b = &b, .stop = &stop }});
    var readers: [4]std.Thread = undefined;
    for (&readers) |*rd| rd.* = try std.Thread.spawn(.{}, Reader.run, .{Reader{ .rw = &rw, .a = &a, .b = &b, .ok = &ok }});
    for (readers) |rd| rd.join();
    stop.store(true, .release);
    for (writers) |w| w.join();
    try std.testing.expect(ok.load(.acquire));
}

test "WaitGroup blocks until every started unit finishes" {
    var wg: Thread.WaitGroup = .{};
    var done = std.atomic.Value(u32).init(0);
    const n = 6;
    const Ctx = struct {
        wg: *Thread.WaitGroup,
        done: *std.atomic.Value(u32),
        fn run(c: @This()) void {
            _ = c.done.fetchAdd(1, .monotonic);
            c.wg.finish();
        }
    };
    var i: usize = 0;
    while (i < n) : (i += 1) wg.start();
    var handles: [n]std.Thread = undefined;
    for (&handles) |*h| h.* = try std.Thread.spawn(.{}, Ctx.run, .{Ctx{ .wg = &wg, .done = &done }});
    wg.wait(); // must not return until all n finished
    try std.testing.expectEqual(@as(u32, n), done.load(.acquire));
    for (handles) |h| h.join();
}