480 lines
14 KiB
Zig
480 lines
14 KiB
Zig
const std = @import("std");
|
|
const Allocator = std.mem.Allocator;
|
|
const Node = @import("node.zig");
|
|
const Trace = @import("trace.zig").Trace;
|
|
|
|
const Scheduler = @This();
|
|
|
|
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,
|
|
rbp: usize,
|
|
rip: usize,
|
|
rdi: usize,
|
|
};
|
|
|
|
const ContextSwitch = extern struct {
|
|
old: *Registers,
|
|
new: *Registers,
|
|
};
|
|
|
|
const STACK_CANARY_SIZE = 256;
|
|
const STACK_CANARY_BYTE = 0xa5;
|
|
const TASK_RUN_COST_US = 50;
|
|
|
|
pub const TaskID = u64;
|
|
|
|
const State = enum {
|
|
ready,
|
|
running,
|
|
blocked,
|
|
failed,
|
|
returned,
|
|
};
|
|
|
|
const Task = struct {
|
|
id: TaskID,
|
|
regs: Registers,
|
|
stack: []align(16) u8,
|
|
entry: EntryPoint,
|
|
context: ?*const anyopaque,
|
|
state: State,
|
|
node: *Node,
|
|
|
|
// If this is a subtask, parent_id refers to the parent.
|
|
// The cancel flag is set when the parent requests cancellation.
|
|
parent_id: ?TaskID,
|
|
cancel: bool,
|
|
|
|
// 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,
|
|
wakeup_futex: ?*const u32,
|
|
|
|
fn stackCanaryIsIntact(self: *Task) bool {
|
|
const canary = self.stack[0..STACK_CANARY_SIZE];
|
|
for (canary) |byte| {
|
|
if (byte != STACK_CANARY_BYTE)
|
|
return false;
|
|
}
|
|
return true;
|
|
}
|
|
};
|
|
|
|
gpa: Allocator,
|
|
trace: *Trace,
|
|
prng: *std.Random.DefaultPrng,
|
|
tasks: std.ArrayList(Task),
|
|
regs: Registers,
|
|
current_id: ?TaskID,
|
|
current_time: u64,
|
|
next_task_id: u64,
|
|
|
|
pub fn init(self: *Scheduler, gpa: Allocator, trace: *Trace, prng: *std.Random.DefaultPrng) void {
|
|
self.gpa = gpa;
|
|
self.trace = trace;
|
|
self.prng = prng;
|
|
self.tasks = .empty;
|
|
self.current_id = null;
|
|
self.current_time = 0;
|
|
self.next_task_id = 0;
|
|
}
|
|
|
|
pub fn deinit(self: *Scheduler) void {
|
|
for (self.tasks.items) |task| {
|
|
std.heap.page_allocator.free(task.stack);
|
|
}
|
|
self.tasks.deinit(self.gpa);
|
|
}
|
|
|
|
fn traceTaskState(self: *Scheduler, task: *const Task, state: Trace.TaskState, reason: []const u8) void {
|
|
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 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;
|
|
|
|
const stack = try std.heap.page_allocator.alignedAlloc(u8, .fromByteUnits(16), stack_size + STACK_CANARY_SIZE);
|
|
errdefer std.heap.page_allocator.free(stack);
|
|
@memset(stack[0..STACK_CANARY_SIZE], STACK_CANARY_BYTE);
|
|
|
|
var stack_top = @intFromPtr(stack.ptr) + stack.len;
|
|
stack_top &= ~@as(usize, 0xf);
|
|
stack_top -= @sizeOf(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,
|
|
.rip = @intFromPtr(&taskStart),
|
|
.rdi = @intFromPtr(self),
|
|
},
|
|
.stack = stack,
|
|
.entry = entry,
|
|
.context = context,
|
|
.state = .ready,
|
|
.wakeup_time = null,
|
|
.wakeup_tasks = null,
|
|
.wakeup_futex = null,
|
|
.node = node,
|
|
.parent_id = parent_id,
|
|
.cancel = false,
|
|
};
|
|
|
|
try self.tasks.append(self.gpa, task);
|
|
const appended = &self.tasks.items[self.tasks.items.len - 1];
|
|
self.trace.taskSpawned(appended.id, appended.node, appended.parent_id);
|
|
self.traceTaskState(appended, .ready, "spawned");
|
|
return appended;
|
|
}
|
|
|
|
// 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];
|
|
|
|
self.trace.taskRemoved(task.id, task.node);
|
|
std.heap.page_allocator.free(task.stack);
|
|
_ = self.tasks.orderedRemove(index);
|
|
}
|
|
|
|
fn countTasksWithState(self: *Scheduler, state: State) usize {
|
|
var count: usize = 0;
|
|
for (self.tasks.items) |*task| {
|
|
if (task.state == state)
|
|
count += 1;
|
|
}
|
|
return count;
|
|
}
|
|
|
|
fn findNthTaskWithState(self: *Scheduler, state: State, index: usize) ?*Task {
|
|
var remaining = index;
|
|
for (self.tasks.items) |*task| {
|
|
if (task.state != state)
|
|
continue;
|
|
if (remaining == 0)
|
|
return task;
|
|
remaining -= 1;
|
|
}
|
|
return null;
|
|
}
|
|
|
|
fn pickReadyTask(self: *Scheduler) ?*Task {
|
|
const ready_count = self.countTasksWithState(.ready);
|
|
if (ready_count == 0)
|
|
return null;
|
|
|
|
const random = self.prng.random();
|
|
const index = random.uintLessThan(usize, ready_count);
|
|
return self.findNthTaskWithState(.ready, index).?;
|
|
}
|
|
|
|
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: ?*Task = null;
|
|
for (self.tasks.items) |*t| {
|
|
if (t.state != .blocked or t.wakeup_time == null)
|
|
continue;
|
|
if (task == null or t.wakeup_time.? < task.?.wakeup_time.?)
|
|
task = t;
|
|
}
|
|
return task;
|
|
}
|
|
|
|
fn advanceTimeAndUnblockTasks(self: *Scheduler, new_time: u64) void {
|
|
std.debug.assert(self.current_time < new_time);
|
|
const old_time = self.current_time;
|
|
self.current_time = new_time;
|
|
for (self.tasks.items) |task| {
|
|
task.node.local_time = @max(task.node.local_time, new_time);
|
|
}
|
|
self.trace.timeAdvanced(old_time, new_time);
|
|
for (self.tasks.items) |*task| {
|
|
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;
|
|
task.wakeup_futex = null;
|
|
self.traceTaskState(task, .ready, "sleep elapsed");
|
|
}
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
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_time.?);
|
|
return task;
|
|
}
|
|
|
|
pub fn scheduleOne(self: *Scheduler) bool {
|
|
const task = self.pickReadyTask() orelse self.advanceTimeAndPickTask() orelse return false;
|
|
const id = task.id;
|
|
self.current_id = id;
|
|
task.state = .running;
|
|
self.traceTaskState(task, .running, "scheduled");
|
|
contextSwitch(&self.regs, &task.regs);
|
|
self.trace.leaveTask();
|
|
const current = self.findTaskByID(id) orelse return true;
|
|
if (!current.stackCanaryIsIntact())
|
|
@panic("Task stack canary was overwritten");
|
|
return true;
|
|
}
|
|
|
|
fn contextSwitch(old: *Registers, new: *Registers) void {
|
|
asm volatile (
|
|
\\ movq 0(%%rsi), %%rax
|
|
\\ movq 8(%%rsi), %%rcx
|
|
\\ leaq 0f(%%rip), %%rdx
|
|
\\ movq %%rsp, 0(%%rax)
|
|
\\ movq %%rbp, 8(%%rax)
|
|
\\ movq %%rdx, 16(%%rax)
|
|
\\ movq 0(%%rcx), %%rsp
|
|
\\ movq 8(%%rcx), %%rbp
|
|
\\ movq 24(%%rcx), %%rdi
|
|
\\ jmpq *16(%%rcx)
|
|
\\0:
|
|
:
|
|
: [message] "{rsi}" (&ContextSwitch{ .old = old, .new = new }),
|
|
: .{
|
|
.rax = true,
|
|
.rcx = true,
|
|
.rdx = true,
|
|
.rbx = true,
|
|
.rsi = true,
|
|
.rdi = true,
|
|
.r8 = true,
|
|
.r9 = true,
|
|
.r10 = true,
|
|
.r11 = true,
|
|
.r12 = true,
|
|
.r13 = true,
|
|
.r14 = true,
|
|
.r15 = true,
|
|
.memory = true,
|
|
});
|
|
}
|
|
|
|
fn taskStart(self: *Scheduler) callconv(.c) noreturn {
|
|
const id = self.current_id.?;
|
|
const task = self.findTaskByID(id).?;
|
|
self.trace.enterTask(task.id, task.node);
|
|
|
|
var failed = false;
|
|
switch (task.entry) {
|
|
.main => |entry| {
|
|
const node = task.node;
|
|
entry(node.processInit()) catch {
|
|
failed = true;
|
|
};
|
|
},
|
|
.nested => |entry| entry(task.context.?),
|
|
}
|
|
|
|
const current = self.findTaskByID(id).?;
|
|
current.state = if (failed) .failed else .returned;
|
|
current.wakeup_time = null;
|
|
current.wakeup_tasks = null;
|
|
current.wakeup_futex = null;
|
|
current.node.local_time += TASK_RUN_COST_US;
|
|
self.traceTaskState(current, if (failed) .failed else .returned, "task completed");
|
|
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;
|
|
parent.wakeup_futex = null;
|
|
self.traceTaskState(parent, .ready, "child completed");
|
|
}
|
|
}
|
|
}
|
|
contextSwitch(¤t.regs, &self.regs);
|
|
unreachable;
|
|
}
|
|
|
|
// Dummy return address on the task's stack. Should never be reached.
|
|
fn taskReturned() callconv(.c) noreturn {
|
|
@panic("Task returned through the fake return address");
|
|
}
|
|
|
|
// Called by the current task to return control to the scheduler
|
|
pub fn sleep(self: *Scheduler, delta_us: u64) !void {
|
|
const current = self.findTaskByID(self.current_id.?).?;
|
|
const id = current.id;
|
|
current.state = .blocked;
|
|
current.wakeup_time = self.current_time + delta_us;
|
|
current.wakeup_tasks = null;
|
|
current.wakeup_futex = null;
|
|
current.node.local_time += TASK_RUN_COST_US;
|
|
self.traceTaskState(current, .sleeping, "sleep");
|
|
contextSwitch(¤t.regs, &self.regs);
|
|
const resumed = self.findTaskByID(id).?;
|
|
self.trace.enterTask(resumed.id, resumed.node);
|
|
try self.checkCancel();
|
|
}
|
|
|
|
pub fn futexWait(self: *Scheduler, ptr: *const u32, expected: u32) !void {
|
|
try self.checkCancel();
|
|
self.futexWaitUncancelable(ptr, expected);
|
|
try self.checkCancel();
|
|
}
|
|
|
|
pub fn futexWaitUncancelable(self: *Scheduler, ptr: *const u32, expected: u32) void {
|
|
if (@atomicLoad(u32, ptr, .seq_cst) != expected)
|
|
return;
|
|
|
|
const current = self.findTaskByID(self.current_id.?).?;
|
|
const id = current.id;
|
|
current.state = .blocked;
|
|
current.wakeup_time = null;
|
|
current.wakeup_tasks = null;
|
|
current.wakeup_futex = ptr;
|
|
current.node.local_time += TASK_RUN_COST_US;
|
|
self.traceTaskState(current, .waiting_futex, "futex wait");
|
|
contextSwitch(¤t.regs, &self.regs);
|
|
const resumed = self.findTaskByID(id).?;
|
|
self.trace.enterTask(resumed.id, resumed.node);
|
|
}
|
|
|
|
pub fn futexWake(self: *Scheduler, ptr: *const u32, max_waiters: u32) void {
|
|
var woken: u32 = 0;
|
|
for (self.tasks.items) |*task| {
|
|
if (woken == max_waiters)
|
|
break;
|
|
if (task.state == .blocked and task.wakeup_futex == ptr) {
|
|
task.state = .ready;
|
|
task.wakeup_time = null;
|
|
task.wakeup_tasks = null;
|
|
task.wakeup_futex = null;
|
|
self.traceTaskState(task, .ready, "futex wake");
|
|
woken += 1;
|
|
}
|
|
}
|
|
}
|
|
|
|
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;
|
|
task.wakeup_futex = null;
|
|
task.node.local_time += TASK_RUN_COST_US;
|
|
self.traceTaskState(task, .waiting_task, "wait");
|
|
contextSwitch(&task.regs, &self.regs);
|
|
const resumed = self.findTaskByID(id).?;
|
|
self.trace.enterTask(resumed.id, resumed.node);
|
|
try self.checkCancel();
|
|
}
|
|
}
|
|
|
|
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 checkCancel(self: *Scheduler) !void {
|
|
const task = self.findTaskByID(self.current_id.?).?;
|
|
if (!task.cancel)
|
|
return;
|
|
|
|
task.cancel = false;
|
|
return error.Canceled;
|
|
}
|
|
|
|
pub fn recancel(self: *Scheduler) void {
|
|
const task = self.findTaskByID(self.current_id.?).?;
|
|
task.cancel = true;
|
|
}
|