【Golang】——由浅入深分析并发利器sync.WaitGroup

朴灿烈づ我的快乐病毒、 2022-11-15 14:21 179阅读 0赞

### 一、前言 ###

官方文档对`sync.WatiGroup`的描述是：一个`waitGroup`对象可以等待一组协程结束，也就等待一组`goroutine`返回。有了`sync.Waitgroup`我们可以将原本顺序执行的代码在多个`Goroutine`中并发执行，加快程序处理的速度。我们先来感受一下官方给出的样例：(官方给出的样例干的事情其实是访问三个网站，访问完成打印出来)

package main
    
    import (
        "fmt"
        "sync"
        "net/http"
    )
    
    func main() {
        var wg sync.WaitGroup
        var urls = []string{
                "http://www.golang.org/",
                "http://www.google.com/",
                "http://www.baiyuxiong.com/",
        }
        for _, url := range urls {
                // Increment the WaitGroup counter.
                wg.Add(1)
                // Launch a goroutine to fetch the URL.
                go func(url string) {
                        // Decrement the counter when the goroutine completes.
                        defer wg.Done()
                        // Fetch the URL.
                        http.Get(url)
                fmt.Println(url);
                }(url)
        }
        // Wait for all HTTP fetches to complete.
        wg.Wait()
        fmt.Println("over");
    }

看完描述和样例，是不是感觉还有点绕，接着往下看。

### 二、举个栗子 ###

1. 阶段一

我现在有100个小弟在帮我做一个工程，他们分工明确，相互之间没有依赖关系，但是我必须等到所有小弟完工才能算完成工程。但我不知道每个小弟的工作效率怎么样，什么时候能来验收整个工程把我难住了。此时的我三个狗头军师给我出主意，说：咱就回家等着睡觉等就行，睡多久呢，不知道，靠猜。翻译成代码：

func main(){
        for i := 0; i < 100 ; i++{
            go fmt.Println("第 %s 个小弟完成了",i)
        }
        time.Sleep(time.Second) //睡多久，全靠蒙
    }

2. 阶段二

后来我发现不对劲，这玩意睡多久全靠猜肯定不行，赶紧换一批军师吧。正好有个军师深造回来，前几天学习了channel，说我们就可以使用channel来达到我们的目的了：

func main() {
        c := make(chan bool, 100)
        for i := 0; i < 100; i++ {
            go func(i int) {
                fmt.Println("第 %s 个小弟完成了",i)
                c <- true
            }(i)
        }
     
        for i := 0; i < 100; i++ {
            <-c
        }
    }

3. 阶段三

阶段二老感觉有点大材小用了，是不是有更简单的方式呢，别说，还真找到了。

func main() {
       wg := sync.WaitGroup{}
       wg.Add(100)
       for i := 0; i < 100; i++ {
          go func(i int) {
             defer wg.Done()
             fmt.Println("第 %s 个小弟完成了",i)
          }(i)
       }
       wg.Wait()
    }

首先使用 add(100) 把计数器的数量设置为100，每次循环运行完毕都把计数器减1，主函数使用Wait(）阻塞，直到计数器归0结束主线程结束阻塞。这里会阻塞主线程！！！

另外：以上三种方式打印出的数字皆无序。

### 三、原理剖析 ###

**1. 数据结构**

// A WaitGroup must not be copied after first use.
    type WaitGroup struct {
        noCopy noCopy
    
        // 64-bit value: high 32 bits are counter, low 32 bits are waiter count.
        // 64-bit atomic operations require 64-bit alignment, but 32-bit
        // compilers do not ensure it. So we allocate 12 bytes and then use
        // the aligned 8 bytes in them as state, and the other 4 as storage
        // for the sema.
        state1 [3]uint32
    }

总共就有两个字段，`nocopy`是为了保证该结构不会被进行拷贝，这是一种保护机制，会在后面进行介绍；`state1`主要是存储着状态和信号量，这里使用的8字节对齐处理的方式很有意思，我先来一起看看这种处理。

`1.1 state1`状态和信号量处理

`state1`这里总共被分配了`12`个字节，这里被设计了三种状态：

*  其中对齐的`8`个字节作为状态，高`32`位为计数的数量，低`32`位为等待的`goroutine`数量
 *  其中的`4`个字节作为信号量存储

提供了`(wg *WaitGroup) state() (statep *uint64, semap *uint32)`帮助我们从`state1`字段中取出他的状态和信号量，为什么要这样设计呢？那是因为64位原子操作需要64位对齐，但是32位编译器不能保证这一点，所以为了保证`waitGroup`在`32`位平台上使用的话，就必须保证在任何时候，`64位`操作不会报错。所以也就不能分成两个字段来写，考虑到字段顺序不同、平台不同，内存对齐也就不同。因此这里采用动态识别当前我们操作的`64`位数到底是不是在`8`字节对齐的位置上面，我们来分析一下`state`方法：

// state returns pointers to the state and sema fields stored within wg.state1.
    func (wg *WaitGroup) state() (statep *uint64, semap *uint32) {
        if uintptr(unsafe.Pointer(&wg.state1))%8 == 0 {
            return (*uint64)(unsafe.Pointer(&wg.state1)), &wg.state1[2]
        } else {
            return (*uint64)(unsafe.Pointer(&wg.state1[1])), &wg.state1[0]
        }
    }

当数组的首地址是处于一个`8`字节对齐的位置上时，那么就将这个数组的前`8`个字节作为`64`位值使用表示状态，后`4`个字节作为`32`位值表示信号量(`semaphore`)。同理如果首地址没有处于`8`字节对齐的位置上时，那么就将前`4`个字节作为`semaphore`，后`8`个字节作为`64`位数值。

**`2. Add()`、`Done()`方法**

`sync.WaitGroup`提供了`Add()`方法增加一个计数器，`Done()`方法减掉一个计数，`Done()`方法实现比较简单，内部调用的`Add()`方法实现的计数器减一操作，也就是增减逻辑都在`Add()`方法中，所以我们重点看一下`Add()`是如何实现的：

func (wg *WaitGroup) Add(delta int) {
      // 获取状态(Goroutine Counter 和 Waiter Counter)和信号量
        statep, semap := wg.state()
        if race.Enabled {
            _ = *statep // trigger nil deref early
            if delta < 0 {
                // Synchronize decrements with Wait.
                race.ReleaseMerge(unsafe.Pointer(wg))
            }
            race.Disable()
            defer race.Enable()
        }
      // 原子操作，goroutine counter累加delta
        state := atomic.AddUint64(statep, uint64(delta)<<32)
      // 获取当前goroutine counter的值(高32位)
        v := int32(state >> 32)
      // 获取当前waiter counter的值(低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(semap))
        }
      // Goroutine counter是不允许为负数的，否则会发生panic
        if v < 0 {
            panic("sync: negative WaitGroup counter")
        }
      // 当wait的Goroutine不为0时，累加后的counter值和delta相等,说明Add()和Wait()同时调用了,所以发生panic,因为正确的做法是先Add()后Wait()，也就是已经调用了wait()就不允许再添加任务了
        if w != 0 && delta > 0 && v == int32(delta) {
            panic("sync: WaitGroup misuse: Add called concurrently with Wait")
        }
      // 正常`Add()`方法后，`goroutine Counter`计数器大于0或者`waiter Counter`计数器等于0时，不需要释放信号量
        if v > 0 || w == 0 {
            return
        }
        // 能走到这里说明当前Goroutine Counter计数器为0，Waiter Counter计数器大于0, 到这里数据也就是允许发生变动了，如果发生变动了，则出发panic
        if *statep != state {
            panic("sync: WaitGroup misuse: Add called concurrently with Wait")
        }
        // 重置状态，并发出信号量告诉wait所有任务已经完成
        *statep = 0
        for ; w != 0; w-- {
            runtime_Semrelease(semap, false, 0)
        }
    }

上面的代码有一部分是`race`静态检测，下面的分析会省略这一部分，因为它并不是本文的重点。注释我都添加到对应的代码行上了，你是否都看懂了，没看懂不要紧，因为`Add()`是与`Wait()`方法一块使用的，所以有些逻辑与`wait()`里的逻辑是相互照应的，所以当我们看完`wait()`方法的实现在总结一下你们就明白了。

**`3. Wait()`方法**

`sync.Wait()`方法会阻塞主`Goroutine`直到`WaitGroup`计数器变为0。我们一起来看一下`Wait()`方法的源码：

// Wait blocks until the WaitGroup counter is zero.
    func (wg *WaitGroup) Wait() {
      // 获取状态(Goroutine Counter 和 Waiter Counter)和信号量
        statep, semap := wg.state()
        if race.Enabled {
            _ = *statep // trigger nil deref early
            race.Disable()
        }
        for {
        // 使用原子操作读取state，是为了保证Add中的写入操作已经完成
            state := atomic.LoadUint64(statep)
        // 获取当前goroutine counter的值(高32位)
            v := int32(state >> 32)
         // 获取当前waiter counter的值(低32位)
            w := uint32(state)
        // 如果没有任务，或者任务已经在调用`wait`方法前已经执行完成了，就不用阻塞了
            if v == 0 {
                // Counter is 0, no need to wait.
                if race.Enabled {
                    race.Enable()
                    race.Acquire(unsafe.Pointer(wg))
                }
                return
            }
            // 使用CAS操作对`waiter Counter`计数器进行+1操作，外面有for循环保证这里可以进行重试操作
            if atomic.CompareAndSwapUint64(statep, 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(semap))
                }
          // 在这里获取信号量，使线程进入睡眠状态，与Add方法中最后的增加信号量相对应，也就是当最后一个任务调用Done方法
          // 后会调用Add方法对goroutine counter的值减到0，就会走到最后的增加信号量
                runtime_Semacquire(semap)
          // 在Add方法中增加信号量时已经将statep的值设为0了，如果这里不是0，说明在wait之后又调用了Add方法，使用时机不对，触发panic
                if *statep != 0 {
                    panic("sync: WaitGroup is reused before previous Wait has returned")
                }
                if race.Enabled {
                    race.Enable()
                    race.Acquire(unsafe.Pointer(wg))
                }
                return
            }
        }
    }

**4. 源码总结**

分了源码，我们可以总结如下：

*  `Add`方法与`wait`方法不可以并发同时调用，`Add`方法要在`wait`方法之前调用.
 *  `Add()`设置的值必须与实际等待的`goroutine`个数一致，否则会`panic`.
 *  调用了`wait`方法后，必须要在`wait`方法返回以后才能再次重新使用`waitGroup`，也就是`Wait`没有返回之前不要在调用`Add`方法，否则会发生`Panic`.
 *  `Done` 只是对`Add` 方法的简单封装，我们可以向 `Add`方法传入任意负数（需要保证计数器非负）快速将计数器归零以唤醒等待的 `Goroutine`.
 *  `waitGroup`对象只能有一份，不可以拷贝给其他变量，否则会造成意想不到的`Bug`.

### 四、参考资料 ###

1. [https://segmentfault.com/a/1190000039192995?utm\_source=tag-newest][https_segmentfault.com_a_1190000039192995_utm_source_tag-newest]

[https_segmentfault.com_a_1190000039192995_utm_source_tag-newest]: https://segmentfault.com/a/1190000039192995?utm_source=tag-newest