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 {...@@ -1173,12 +1173,20 @@ pub fn checkCancel(io: Io) Cancelable!void {
1173 return io.vtable.checkCancel(io.userdata);1173 return io.vtable.checkCancel(io.userdata);
1174}1174}
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.
1176pub fn Select(comptime U: type) type {1183pub fn Select(comptime U: type) type {
1177 return struct {1184 return struct {
1178 io: Io,1185 io: Io,
1179 group: Group,1186 group: Group,
1187 /// The queue is never closed because there may be live resources
1188 /// inserted into it which would otherwise leak.
1180 queue: Queue(U),1189 queue: Queue(U),
1181 outstanding: usize,
11821190
1183 const S = @This();1191 const S = @This();
11841192
...@@ -1191,7 +1199,6 @@ pub fn Select(comptime U: type) type {...@@ -1191,7 +1199,6 @@ pub fn Select(comptime U: type) type {
1191 .io = io,1199 .io = io,
1192 .queue = .init(buffer),1200 .queue = .init(buffer),
1193 .group = .init,1201 .group = .init,
1194 .outstanding = 0,
1195 };1202 };
1196 }1203 }
11971204
...@@ -1235,7 +1242,6 @@ pub fn Select(comptime U: type) type {...@@ -1235,7 +1242,6 @@ pub fn Select(comptime U: type) type {
1235 }1242 }
1236 };1243 };
1237 const context: Context = .{ .select = s, .args = args };1244 const context: Context = .{ .select = s, .args = args };
1238 _ = @atomicRmw(usize, &s.outstanding, .Add, 1, .monotonic);
1239 s.io.vtable.groupAsync(s.io.userdata, &s.group, @ptrCast(&context), .of(Context), Context.start);1245 s.io.vtable.groupAsync(s.io.userdata, &s.group, @ptrCast(&context), .of(Context), Context.start);
1240 }1246 }
12411247
...@@ -1276,22 +1282,31 @@ pub fn Select(comptime U: type) type {...@@ -1276,22 +1282,31 @@ pub fn Select(comptime U: type) type {
1276 };1282 };
1277 const context: Context = .{ .select = s, .args = args };1283 const context: Context = .{ .select = s, .args = args };
1278 try s.io.vtable.groupConcurrent(s.io.userdata, &s.group, @ptrCast(&context), .of(Context), Context.start);1284 try s.io.vtable.groupConcurrent(s.io.userdata, &s.group, @ptrCast(&context), .of(Context), Context.start);
1279 _ = @atomicRmw(usize, &s.outstanding, .Add, 1, .monotonic);
1280 }1285 }
12811286
1282 /// Blocks until another task of the select finishes.1287 /// Blocks until another task of the select finishes.
1283 ///1288 ///
1284 /// Asserts there is at least one more `outstanding` task.1289 /// Threadsafe.
1285 ///
1286 /// Not threadsafe.
1287 pub fn await(s: *S) Cancelable!U {1290 pub fn await(s: *S) Cancelable!U {
1288 s.outstanding -= 1;
1289 return s.queue.getOne(s.io) catch |err| switch (err) {1291 return s.queue.getOne(s.io) catch |err| switch (err) {
1290 error.Canceled => |e| return e,1292 error.Canceled => |e| return e,
1291 error.Closed => unreachable,1293 error.Closed => unreachable,
1292 };1294 };
1293 }1295 }
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
1295 /// Equivalent to `await` but requests cancelation on all remaining1310 /// Equivalent to `await` but requests cancelation on all remaining
1296 /// tasks owned by the select.1311 /// tasks owned by the select.
1297 ///1312 ///
...@@ -1299,9 +1314,8 @@ pub fn Select(comptime U: type) type {...@@ -1299,9 +1314,8 @@ pub fn Select(comptime U: type) type {
1299 ///1314 ///
1300 /// It is illegal to call `await` after this.1315 /// It is illegal to call `await` after this.
1301 ///1316 ///
1302 /// Idempotent. Not threadsafe.1317 /// Idempotent. Threadsafe.
1303 pub fn cancel(s: *S) void {1318 pub fn cancel(s: *S) void {
1304 s.outstanding = 0;
1305 s.group.cancel(s.io);1319 s.group.cancel(s.io);
1306 }1320 }
1307 };1321 };
...@@ -1731,7 +1745,7 @@ pub const TypeErasedQueue = struct {...@@ -1731,7 +1745,7 @@ pub const TypeErasedQueue = struct {
1731 return if (slice.len > 0) slice else null;1745 return if (slice.len > 0) slice else null;
1732 }1746 }
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 {
1735 // A closed queue cannot be added to, even if there is space in the buffer.1749 // A closed queue cannot be added to, even if there is space in the buffer.
1736 if (q.closed) return error.Closed;1750 if (q.closed) return error.Closed;
17371751
...@@ -1767,12 +1781,12 @@ pub const TypeErasedQueue = struct {...@@ -1767,12 +1781,12 @@ pub const TypeErasedQueue = struct {
1767 if (n == elements.len) return elements.len;1781 if (n == elements.len) return elements.len;
1768 }1782 }
17691783
1770 // Don't block if we hit the target.1784 // Don't block if we hit the min.
1771 if (n >= target) return n;1785 if (n >= min) return n;
17721786
1773 var pending: Put = .{1787 var pending: Put = .{
1774 .remaining = elements[n..],1788 .remaining = elements[n..],
1775 .needed = target - n,1789 .needed = min - n,
1776 .condition = .init,1790 .condition = .init,
1777 .node = .{},1791 .node = .{},
1778 };1792 };
...@@ -1831,7 +1845,7 @@ pub const TypeErasedQueue = struct {...@@ -1831,7 +1845,7 @@ pub const TypeErasedQueue = struct {
1831 return if (slice.len > 0) slice else null;1845 return if (slice.len > 0) slice else null;
1832 }1846 }
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 {
1835 // The ring buffer gets first priority, then data should come from any1849 // The ring buffer gets first priority, then data should come from any
1836 // queued putters, then finally the ring buffer should be filled with1850 // queued putters, then finally the ring buffer should be filled with
1837 // data from putters so they can be resumed.1851 // data from putters so they can be resumed.
...@@ -1877,15 +1891,15 @@ pub const TypeErasedQueue = struct {...@@ -1877,15 +1891,15 @@ pub const TypeErasedQueue = struct {
1877 // No need to call `fillRingBufferFromPutters` from this point onwards,1891 // No need to call `fillRingBufferFromPutters` from this point onwards,
1878 // because we emptied the ring buffer *and* the putter queue!1892 // 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 how1894 // Don't block if we hit the min or if the queue is closed. Return how
1881 // many elements we could get immediately, unless the queue was closed and1895 // many elements we could get immediately, unless the queue was closed and
1882 // empty, in which case report `error.Closed`.1896 // empty, in which case report `error.Closed`.
1883 if (n == 0 and q.closed) return error.Closed;1897 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
1886 var pending: Get = .{1900 var pending: Get = .{
1887 .remaining = buffer[n..],1901 .remaining = buffer[n..],
1888 .needed = target - n,1902 .needed = min - n,
1889 .condition = .init,1903 .condition = .init,
1890 .node = .{},1904 .node = .{},
1891 };1905 };
...@@ -1961,7 +1975,7 @@ pub fn Queue(Elem: type) type {...@@ -1961,7 +1975,7 @@ pub fn Queue(Elem: type) type {
1961 /// there is insufficient capacity. Returns when any one of the1975 /// there is insufficient capacity. Returns when any one of the
1962 /// following conditions is satisfied:1976 /// following conditions is satisfied:
1963 ///1977 ///
1964 /// * At least `target` elements have been added to the queue1978 /// * At least `min` elements have been added to the queue
1965 /// * The queue is closed1979 /// * The queue is closed
1966 /// * The current task is canceled1980 /// * The current task is canceled
1967 ///1981 ///
...@@ -1970,16 +1984,16 @@ pub fn Queue(Elem: type) type {...@@ -1970,16 +1984,16 @@ pub fn Queue(Elem: type) type {
1970 ///1984 ///
1971 /// If the queue is closed or the task is canceled, but some items were1985 /// If the queue is closed or the task is canceled, but some items were
1972 /// already added before the closure or cancelation, then `put` may1986 /// already added before the closure or cancelation, then `put` may
1973 /// return a number lower than `target`, in which case future calls are1987 /// return a number lower than `min`, in which case future calls are
1974 /// guaranteed to return `error.Canceled` or `error.Closed`.1988 /// guaranteed to return `error.Canceled` or `error.Closed`.
1975 ///1989 ///
1976 /// A return value of 0 is only possible if `target` is 0, in which case1990 /// A return value of 0 is only possible if `min` is 0, in which case
1977 /// the call is guaranteed to queue as many of `elements` as is possible1991 /// the call is guaranteed to queue as many of `elements` as is possible
1978 /// *without* blocking.1992 /// *without* blocking.
1979 ///1993 ///
1980 /// Asserts that `elements.len >= target`.1994 /// Asserts that `elements.len >= min`.
1981 pub fn put(q: *@This(), io: Io, elements: []const Elem, target: usize) (QueueClosedError || Cancelable)!usize {1995 pub fn put(q: *@This(), io: Io, elements: []const Elem, min: usize) (QueueClosedError || Cancelable)!usize {
1982 return @divExact(try q.type_erased.put(io, @ptrCast(elements), target * @sizeOf(Elem)), @sizeOf(Elem));1996 return @divExact(try q.type_erased.put(io, @ptrCast(elements), min * @sizeOf(Elem)), @sizeOf(Elem));
1983 }1997 }
19841998
1985 /// Same as `put` but blocks until all elements have been added to the queue.1999 /// Same as `put` but blocks until all elements have been added to the queue.
...@@ -2018,7 +2032,7 @@ pub fn Queue(Elem: type) type {...@@ -2018,7 +2032,7 @@ pub fn Queue(Elem: type) type {
2018 /// if there are insufficient elements currently in the queue. Returns when2032 /// if there are insufficient elements currently in the queue. Returns when
2019 /// any one of the following conditions is satisfied:2033 /// any one of the following conditions is satisfied:
2020 ///2034 ///
2021 /// * At least `target` elements have been received from the queue2035 /// * At least `min` elements have been received from the queue
2022 /// * The queue is closed and contains no buffered elements2036 /// * The queue is closed and contains no buffered elements
2023 /// * The current task is canceled2037 /// * The current task is canceled
2024 ///2038 ///
...@@ -2027,16 +2041,16 @@ pub fn Queue(Elem: type) type {...@@ -2027,16 +2041,16 @@ pub fn Queue(Elem: type) type {
2027 ///2041 ///
2028 /// If the queue is closed or the task is canceled, but some items were2042 /// If the queue is closed or the task is canceled, but some items were
2029 /// already received before the closure or cancelation, then `get` may2043 /// already received before the closure or cancelation, then `get` may
2030 /// return a number lower than `target`, in which case future calls are2044 /// return a number lower than `min`, in which case future calls are
2031 /// guaranteed to return `error.Canceled` or `error.Closed`.2045 /// guaranteed to return `error.Canceled` or `error.Closed`.
2032 ///2046 ///
2033 /// A return value of 0 is only possible if `target` is 0, in which case2047 /// A return value of 0 is only possible if `min` is 0, in which case
2034 /// the call is guaranteed to fill as much of `buffer` as is possible2048 /// the call is guaranteed to fill as much of `buffer` as is possible
2035 /// *without* blocking.2049 /// *without* blocking.
2036 ///2050 ///
2037 /// Asserts that `buffer.len >= target`.2051 /// Asserts that `buffer.len >= min`.
2038 pub fn get(q: *@This(), io: Io, buffer: []Elem, target: usize) (QueueClosedError || Cancelable)!usize {2052 pub fn get(q: *@This(), io: Io, buffer: []Elem, min: usize) (QueueClosedError || Cancelable)!usize {
2039 return @divExact(try q.type_erased.get(io, @ptrCast(buffer), target * @sizeOf(Elem)), @sizeOf(Elem));2053 return @divExact(try q.type_erased.get(io, @ptrCast(buffer), min * @sizeOf(Elem)), @sizeOf(Elem));
2040 }2054 }
20412055
2042 /// Same as `get`, except does not introduce a cancelation point.2056 /// Same as `get`, except does not introduce a cancelation point.
lib/std/Io/test.zig+38
...@@ -810,3 +810,41 @@ test "Event broadcast" {...@@ -810,3 +810,41 @@ test "Event broadcast" {
810810
811 try ctx.run();811 try ctx.run();
812}812}
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(...@@ -883,6 +883,7 @@ fn ktMultiThreaded(
883 defer allocator.free(pending_cv_buf);883 defer allocator.free(pending_cv_buf);
884 var pending_cv_lens: [256]usize = .{0} ** 256;884 var pending_cv_lens: [256]usize = .{0} ** 256;
885885
886 var select_outstanding: usize = 0;
886 var select: Select = .init(io, select_buf);887 var select: Select = .init(io, select_buf);
887 defer select.cancel();888 defer select.cancel();
888 var batches_spawned: usize = 0;889 var batches_spawned: usize = 0;
...@@ -894,6 +895,7 @@ fn ktMultiThreaded(...@@ -894,6 +895,7 @@ fn ktMultiThreaded(
894 const batch_leaves = @min(leaves_per_batch, full_leaves - batch_start_leaf);895 const batch_leaves = @min(leaves_per_batch, full_leaves - batch_start_leaf);
895 const start_offset = chunk_size + batch_start_leaf * chunk_size;896 const start_offset = chunk_size + batch_start_leaf * chunk_size;
896897
898 select_outstanding += 1;
897 select.async(.batch, SelectLeafContext(Variant).process, .{SelectLeafContext(Variant){899 select.async(.batch, SelectLeafContext(Variant).process, .{SelectLeafContext(Variant){
898 .view = view,900 .view = view,
899 .batch_idx = batches_spawned,901 .batch_idx = batches_spawned,
...@@ -903,6 +905,7 @@ fn ktMultiThreaded(...@@ -903,6 +905,7 @@ fn ktMultiThreaded(
903 batches_spawned += 1;905 batches_spawned += 1;
904 }906 }
905907
908 select_outstanding -= 1;
906 const result = try select.await();909 const result = try select.await();
907 const batch = result.batch;910 const batch = result.batch;
908 const slot = batch.batch_idx % max_concurrent;911 const slot = batch.batch_idx % max_concurrent;
...@@ -927,7 +930,7 @@ fn ktMultiThreaded(...@@ -927,7 +930,7 @@ fn ktMultiThreaded(
927 }930 }
928 }931 }
929932
930 assert(select.outstanding == 0);933 assert(select_outstanding == 0);
931 }934 }
932935
933 if (has_partial_leaf) {936 if (has_partial_leaf) {