Go并发编程-WaitGroup

Go并发编程-WaitGroup

​ 协同等待，任务编排利器

简单用

​ WaitGroup解决了并发等待的问题。在使用groutine执行任务时，经常需要等待goroutine全部执行完成后再执行下一步。WaitGroup并发原语非常容易的解决了这个问题。

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

	wg.Add(10)
	for i := 0; i < 10; i++ {
		go func() {
			defer wg.Done()
			fmt.Println("task")
		}()
	}
	wg.Wait()
	fmt.Println("end")
}

​ 分析以上代码：创建一个WaitGroup对象，初始值是0。通过Add(10)将他的计数器设置为10。创建10个goroutine执行任务，并在任务结束时通过Done()将计数器减1。通过Wait阻塞，等待所有的goroutine执行完成。打印结束。

看实现

type WaitGroup struct {
  // 避免复制使用的技巧。vet工具可以检查实现Locker接口，使用noCopy可以告诉vet工具是否复制使用
	noCopy noCopy
	// 复合意义的字段，包含WaitGroup的计数、阻塞在检查点的waiter数和信号量
  // 不同的操作系统组成不同，64位整数的原子操作要求整数的地址是64位对齐的，所以32位和64位操作系统值不同
  // 64位：第一个元素是waiter的数量,第二个元素WaitGroup的计数值，第三个元素是信号量
  // 32位：第一个元素是信号量并且是64bit对齐的，第二个元素是waiter的数量，第三个元素是WaitGroup的计数值
	state1 [3]uint32
}

核心功能实现

func (wg *WaitGroup) Add(delta int) {
	statep, semap := wg.state()
  //高32位是计数值v，所以delta左移32，增加到计数上
	state := atomic.AddUint64(statep, uint64(delta)<<32)
	v := int32(state >> 32) //当前计数值
	w := uint32(state) //waiter数量
	if v > 0 || w == 0 {
		return
	}
  // 如果计数值v为0并且waiter部位0，那么state的值是waiter的数量
  // 将waiter的数量设置为0，因为计数值也是0，座椅他们的组合*statep设置为0即可
	*statep = 0
	for ; w != 0; w-- {
		runtime_Semrelease(semap, false, 0)
	}
}

func (wg *WaitGroup) Done() {
	wg.Add(-1)
}
// 不断检查state的值，计数值变为0时则所有任务完成不需要阻塞，如果大于0则任务没完成，阻塞
func (wg *WaitGroup) Wait() {
	statep, semap := wg.state()
	for {
		state := atomic.LoadUint64(statep)
		v := int32(state >> 32)
		w := uint32(state)
		if v == 0 {
			return
		}
		// Increment waiters count.
		if atomic.CompareAndSwapUint64(statep, state, state+1) {
			runtime_Semacquire(semap)
			if *statep != 0 {
				panic("sync: WaitGroup is reused before previous Wait has returned")
			}
			return
		}
	}
}

别踩坑

1. 记数器设置为负值

   WaitGroup计数值必须大于等于0，否则会panic

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

   调用Done的次数过多，导致计数值小于0，引发panic

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

2. 不期望的Add的时机

   使用WaitGroup应该在所有的Add方法后调用Wait，否则会有panic或者不期望的结果

   func main() {
   	wg := sync.WaitGroup{}
   	go doSomething(&wg)
   	go doSomething(&wg)
   
   	wg.Wait()
   }
   
   func doSomething(wg *sync.WaitGroup) {
   	wg.Add(1)
   	fmt.Println("do something")
   	wg.Done()
   }
   

   上述情况，期望时两个doSomething之后再结束，但是Add是在goroutine中去add的，没有得到想要的结果。一般使用WaitGroup先Add再启动groutine

3. 重用WaitGroup

   WaitGroup是允许重用的，但是第二次使用的时候要确保计数值恢复到零值.

   func main() {
   	wg := sync.WaitGroup{}
   	wg.Add(1)
   	go func() {
   		time.Sleep(time.Millisecond)
   		wg.Done()
   		wg.Add(1)
   	}()
   
   	wg.Wait()
   }