Question

允许有多个等待。

源码剖析

通过 Add/Done 和 Wait 两个计数器来判断状态。

 64                  32                   0
  +-------------------+-------------------+-------------------+
  |  Add/Done counter |    Wait counter   |       Sema        |
  +-------------------+-------------------+-------------------+
  |<-------------- state1 --------------->|<---- state2 ----->|
  

• 正确逻辑： Add 在 Wait 之前，不应并发。

// sync/waitgroup.go

type WaitGroup struct {
	noCopy noCopy

	// 64-bit value: high 32 bits are counter, 
    // low 32 bits are waiter count.
    
	state1 uint64   // [Add/Done, Wait] counter。
	state2 uint32   // Sema
}

func (wg *WaitGroup) state() (statep *uint64, semap *uint32) {
	// state1 is 64-bit aligned: nothing to do.
	return &wg.state1, &wg.state2
}

累加或递减计数，归零时唤醒所有等待。

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

func (wg *WaitGroup) Add(delta int) {
    
    // 累加 Add 计数。
	statep, semap := wg.state()
	state := atomic.AddUint64(statep, uint64(delta)<<32)
    
	v := int32(state >> 32)  // add.count
	w := uint32(state)       // wait.count
    
    // 不可能小于零。
	if v < 0 {
		panic("sync: negative WaitGroup counter")
	}
    
    // add.count v == delta，第一次 Add 操作，
    // wait.count w != 0，先有等待，显然不行！
	if w != 0 && delta > 0 && v == int32(delta) {
		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 {
		panic("sync: WaitGroup misuse: Add called concurrently with Wait")
	}

    // 上面拦截了 v > 0 的情形。
    // 此处只能是最后一次 Done/Add(-1) == 0，
    // 唤醒所有等待。
	*statep = 0
	for ; w != 0; w-- {
		runtime_Semrelease(semap, false, 0)
	}
}

如果计数为零，自然不需要等待。否则，累加等待计数，休眠。

// Wait blocks until the WaitGroup counter is zero.

func (wg *WaitGroup) Wait() {
	statep, semap := wg.state()
	for {
		state := atomic.LoadUint64(statep)
        
		v := int32(state >> 32)  // add.count
		w := uint32(state)       // wait.count
        
		if v == 0 {
			// Counter is 0, no need to wait.
			return
		}
        
		// Increment waiters count.
		if atomic.CompareAndSwapUint64(statep, state, state+1) {
			runtime_Semacquire(semap)
			return
		}
	}
}