面试题：Go中WaitGroup的基本使用及原理简述

使用WaitGroup等待一组goroutine完成

引入包：首先需要引入 sync 包，因为 WaitGroup 是在这个包中定义的。

package main

import (
    "fmt"
    "sync"
)

创建WaitGroup实例：在主函数或其他需要等待goroutine完成的地方，创建一个 WaitGroup 实例。

var wg sync.WaitGroup

添加等待计数：在启动每个goroutine之前，调用 wg.Add(1) 来增加等待计数。例如：

for i := 0; i < 5; i++ {
    wg.Add(1)
    go func(num int) {
        defer wg.Done()
        // 模拟goroutine执行任务
        fmt.Printf("Goroutine %d is working\n", num)
    }(i)
}

等待所有goroutine完成：在所有goroutine启动后，调用 wg.Wait() 来阻塞当前goroutine，直到所有已添加的等待计数都被 wg.Done() 减为0。

wg.Wait()
fmt.Println("All goroutines have finished")

WaitGroup内部实现原理

结构体定义：WaitGroup 本质上是一个结构体，其内部维护了一个计数器（state1）和一个信号量（state2）。在64位系统上，这两个部分被合并在一个64位的字中。

// src/sync/waitgroup.go
type WaitGroup struct {
    noCopy noCopy
    state1 uint64
}

Add方法：Add 方法用于增加等待计数，它通过原子操作修改计数器的值。如果计数器变为0，会释放所有等待的goroutine。

func (wg *WaitGroup) Add(delta int) {
    statep, semap := wg.state()
    state := atomic.AddUint64(statep, uint64(delta)<<32)
    v := int32(state >> 32)
    if v < 0 {
        panic("sync: negative WaitGroup counter")
    }
    if v > 0 || int32(state)&^uint64(0xffffffff) != 0 {
        return
    }
    // 计数器变为0，释放所有等待的goroutine
    for ; ; {
        s := atomic.LoadUint64(semap)
        if s == 0 {
            break
        }
        if atomic.CompareAndSwapUint64(semap, s, s-1) {
            runtime_Semrelease(semap, false, 0)
            break
        }
    }
}

Done方法：Done 方法实际上是 Add(-1) 的便捷调用，它减少等待计数。

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

Wait方法：Wait 方法会阻塞当前goroutine，直到等待计数变为0。它通过原子操作检查计数器，并在计数器不为0时，使用信号量等待。

func (wg *WaitGroup) Wait() {
    statep, semap := wg.state()
    for {
        state := atomic.LoadUint64(statep)
        v := int32(state >> 32)
        if v == 0 {
            return
        }
        if atomic.CompareAndSwapUint64(statep, state, state+1) {
            runtime_Semacquire(semap)
            if atomic.LoadUint64(statep)&1 != 0 {
                panic("sync: WaitGroup is reused before previous Wait has returned")
            }
            return
        }
    }
}

通过这种方式，WaitGroup 实现了在Go语言中高效地等待一组goroutine完成的功能。

星途面试题库

面试题：Go中WaitGroup的基本使用及原理简述

知识考点

面试题答案

使用WaitGroup等待一组goroutine完成

WaitGroup内部实现原理