golang sync.WaitGroup 源码分析

最近在学习golang源码，学习golang源码是学习golang的非常好的途径。

先来记录一波sync包的学习。版本 go1.14.2 darwin/amd64

sync.WaitGroup

我们一般使用sync.waitGroup 做并发控制，使用方式一般如下

func main() {
    wg := sync.WaitGroup{}

    for i := 0; i < 10; i++ {
        wg.Add(1)
        go func(index int) {
            defer wg.Done()
            fmt.Println(index)
        }(i)
    }
    
    wg.Wait()
}

所以我们先看下waitGroup 主要有是哪个函数，Add, Done, Wait 函数，在看函数源码前，我们先看下WaitGroup 结构

type WaitGroup struct {
    noCopy noCopy          //禁止拷贝，如果有拷贝构建不会报错，可以用go vet 或 go tool vet 检测是否有拷贝错误，
    state1 [3]uint32      //重要，存有 计数器，等待数，信号量的值
}

其中 state1 成员存放的值在64 位系统中如下：

接下来看下Add方法

func (wg *WaitGroup) Add(delta int) {
  statep, semap := wg.state()   //获取wg状态，statep地址(高32位的计数器值和低32位的等待数量值), semap 信号量
    if race.Enabled {  //数据竞态检测，默认是false,开启消耗cpu性能 ,先不管
        _ = *statep
        if delta < 0 { 
            race.ReleaseMerge(unsafe.Pointer(wg))
        }
        race.Disable()
        defer race.Enable()
    }
    state := atomic.AddUint64(statep, uint64(delta)<<32) //原子操作给计数器加上delta的值
    v := int32(state >> 32)     //高32位 计数器
    w := uint32(state)   //低32位 等待数
    if race.Enabled && delta > 0 && v == int32(delta) {  //数据竞态检测,先不管
        race.Read(unsafe.Pointer(semap))
    }
    if v < 0 {                                  //计数器 <0 
        panic("sync: negative WaitGroup counter")
    }
    if w != 0 && delta > 0 && v == int32(delta) {    //在add之前，已经调用过wait函数？？（没太看明白）
        panic("sync: WaitGroup misuse: Add called concurrently with Wait")
    }
    if v > 0 || w == 0 { // wait数 为0 或 计数>0 直接返回
        return
    }

    if *statep != state {   // ？ 理论上应该相等
        panic("sync: WaitGroup misuse: Add called concurrently with Wait")
    }
    
    *statep = 0
    for ; w != 0; w-- {  //计数器 为0， 释放所有wait 信号量，
        runtime_Semrelease(semap, false, 0)
    }
}

Done函数, 很简单了：

// Done decrements the WaitGroup counter by one.
func (wg *WaitGroup) Done() {
    wg.Add(-1)
}

Wait函数, 把数据竞态检测 部分去了，为代码看起来简洁

func (wg *WaitGroup) Wait() {
    statep, semap := wg.state()  //获取wg状态，statep地址(高32位的计数器值和低32位的等待数量值), semap 信号量
    for {
        state := atomic.LoadUint64(statep) 
        v := int32(state >> 32)  //如上
        w := uint32(state)
        if v == 0 {  //计数为0，不用wait
            return
        }

        if atomic.CompareAndSwapUint64(statep, state, state+1) { //等待数加1
            runtime_Semacquire(semap)    //获取信号量
            if *statep != 0 {
                panic("sync: WaitGroup is reused before previous Wait has returned")
            }
            return
        }
    }
}

runtime_Semacquire 和 runtime_Semrelease 获取和释放信号量，用于休眠和唤醒 协程，具体实现等下次深入研究下