authorgravatar for andrew@ziglang.orgAndrew Kelley <andrew@ziglang.org> 2026-02-20 17:28:33-08:00
committergravatar for andrew@ziglang.orgAndrew Kelley <andrew@ziglang.org> 2026-02-20 17:28:33-08:00
log311bba4af0985241601fed4cf6aba2495fd2912f
tree68d1c531f78f7b1fb8718e6fa4da7816cbab94c9
parentf9053f38e5ca29c23d2328ed06293727cd1fe8d4

std.Io.Select: add awaitMany function and unit test

and fix documentation. these functions are in fact threadsafe.

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

lib/std/Io.zig+35-22
......@@ -1286,7 +1286,7 @@ pub fn Select(comptime U: type) type {
12861286
12871287 /// Blocks until another task of the select finishes.
12881288 ///
1289 /// Not threadsafe.
1289 /// Threadsafe.
12901290 pub fn await(s: *S) Cancelable!U {
12911291 return s.queue.getOne(s.io) catch |err| switch (err) {
12921292 error.Canceled => |e| return e,
......@@ -1294,6 +1294,19 @@ pub fn Select(comptime U: type) type {
12941294 };
12951295 }
12961296
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
12971310 /// Equivalent to `await` but requests cancelation on all remaining
12981311 /// tasks owned by the select.
12991312 ///
......@@ -1301,7 +1314,7 @@ pub fn Select(comptime U: type) type {
13011314 ///
13021315 /// It is illegal to call `await` after this.
13031316 ///
1304 /// Idempotent. Not threadsafe.
1317 /// Idempotent. Threadsafe.
13051318 pub fn cancel(s: *S) void {
13061319 s.group.cancel(s.io);
13071320 }
......@@ -1732,7 +1745,7 @@ pub const TypeErasedQueue = struct {
17321745 return if (slice.len > 0) slice else null;
17331746 }
17341747
1735 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 {
17361749 // A closed queue cannot be added to, even if there is space in the buffer.
17371750 if (q.closed) return error.Closed;
17381751
......@@ -1768,12 +1781,12 @@ pub const TypeErasedQueue = struct {
17681781 if (n == elements.len) return elements.len;
17691782 }
17701783
1771 // Don't block if we hit the target.
1772 if (n >= target) return n;
1784 // Don't block if we hit the min.
1785 if (n >= min) return n;
17731786
17741787 var pending: Put = .{
17751788 .remaining = elements[n..],
1776 .needed = target - n,
1789 .needed = min - n,
17771790 .condition = .init,
17781791 .node = .{},
17791792 };
......@@ -1832,7 +1845,7 @@ pub const TypeErasedQueue = struct {
18321845 return if (slice.len > 0) slice else null;
18331846 }
18341847
1835 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 {
18361849 // The ring buffer gets first priority, then data should come from any
18371850 // queued putters, then finally the ring buffer should be filled with
18381851 // data from putters so they can be resumed.
......@@ -1878,15 +1891,15 @@ pub const TypeErasedQueue = struct {
18781891 // No need to call `fillRingBufferFromPutters` from this point onwards,
18791892 // because we emptied the ring buffer *and* the putter queue!
18801893
1881 // 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
18821895 // many elements we could get immediately, unless the queue was closed and
18831896 // empty, in which case report `error.Closed`.
18841897 if (n == 0 and q.closed) return error.Closed;
1885 if (n >= target or q.closed) return n;
1898 if (n >= min or q.closed) return n;
18861899
18871900 var pending: Get = .{
18881901 .remaining = buffer[n..],
1889 .needed = target - n,
1902 .needed = min - n,
18901903 .condition = .init,
18911904 .node = .{},
18921905 };
......@@ -1962,7 +1975,7 @@ pub fn Queue(Elem: type) type {
19621975 /// there is insufficient capacity. Returns when any one of the
19631976 /// following conditions is satisfied:
19641977 ///
1965 /// * At least `target` elements have been added to the queue
1978 /// * At least `min` elements have been added to the queue
19661979 /// * The queue is closed
19671980 /// * The current task is canceled
19681981 ///
......@@ -1971,16 +1984,16 @@ pub fn Queue(Elem: type) type {
19711984 ///
19721985 /// If the queue is closed or the task is canceled, but some items were
19731986 /// already added before the closure or cancelation, then `put` may
1974 /// 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
19751988 /// guaranteed to return `error.Canceled` or `error.Closed`.
19761989 ///
1977 /// 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
19781991 /// the call is guaranteed to queue as many of `elements` as is possible
19791992 /// *without* blocking.
19801993 ///
1981 /// Asserts that `elements.len >= target`.
1982 pub fn put(q: *@This(), io: Io, elements: []const Elem, target: usize) (QueueClosedError || Cancelable)!usize {
1983 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));
19841997 }
19851998
19861999 /// Same as `put` but blocks until all elements have been added to the queue.
......@@ -2019,7 +2032,7 @@ pub fn Queue(Elem: type) type {
20192032 /// if there are insufficient elements currently in the queue. Returns when
20202033 /// any one of the following conditions is satisfied:
20212034 ///
2022 /// * At least `target` elements have been received from the queue
2035 /// * At least `min` elements have been received from the queue
20232036 /// * The queue is closed and contains no buffered elements
20242037 /// * The current task is canceled
20252038 ///
......@@ -2028,16 +2041,16 @@ pub fn Queue(Elem: type) type {
20282041 ///
20292042 /// If the queue is closed or the task is canceled, but some items were
20302043 /// already received before the closure or cancelation, then `get` may
2031 /// 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
20322045 /// guaranteed to return `error.Canceled` or `error.Closed`.
20332046 ///
2034 /// 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
20352048 /// the call is guaranteed to fill as much of `buffer` as is possible
20362049 /// *without* blocking.
20372050 ///
2038 /// Asserts that `buffer.len >= target`.
2039 pub fn get(q: *@This(), io: Io, buffer: []Elem, target: usize) (QueueClosedError || Cancelable)!usize {
2040 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));
20412054 }
20422055
20432056 /// 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}