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 {...@@ -1286,7 +1286,7 @@ pub fn Select(comptime U: type) type {
12861286
1287 /// Blocks until another task of the select finishes.1287 /// Blocks until another task of the select finishes.
1288 ///1288 ///
1289 /// Not threadsafe.1289 /// Threadsafe.
1290 pub fn await(s: *S) Cancelable!U {1290 pub fn await(s: *S) Cancelable!U {
1291 return s.queue.getOne(s.io) catch |err| switch (err) {1291 return s.queue.getOne(s.io) catch |err| switch (err) {
1292 error.Canceled => |e| return e,1292 error.Canceled => |e| return e,
...@@ -1294,6 +1294,19 @@ pub fn Select(comptime U: type) type {...@@ -1294,6 +1294,19 @@ pub fn Select(comptime U: type) type {
1294 };1294 };
1295 }1295 }
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
1297 /// Equivalent to `await` but requests cancelation on all remaining1310 /// Equivalent to `await` but requests cancelation on all remaining
1298 /// tasks owned by the select.1311 /// tasks owned by the select.
1299 ///1312 ///
...@@ -1301,7 +1314,7 @@ pub fn Select(comptime U: type) type {...@@ -1301,7 +1314,7 @@ pub fn Select(comptime U: type) type {
1301 ///1314 ///
1302 /// It is illegal to call `await` after this.1315 /// It is illegal to call `await` after this.
1303 ///1316 ///
1304 /// Idempotent. Not threadsafe.1317 /// Idempotent. Threadsafe.
1305 pub fn cancel(s: *S) void {1318 pub fn cancel(s: *S) void {
1306 s.group.cancel(s.io);1319 s.group.cancel(s.io);
1307 }1320 }
...@@ -1732,7 +1745,7 @@ pub const TypeErasedQueue = struct {...@@ -1732,7 +1745,7 @@ pub const TypeErasedQueue = struct {
1732 return if (slice.len > 0) slice else null;1745 return if (slice.len > 0) slice else null;
1733 }1746 }
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 {
1736 // 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.
1737 if (q.closed) return error.Closed;1750 if (q.closed) return error.Closed;
17381751
...@@ -1768,12 +1781,12 @@ pub const TypeErasedQueue = struct {...@@ -1768,12 +1781,12 @@ pub const TypeErasedQueue = struct {
1768 if (n == elements.len) return elements.len;1781 if (n == elements.len) return elements.len;
1769 }1782 }
17701783
1771 // Don't block if we hit the target.1784 // Don't block if we hit the min.
1772 if (n >= target) return n;1785 if (n >= min) return n;
17731786
1774 var pending: Put = .{1787 var pending: Put = .{
1775 .remaining = elements[n..],1788 .remaining = elements[n..],
1776 .needed = target - n,1789 .needed = min - n,
1777 .condition = .init,1790 .condition = .init,
1778 .node = .{},1791 .node = .{},
1779 };1792 };
...@@ -1832,7 +1845,7 @@ pub const TypeErasedQueue = struct {...@@ -1832,7 +1845,7 @@ pub const TypeErasedQueue = struct {
1832 return if (slice.len > 0) slice else null;1845 return if (slice.len > 0) slice else null;
1833 }1846 }
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 {
1836 // The ring buffer gets first priority, then data should come from any1849 // The ring buffer gets first priority, then data should come from any
1837 // queued putters, then finally the ring buffer should be filled with1850 // queued putters, then finally the ring buffer should be filled with
1838 // data from putters so they can be resumed.1851 // data from putters so they can be resumed.
...@@ -1878,15 +1891,15 @@ pub const TypeErasedQueue = struct {...@@ -1878,15 +1891,15 @@ pub const TypeErasedQueue = struct {
1878 // No need to call `fillRingBufferFromPutters` from this point onwards,1891 // No need to call `fillRingBufferFromPutters` from this point onwards,
1879 // because we emptied the ring buffer *and* the putter queue!1892 // 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 how1894 // Don't block if we hit the min or if the queue is closed. Return how
1882 // many elements we could get immediately, unless the queue was closed and1895 // many elements we could get immediately, unless the queue was closed and
1883 // empty, in which case report `error.Closed`.1896 // empty, in which case report `error.Closed`.
1884 if (n == 0 and q.closed) return error.Closed;1897 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
1887 var pending: Get = .{1900 var pending: Get = .{
1888 .remaining = buffer[n..],1901 .remaining = buffer[n..],
1889 .needed = target - n,1902 .needed = min - n,
1890 .condition = .init,1903 .condition = .init,
1891 .node = .{},1904 .node = .{},
1892 };1905 };
...@@ -1962,7 +1975,7 @@ pub fn Queue(Elem: type) type {...@@ -1962,7 +1975,7 @@ pub fn Queue(Elem: type) type {
1962 /// there is insufficient capacity. Returns when any one of the1975 /// there is insufficient capacity. Returns when any one of the
1963 /// following conditions is satisfied:1976 /// following conditions is satisfied:
1964 ///1977 ///
1965 /// * At least `target` elements have been added to the queue1978 /// * At least `min` elements have been added to the queue
1966 /// * The queue is closed1979 /// * The queue is closed
1967 /// * The current task is canceled1980 /// * The current task is canceled
1968 ///1981 ///
...@@ -1971,16 +1984,16 @@ pub fn Queue(Elem: type) type {...@@ -1971,16 +1984,16 @@ pub fn Queue(Elem: type) type {
1971 ///1984 ///
1972 /// 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
1973 /// already added before the closure or cancelation, then `put` may1986 /// already added before the closure or cancelation, then `put` may
1974 /// return a number lower than `target`, in which case future calls are1987 /// return a number lower than `min`, in which case future calls are
1975 /// guaranteed to return `error.Canceled` or `error.Closed`.1988 /// guaranteed to return `error.Canceled` or `error.Closed`.
1976 ///1989 ///
1977 /// 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
1978 /// 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
1979 /// *without* blocking.1992 /// *without* blocking.
1980 ///1993 ///
1981 /// Asserts that `elements.len >= target`.1994 /// Asserts that `elements.len >= min`.
1982 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 {
1983 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));
1984 }1997 }
19851998
1986 /// 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.
...@@ -2019,7 +2032,7 @@ pub fn Queue(Elem: type) type {...@@ -2019,7 +2032,7 @@ pub fn Queue(Elem: type) type {
2019 /// if there are insufficient elements currently in the queue. Returns when2032 /// if there are insufficient elements currently in the queue. Returns when
2020 /// any one of the following conditions is satisfied:2033 /// any one of the following conditions is satisfied:
2021 ///2034 ///
2022 /// * At least `target` elements have been received from the queue2035 /// * At least `min` elements have been received from the queue
2023 /// * The queue is closed and contains no buffered elements2036 /// * The queue is closed and contains no buffered elements
2024 /// * The current task is canceled2037 /// * The current task is canceled
2025 ///2038 ///
...@@ -2028,16 +2041,16 @@ pub fn Queue(Elem: type) type {...@@ -2028,16 +2041,16 @@ pub fn Queue(Elem: type) type {
2028 ///2041 ///
2029 /// 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
2030 /// already received before the closure or cancelation, then `get` may2043 /// already received before the closure or cancelation, then `get` may
2031 /// return a number lower than `target`, in which case future calls are2044 /// return a number lower than `min`, in which case future calls are
2032 /// guaranteed to return `error.Canceled` or `error.Closed`.2045 /// guaranteed to return `error.Canceled` or `error.Closed`.
2033 ///2046 ///
2034 /// 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
2035 /// 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
2036 /// *without* blocking.2049 /// *without* blocking.
2037 ///2050 ///
2038 /// Asserts that `buffer.len >= target`.2051 /// Asserts that `buffer.len >= min`.
2039 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 {
2040 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));
2041 }2054 }
20422055
2043 /// 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}