golang源代码阅读，sync系列-WaitGroup

概览

使用场景：管理多个goroutine，用于等待多个goroutine都完成后，主goroutine继续运行

例子

wg  := new(sync.WaitGroup)
wg.Add(n) //n为要等待的goroutine个数
go func() {
    wg.Done() //对应wg.Add(-1)
}
...
wg.Wait()

源码解读

总结

通过state1一个数组来表示当前的delta，waiter、semap信号量大小，如果调用Add，会delta会加上输入的大小，然后返回，如果多个Wait调用，每个Wait都会阻塞，每个进入waiter都会加1，因为semap信号量为0，所以每个waiter都会阻塞，直到delta为0（关键点，Add加负数的时候，如果delta不为0，直接返回，为0，则循环释放信号量），add会循环waiter，每回减1，释放信号量，每回减1，直到为0，每释放一个信号量，一个Wait返回，直到信号量为0，每个Wait都返回了。

结构体字段介绍

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
}

1. noCopy 标记位，用作go vet检查，如果当前Struct被copy了，则会报错
   实例代码：

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

	copyOne := wg
	copyOne.Add(1)
	go func() {
		fmt.Println("add -1")
		wg.Done()
	}()
	wg.Wait()
}

上面代码可以正常运行
goland会提示如下错误

运行go vet会报下面错误

2. state1存储信号量、waiter、delta
state1[0]:semap，信号量
state1[1]:waiter个数
state1[2]:delta

函数介绍

Add函数

func (wg *WaitGroup) Add(delta int) {
	statep, semap := wg.state() //对state1进行拆分，获取statep和semap，statep对应state1[1]state1[2]，semap的对应statep[0]
	state := atomic.AddUint64(statep, uint64(delta)<<32) //stetep[1]即state1[2]存储delta
	v := int32(state >> 32) //v为state[1],当前delta大小
	w := uint32(state) // w为state[0],当前waiter个数

	if v < 0 {
		panic("sync: negative WaitGroup counter")
	}
	if w != 0 && delta > 0 && v == int32(delta) {
		panic("sync: WaitGroup misuse: Add called concurrently with Wait")
	}
	if v > 0 || w == 0 { //关键点，当前还未调用Wait()或者当前的delta总数还是大于0
		return
	}
	
	if *statep != state {
		panic("sync: WaitGroup misuse: Add called concurrently with Wait")
	}
	// Reset waiters count to 0.
	// 当前v即delta为0时，则一个个通知waiter，通过进行release信号量，实际即semap+1
	*statep = 0
	for ; w != 0; w-- { //如果有多个waiter，semap为waiter个数
		runtime_Semrelease(semap, false, 0)
	}
}

Wait函数

func (wg *WaitGroup) Wait() {
	statep, semap := wg.state() //同Add方法第一行，获取statep和semap
	for {
		state := atomic.LoadUint64(statep) 
		v := int32(state >> 32) // 当前delta大小
		w := uint32(state) //当前waiter个数
		if v == 0 { // 如果当前delta大小为0，则直接返回，没有Add过，或者已经被Done掉了
			return
		}
		// Increment waiters count. 如果当前state没有被修改过，则statep加1，实际是w+1了，即多了一个waiter
		if atomic.CompareAndSwapUint64(statep, state, state+1) {
			runtime_Semacquire(semap) //这是acquire一个实际为semap减1，因为semap当前还是0，这时会被阻塞
			if *statep != 0 {
				panic("sync: WaitGroup is reused before previous Wait has returned")
			}
			if race.Enabled {
				race.Enable()
				race.Acquire(unsafe.Pointer(wg))
			}
			return
		}
	}
}

Done

实际为Add -1

func (wg *WaitGroup) Done() {
	wg.Add(-1)
}