authorgravatar for andrew@ziglang.orgAndrew Kelley <andrew@ziglang.org> 2026-02-21 05:22:20+01:00
committergravatar for andrew@ziglang.orgAndrew Kelley <andrew@ziglang.org> 2026-02-21 05:22:20+01:00
logcd02b1703b6acad7e76f3592b0e637ff6a7c8c99
tree68d1c531f78f7b1fb8718e6fa4da7816cbab94c9
parent5ac6ff43d41f23d7d215c3164848bb4ffcf00d59
parent311bba4af0985241601fed4cf6aba2495fd2912f

Merge pull request 'std.Io.Select: add `awaitMany`, documentation, and unit test; remove `outstanding` field' (#31296) from select-enhancement into master

Reviewed-on: https://codeberg.org/ziglang/zig/pulls/31296

3 files changed, 86 insertions(+), 31 deletions(-)

lib/std/Io.zig+44-30
......@@ -1173,12 +1173,20 @@ pub fn checkCancel(io: Io) Cancelable!void {
11731173 return io.vtable.checkCancel(io.userdata);
11741174}
11751175
1176/// Executes tasks together, providing a mechanism to wait until one or more
1177/// tasks complete. Similar to `Batch` but operates at the higher level task
1178/// abstraction layer rather than lower level `Operation` abstraction layer.
1179///
1180/// The provided tagged union will be used as the return type of the await
1181/// function. When calling async or concurrent, one specifies which union field
1182/// the called function's result will be placed into upon completion.
11761183pub fn Select(comptime U: type) type {
11771184 return struct {
11781185 io: Io,
11791186 group: Group,
1187 /// The queue is never closed because there may be live resources
1188 /// inserted into it which would otherwise leak.
11801189 queue: Queue(U),
1181 outstanding: usize,
11821190
11831191 const S = @This();
11841192
......@@ -1191,7 +1199,6 @@ pub fn Select(comptime U: type) type {
11911199 .io = io,
11921200 .queue = .init(buffer),
11931201 .group = .init,
1194 .outstanding = 0,
11951202 };
11961203 }
11971204
......@@ -1235,7 +1242,6 @@ pub fn Select(comptime U: type) type {
12351242 }
12361243 };
12371244 const context: Context = .{ .select = s, .args = args };
1238 _ = @atomicRmw(usize, &s.outstanding, .Add, 1, .monotonic);
12391245 s.io.vtable.groupAsync(s.io.userdata, &s.group, @ptrCast(&context), .of(Context), Context.start);
12401246 }
12411247
......@@ -1276,22 +1282,31 @@ pub fn Select(comptime U: type) type {
12761282 };
12771283 const context: Context = .{ .select = s, .args = args };
12781284 try s.io.vtable.groupConcurrent(s.io.userdata, &s.group, @ptrCast(&context), .of(Context), Context.start);
1279 _ = @atomicRmw(usize, &s.outstanding, .Add, 1, .monotonic);
12801285 }
12811286
12821287 /// Blocks until another task of the select finishes.
12831288 ///
1284 /// Asserts there is at least one more `outstanding` task.
1285 ///
1286 /// Not threadsafe.
1289 /// Threadsafe.
12871290 pub fn await(s: *S) Cancelable!U {
1288 s.outstanding -= 1;
12891291 return s.queue.getOne(s.io) catch |err| switch (err) {
12901292 error.Canceled => |e| return e,
12911293 error.Closed => unreachable,
12921294 };
12931295 }
12941296
1297 /// Blocks until at least `min` number of results have been copied
1298 /// into `buffer`.
1299 ///
1300 /// Asserts that `buffer.len >= min`.
1301 ///
1302 /// Threadsafe.
1303 pub fn awaitMany(s: *S, buffer: []U, min: usize) Cancelable!usize {
1304 return s.queue.get(s.io, buffer, min) catch |err| switch (err) {
1305 error.Canceled => |e| return e,
1306 error.Closed => unreachable,
1307 };
1308 }
1309
12951310 /// Equivalent to `await` but requests cancelation on all remaining
12961311 /// tasks owned by the select.
12971312 ///
......@@ -1299,9 +1314,8 @@ pub fn Select(comptime U: type) type {
12991314 ///
13001315 /// It is illegal to call `await` after this.
13011316 ///
1302 /// Idempotent. Not threadsafe.
1317 /// Idempotent. Threadsafe.
13031318 pub fn cancel(s: *S) void {
1304 s.outstanding = 0;
13051319 s.group.cancel(s.io);
13061320 }
13071321 };
......@@ -1731,7 +1745,7 @@ pub const TypeErasedQueue = struct {
17311745 return if (slice.len > 0) slice else null;
17321746 }
17331747
1734 fn putLocked(q: *TypeErasedQueue, io: Io, elements: []const u8, target: usize, uncancelable: bool) (QueueClosedError || Cancelable)!usize {
1748 fn putLocked(q: *TypeErasedQueue, io: Io, elements: []const u8, min: usize, uncancelable: bool) (QueueClosedError || Cancelable)!usize {
17351749 // A closed queue cannot be added to, even if there is space in the buffer.
17361750 if (q.closed) return error.Closed;
17371751
......@@ -1767,12 +1781,12 @@ pub const TypeErasedQueue = struct {
17671781 if (n == elements.len) return elements.len;
17681782 }
17691783
1770 // Don't block if we hit the target.
1771 if (n >= target) return n;
1784 // Don't block if we hit the min.
1785 if (n >= min) return n;
17721786
17731787 var pending: Put = .{
17741788 .remaining = elements[n..],
1775 .needed = target - n,
1789 .needed = min - n,
17761790 .condition = .init,
17771791 .node = .{},
17781792 };
......@@ -1831,7 +1845,7 @@ pub const TypeErasedQueue = struct {
18311845 return if (slice.len > 0) slice else null;
18321846 }
18331847
1834 fn getLocked(q: *TypeErasedQueue, io: Io, buffer: []u8, target: usize, uncancelable: bool) (QueueClosedError || Cancelable)!usize {
1848 fn getLocked(q: *TypeErasedQueue, io: Io, buffer: []u8, min: usize, uncancelable: bool) (QueueClosedError || Cancelable)!usize {
18351849 // The ring buffer gets first priority, then data should come from any
18361850 // queued putters, then finally the ring buffer should be filled with
18371851 // data from putters so they can be resumed.
......@@ -1877,15 +1891,15 @@ pub const TypeErasedQueue = struct {
18771891 // No need to call `fillRingBufferFromPutters` from this point onwards,
18781892 // because we emptied the ring buffer *and* the putter queue!
18791893
1880 // Don't block if we hit the target or if the queue is closed. Return how
1894 // Don't block if we hit the min or if the queue is closed. Return how
18811895 // many elements we could get immediately, unless the queue was closed and
18821896 // empty, in which case report `error.Closed`.
18831897 if (n == 0 and q.closed) return error.Closed;
1884 if (n >= target or q.closed) return n;
1898 if (n >= min or q.closed) return n;
18851899
18861900 var pending: Get = .{
18871901 .remaining = buffer[n..],
1888 .needed = target - n,
1902 .needed = min - n,
18891903 .condition = .init,
18901904 .node = .{},
18911905 };
......@@ -1961,7 +1975,7 @@ pub fn Queue(Elem: type) type {
19611975 /// there is insufficient capacity. Returns when any one of the
19621976 /// following conditions is satisfied:
19631977 ///
1964 /// * At least `target` elements have been added to the queue
1978 /// * At least `min` elements have been added to the queue
19651979 /// * The queue is closed
19661980 /// * The current task is canceled
19671981 ///
......@@ -1970,16 +1984,16 @@ pub fn Queue(Elem: type) type {
19701984 ///
19711985 /// If the queue is closed or the task is canceled, but some items were
19721986 /// already added before the closure or cancelation, then `put` may
1973 /// return a number lower than `target`, in which case future calls are
1987 /// return a number lower than `min`, in which case future calls are
19741988 /// guaranteed to return `error.Canceled` or `error.Closed`.
19751989 ///
1976 /// A return value of 0 is only possible if `target` is 0, in which case
1990 /// A return value of 0 is only possible if `min` is 0, in which case
19771991 /// the call is guaranteed to queue as many of `elements` as is possible
19781992 /// *without* blocking.
19791993 ///
1980 /// Asserts that `elements.len >= target`.
1981 pub fn put(q: *@This(), io: Io, elements: []const Elem, target: usize) (QueueClosedError || Cancelable)!usize {
1982 return @divExact(try q.type_erased.put(io, @ptrCast(elements), target * @sizeOf(Elem)), @sizeOf(Elem));
1994 /// Asserts that `elements.len >= min`.
1995 pub fn put(q: *@This(), io: Io, elements: []const Elem, min: usize) (QueueClosedError || Cancelable)!usize {
1996 return @divExact(try q.type_erased.put(io, @ptrCast(elements), min * @sizeOf(Elem)), @sizeOf(Elem));
19831997 }
19841998
19851999 /// Same as `put` but blocks until all elements have been added to the queue.
......@@ -2018,7 +2032,7 @@ pub fn Queue(Elem: type) type {
20182032 /// if there are insufficient elements currently in the queue. Returns when
20192033 /// any one of the following conditions is satisfied:
20202034 ///
2021 /// * At least `target` elements have been received from the queue
2035 /// * At least `min` elements have been received from the queue
20222036 /// * The queue is closed and contains no buffered elements
20232037 /// * The current task is canceled
20242038 ///
......@@ -2027,16 +2041,16 @@ pub fn Queue(Elem: type) type {
20272041 ///
20282042 /// If the queue is closed or the task is canceled, but some items were
20292043 /// already received before the closure or cancelation, then `get` may
2030 /// return a number lower than `target`, in which case future calls are
2044 /// return a number lower than `min`, in which case future calls are
20312045 /// guaranteed to return `error.Canceled` or `error.Closed`.
20322046 ///
2033 /// A return value of 0 is only possible if `target` is 0, in which case
2047 /// A return value of 0 is only possible if `min` is 0, in which case
20342048 /// the call is guaranteed to fill as much of `buffer` as is possible
20352049 /// *without* blocking.
20362050 ///
2037 /// Asserts that `buffer.len >= target`.
2038 pub fn get(q: *@This(), io: Io, buffer: []Elem, target: usize) (QueueClosedError || Cancelable)!usize {
2039 return @divExact(try q.type_erased.get(io, @ptrCast(buffer), target * @sizeOf(Elem)), @sizeOf(Elem));
2051 /// Asserts that `buffer.len >= min`.
2052 pub fn get(q: *@This(), io: Io, buffer: []Elem, min: usize) (QueueClosedError || Cancelable)!usize {
2053 return @divExact(try q.type_erased.get(io, @ptrCast(buffer), min * @sizeOf(Elem)), @sizeOf(Elem));
20402054 }
20412055
20422056 /// Same as `get`, except does not introduce a cancelation point.
lib/std/Io/test.zig+38
......@@ -810,3 +810,41 @@ test "Event broadcast" {
810810
811811 try ctx.run();
812812}
813
814test "Select" {
815 const S = struct {
816 fn foo() bool {
817 return true;
818 }
819
820 fn bar(io: Io) Io.Cancelable!void {
821 try io.sleep(.fromSeconds(300), .awake);
822 }
823 };
824
825 const io = testing.io;
826
827 const U = union(enum) {
828 foo: bool,
829 bar: Io.Cancelable!void,
830 };
831 var buffer: [4]U = undefined;
832 var select: Io.Select(U) = .init(io, &buffer);
833 defer select.cancel();
834
835 select.async(.foo, S.foo, .{});
836 select.concurrent(.bar, S.bar, .{io}) catch |err| switch (err) {
837 error.ConcurrencyUnavailable => return error.SkipZigTest,
838 };
839
840 switch (try select.await()) {
841 .foo => {},
842 .bar => return error.TestFailed, // should be sleeping
843 }
844 select.async(.foo, S.foo, .{});
845 select.async(.foo, S.foo, .{});
846
847 var finished_buffer: [3]U = undefined;
848 const finished = finished_buffer[0..try select.awaitMany(&finished_buffer, 2)];
849 try testing.expectEqualSlices(U, &.{ .{ .foo = true }, .{ .foo = true } }, finished);
850}
lib/std/crypto/kangarootwelve.zig+4-1
......@@ -883,6 +883,7 @@ fn ktMultiThreaded(
883883 defer allocator.free(pending_cv_buf);
884884 var pending_cv_lens: [256]usize = .{0} ** 256;
885885
886 var select_outstanding: usize = 0;
886887 var select: Select = .init(io, select_buf);
887888 defer select.cancel();
888889 var batches_spawned: usize = 0;
......@@ -894,6 +895,7 @@ fn ktMultiThreaded(
894895 const batch_leaves = @min(leaves_per_batch, full_leaves - batch_start_leaf);
895896 const start_offset = chunk_size + batch_start_leaf * chunk_size;
896897
898 select_outstanding += 1;
897899 select.async(.batch, SelectLeafContext(Variant).process, .{SelectLeafContext(Variant){
898900 .view = view,
899901 .batch_idx = batches_spawned,
......@@ -903,6 +905,7 @@ fn ktMultiThreaded(
903905 batches_spawned += 1;
904906 }
905907
908 select_outstanding -= 1;
906909 const result = try select.await();
907910 const batch = result.batch;
908911 const slot = batch.batch_idx % max_concurrent;
......@@ -927,7 +930,7 @@ fn ktMultiThreaded(
927930 }
928931 }
929932
930 assert(select.outstanding == 0);
933 assert(select_outstanding == 0);
931934 }
932935
933936 if (has_partial_leaf) {