From d6e098cf40dbbc36f7e5dd024424f72749d869e0 Mon Sep 17 00:00:00 2001 From: cozis Date: Sat, 20 Jun 2026 14:14:19 +0200 Subject: [PATCH] scheduler: support crashing and restarting tasks --- src/network.zig | 11 +++++ src/node.zig | 2 +- src/scheduler.zig | 121 +++++++++++++++++++++++++++++++++++++++++----- src/simulator.zig | 2 +- src/trace.zig | 28 +++++------ 5 files changed, 135 insertions(+), 29 deletions(-) diff --git a/src/network.zig b/src/network.zig index c396d1a..af96bb8 100644 --- a/src/network.zig +++ b/src/network.zig @@ -91,6 +91,8 @@ pub const Host = struct { while (self.listen_list) |s| { self.closeListenSocket(s); } + + self.network.unregisterHost(self); } fn linkConnSocket(self: *Host, socket: *ConnSocket) void { @@ -313,6 +315,15 @@ pub fn registerHost(self: *Network, host: *Host) Allocator.Error!void { host.network = self; } +pub fn unregisterHost(self: *Network, host: *Host) void { + for (self.hosts.items, 0..) |item, i| { + if (item == host) { + _ = self.hosts.swapRemove(i); + return; + } + } +} + pub fn findHostByIPv4(self: *Network, ipv4: u32) ?*Host { for (self.hosts.items) |host| { if (host.isAddressAvailable(ipv4)) diff --git a/src/node.zig b/src/node.zig index 0144f21..d901e73 100644 --- a/src/node.zig +++ b/src/node.zig @@ -194,7 +194,7 @@ pub fn despawn(self: *Node, id: TaskID) void { } pub fn cancel(self: *Node, id: TaskID) void { - self.scheduler.cancel(id); + _ = self.scheduler.cancel(id); } pub fn checkCancel(self: *Node) !void { diff --git a/src/scheduler.zig b/src/scheduler.zig index 29f502b..2b70095 100644 --- a/src/scheduler.zig +++ b/src/scheduler.zig @@ -37,6 +37,11 @@ const State = enum { blocked, failed, returned, + crashed, + + pub fn isCompleted(self: State) bool { + return self == .failed or self == .crashed or self == .returned; + } }; const Task = struct { @@ -100,8 +105,8 @@ fn traceTaskState(self: *Scheduler, task: *const Task, state: Trace.TaskState, r self.trace.taskState(task.id, task.node, state, reason); } -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 spawn(self: *Scheduler, node: *Node, entry: MainEntryPoint, stack_size: usize) !TaskID { + return (try self.spawnInner(node, .{ .main = entry }, stack_size, null, null)).id; } pub fn spawnNested(self: *Scheduler, node: *Node, entry: NestedEntryPoint, context: *const anyopaque) !TaskID { @@ -166,14 +171,81 @@ fn spawnInner( // The intended use-case is undoing a previous call .spawnNested to // in case of errors. pub fn despawnNested(self: *Scheduler, id: TaskID) void { + self.removeTask(id); +} + +fn removeTask(self: *Scheduler, id: TaskID) void { const index = self.findTaskIndexByID(id) orelse return; const task = &self.tasks.items[index]; - + if (self.current_id != null and self.current_id.? == task.id) + self.current_id = null; self.trace.taskRemoved(task.id, task.node); std.heap.page_allocator.free(task.stack); _ = self.tasks.orderedRemove(index); } +fn findChildOf(self: *Scheduler, parent_id: TaskID) ?TaskID { + for (self.tasks.items) |task| { + if (task.parent_id == parent_id) + return task.id; + } + return null; +} + +fn removeDescendants(self: *Scheduler, parent_id: TaskID) void { + while (self.findChildOf(parent_id)) |child_id| { + self.removeDescendants(child_id); + self.removeTask(child_id); + } +} + +fn markTaskCrashed(self: *Scheduler, id: TaskID) void { + const task = self.findTaskByID(id) orelse return; + + if (self.current_id != null and self.current_id.? == task.id) + self.current_id = null; + + task.state = .crashed; + task.wakeup_time = null; + task.wakeup_tasks = null; + task.wakeup_futex = null; + task.cancel = false; + self.traceTaskState(task, .crashed, "node crashed"); +} + +// Tears down a task tree. Subtasks are always removed; the root is either +// removed outright or kept and marked crashed as a record that the process died. +pub fn removeTaskTree(self: *Scheduler, root_id: TaskID, keep_root_but_mark_crashed: bool) void { + self.removeDescendants(root_id); + if (keep_root_but_mark_crashed) + self.markTaskCrashed(root_id) + else + self.removeTask(root_id); +} + +pub fn terminate(self: *Scheduler, id: TaskID) void { + _ = self.cancel(id); + while (true) { + const task = self.findTaskByID(id) orelse break; + switch (task.state) { + .ready => {}, + .blocked => { + // Don't wait for the task to wake up on its own; force the + // blocking operation to return Canceled instead. + _ = self.cancel(id); + }, + .running, .failed, .returned, .crashed => break, + } + + const ready = self.findTaskByID(id) orelse break; + if (ready.state != .ready) + break; + + if (!self.runTask(id)) + break; + } +} + fn countTasksWithState(self: *Scheduler, state: State) usize { var count: usize = 0; for (self.tasks.items) |*task| { @@ -252,6 +324,14 @@ fn advanceTimeAndUnblockTasks(self: *Scheduler, new_time: u64) void { } } +// Jumps the clock forward to `new_time`, waking any tasks whose sleep elapsed. +// Unlike the task-driven time advance, this lets an external timer (e.g. the +// fault injector) move time when no task is scheduled to do so. +pub fn advanceTimeTo(self: *Scheduler, new_time: u64) void { + if (new_time > self.current_time) + self.advanceTimeAndUnblockTasks(new_time); +} + fn taskIsWaitingFor(task: *const Task, id: TaskID) bool { if (task.state != .blocked) return false; @@ -273,11 +353,19 @@ fn advanceTimeAndPickTask(self: *Scheduler) ?*Task { pub fn scheduleOne(self: *Scheduler) bool { const task = self.pickReadyTask() orelse self.advanceTimeAndPickTask() orelse return false; const id = task.id; + return self.runTask(id); +} + +fn runTask(self: *Scheduler, id: TaskID) bool { + const task = self.findTaskByID(id) orelse return false; + std.debug.assert(task.state == .ready); + self.current_id = id; task.state = .running; self.traceTaskState(task, .running, "scheduled"); contextSwitch(&self.regs, &task.regs); self.trace.leaveTask(); + self.current_id = null; const current = self.findTaskByID(id) orelse return true; if (!current.stackCanaryIsIntact()) @panic("Task stack canary was overwritten"); @@ -351,6 +439,8 @@ fn taskStart(self: *Scheduler) callconv(.c) noreturn { self.traceTaskState(parent, .ready, "child completed"); } } + } else { + self.removeDescendants(current.id); } contextSwitch(¤t.regs, &self.regs); unreachable; @@ -451,17 +541,22 @@ pub fn wait(self: *Scheduler, ids: []const TaskID) !TaskID { } } -pub fn cancel(self: *Scheduler, id: TaskID) void { - const child = self.findTaskByID(id) orelse return; - std.debug.assert(child.parent_id == self.current_id); - child.cancel = true; - if (child.state == .blocked) { - child.state = .ready; - child.wakeup_time = null; - child.wakeup_tasks = null; - child.wakeup_futex = null; - self.traceTaskState(child, .ready, "cancel"); +pub fn cancel(self: *Scheduler, id: TaskID) bool { + const task = self.findTaskByID(id) orelse return false; + if (self.current_id) |current_id| + std.debug.assert(task.parent_id == current_id); + if (task.state.isCompleted()) + return false; + + task.cancel = true; + if (task.state == .blocked) { + task.state = .ready; + task.wakeup_time = null; + task.wakeup_tasks = null; + task.wakeup_futex = null; + self.traceTaskState(task, .ready, "cancel"); } + return true; } pub fn checkCancel(self: *Scheduler) !void { diff --git a/src/simulator.zig b/src/simulator.zig index 3fcadb5..c692b28 100644 --- a/src/simulator.zig +++ b/src/simulator.zig @@ -304,7 +304,7 @@ pub fn spawn(self: *Simulator, command: []const u8, options: SpawnOptions) Spawn try self.nodes.append(self.gpa, node); errdefer _ = self.nodes.swapRemove(self.nodes.items.len - 1); - try self.scheduler.spawn(node, entry, options.stack_size); + _ = try self.scheduler.spawn(node, entry, options.stack_size); } pub fn scheduleOne(self: *Simulator) bool { diff --git a/src/trace.zig b/src/trace.zig index 49a2791..c8c7b4b 100644 --- a/src/trace.zig +++ b/src/trace.zig @@ -5,7 +5,6 @@ const Scheduler = @import("scheduler.zig"); const TaskID = Scheduler.TaskID; pub const Trace = struct { - pub const TaskState = enum { ready, running, @@ -14,16 +13,18 @@ pub const Trace = struct { waiting_task, failed, returned, + crashed, pub fn toString(self: TaskState) []const u8 { return switch (self) { - .ready => "ready", - .running => "running", - .sleeping => "sleeping", + .ready => "ready", + .running => "running", + .sleeping => "sleeping", .waiting_futex => "waiting-futex", - .waiting_task => "waiting-task", - .failed => "failed", - .returned => "returned", + .waiting_task => "waiting-task", + .failed => "failed", + .returned => "returned", + .crashed => "crashed", }; } }; @@ -99,7 +100,7 @@ pub const Trace = struct { if (self.file == null) return; const global_time = self.scheduler.current_time; - const local_time = node.local_time; + const local_time = node.local_time; var line_buffer: [256]u8 = undefined; const line = std.fmt.bufPrint(&line_buffer, @@ -135,7 +136,6 @@ pub const Trace = struct { \\ , .{ to, to, from, to }) catch return; self.write(line); - } pub fn beginIO(self: *Trace, disk: bool, source: std.builtin.SourceLocation) Pending { @@ -156,7 +156,7 @@ pub const Trace = struct { const task = self.task.?; const global_time = self.scheduler.current_time; - const local_time = node.local_time; + const local_time = node.local_time; const trace_id = self.next_trace_id; self.next_trace_id += 1; @@ -200,14 +200,14 @@ pub const Trace = struct { var line_buffer: [256]u8 = undefined; const prefix = std.fmt.bufPrint(&line_buffer, \\{{"time":{},"global_time":{},"event":"io_complete","id":{},"node":{},"task":{},"op": - , .{ local_time, global_time, pending_trace.trace_id, pending_trace.node_id, pending_trace.task_id }) catch return; + , .{ local_time, global_time, pending_trace.trace_id, pending_trace.node_id, pending_trace.task_id }) catch return; self.write(prefix); self.writeJsonString(pending_trace.op); self.write(",\"resource\":"); self.writeJsonString(if (pending_trace.disk) "disk" else "io"); const suffix = std.fmt.bufPrint(&line_buffer, \\,"start":{},"end":{},"result": - , .{ pending_trace.local_time, local_time }) catch return; + , .{ pending_trace.local_time, local_time }) catch return; self.write(suffix); var result_buffer: [96]u8 = undefined; self.writeJsonString(resultText(&result_buffer, result)); @@ -217,12 +217,12 @@ pub const Trace = struct { var disk_line_buffer: [256]u8 = undefined; const disk_prefix = std.fmt.bufPrint(&disk_line_buffer, \\{{"time":{},"global_time":{},"event":"disk","node":{},"op": - , .{ pending_trace.local_time, global_time, pending_trace.node_id }) catch return; + , .{ pending_trace.local_time, global_time, pending_trace.node_id }) catch return; self.write(disk_prefix); self.writeJsonString(pending_trace.op); const middle = std.fmt.bufPrint(&disk_line_buffer, \\,"start":{},"end":{},"detail": - , .{ pending_trace.local_time, local_time }) catch return; + , .{ pending_trace.local_time, local_time }) catch return; self.write(middle); self.writeJsonString(resultText(&result_buffer, result)); self.write("}\n");