1 // Copyright 2011 The Go Authors. All rights reserved. 2 // Use of this source code is governed by a BSD-style 3 // license that can be found in the LICENSE file. 4 5 package sync 6 7 import ( 8 "internal/race" 9 "sync/atomic" 10 "unsafe" 11 ) 12 13 // A WaitGroup waits for a collection of goroutines to finish. 14 // The main goroutine calls [WaitGroup.Add] to set the number of 15 // goroutines to wait for. Then each of the goroutines 16 // runs and calls [WaitGroup.Done] when finished. At the same time, 17 // [WaitGroup.Wait] can be used to block until all goroutines have finished. 18 // 19 // A WaitGroup must not be copied after first use. 20 // 21 // In the terminology of [the Go memory model], a call to [WaitGroup.Done] 22 // “synchronizes before” the return of any Wait call that it unblocks. 23 // 24 // [the Go memory model]: https://go.dev/ref/mem 25 type WaitGroup struct { 26 noCopy noCopy 27 28 state atomic.Uint64 // high 32 bits are counter, low 32 bits are waiter count. 29 sema uint32 30 } 31 32 // Add adds delta, which may be negative, to the [WaitGroup] counter. 33 // If the counter becomes zero, all goroutines blocked on [WaitGroup.Wait] are released. 34 // If the counter goes negative, Add panics. 35 // 36 // Note that calls with a positive delta that occur when the counter is zero 37 // must happen before a Wait. Calls with a negative delta, or calls with a 38 // positive delta that start when the counter is greater than zero, may happen 39 // at any time. 40 // Typically this means the calls to Add should execute before the statement 41 // creating the goroutine or other event to be waited for. 42 // If a WaitGroup is reused to wait for several independent sets of events, 43 // new Add calls must happen after all previous Wait calls have returned. 44 // See the WaitGroup example. 45 func (wg *WaitGroup) Add(delta int) { 46 if race.Enabled { 47 if delta < 0 { 48 // Synchronize decrements with Wait. 49 race.ReleaseMerge(unsafe.Pointer(wg)) 50 } 51 race.Disable() 52 defer race.Enable() 53 } 54 state := wg.state.Add(uint64(delta) << 32) 55 v := int32(state >> 32) 56 w := uint32(state) 57 if race.Enabled && delta > 0 && v == int32(delta) { 58 // The first increment must be synchronized with Wait. 59 // Need to model this as a read, because there can be 60 // several concurrent wg.counter transitions from 0. 61 race.Read(unsafe.Pointer(&wg.sema)) 62 } 63 if v < 0 { 64 panic("sync: negative WaitGroup counter") 65 } 66 if w != 0 && delta > 0 && v == int32(delta) { 67 panic("sync: WaitGroup misuse: Add called concurrently with Wait") 68 } 69 if v > 0 || w == 0 { 70 return 71 } 72 // This goroutine has set counter to 0 when waiters > 0. 73 // Now there can't be concurrent mutations of state: 74 // - Adds must not happen concurrently with Wait, 75 // - Wait does not increment waiters if it sees counter == 0. 76 // Still do a cheap sanity check to detect WaitGroup misuse. 77 if wg.state.Load() != state { 78 panic("sync: WaitGroup misuse: Add called concurrently with Wait") 79 } 80 // Reset waiters count to 0. 81 wg.state.Store(0) 82 for ; w != 0; w-- { 83 runtime_Semrelease(&wg.sema, false, 0) 84 } 85 } 86 87 // Done decrements the [WaitGroup] counter by one. 88 func (wg *WaitGroup) Done() { 89 wg.Add(-1) 90 } 91 92 // Wait blocks until the [WaitGroup] counter is zero. 93 func (wg *WaitGroup) Wait() { 94 if race.Enabled { 95 race.Disable() 96 } 97 for { 98 state := wg.state.Load() 99 v := int32(state >> 32) 100 w := uint32(state) 101 if v == 0 { 102 // Counter is 0, no need to wait. 103 if race.Enabled { 104 race.Enable() 105 race.Acquire(unsafe.Pointer(wg)) 106 } 107 return 108 } 109 // Increment waiters count. 110 if wg.state.CompareAndSwap(state, state+1) { 111 if race.Enabled && w == 0 { 112 // Wait must be synchronized with the first Add. 113 // Need to model this is as a write to race with the read in Add. 114 // As a consequence, can do the write only for the first waiter, 115 // otherwise concurrent Waits will race with each other. 116 race.Write(unsafe.Pointer(&wg.sema)) 117 } 118 runtime_Semacquire(&wg.sema) 119 if wg.state.Load() != 0 { 120 panic("sync: WaitGroup is reused before previous Wait has returned") 121 } 122 if race.Enabled { 123 race.Enable() 124 race.Acquire(unsafe.Pointer(wg)) 125 } 126 return 127 } 128 } 129 } 130