WaitGroup实现机制 - 元素码农

发布时间: 2025-04-24 23:09

↑

# WaitGroup实现机制

## 简介

WaitGroup是Go语言标准库sync包中提供的一种同步原语，用于等待一组goroutine完成执行。它的主要用途是协调多个并发任务，确保所有任务都完成后再继续执行后续操作。WaitGroup就像一个计数器，可以增加计数（Add方法）、减少计数（Done方法）以及等待计数归零（Wait方法）。

## 基本用法

```go
package main

import (
    "fmt"
    "sync"
    "time"
)

func main() {
    var wg sync.WaitGroup
    
    // 设置计数器为2，表示有两个任务需要等待
    wg.Add(2)
    
    // 启动第一个goroutine
    go func() {
        defer wg.Done() // 任务完成时减少计数
        fmt.Println("任务1执行中...")
        time.Sleep(1 * time.Second)
        fmt.Println("任务1完成")
    }()
    
    // 启动第二个goroutine
    go func() {
        defer wg.Done() // 任务完成时减少计数
        fmt.Println("任务2执行中...")
        time.Sleep(2 * time.Second)
        fmt.Println("任务2完成")
    }()
    
    // 等待所有任务完成
    wg.Wait()
    fmt.Println("所有任务已完成")
}
```

## 内部结构

WaitGroup的内部结构相对简单，但实现非常精巧。以下是Go 1.17版本中WaitGroup的定义：

```go
type WaitGroup struct {
    noCopy noCopy
    state1 [3]uint32
}
```

这个结构看起来很简洁，但实际上包含了几个重要的字段：

- `noCopy`：一个空结构体，用于防止WaitGroup被复制使用
- `state1`：一个长度为3的uint32数组，实际上包含了两个概念上的字段：
  - 计数器（counter）：记录当前等待完成的goroutine数量
  - 等待者计数（waiter count）：记录有多少个goroutine在等待
  - 信号量（semaphore）：用于阻塞等待的goroutine

为了提高内存效率，WaitGroup将这些字段打包在一个12字节的数组中，并根据CPU架构的不同，以不同方式解释这些字节。

## 核心方法实现

### Add方法

Add方法用于增加或减少WaitGroup的计数器：

```go
func (wg *WaitGroup) Add(delta int) {
    statep, semap := wg.state()
    state := atomic.AddUint64(statep, uint64(delta)<<32)
    v := int32(state >> 32) // 提取计数器值
    w := uint32(state) // 提取等待者数量
    
    if v < 0 {
        panic("sync: negative WaitGroup counter")
    }
    
    // 如果计数器变为0，且有goroutine在等待，则唤醒所有等待者
    if v == 0 && w != 0 {
        atomic.AddUint64(statep, ^uint64(0)>>32) // 清零等待者计数
        for ; w != 0; w-- {
            runtime_Semrelease(semap, false, 0) // 释放信号量，唤醒等待的goroutine
        }
    }
}
```

主要逻辑：
1. 原子地增加计数器值
2. 检查计数器是否为负，如果是则触发panic
3. 如果计数器变为0且有goroutine在等待，则唤醒所有等待的goroutine

### Done方法

Done方法实际上是对Add方法的简单封装，相当于`Add(-1)`：

```go
func (wg *WaitGroup) Done() {
    wg.Add(-1)
}
```

### Wait方法

Wait方法会阻塞当前goroutine，直到WaitGroup的计数器变为0：

```go
func (wg *WaitGroup) Wait() {
    statep, semap := wg.state()
    for {
        state := atomic.LoadUint64(statep)
        v := int32(state >> 32) // 提取计数器值
        w := uint32(state) // 提取等待者数量
        
        if v == 0 {
            // 如果计数器已经为0，直接返回
            return
        }
        
        // 原子地增加等待者计数
        if atomic.CompareAndSwapUint64(statep, state, state+1) {
            // 等待信号量
            runtime_Semacquire(semap)
            // 被唤醒后直接返回，不需要再检查计数器
            // 因为只有在计数器变为0时才会被唤醒
            return
        }
    }
}
```

主要逻辑：
1. 检查计数器是否为0，如果是则直接返回
2. 如果计数器不为0，则原子地增加等待者计数
3. 通过信号量阻塞当前goroutine，直到被Add/Done方法唤醒

## 内存布局优化

WaitGroup的内存布局经过了精心设计，以适应不同的CPU架构。在64位系统上，state1数组的前8个字节用于存储计数器和等待者计数（各占4字节），后4个字节用于存储信号量。而在32位系统上，为了保证64位原子操作的对齐要求，这个布局会有所不同。

```go
func (wg *WaitGroup) state() (statep *uint64, semap *uint32) {
    if uintptr(unsafe.Pointer(&wg.state1))%8 == 0 {
        // 64位对齐的情况
        return (*uint64)(unsafe.Pointer(&wg.state1)), &wg.state1[2]
    } else {
        // 32位对齐的情况
        return (*uint64)(unsafe.Pointer(&wg.state1[1])), &wg.state1[0]
    }
}
```

这种设计确保了在任何架构上都能高效地进行原子操作。

## 使用注意事项

1. **计数器不能为负**：如果调用Done()的次数超过了Add()设置的计数值，WaitGroup会panic

2. **不要复制WaitGroup**：WaitGroup包含了noCopy字段，通过go vet工具可以检测到复制行为

3. **Add必须在Wait之前调用**：确保在调用Wait()方法前设置好计数器的初始值

4. **可重用性**：WaitGroup在计数器归零后可以被重新使用，但要确保没有goroutine仍在等待

## 实际应用场景

### 并行任务处理

```go
func processItems(items []Item) error {
    var wg sync.WaitGroup
    errs := make(chan error, len(items))
    
    for _, item := range items {
        wg.Add(1)
        go func(item Item) {
            defer wg.Done()
            if err := processItem(item); err != nil {
                errs <- err
            }
        }(item)
    }
    
    // 等待所有处理完成
    wg.Wait()
    close(errs)
    
    // 检查是否有错误
    for err := range errs {
        if err != nil {
            return err
        }
    }
    
    return nil
}
```

### 优雅关闭

```go
func main() {
    var wg sync.WaitGroup
    
    // 启动工作goroutine
    for i := 0; i < 10; i++ {
        wg.Add(1)
        go worker(&wg, i)
    }
    
    // 接收关闭信号
    sigCh := make(chan os.Signal, 1)
    signal.Notify(sigCh, syscall.SIGINT, syscall.SIGTERM)
    <-sigCh
    
    // 发送停止信号给所有worker
    stopWorkers()
    
    // 等待所有worker优雅退出
    fmt.Println("等待所有worker退出...")
    wg.Wait()
    fmt.Println("所有worker已退出，程序关闭")
}
```

## 与其他同步原语的比较

| 同步原语 | 主要用途 | 特点 |
|---------|---------|------|
| WaitGroup | 等待一组goroutine完成 | 简单易用，适合一次性等待场景 |
| Channel | 通信和同步 | 更灵活，可传递数据，适合复杂协作 |
| Mutex | 互斥访问共享资源 | 保护临界区，一次只允许一个goroutine访问 |
| RWMutex | 读写锁 | 允许多个读操作并发，写操作互斥 |
| Cond | 条件变量 | 等待特定条件满足，可以唤醒一个或所有等待者 |

## 总结

WaitGroup是Go语言中一个简单而强大的同步原语，特别适合于等待一组并发任务完成的场景。它的实现虽然简洁，但包含了许多精巧的设计，如内存布局优化、原子操作和信号量使用等。理解WaitGroup的内部实现有助于我们更好地使用它，并在适当的场景选择合适的同步机制。

在实际开发中，WaitGroup通常与其他并发原语（如Channel、Context等）结合使用，以构建更复杂、更健壮的并发程序。