WaitGroup

数据结构

WaitGroup 的数据结构如下，包含长度为 3 的 uint32 数组。

1
2
3

type WaitGroup struct {
    state1 [3]uint32
}

其中，数组包含了一个 state（由两个计数器组成）和一个信号量：

counter： 当前还未执行结束的 goroutine 计数器。
waiter：等候者的数量。
semaphore：信号量。

也可以认为是 state 的类型为 uint64，而信号量的类型为 uint32。

Add Wait Done 实现

WaitGroup 实现了三个方法：

Add(delta int)：把 delta 的值加到 counter 中。
Wait()：waiter 的值增加 1，并阻塞等待信号量。
Done()：counter 递减 1，当 counter 的值等于 0 时，按照 waiter 数值释放相应次数信号量。

Add

Add 函数主要的流程如下：

首先把 delta 加到 counter 中，因为 delta 可能是负值，所以 counter 可能小于 0，当 counter 小于 0 时发出 panic。
接着如果 counter 等于 0，则释放 waiter 个信号量。

func (wg *WaitGroup) Add(delta int) {
    statep, semap := wg.state() //获取state和semaphore地址指针

    state := atomic.AddUint64(statep, uint64(delta)<<32) //把delta左移32位累加到state，即累加到counter中
    v := int32(state >> 32) //获取counter值
    w := uint32(state)      //获取waiter值

    if v < 0 {              //经过累加后counter值变为负值，panic
        panic("sync: negative WaitGroup counter")
    }

    //经过累加后，此时，counter >= 0
    //如果counter为正，说明不需要释放信号量，直接退出
    //如果waiter为零，说明没有等待者，也不需要释放信号量，直接退出
    if v > 0 || w == 0 {
        return
    }

    //此时，counter一定等于0，而waiter一定大于0（内部维护waiter，不会出现小于0的情况），
    //先把counter置为0，再释放waiter个数的信号量
    *statep = 0
    for ; w != 0; w-- {
        runtime_Semrelease(semap, false) //释放信号量，执行一次释放一个，唤醒一个等待者
    }
}

Wait

Wait 函数主要的流程如下：

首先，累加 waiter 的数量。
接着，如果 counter 大于 0，则阻塞当前协程等待信号量。

func (wg *WaitGroup) Wait() {
    statep, semap := wg.state() //获取state和semaphore地址指针
    for {
        state := atomic.LoadUint64(statep) //获取state值
        v := int32(state >> 32)            //获取counter值
        w := uint32(state)                 //获取waiter值
        if v == 0 {                        //如果counter值为0，说明所有goroutine都退出了，不需要待待，直接返回
            return
        }

        // 使用CAS（比较交换算法）累加waiter，累加可能会失败，失败后通过for loop下次重试
        if atomic.CompareAndSwapUint64(statep, state, state+1) {
            runtime_Semacquire(semap) //累加成功后，等待信号量唤醒自己
            return
        }
    }
}

这里用到了 CAS 算法，保证有多个 goroutine 同时执行 Wait() 时，也能正确累加 waiter。

Done

Done 只做了 counter 减一这一件事，这样当 counter 等于 0 时由 Add 函数唤醒等待的协程。

1
2
3

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

Dawn's Blogs

GO专家编程读书笔记 (4) 常见控制结构实现原理之WaitGroup

WaitGroup

数据结构

Add Wait Done 实现

Add

Wait

Done