go WaitGroup源码分析

适用场景

等待若干个任务执行完成。

实现原理

字段

type WaitGroup struct { 
    省略
    
    // 存储waiter数、WaitGroup计数和信号量
    state1 [3]uint32
}

添加任务Add函数

设置WaitGroup计数

func (wg *WaitGroup) Add(delta int) { 
    statep, semap := wg.state()
    // 通过原子操作来增加WaitGroup计数，高32位是WaitGroup计数
    state := atomic.AddUint64(statep, uint64(delta)<<32)
    // 获取当前WaitGroup计数
    v := int32(state >> 32)
    // 获取当前waiter数
    w := uint32(state)
 
    // 如果还有等待任务或者等待者数量是0，那么不需要唤醒等待者
    if v > 0 || w == 0 { 
        return 
    } 
 
    // WaitGroup计数是0，waiter数设为0
    *statep = 0
    // 唤醒所有的waiter
    for ; w != 0; w-- {
        runtime_Semrelease(semap, false, 0)
    }
}

任务结束Done函数

WaitGroup计数减1

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

等待任务结束Wait函数

如果WaitGroup计数是0，那么调用者不等待，直接返回；如果WaitGroup计数大于0，那么调用者加入waiter队列，阻塞自己。

func (wg *WaitGroup) Wait() { 
    statep, semap := wg.state() 
     
    for { 
        state := atomic.LoadUint64(statep) 
        // 获取当前WaitGroup计数
        v := int32(state >> 32)
        // 获取当前waiter数
        w := uint32(state)
        if v == 0 { 
            // 如果WaitGroup计数为0，那么调用这个方法的goroutine不再等待 
            return 
        } 
        // waiter数加1，cas失败后通过外层for循环来继续尝试cas 
        if atomic.CompareAndSwapUint64(statep, state, state+1) { 
            // 阻塞等待 
            runtime_Semacquire(semap) 
            // 被唤醒，不再阻塞，返回 
            return 
        } 
    } 
}