GO sync包之WaitGroup详解

WaitGroup

WaitGroup可以用于等待多个并发操作的完成。如果你不考虑并发操作的结果或者你有其他方式收集并发操作的结果，还是建议使用channel和select关键字。

一个var wg sync.WaitGroup
wg.Add(1) // 调用Add方法并传入参数1，表示开启一个goroutine
go func() {
   defer wg.Done() // 使用defer关键字确保goroutine退出时调用Done方法，Done方法表示退出 了一个goroutine
   fmt.Println("1st goroutine sleeping...")
   time.Sleep(1)
}()

wg.Add(1)
go func() {
   defer wg.Done()
   fmt.Println("2nd goroutine sleeping...")
   time.Sleep(2)
}()

wg.Wait() //调用Wait方法阻塞main goroutine直到其余子goroutine全部退出
fmt.Println("All goroutines complete")

结果：

2nd goroutine sleeping...
1st goroutine sleeping...
All goroutines complete

WaitGroup相当于一个线程安全的计数器，调用Add方法表明计数器增加多少，调用Done表示计数器减一，Wait方法表示阻塞直到计数器等于。注意Add方法要在goroutine外部调用，因为如果不这样做，会产生数据竞争的问题。调度器不保证什么时候执行goroutine，有可能已经执行到Wait方法时，goroutine还没有执行，这时Wait方法并不会阻塞程序运行。同时，不建议将WaitGroup做为方法的参数进行传递，因为值传递的问题，方法外和方法内属于两个WaitGroup，如果非要当作方法参数进行传递，可以使用指针。

WaitGroup的结构

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
   // 64位系统中state1[0]为counter，state1[1]为waiter的个数，state1[2]为信号量
   // 32位系统中state[1]为counter，state[2]为waiter的个数，state1[0]为信号量
}

// 返回counter和waiter的指针和信号量
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 {
      return (*uint64)(unsafe.Pointer(&wg.state1[1])), &wg.state1[0]
   }
}

Add方法

func (wg *WaitGroup) Add(delta int) {
   statep, semap := wg.state()
   
   state := atomic.AddUint64(statep, uint64(delta)<<32) //通过原子操作对counter进行增加操作
   v := int32(state >> 32) // 获取到添加后的counter
   w := uint32(state) // 获取waiter
  
   if v < 0 {
      // 如果counter小于0，抛出panic
      panic("sync: negative WaitGroup counter")
   }
   if w != 0 && delta > 0 && v == int32(delta) {
      // wait方法释放阻塞期间，调用了Add方法，如果想重用WaitGroup，需要确保上一次的wait方法调用完毕后，再调用Add方法
      panic("sync: WaitGroup misuse: Add called concurrently with Wait")
   }
   if v > 0 || w == 0 {
      return
   }
   // This goroutine has set counter to 0 when waiters > 0.
   // Now there can't be concurrent mutations of state:
   // - Adds must not happen concurrently with Wait,
   // - Wait does not increment waiters if it sees counter == 0.
   // Still do a cheap sanity check to detect WaitGroup misuse.
   if *statep != state {
      // 并发情况下，如果waiter大于0时Add在并发情况下调用会出现该情况
      panic("sync: WaitGroup misuse: Add called concurrently with Wait")
   }
   // Reset waiters count to 0.
   // 重置waiter 和count，并且释放waiters
   *statep = 0
   for ; w != 0; w-- {
      runtime_Semrelease(semap, false, 0)
   }
}

Wait方法

func (wg *WaitGroup) Wait() {
   statep, semap := wg.state()

// 循环
   for {
      state := atomic.LoadUint64(statep) //原子操作获取counter和waiter
      v := int32(state >> 32)
      w := uint32(state)
      if v == 0 {
         // 值等于0时退出
         return
      }
      // Increment waiters count.
      // 原子操作获取状态值，如果获取到了在waiter位加1
      if atomic.CompareAndSwapUint64(statep, state, state+1) {
         //等待Add方法中调用runtime_Semrelease方法
         runtime_Semacquire(semap)
         if *statep != 0 {
            panic("sync: WaitGroup is reused before previous Wait has returned")
         }

         return
      }
   }
}