golang 源码学习之WaitGroup

数据结构

//sync/waitgroup.go
type WaitGroup struct {
    noCopy noCopy // 不可复制，只能指针传递

    // 对齐的8byte存储状态， 另外4byte存储信号。状态的高32位counter计数器，低32位waiter计数器
    // counter未执行完的goroutine。waiter正在等待的goroutine。那么WaitGroup可以放多少goroutine呢？2^32
    state1 [3]uint32
}

Add

//sync/waitgroup.go
func (wg *WaitGroup) Add(delta int) {

    // statep状态（包含counter和waiter计数器），semap信号
    statep, semap := wg.state() 
    
    // uint64(delta)<<32 获取counter计数器， 这里是counter + delta
    state := atomic.AddUint64(statep, uint64(delta)<<32) 
    v := int32(state >> 32) // counter 计数器
    w := uint32(state) // waiter 计数器

    // counter 计数器 < 0 恐慌
    if v < 0 {
        panic("sync: negative WaitGroup counter")
    }
    //当 wait 和 add并发执行时，可能会发生恐慌。注意这里并不是说wait之后就不能添加Add了
    if w != 0 && delta > 0 && v == int32(delta) {
        panic("sync: WaitGroup misuse: Add called concurrently with Wait")
    }

    // add 完成
    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")
    }
    // Reset waiters count to 0.
    *statep = 0 // 计数器置零
    for ; w != 0; w-- {
        runtime_Semrelease(semap, false)  // 一个一个唤醒waiter
    }
}

Done

//sync/waitgroup.go
func (wg *WaitGroup) Done() {
    wg.Add(-1) // Done其实也是执行的Add
}

Wait

//sync/waitgroup.go
func (wg *WaitGroup) Wait() {
    statep, semap := wg.state()

    for {
        state := atomic.LoadUint64(statep)
        v := int32(state >> 32) //
        w := uint32(state)

        // 如果counter计数器 == 0 ， 直接返回
        if v == 0 {
            // Counter is 0, no need to wait.
            if race.Enabled {
                race.Enable()
                race.Acquire(unsafe.Pointer(wg))
            }
            return
        }
        // Increment waiters count.
        if atomic.CompareAndSwapUint64(statep, state, state+1) { //waiter计数器+1
            if race.Enabled && w == 0 {
                race.Write(unsafe.Pointer(semap))
            }
            runtime_Semacquire(semap)  //等待信号唤醒自己
            if *statep != 0 {
                panic("sync: WaitGroup is reused before previous Wait has returned")
            }

            return
        }
    }
}

小结

• 可以复制，只能通过指针传递
• 最多可存放2^32个未执行完的goroutine
• wait和add不要并发执行，有可能发生恐慌
• counter计数器负数时 恐慌。Add负数 或者 过多的Done都有可能恐慌

有疑问加站长微信联系（非本文作者）