diff --git a/src/scheduler.zig b/src/scheduler.zig index 56340fb..25b2df1 100644 --- a/src/scheduler.zig +++ b/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 - std.debug.assert(t.state != .blocked or t.wakeup != null); - if (t.state == .blocked and t.wakeup.? < task.wakeup.?) { + if (t.state != .blocked or t.wakeup_time == null) + continue; + 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) { - task.state = .ready; - task.wakeup = null; + if (task.state == .blocked) { + if (task.wakeup_time) |wakeup_time| { + 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 { - task.state = .failed; - contextSwitch(&task.regs, &self.regs); - unreachable; - }; + var failed = false; + switch (task.entry) { + .main => |entry| { + const node = task.node; + entry(node.processInit()) catch { + failed = true; + }; + }, + .nested => |entry| entry(task.context.?), + } - task.state = .returned; - contextSwitch(&task.regs, &self.regs); + const current = self.findTaskByID(id).?; + 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(¤t.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.?; - task.state = .blocked; - task.wakeup = self.current_time + delta_us; - contextSwitch(&task.regs, &self.regs); + const current = self.findTaskByID(self.current_id.?).?; + current.state = .blocked; + current.wakeup_time = self.current_time + delta_us; + current.wakeup_tasks = null; + contextSwitch(¤t.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); + } } diff --git a/src/simulator.zig b/src/simulator.zig index 9cfc6fd..4879286 100644 --- a/src/simulator.zig +++ b/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,