Go语言之sync包 WaitGroup的使用和底层实现

在 Go 语言里，sync 包中的 WaitGroup 是一个实用工具，用于等待一组 goroutine 完成任务。其核心原理是通过内部维护一个计数器，该计数器初始值为 0，每启动一个新的 goroutine 就将计数器加 1，每个 goroutine 完成任务后会将计数器减 1，当计数器变为 0 时，意味着所有 goroutine 都已完成任务。

下面为你展示 WaitGroup 的使用示例：

package mainimport ("fmt""sync"
)var (counter intmutex   sync.Mutex
)func increment() {mutex.Lock()defer mutex.Unlock()counter++
}func main() {var wg sync.WaitGroupnumGoroutines := 1000wg.Add(numGoroutines)for i := 0; i < numGoroutines; i++ {go func() {defer wg.Done()increment()}()}wg.Wait()wg.Add(1)fmt.Println("Counter:", counter)
}

代码解释

1. sync.WaitGroup 的声明：在 main 函数里，借助 var wg sync.WaitGroup 声明了一个 WaitGroup 实例。
2. wg.Add 方法：在启动每个 worker goroutine 之前，调用 wg.Add(1) 把计数器的值加 1，以此表明有一个新的 goroutine 开始工作。
3. defer wg.Done()：在 worker 函数中，使用 defer wg.Done() 保证当该函数执行结束时，计数器的值会减 1。
4. wg.Wait 方法：在 main 函数里调用 wg.Wait()，此操作会让 main 函数阻塞，直至计数器的值变为 0，也就是所有的 goroutine 都已完成任务。

注意事项

• 要保证 wg.Add 方法在 goroutine 启动前调用，wg.Done 方法在 goroutine 结束时调用。
• 不能在 wg.Wait 调用之后再调用 wg.Add，否则会引发恐慌（panic），实际上高版本并没有这个问题。
• WaitGroup 是按值传递的，若要在多个 goroutine 中使用，需传递其指针。

底层原理实现

// Copyright 2011 The Go Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.package syncimport ("internal/race""sync/atomic""unsafe"
)// A WaitGroup waits for a collection of goroutines to finish.
// The main goroutine calls [WaitGroup.Add] to set the number of
// goroutines to wait for. Then each of the goroutines
// runs and calls [WaitGroup.Done] when finished. At the same time,
// [WaitGroup.Wait] can be used to block until all goroutines have finished.
//
// A WaitGroup must not be copied after first use.
//
// In the terminology of [the Go memory model], a call to [WaitGroup.Done]
// “synchronizes before” the return of any Wait call that it unblocks.
//
// [the Go memory model]: https://go.dev/ref/mem
type WaitGroup struct {noCopy noCopystate atomic.Uint64 // high 32 bits are counter, low 32 bits are waiter count.sema  uint32
}// Add adds delta, which may be negative, to the [WaitGroup] counter.
// If the counter becomes zero, all goroutines blocked on [WaitGroup.Wait] are released.
// If the counter goes negative, Add panics.
//
// Note that calls with a positive delta that occur when the counter is zero
// must happen before a Wait. Calls with a negative delta, or calls with a
// positive delta that start when the counter is greater than zero, may happen
// at any time.
// Typically this means the calls to Add should execute before the statement
// creating the goroutine or other event to be waited for.
// If a WaitGroup is reused to wait for several independent sets of events,
// new Add calls must happen after all previous Wait calls have returned.
// See the WaitGroup example.
func (wg *WaitGroup) Add(delta int) {if race.Enabled {if delta < 0 {// Synchronize decrements with Wait.race.ReleaseMerge(unsafe.Pointer(wg))}race.Disable()defer race.Enable()}state := wg.state.Add(uint64(delta) << 32)v := int32(state >> 32)w := uint32(state)if race.Enabled && delta > 0 && v == int32(delta) {// The first increment must be synchronized with Wait.// Need to model this as a read, because there can be// several concurrent wg.counter transitions from 0.race.Read(unsafe.Pointer(&wg.sema))}if v < 0 {panic("sync: negative WaitGroup counter")}if w != 0 && delta > 0 && v == int32(delta) {panic("sync: WaitGroup misuse: Add called concurrently with Wait")}if v > 0 || w == 0 {return}// This goroutine has set counter to 0 when waiters > 0.// Now there can't be concurrent mutations of state:// - Adds must not happen concurrently with Wait,// - Wait does not increment waiters if it sees counter == 0.// Still do a cheap sanity check to detect WaitGroup misuse.if wg.state.Load() != state {panic("sync: WaitGroup misuse: Add called concurrently with Wait")}// Reset waiters count to 0.wg.state.Store(0)for ; w != 0; w-- {runtime_Semrelease(&wg.sema, false, 0)}
}// Done decrements the [WaitGroup] counter by one.
func (wg *WaitGroup) Done() {wg.Add(-1)
}// Wait blocks until the [WaitGroup] counter is zero.
func (wg *WaitGroup) Wait() {if race.Enabled {race.Disable()}for {state := wg.state.Load()v := int32(state >> 32)w := uint32(state)if v == 0 {// Counter is 0, no need to wait.if race.Enabled {race.Enable()race.Acquire(unsafe.Pointer(wg))}return}// Increment waiters count.if wg.state.CompareAndSwap(state, state+1) {if race.Enabled && w == 0 {// Wait must be synchronized with the first Add.// Need to model this is as a write to race with the read in Add.// As a consequence, can do the write only for the first waiter,// otherwise concurrent Waits will race with each other.race.Write(unsafe.Pointer(&wg.sema))}runtime_Semacquire(&wg.sema)if wg.state.Load() != 0 {panic("sync: WaitGroup is reused before previous Wait has returned")}if race.Enabled {race.Enable()race.Acquire(unsafe.Pointer(wg))}return}}
}