1 package sync
2 3 import "internal/task"
4 5 type WaitGroup struct {
6 futex task.Futex
7 }
8 9 func (wg *WaitGroup) Add(delta int) {
10 switch {
11 case delta > 0:
12 // Delta is positive.
13 for {
14 // Check for overflow.
15 counter := wg.futex.Load()
16 if uint32(delta) > (^uint32(0))-counter {
17 panic("sync: WaitGroup counter overflowed")
18 }
19 20 // Add to the counter.
21 if wg.futex.CompareAndSwap(counter, counter+uint32(delta)) {
22 // Successfully added.
23 return
24 }
25 }
26 default:
27 // Delta is negative (or zero).
28 for {
29 counter := wg.futex.Load()
30 31 // Check for underflow.
32 if uint32(-delta) > counter {
33 panic("sync: negative WaitGroup counter")
34 }
35 36 // Subtract from the counter.
37 if !wg.futex.CompareAndSwap(counter, counter-uint32(-delta)) {
38 // Could not swap, trying again.
39 continue
40 }
41 42 // If the counter is zero, everything is done and the waiters should
43 // be resumed.
44 // When there are multiple thread, there is a chance for the counter
45 // to go to zero, WakeAll to be called, and then the counter to be
46 // incremented again before a waiting goroutine has a chance to
47 // check the new (zero) value. However the last increment is
48 // explicitly given in the docs as something that should not be
49 // done:
50 //
51 // > Note that calls with a positive delta that occur when the
52 // > counter is zero must happen before a Wait.
53 //
54 // So we're fine here.
55 if counter-uint32(-delta) == 0 {
56 // TODO: this is not the most efficient implementation possible
57 // because we wake up all waiters unconditionally, even if there
58 // might be none. Though since the common usage is for this to
59 // be called with at least one waiter, it's probably fine.
60 wg.futex.WakeAll()
61 }
62 63 // Successfully swapped (and woken all waiting tasks if needed).
64 return
65 }
66 }
67 }
68 69 func (wg *WaitGroup) Done() {
70 wg.Add(-1)
71 }
72 73 func (wg *WaitGroup) Wait() {
74 for {
75 counter := wg.futex.Load()
76 if counter == 0 {
77 return // everything already finished
78 }
79 80 if wg.futex.Wait(counter) {
81 // Successfully woken by WakeAll (in wg.Add).
82 break
83 }
84 }
85 }
86