| ... | ... | @@ -0,0 +1,210 @@ |
| 1 | const std = @import("../std.zig"); |
| 2 | const assert = std.debug.assert; |
| 3 | const Allocator = std.mem.Allocator; |
| 4 | const Io = std.Io; |
| 5 | const EventLoop = @This(); |
| 6 | |
| 7 | gpa: Allocator, |
| 8 | queue: std.DoublyLinkedList(void), |
| 9 | free: std.DoublyLinkedList(void), |
| 10 | main_fiber_buffer: [@sizeOf(Fiber) + max_result_len]u8 align(@alignOf(Fiber)), |
| 11 | |
| 12 | threadlocal var current_fiber: *Fiber = undefined; |
| 13 | |
| 14 | const max_result_len = 64; |
| 15 | const min_stack_size = 4 * 1024 * 1024; |
| 16 | |
| 17 | const Fiber = struct { |
| 18 | regs: Regs, |
| 19 | awaiter: ?*Fiber, |
| 20 | queue_node: std.DoublyLinkedList(void).Node, |
| 21 | |
| 22 | const finished: ?*Fiber = @ptrFromInt(std.mem.alignBackward(usize, std.math.maxInt(usize), @alignOf(Fiber))); |
| 23 | |
| 24 | fn resultPointer(f: *Fiber) [*]u8 { |
| 25 | const base: [*]u8 = @ptrCast(f); |
| 26 | return base + @sizeOf(Fiber); |
| 27 | } |
| 28 | |
| 29 | fn stackEndPointer(f: *Fiber) [*]u8 { |
| 30 | const base: [*]u8 = @ptrCast(f); |
| 31 | return base + std.mem.alignForward( |
| 32 | usize, |
| 33 | @sizeOf(Fiber) + max_result_len + min_stack_size, |
| 34 | std.heap.page_size_max, |
| 35 | ); |
| 36 | } |
| 37 | }; |
| 38 | |
| 39 | pub fn init(el: *EventLoop, gpa: Allocator) void { |
| 40 | el.* = .{ |
| 41 | .gpa = gpa, |
| 42 | .queue = .{}, |
| 43 | .free = .{}, |
| 44 | .main_fiber_buffer = undefined, |
| 45 | }; |
| 46 | current_fiber = @ptrCast(&el.main_fiber_buffer); |
| 47 | } |
| 48 | |
| 49 | fn allocateFiber(el: *EventLoop, result_len: usize) error{OutOfMemory}!*Fiber { |
| 50 | assert(result_len <= max_result_len); |
| 51 | const free_node = el.free.pop() orelse { |
| 52 | const n = std.mem.alignForward( |
| 53 | usize, |
| 54 | @sizeOf(Fiber) + max_result_len + min_stack_size, |
| 55 | std.heap.page_size_max, |
| 56 | ); |
| 57 | return @alignCast(@ptrCast(try el.gpa.alignedAlloc(u8, @alignOf(Fiber), n))); |
| 58 | }; |
| 59 | return @fieldParentPtr("queue_node", free_node); |
| 60 | } |
| 61 | |
| 62 | fn yield(el: *EventLoop, optional_fiber: ?*Fiber) void { |
| 63 | if (optional_fiber) |fiber| { |
| 64 | const old = &current_fiber.regs; |
| 65 | current_fiber = fiber; |
| 66 | contextSwitch(old, &fiber.regs); |
| 67 | return; |
| 68 | } |
| 69 | if (el.queue.pop()) |node| { |
| 70 | const fiber: *Fiber = @fieldParentPtr("queue_node", node); |
| 71 | const old = &current_fiber.regs; |
| 72 | current_fiber = fiber; |
| 73 | contextSwitch(old, &fiber.regs); |
| 74 | return; |
| 75 | } |
| 76 | @panic("everything is done"); |
| 77 | } |
| 78 | |
| 79 | /// Equivalent to calling `yield` and then giving the fiber back to the event loop. |
| 80 | fn exit(el: *EventLoop, optional_fiber: ?*Fiber) noreturn { |
| 81 | yield(el, optional_fiber); |
| 82 | @panic("TODO recycle the fiber"); |
| 83 | } |
| 84 | |
| 85 | fn schedule(el: *EventLoop, fiber: *Fiber) void { |
| 86 | el.queue.append(&fiber.queue_node); |
| 87 | } |
| 88 | |
| 89 | fn myFiber(el: *EventLoop) *Fiber { |
| 90 | _ = el; |
| 91 | return current_fiber; |
| 92 | } |
| 93 | |
| 94 | const Regs = extern struct { |
| 95 | rsp: usize, |
| 96 | r15: usize, |
| 97 | r14: usize, |
| 98 | r13: usize, |
| 99 | r12: usize, |
| 100 | rbx: usize, |
| 101 | rbp: usize, |
| 102 | }; |
| 103 | |
| 104 | const contextSwitch: *const fn (old: *Regs, new: *Regs) callconv(.c) void = @ptrCast(&contextSwitch_naked); |
| 105 | |
| 106 | noinline fn contextSwitch_naked() callconv(.naked) void { |
| 107 | asm volatile ( |
| 108 | \\movq %%rsp, 0x00(%%rdi) |
| 109 | \\movq %%r15, 0x08(%%rdi) |
| 110 | \\movq %%r14, 0x10(%%rdi) |
| 111 | \\movq %%r13, 0x18(%%rdi) |
| 112 | \\movq %%r12, 0x20(%%rdi) |
| 113 | \\movq %%rbx, 0x28(%%rdi) |
| 114 | \\movq %%rbp, 0x30(%%rdi) |
| 115 | \\ |
| 116 | \\movq 0x00(%%rsi), %%rsp |
| 117 | \\movq 0x08(%%rsi), %%r15 |
| 118 | \\movq 0x10(%%rsi), %%r14 |
| 119 | \\movq 0x18(%%rsi), %%r13 |
| 120 | \\movq 0x20(%%rsi), %%r12 |
| 121 | \\movq 0x28(%%rsi), %%rbx |
| 122 | \\movq 0x30(%%rsi), %%rbp |
| 123 | \\ |
| 124 | \\ret |
| 125 | ); |
| 126 | } |
| 127 | |
| 128 | fn popRet() callconv(.naked) void { |
| 129 | asm volatile ( |
| 130 | \\pop %%rdi |
| 131 | \\ret |
| 132 | ); |
| 133 | } |
| 134 | |
| 135 | pub fn @"async"( |
| 136 | userdata: ?*anyopaque, |
| 137 | eager_result: []u8, |
| 138 | context: ?*anyopaque, |
| 139 | start: *const fn (context: ?*anyopaque, result: *anyopaque) void, |
| 140 | ) ?*std.Io.AnyFuture { |
| 141 | const event_loop: *EventLoop = @alignCast(@ptrCast(userdata)); |
| 142 | const fiber = event_loop.allocateFiber(eager_result.len) catch { |
| 143 | start(context, eager_result.ptr); |
| 144 | return null; |
| 145 | }; |
| 146 | fiber.awaiter = null; |
| 147 | fiber.queue_node = .{ .data = {} }; |
| 148 | |
| 149 | const closure: *AsyncClosure = @ptrFromInt(std.mem.alignBackward( |
| 150 | usize, |
| 151 | @intFromPtr(fiber.stackEndPointer() - @sizeOf(AsyncClosure)), |
| 152 | @alignOf(AsyncClosure), |
| 153 | )); |
| 154 | closure.* = .{ |
| 155 | .event_loop = event_loop, |
| 156 | .context = context, |
| 157 | .fiber = fiber, |
| 158 | .start = start, |
| 159 | }; |
| 160 | const stack_end_ptr: [*]align(16) usize = @alignCast(@ptrCast(closure)); |
| 161 | (stack_end_ptr - 1)[0] = 0; |
| 162 | (stack_end_ptr - 2)[0] = @intFromPtr(&AsyncClosure.call); |
| 163 | (stack_end_ptr - 3)[0] = @intFromPtr(closure); |
| 164 | (stack_end_ptr - 4)[0] = @intFromPtr(&popRet); |
| 165 | |
| 166 | fiber.regs = .{ |
| 167 | .rsp = @intFromPtr(stack_end_ptr - 4), |
| 168 | .r15 = 0, |
| 169 | .r14 = 0, |
| 170 | .r13 = 0, |
| 171 | .r12 = 0, |
| 172 | .rbx = 0, |
| 173 | .rbp = 0, |
| 174 | }; |
| 175 | |
| 176 | event_loop.schedule(fiber); |
| 177 | return @ptrCast(fiber); |
| 178 | } |
| 179 | |
| 180 | const AsyncClosure = struct { |
| 181 | _: void align(16) = {}, |
| 182 | event_loop: *EventLoop, |
| 183 | context: ?*anyopaque, |
| 184 | fiber: *EventLoop.Fiber, |
| 185 | start: *const fn (context: ?*anyopaque, result: *anyopaque) void, |
| 186 | |
| 187 | fn call(closure: *AsyncClosure) callconv(.c) void { |
| 188 | std.log.debug("wrap called in async", .{}); |
| 189 | closure.start(closure.context, closure.fiber.resultPointer()); |
| 190 | const awaiter = @atomicRmw(?*EventLoop.Fiber, &closure.fiber.awaiter, .Xchg, EventLoop.Fiber.finished, .seq_cst); |
| 191 | closure.event_loop.exit(awaiter); |
| 192 | } |
| 193 | }; |
| 194 | |
| 195 | pub fn @"await"(userdata: ?*anyopaque, any_future: *std.Io.AnyFuture, result: []u8) void { |
| 196 | const event_loop: *EventLoop = @alignCast(@ptrCast(userdata)); |
| 197 | const future_fiber: *EventLoop.Fiber = @alignCast(@ptrCast(any_future)); |
| 198 | const result_src = future_fiber.resultPointer()[0..result.len]; |
| 199 | const my_fiber = event_loop.myFiber(); |
| 200 | |
| 201 | const prev = @atomicRmw(?*EventLoop.Fiber, &future_fiber.awaiter, .Xchg, my_fiber, .seq_cst); |
| 202 | if (prev == EventLoop.Fiber.finished) { |
| 203 | @memcpy(result, result_src); |
| 204 | return; |
| 205 | } |
| 206 | event_loop.yield(prev); |
| 207 | // Resumed when the value is available. |
| 208 | std.log.debug("yield returned in await", .{}); |
| 209 | @memcpy(result, result_src); |
| 210 | } |