scheduler: Implement spawning of nested tasks

src/scheduler.zig

@@ -3,7 +3,14 @@ const Allocator = std.mem.Allocator;
 const Node = @import("node.zig");
 
 const Scheduler = @This();
-pub const EntryPoint = *const fn(std.process.Init) anyerror!void;
+
+pub const MainEntryPoint = *const fn(std.process.Init) anyerror!void;
+pub const NestedEntryPoint = *const fn(context: *const anyopaque) void;
+
+pub const EntryPoint = union(enum) {
+    main: MainEntryPoint,
+    nested: NestedEntryPoint,
+};
 
 const Registers = struct {
     rsp: usize,
@@ -20,6 +27,8 @@ const ContextSwitch = extern struct {
 const STACK_CANARY_SIZE = 256;
 const STACK_CANARY_BYTE = 0xa5;
 
+pub const TaskID = u64;
+
 const State = enum {
     ready,
     running,
@@ -29,12 +38,20 @@ const State = enum {
 };
 
 const Task = struct {
+    id    : TaskID,
     regs  : Registers,
     stack : []align(16) u8,
     entry : EntryPoint,
+    context: ?*const anyopaque,
     state : State,
-    wakeup: ?u64,
     node  : *Node,
+    parent_id: ?TaskID,
+
+    // These fields are only used when state=.blocked and
+    // are not mutually exclusive. Each represents a different
+    // wakeup condition.
+    wakeup_time: ?u64,
+    wakeup_tasks: ?[]const TaskID,
 
     fn stackCanaryIsIntact(self: *Task) bool {
         const canary = self.stack[0..STACK_CANARY_SIZE];
@@ -49,14 +66,16 @@ const Task = struct {
 gpa: Allocator,
 tasks: std.ArrayList(Task),
 regs: Registers,
-current: ?*Task,
+current_id: ?TaskID,
 current_time: u64,
+next_task_id: u64,
 
 pub fn init(self: *Scheduler, gpa: Allocator) void {
     self.gpa = gpa;
     self.tasks = .empty;
-    self.current = null;
+    self.current_id = null;
     self.current_time = 0;
+    self.next_task_id = 0;
 }
 
 pub fn deinit(self: *Scheduler) void {
@@ -66,7 +85,23 @@ pub fn deinit(self: *Scheduler) void {
     self.tasks.deinit(self.gpa);
 }
 
-pub fn spawn(self: *Scheduler, node: *Node, entry: EntryPoint, stack_size: usize) !void {
+pub fn spawn(self: *Scheduler, node: *Node, entry: MainEntryPoint, stack_size: usize) !void {
+    _ = try self.spawnInner(node, .{ .main = entry }, stack_size, null, null);
+}
+
+pub fn spawnNested(self: *Scheduler, node: *Node, entry: NestedEntryPoint, context: *const anyopaque) !TaskID {
+    const parent_id = self.current_id.?;
+    return (try self.spawnInner(node, .{ .nested = entry }, 64 * 1024, parent_id, context)).id;
+}
+
+fn spawnInner(
+    self: *Scheduler,
+    node: *Node,
+    entry: EntryPoint,
+    stack_size: usize,
+    parent_id: ?TaskID,
+    context: ?*const anyopaque,
+) !*Task {
 
     if (stack_size > std.math.maxInt(usize) - STACK_CANARY_SIZE)
         return Allocator.Error.OutOfMemory;
@@ -82,7 +117,11 @@ pub fn spawn(self: *Scheduler, node: *Node, entry: EntryPoint, stack_size: usize
     // If the entry point returns, `ret` jumps here instead of into nowhere.
     @as(*usize, @ptrFromInt(stack_top)).* = @intFromPtr(&taskReturned);
 
+    const id = self.next_task_id;
+    defer self.next_task_id += 1;
+
     const task = Task {
+        .id = id,
         .regs = .{
             .rsp = stack_top,
             .rbp = 0,
@@ -91,12 +130,28 @@ pub fn spawn(self: *Scheduler, node: *Node, entry: EntryPoint, stack_size: usize
         },
         .stack  = stack,
         .entry  = entry,
+        .context = context,
         .state  = .ready,
-        .wakeup = null,
+        .wakeup_time = null,
+        .wakeup_tasks = null,
         .node   = node,
+        .parent_id = parent_id,
     };
 
     try self.tasks.append(self.gpa, task);
+    return &self.tasks.items[self.tasks.items.len - 1];
+}
+
+// Removes from the scheduler a nested task. This should only be used
+// on a task that has never ran yet or the cleanup won't be clean.
+// The intended use-case is undoing a previous call .spawnNested to
+// in case of errors.
+pub fn despawnNested(self: *Scheduler, id: TaskID) void {
+    const index = self.findTaskIndexByID(id) orelse return;
+    const task = &self.tasks.items[index];
+
+    std.heap.page_allocator.free(task.stack);
+    _ = self.tasks.orderedRemove(index);
 }
 
 fn findTaskWithState(self: *Scheduler, state: State) ?*Task {
@@ -107,14 +162,26 @@ fn findTaskWithState(self: *Scheduler, state: State) ?*Task {
     return null;
 }
 
+fn findTaskIndexByID(self: *Scheduler, id: TaskID) ?usize {
+    for (self.tasks.items, 0..) |*task, i| {
+        if (task.id == id)
+            return i;
+    }
+    return null;
+}
+
+fn findTaskByID(self: *Scheduler, id: TaskID) ?*Task {
+    const index = self.findTaskIndexByID(id) orelse return null;
+    return &self.tasks.items[index];
+}
+
 fn findBlockedTaskWithLowestWakeupTime(self: *Scheduler) ?*Task {
-    var task = self.findTaskWithState(.blocked) orelse return null;
+    var task: ?*Task = null;
     for (self.tasks.items) |*t| {
-        // If the state is blocked, wakeup MUST be set
+        if (t.state != .blocked or t.wakeup_time == null)
-        std.debug.assert(t.state != .blocked or t.wakeup != null);
+            continue;
-        if (t.state == .blocked and t.wakeup.? < task.wakeup.?) {
+        if (task == null or t.wakeup_time.? < task.?.wakeup_time.?)
             task = t;
-        }
     }
     return task;
 }
@@ -123,16 +190,33 @@ fn advanceTimeAndUnblockTasks(self: *Scheduler, new_time: u64) void {
     std.debug.assert(self.current_time < new_time);
     self.current_time = new_time;
     for (self.tasks.items) |*task| {
-        if (task.state == .blocked and task.wakeup.? <= new_time) {
+        if (task.state == .blocked) {
-            task.state = .ready;
+            if (task.wakeup_time) |wakeup_time| {
-            task.wakeup = null;
+                if (wakeup_time <= new_time) {
+                    task.state = .ready;
+                    task.wakeup_time = null;
+                    task.wakeup_tasks = null;
+                }
+            }
         }
     }
 }
 
+fn taskIsWaitingFor(task: *const Task, id: TaskID) bool {
+    if (task.state != .blocked)
+        return false;
+    if (task.wakeup_tasks) |wakeup_tasks| {
+        for (wakeup_tasks) |wakeup_id| {
+            if (wakeup_id == id)
+                return true;
+        }
+    }
+    return false;
+}
+
 fn advanceTimeAndPickTask(self: *Scheduler) ?*Task {
     const task = self.findBlockedTaskWithLowestWakeupTime() orelse return null;
-    self.advanceTimeAndUnblockTasks(task.wakeup.?);
+    self.advanceTimeAndUnblockTasks(task.wakeup_time.?);
     return task;
 }
 
@@ -140,10 +224,12 @@ pub fn scheduleOne(self: *Scheduler) bool {
     const task = self.findTaskWithState(.ready)
         orelse self.advanceTimeAndPickTask()
         orelse return false;
-    self.current = task;
+    const id = task.id;
+    self.current_id = id;
     task.state = .running;
     contextSwitch(&self.regs, &task.regs);
-    if (!task.stackCanaryIsIntact())
+    const current = self.findTaskByID(id) orelse return true;
+    if (!current.stackCanaryIsIntact())
         @panic("Task stack canary was overwritten");
     return true;
 }
@@ -183,16 +269,34 @@ fn contextSwitch(old: *Registers, new: *Registers) void {
 }
 
 fn taskStart(self: *Scheduler) callconv(.c) noreturn {
-    const task = self.current.?;
+    const id = self.current_id.?;
+    const task = self.findTaskByID(id).?;
 
-    task.entry(task.node.processInit()) catch {
+    var failed = false;
-        task.state = .failed;
+    switch (task.entry) {
-        contextSwitch(&task.regs, &self.regs);
+        .main => |entry| {
-        unreachable;
+            const node = task.node;
-    };
+            entry(node.processInit()) catch {
+                failed = true;
+            };
+        },
+        .nested => |entry| entry(task.context.?),
+    }
 
-    task.state = .returned;
+    const current = self.findTaskByID(id).?;
-    contextSwitch(&task.regs, &self.regs);
+    current.state = if (failed) .failed else .returned;
+    current.wakeup_time = null;
+    current.wakeup_tasks = null;
+    if (current.parent_id) |parent_id| {
+        if (self.findTaskByID(parent_id)) |parent| {
+            if (taskIsWaitingFor(parent, current.id)) {
+                parent.state = .ready;
+                parent.wakeup_time = null;
+                parent.wakeup_tasks = null;
+            }
+        }
+    }
+    contextSwitch(&current.regs, &self.regs);
     unreachable;
 }
 
@@ -203,8 +307,38 @@ fn taskReturned() callconv(.c) noreturn {
 
 // Called by the current task to return control to the scheduler
 pub fn sleep(self: *Scheduler, delta_us: u64) void {
-    const task = self.current.?;
+    const current = self.findTaskByID(self.current_id.?).?;
-    task.state = .blocked;
+    current.state = .blocked;
-    task.wakeup = self.current_time + delta_us;
+    current.wakeup_time = self.current_time + delta_us;
-    contextSwitch(&task.regs, &self.regs);
+    current.wakeup_tasks = null;
+    contextSwitch(&current.regs, &self.regs);
+}
+
+fn findCompletedTaskInSet(self: *Scheduler, ids: []const TaskID) !?*Task {
+    for (ids) |id| {
+        const child = self.findTaskByID(id) orelse return error.InvalidHandle;
+        switch (child.state) {
+            .returned, .failed => return child,
+            else => {},
+        }
+    }
+    return null;
+}
+
+pub fn wait(self: *Scheduler, ids: []const TaskID) !TaskID {
+    const id = self.current_id.?;
+
+    while (true) {
+        const child = try self.findCompletedTaskInSet(ids);
+        if (child) |c| {
+            std.debug.assert(c.parent_id == id);
+            return c.id;
+        }
+
+        const task = self.findTaskByID(self.current_id.?);
+        task.state = .blocked;
+        task.wakeup_time = null;
+        task.wakeup_tasks = ids;
+        contextSwitch(&task.regs, &self.regs);
+    }
 }

src/simulator.zig

@@ -7,7 +7,7 @@ const Scheduler = @import("scheduler.zig");
 const Network   = @import("network.zig");
 const Node      = @import("node.zig");
 
-pub const EntryPoint = Scheduler.EntryPoint;
+pub const EntryPoint = Scheduler.MainEntryPoint;
 
 const SpawnOptions = struct {
     stack_size: usize = 64 * 1024,