cozis/Zigmulator
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

scheduler: Implement spawning of nested tasks

Browse Source
This commit is contained in:
cozis
2026-06-04 13:32:12 +02:00
parent aec8ff5154
commit 1f5e185bec
2 changed files with 164 additions and 30 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/scheduler.zig
+163 -29
View File
@@ -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(&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.?;
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(&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);
}
}