浅谈Golang并发控制WaitGroup

package main
import (
  "fmt"
  "sync"
)
func wg() {
  var wg sync.WaitGroup
  for i := 0; i < 10; i++ {
    wg.Add(1) //Add放在Goroutine外
    go func(i int) {
      //TODO
      defer wg.Done() //建议用defer保证调用
    }(i)
  }
  wg.Wait()
  fmt.Println("finished")
}
// Tip: waitGroup 不要进行copy
func errWg1() {
  var wg sync.WaitGroup
  for i := 0; i < 10; i++ {
    wg.Add(1)
    go func(i int, wg sync.WaitGroup) {
      fmt.Println(i)
      defer wg.Done()
    }(i, wg)
    /*
      type WaitGroup struct {
        noCopy noCopy
        state1 [3]uint32
      }
      sync.WaitGroup是一个结构体类型，当参数传递是值拷贝
    */
  }
  wg.Wait()
  fmt.Println("finished")
}
// Tip: waitGroup 的 Add 要在goroutine前执行
func errWg2() {
  var wg sync.WaitGroup
  for i := 0; i < 10; i++ {
    go func(i int) {
      wg.Add(1)
      fmt.Println(i)
      defer wg.Done()
    }(i)
  }
  wg.Wait()
  /*
    极端情况，wg.Add放在Goroutine里
    可能for循环完了，还没执行到wg.Add(1)
    就到了wg.Wait()这一步，此时直接wait掉结束了
  */
  fmt.Println("finished")
}
// Tip: waitGroup 的 Add 很大会有影响吗？
func errWg3() {
  var wg sync.WaitGroup
  for i := 0; i < 10; i++ {
    wg.Add(100)
    go func(i int) {
      fmt.Println(i)
      defer wg.Done()
      wg.Add(-100)
      /*
        使用是没问题的，查阅wg.Done()源码，其对应的是wg.Add(-1)，查阅wg.Add源码，发现其内部
        state := atomic.AddUint64(statep, uint64(delta)<<32)
      */
    }(i)
  }
  fmt.Println("finished")
}

package sync
import (
  "internal/race"
  "sync/atomic"
  "unsafe"
)
type WaitGroup struct {
  noCopy noCopy
  //64位值：高32位为计数器，低32位为waiter计数,另外32bit是用作信号量
  //64位原子操作需要64位对齐
  //但需要32位对齐编译器并不能保证这一点
  //所以我们分配12个字节，然后使用其中对齐的8个字节作为状态，其他4个字节作为存储
  //state1：具有复合意义的字段，包含WaitGroup计数值，waiter计数和信号量
  state1 [3]uint32
}
// state返回指向存储在wg.state1中的state和sema字段的指针。
func (wg *WaitGroup) state() (statep *uint64, semap *uint32) {
  if uintptr(unsafe.Pointer(&wg.state1))%8 == 0 {
    //如果地址是64bit对齐的，数组前两个元素做state，后一个元素做信号量
    return (*uint64)(unsafe.Pointer(&wg.state1)), &wg.state1[2]
  } else {
    //如果地址是32bit对齐的，数组后两个元素用来做state，它可以用来做64bit的原子操作，第一个元素32bit用来做信号量
    return (*uint64)(unsafe.Pointer(&wg.state1[1])), &wg.state1[0]
  }
}
// Add方法主要操作的state1字段中计数值部分
// 当Add方法被调用时，首先会将delta参数值左移32位(计数值在高32位)
// 然后内部通过原子操作将这个值加到计数值上
// 需要注意的是，delta的取值范围可正可负，因为调用Done()方法时，内部通过Add(-1)方法实现的
func (wg *WaitGroup) Add(delta int) {
  // statep表示wait数和计数值
  // 低32位表示wait数，高32位表示计数值
  statep, semap := wg.state()
  //下面的代码去掉了源码里的race判断，race判断是竞争检测相关的
  //只有当运行程序时带-race参数才起作用，不是WaitGroup的核心功能，去掉后方便看清程序主体
  // uint64(delta)<<32 将delta左移32位
  // 因为高32位表示计数值，所以将delta左移32，v = v + delta
  state := atomic.AddUint64(statep, uint64(delta)<<32)
  //state右移32位得到v的值
  v := int32(state >> 32)
  //state的低32位是w的值
  w := uint32(state)
  //v变成负数时就panic。比如第一次调用Add就传负数，就会出现这种panic
  if v < 0 {
    panic("sync: negative WaitGroup counter")
  }
  //正常Add在Wait之前调用，w是0才对
  if w != 0 && delta > 0 && v == int32(delta) {
    panic("sync: WaitGroup misuse: Add called concurrently with Wait")
  }
  //正常Add之后程序return
  if v > 0 || w == 0 {
    return
  }
  //执行到这,此时计数器V=0；那么等待计数器肯定和整个state的值相等，
  //不然只有一个情况：有人调了Add（），并且是并发调用的。
  if *statep != state {
    panic("sync: WaitGroup misuse: Add called concurrently with Wait")
  }
  //把计数器v和w都置0
  //如果计数值v为0并且waiter的数量w不为0，那么state的值就是waiter的数量
  //将waiter的数量设置为0，因为计数值v也是0,所以它们俩的组合*statep直接设置为0即可。此时需要并唤醒所有的waiter
  *statep = 0
  //发送w个信号量，也就是有多少goroutine调用了Wait()就发送多少信号量，让他们都解除阻塞
  for ; w != 0; w-- {
    runtime_Semrelease(semap, false, 0)
  }
}
//内部就是调用Add(-1)方法
func (wg *WaitGroup) Done() {
  wg.Add(-1)
}
//不断检查state值。如果其中的计数值为零，则说明所有的子goroutine已全部执行完毕，调用者不必等待，直接返回。
//如果计数值大于零，说明此时还有任务没有完成，那么调用者变成等待者，需要加入wait队列，并且阻塞自己。
func (wg *WaitGroup) Wait() {
  //statep表示wait数和计数值
  //低32位表示wait数，高32位表示计数值
  statep, semap := wg.state()
  //循环检查计数器V啥时候等于0
  for {
    //原子操作的方式获取state的值
    state := atomic.LoadUint64(statep)
    //将state右移32位，表示当前计数值
    v := int32(state >> 32)
    //w表示waiter等待值
    w := uint32(state)
    if v == 0 {
      //如果当前计数值为零，表示当前子goroutine已全部执行完毕，则直接返回
      return
    }
    //否则使用原子操作将state值加一
    //调用Wait()的goroutine的数目加1，Add()发送信号量的时候知道多少goroutine在等待
    if atomic.CompareAndSwapUint64(statep, state, state+1) {
      //获取信号量，如果拿不到就阻塞，如果拿到了就继续执行
      runtime_Semacquire(semap)
      //程序运行到这里说明Add()已经发送了信号量，Wait()解除阻塞了，这个时间state的值是0才正常。非0是个异常情况，这里panic
      if *statep != 0 {
        panic("sync: WaitGroup is reused before previous Wait has returned")
      }
      return
    }
  }
}

//wg.Add
state := atomic.AddUint64(statep, uint64(delta)<<32)
//wg.Wait
state := atomic.LoadUint64(statep)
atomic.CompareAndSwapUint64(statep, state, state+1)
//CAS一般情况下会比锁的性能更高