【Golang】WaitGroup 实现原理

前言

在并发编程中，协调多个 goroutine 的执行顺序和同步是一个常见的需求。Golang 提供了 sync.WaitGroup 来简化这一过程。WaitGroup 允许主 goroutine 等待一组 goroutine 完成工作。本文将详细介绍 sync.WaitGroup 的实现原理及其使用方法，帮助读者更好地理解和应用这一工具。

一、介绍

1. 什么是 sync.WaitGroup
sync.WaitGroup 是 Golang 标准库中的一个同步原语，用于等待一组 goroutine 完成。它提供了一种简单的方式来协调多个 goroutine 的执行顺序，确保主 goroutine 在所有子 goroutine 完成后再继续执行。

2. sync.WaitGroup 的基本结构
sync.WaitGroup 的核心是一个计数器和一个条件变量。计数器记录需要等待的 goroutine 数量，条件变量用于在计数器变为零时通知等待的 goroutine。

WaitGroup 的结构体定义如下：

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

state1 数组包含了三个字段：

• counter：计数器，记录需要等待的 goroutine 数量。
• waiters：等待的 goroutine 数量。
• sema：信号量，用于通知等待的 goroutine。

二、实现原理

1. Add 方法
Add 方法用于增加或减少计数器的值。正数表示增加需要等待的 goroutine 数量，负数表示减少。

示例：

func (wg *WaitGroup) Add(delta int) {
    statep := &wg.state1[0]
    state := atomic.AddUint32(statep, uint32(delta)<<1)
    v := int32(state >> 1)
    if v < 0 {
        panic("sync: negative WaitGroup counter")
    }
    if delta > 0 && v == int32(delta) {
        return
    }
    if delta < 0 && v == 0 {
        for i := 0; i < int(state&1); i++ {
            runtime_Semrelease(&wg.state1[2], false, 0)
        }
    }
}

2. Done 方法
Done 方法用于减少计数器的值，相当于 Add(-1)。当计数器变为零时，通知所有等待的 goroutine。

示例：

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

3. Wait 方法
Wait 方法用于等待计数器变为零。当计数器为零时，Wait 方法立即返回；否则，当前 goroutine 会被阻塞，直到计数器变为零。

示例：

func (wg *WaitGroup) Wait() {
    statep := &wg.state1[0]
    if atomic.LoadUint32(statep)>>1 == 0 {
        return
    }
    runtime_Semacquire(&wg.state1[2])
}

三、使用方式

1. 基本使用
sync.WaitGroup 的基本使用方法包括以下几个步骤：

• 创建一个 WaitGroup 实例。
• 使用 Add 方法设置需要等待的 goroutine 数量。
• 在每个 goroutine 中调用 Done 方法，表示完成工作。
• 在主 goroutine 中调用 Wait 方法，等待所有子 goroutine 完成。
  示例：

package main

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

func worker(id int, wg *sync.WaitGroup) {
    defer wg.Done()
    fmt.Printf("Worker %d starting\n", id)
    time.Sleep(time.Second)
    fmt.Printf("Worker %d done\n", id)
}

func main() {
    var wg sync.WaitGroup

    for i := 1; i <= 3; i++ {
        wg.Add(1)
        go worker(i, &wg)
    }

    wg.Wait()
    fmt.Println("All workers done")
}

2. 注意事项

• WaitGroup 的计数器不能为负数。如果调用 Add 方法时传入的负数使计数器变为负数，会引发 panic。
• WaitGroup 不能被复制。应始终传递指向 WaitGroup 的指针，以确保所有 goroutine 操作的是同一个 WaitGroup 实例。

四、总结

sync.WaitGroup 是 Golang 中用于协调多个 goroutine 执行顺序的强大工具。通过理解 WaitGroup 的实现原理和使用方法，开发者可以更好地管理并发任务，确保主 goroutine 在所有子 goroutine 完成后再继续执行。希望通过本文的介绍，读者能够深入了解 sync.WaitGroup 的实现原理及其使用方法，并在实际项目中灵活运用。