authorgravatar for andrew@ziglang.orgAndrew Kelley <andrew@ziglang.org> 2025-03-26 19:24:37-07:00
committergravatar for andrew@ziglang.orgAndrew Kelley <andrew@ziglang.org> 2025-07-20 10:38:38-07:00
log07ee4977dab9696a511fab69fe136280c9bfa0fc
treefdc7fe2c476d08c78f9d234d0cdff5a49c767ec1
parent6ae18714b74a5267608e931b46db20be637ccb03

demo: single-threaded green threads implementation

2 files changed, 213 insertions(+), 0 deletions(-)

lib/std/Io.zig+3
......@@ -914,6 +914,8 @@ test {
914914
915915const Io = @This();
916916
917pub const EventLoop = @import("Io/EventLoop.zig");
918
917919userdata: ?*anyopaque,
918920vtable: *const VTable,
919921
......@@ -925,6 +927,7 @@ pub const VTable = struct {
925927 userdata: ?*anyopaque,
926928 /// The pointer of this slice is an "eager" result value.
927929 /// The length is the size in bytes of the result type.
930 /// This pointer's lifetime expires directly after the call to this function.
928931 eager_result: []u8,
929932 /// Passed to `start`.
930933 context: ?*anyopaque,
lib/std/Io/EventLoop.zig created+210
......@@ -0,0 +1,210 @@
1const std = @import("../std.zig");
2const assert = std.debug.assert;
3const Allocator = std.mem.Allocator;
4const Io = std.Io;
5const EventLoop = @This();
6
7gpa: Allocator,
8queue: std.DoublyLinkedList(void),
9free: std.DoublyLinkedList(void),
10main_fiber_buffer: [@sizeOf(Fiber) + max_result_len]u8 align(@alignOf(Fiber)),
11
12threadlocal var current_fiber: *Fiber = undefined;
13
14const max_result_len = 64;
15const min_stack_size = 4 * 1024 * 1024;
16
17const 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
39pub 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
49fn 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
62fn 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.
80fn exit(el: *EventLoop, optional_fiber: ?*Fiber) noreturn {
81 yield(el, optional_fiber);
82 @panic("TODO recycle the fiber");
83}
84
85fn schedule(el: *EventLoop, fiber: *Fiber) void {
86 el.queue.append(&fiber.queue_node);
87}
88
89fn myFiber(el: *EventLoop) *Fiber {
90 _ = el;
91 return current_fiber;
92}
93
94const 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
104const contextSwitch: *const fn (old: *Regs, new: *Regs) callconv(.c) void = @ptrCast(&contextSwitch_naked);
105
106noinline 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
128fn popRet() callconv(.naked) void {
129 asm volatile (
130 \\pop %%rdi
131 \\ret
132 );
133}
134
135pub 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
180const 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
195pub 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}