authorgravatar for kappaloris@gmail.comLoris Cro <kappaloris@gmail.com> 2020-10-02 19:15:26+02:00
committergravatar for andrew@ziglang.orgAndrew Kelley <andrew@ziglang.org> 2020-10-07 04:34:09-04:00
log53dee08af99dd334b0d227afb5ce2a0f92c35a5d
treee6045bccd1eb7e0ec828980fe08882f33468cc9c
parent0a6863a267d88e73ec62aca2c157654020682b00

add WaitGroup to std.event

Signed-off-by: Loris Cro <kappaloris@gmail.com>

3 files changed, 124 insertions(+), 0 deletions(-)

lib/std/event.zig+2
......@@ -12,6 +12,7 @@ pub const Locked = @import("event/locked.zig").Locked;
1212pub const RwLock = @import("event/rwlock.zig").RwLock;
1313pub const RwLocked = @import("event/rwlocked.zig").RwLocked;
1414pub const Loop = @import("event/loop.zig").Loop;
15pub const WaitGroup = @import("event/WaitGroup.zig").WaitGroup;
1516
1617test "import event tests" {
1718 _ = @import("event/channel.zig");
......@@ -23,4 +24,5 @@ test "import event tests" {
2324 _ = @import("event/rwlock.zig");
2425 _ = @import("event/rwlocked.zig");
2526 _ = @import("event/loop.zig");
27 _ = @import("event/wait_group.zig");
2628}
lib/std/event/loop.zig+2
......@@ -660,9 +660,11 @@ pub const Loop = struct {
660660 const Wrapper = struct {
661661 const Args = @TypeOf(args);
662662 fn run(func_args: Args, loop: *Loop, allocator: *mem.Allocator) void {
663 loop.beginOneEvent();
663664 loop.yield();
664665 const result = @call(.{}, func, func_args);
665666 suspend {
667 loop.finishOneEvent();
666668 allocator.destroy(@frame());
667669 }
668670 }
lib/std/event/wait_group.zig created+120
......@@ -0,0 +1,120 @@
1// SPDX-License-Identifier: MIT
2// Copyright (c) 2015-2020 Zig Contributors
3// This file is part of [zig](https://ziglang.org/), which is MIT licensed.
4// The MIT license requires this copyright notice to be included in all copies
5// and substantial portions of the software.
6const std = @import("../std.zig");
7const builtin = @import("builtin");
8const Loop = std.event.Loop;
9
10/// A WaitGroup keeps track and waits for a group of async tasks to finish.
11/// Call `begin` when creating new tasks, and have tasks call `finish` when done.
12/// You can provide a count for both operations to perform them in bulk.
13/// Call `wait` to suspend until all tasks are completed.
14/// Multiple waiters are supported.
15///
16/// WaitGroup is an instance of WaitGroupGeneric, which takes in a bitsize
17/// for the internal counter. WaitGroup defaults to a `usize` counter.
18/// It's also possible to define a max value for the counter so that
19/// `begin` will return error.Overflow when the limit is reached, even
20/// if the integer type has not has not overflowed.
21/// By default `max_value` is set to std.math.maxInt(CounterType).
22pub const WaitGroup = WaitGroupGeneric(std.meta.bitCount(usize));
23
24pub fn WaitGroupGeneric(comptime counter_size: u16) type {
25 const CounterType = std.meta.Int(false, counter_size);
26
27 const global_event_loop = Loop.instance orelse
28 @compileError("std.event.WaitGroup currently only works with event-based I/O");
29
30 return struct {
31 counter: CounterType = 0,
32 max_counter: CounterType = std.math.maxInt(CounterType),
33 mutex: std.Mutex = .{},
34 waiters: ?*Waiter = null,
35 const Waiter = struct {
36 next: ?*Waiter,
37 tail: *Waiter,
38 node: Loop.NextTickNode,
39 };
40
41 const Self = @This();
42 pub fn begin(self: *Self, count: CounterType) error{Overflow}!void {
43 const held = self.mutex.acquire();
44 defer held.release();
45
46 const new_counter = try std.math.add(CounterType, self.counter, count);
47 if (new_counter > self.max_counter) return error.Overflow;
48 self.counter = new_counter;
49 }
50
51 pub fn finish(self: *Self, count: CounterType) void {
52 var waiters = blk: {
53 const held = self.mutex.acquire();
54 defer held.release();
55 self.counter = std.math.sub(CounterType, self.counter, count) catch unreachable;
56 if (self.counter == 0) {
57 const temp = self.waiters;
58 self.waiters = null;
59 break :blk temp;
60 }
61 break :blk null;
62 };
63
64 // We don't need to hold the lock to reschedule any potential waiter.
65 while (waiters) |w| {
66 const temp_w = w;
67 waiters = w.next;
68 global_event_loop.onNextTick(&temp_w.node);
69 }
70 }
71
72 pub fn wait(self: *Self) void {
73 const held = self.mutex.acquire();
74
75 if (self.counter == 0) {
76 held.release();
77 return;
78 }
79
80 var self_waiter: Waiter = undefined;
81 self_waiter.node.data = @frame();
82 if (self.waiters) |head| {
83 head.tail.next = &self_waiter;
84 head.tail = &self_waiter;
85 } else {
86 self.waiters = &self_waiter;
87 self_waiter.tail = &self_waiter;
88 self_waiter.next = null;
89 }
90 suspend {
91 held.release();
92 }
93 }
94 };
95}
96
97test "basic WaitGroup usage" {
98 if (!std.io.is_async) return error.SkipZigTest;
99
100 // TODO https://github.com/ziglang/zig/issues/1908
101 if (builtin.single_threaded) return error.SkipZigTest;
102
103 // TODO https://github.com/ziglang/zig/issues/3251
104 if (builtin.os.tag == .freebsd) return error.SkipZigTest;
105
106 var initial_wg = WaitGroup{};
107 var final_wg = WaitGroup{};
108
109 try initial_wg.begin(1);
110 try final_wg.begin(1);
111 var task_frame = async task(&initial_wg, &final_wg);
112 initial_wg.finish(1);
113 final_wg.wait();
114 await task_frame;
115}
116
117fn task(wg_i: *WaitGroup, wg_f: *WaitGroup) void {
118 wg_i.wait();
119 wg_f.finish(1);
120}