authorgravatar for andrew@ziglang.orgAndrew Kelley <andrew@ziglang.org> 2025-04-03 21:15:19-07:00
committergravatar for andrew@ziglang.orgAndrew Kelley <andrew@ziglang.org> 2025-10-02 16:30:59-07:00
log0c1e102e975f062fdb7b513256598a0f9846a543
treece58f2e2cab500d0287f5337d46a343bf44f76f0
parentb4fb7df1404f6fb20405ef31abf68f61dd13e4cb

std.Io.EventLoop: implement select

2 files changed, 72 insertions(+), 22 deletions(-)

lib/std/Io/EventLoop.zig+71-21
......@@ -485,6 +485,7 @@ const SwitchMessage = struct {
485485 reschedule,
486486 recycle,
487487 register_awaiter: *?*Fiber,
488 register_select: []const *Io.AnyFuture,
488489 mutex_lock: struct {
489490 prev_state: Io.Mutex.State,
490491 mutex: *Io.Mutex,
......@@ -514,13 +515,21 @@ const SwitchMessage = struct {
514515 .register_awaiter => |awaiter| {
515516 const prev_fiber: *Fiber = @alignCast(@fieldParentPtr("context", message.contexts.prev));
516517 assert(prev_fiber.queue_next == null);
517 if (@atomicRmw(
518 ?*Fiber,
519 awaiter,
520 .Xchg,
521 prev_fiber,
522 .acq_rel,
523 ) == Fiber.finished) el.schedule(thread, .{ .head = prev_fiber, .tail = prev_fiber });
518 if (@atomicRmw(?*Fiber, awaiter, .Xchg, prev_fiber, .acq_rel) == Fiber.finished)
519 el.schedule(thread, .{ .head = prev_fiber, .tail = prev_fiber });
520 },
521 .register_select => |futures| {
522 const prev_fiber: *Fiber = @alignCast(@fieldParentPtr("context", message.contexts.prev));
523 assert(prev_fiber.queue_next == null);
524 for (futures) |any_future| {
525 const future_fiber: *Fiber = @alignCast(@ptrCast(any_future));
526 if (@atomicRmw(?*Fiber, &future_fiber.awaiter, .Xchg, prev_fiber, .acq_rel) == Fiber.finished) {
527 const closure: *AsyncClosure = .fromFiber(future_fiber);
528 if (!@atomicRmw(bool, &closure.already_awaited, .Xchg, true, .seq_cst)) {
529 el.schedule(thread, .{ .head = prev_fiber, .tail = prev_fiber });
530 }
531 }
532 }
524533 },
525534 .mutex_lock => |mutex_lock| {
526535 const prev_fiber: *Fiber = @alignCast(@fieldParentPtr("context", message.contexts.prev));
......@@ -661,6 +670,7 @@ const AsyncClosure = struct {
661670 fiber: *Fiber,
662671 start: *const fn (context: *const anyopaque, result: *anyopaque) void,
663672 result_align: Alignment,
673 already_awaited: bool,
664674
665675 fn contextPointer(closure: *AsyncClosure) [*]align(Fiber.max_context_align.toByteUnits()) u8 {
666676 return @alignCast(@as([*]u8, @ptrCast(closure)) + @sizeOf(AsyncClosure));
......@@ -668,12 +678,24 @@ const AsyncClosure = struct {
668678
669679 fn call(closure: *AsyncClosure, message: *const SwitchMessage) callconv(.withStackAlign(.c, @alignOf(AsyncClosure))) noreturn {
670680 message.handle(closure.event_loop);
671 std.log.debug("{*} performing async", .{closure.fiber});
672 closure.start(closure.contextPointer(), closure.fiber.resultBytes(closure.result_align));
673 const awaiter = @atomicRmw(?*Fiber, &closure.fiber.awaiter, .Xchg, Fiber.finished, .acq_rel);
674 closure.event_loop.yield(awaiter, .nothing);
681 const fiber = closure.fiber;
682 std.log.debug("{*} performing async", .{fiber});
683 closure.start(closure.contextPointer(), fiber.resultBytes(closure.result_align));
684 const awaiter = @atomicRmw(?*Fiber, &fiber.awaiter, .Xchg, Fiber.finished, .acq_rel);
685 const ready_awaiter = r: {
686 const a = awaiter orelse break :r null;
687 if (@atomicRmw(bool, &closure.already_awaited, .Xchg, true, .acq_rel)) break :r null;
688 break :r a;
689 };
690 closure.event_loop.yield(ready_awaiter, .nothing);
675691 unreachable; // switched to dead fiber
676692 }
693
694 fn fromFiber(fiber: *Fiber) *AsyncClosure {
695 return @ptrFromInt(Fiber.max_context_align.max(.of(AsyncClosure)).backward(
696 @intFromPtr(fiber.allocatedEnd()) - Fiber.max_context_size,
697 ) - @sizeOf(AsyncClosure));
698 }
677699};
678700
679701fn @"async"(
......@@ -696,9 +718,7 @@ fn @"async"(
696718 };
697719 std.log.debug("allocated {*}", .{fiber});
698720
699 const closure: *AsyncClosure = @ptrFromInt(Fiber.max_context_align.max(.of(AsyncClosure)).backward(
700 @intFromPtr(fiber.allocatedEnd()) - Fiber.max_context_size,
701 ) - @sizeOf(AsyncClosure));
721 const closure: *AsyncClosure = .fromFiber(fiber);
702722 const stack_end: [*]usize = @alignCast(@ptrCast(closure));
703723 (stack_end - 1)[0..1].* = .{@intFromPtr(&AsyncClosure.call)};
704724 fiber.* = .{
......@@ -721,6 +741,7 @@ fn @"async"(
721741 .fiber = fiber,
722742 .start = start,
723743 .result_align = result_alignment,
744 .already_awaited = false,
724745 };
725746 @memcpy(closure.contextPointer(), context);
726747
......@@ -728,13 +749,6 @@ fn @"async"(
728749 return @ptrCast(fiber);
729750}
730751
731fn select(userdata: ?*anyopaque, futures: []const *Io.AnyFuture) usize {
732 const el: *EventLoop = @alignCast(@ptrCast(userdata));
733 _ = el;
734 _ = futures;
735 @panic("TODO");
736}
737
738752const DetachedClosure = struct {
739753 event_loop: *EventLoop,
740754 fiber: *Fiber,
......@@ -836,6 +850,42 @@ fn @"await"(
836850 event_loop.recycle(future_fiber);
837851}
838852
853fn select(userdata: ?*anyopaque, futures: []const *Io.AnyFuture) usize {
854 const el: *EventLoop = @alignCast(@ptrCast(userdata));
855
856 // Optimization to avoid the yield below.
857 for (futures, 0..) |any_future, i| {
858 const future_fiber: *Fiber = @alignCast(@ptrCast(any_future));
859 if (@atomicLoad(?*Fiber, &future_fiber.awaiter, .acquire) == Fiber.finished)
860 return i;
861 }
862
863 el.yield(null, .{ .register_select = futures });
864
865 std.log.debug("back from select yield", .{});
866
867 const my_thread: *Thread = .current();
868 const my_fiber = my_thread.currentFiber();
869 var result: ?usize = null;
870
871 for (futures, 0..) |any_future, i| {
872 const future_fiber: *Fiber = @alignCast(@ptrCast(any_future));
873 if (@cmpxchgStrong(?*Fiber, &future_fiber.awaiter, my_fiber, null, .seq_cst, .seq_cst)) |awaiter| {
874 if (awaiter == Fiber.finished) {
875 if (result == null) result = i;
876 } else if (awaiter) |a| {
877 const closure: *AsyncClosure = .fromFiber(a);
878 closure.already_awaited = false;
879 }
880 } else {
881 const closure: *AsyncClosure = .fromFiber(my_fiber);
882 closure.already_awaited = false;
883 }
884 }
885
886 return result.?;
887}
888
839889fn cancel(
840890 userdata: ?*anyopaque,
841891 any_future: *std.Io.AnyFuture,
lib/std/Thread/Pool.zig+1-1
......@@ -364,7 +364,7 @@ const AsyncClosure = struct {
364364 context_offset: usize,
365365 result_offset: usize,
366366
367 const done_reset_event: *std.Thread.ResetEvent = @ptrFromInt(std.mem.alignBackward(usize, std.math.maxInt(usize), @alignOf(std.Thread.ResetEvent)));
367 const done_reset_event: *std.Thread.ResetEvent = @ptrFromInt(@alignOf(std.Thread.ResetEvent));
368368
369369 const canceling_tid: std.Thread.Id = switch (@typeInfo(std.Thread.Id)) {
370370 .int => |int_info| switch (int_info.signedness) {