Go語言學(xué)習(xí)之WaitGroup用法詳解

更新時間：2022年06月11日 08:23:30 作者：OOMO1Edc

Go語言中的?WaitGroup?和?Java?中的?CyclicBarrier、CountDownLatch?非常類似。本文將詳細(xì)為大家講講WaitGroup的用法，感興趣的小伙伴可以跟隨小編一起學(xué)習(xí)一下

前言

在前面的文章中，我們使用過 WaitGroup 進(jìn)行任務(wù)編排，Go語言中的 WaitGroup 和 Java 中的 CyclicBarrier、CountDownLatch 非常類似。比如我們有一個主任務(wù)在執(zhí)行，執(zhí)行到某一點(diǎn)時需要并行執(zhí)行三個子任務(wù)，并且需要等到三個子任務(wù)都執(zhí)行完后，再繼續(xù)執(zhí)行主任務(wù)。那我們就需要設(shè)置一個檢查點(diǎn)，使主任務(wù)一直阻塞在這，等三個子任務(wù)執(zhí)行完后再放行。

說明：本文中的示例，均是基于Go1.17 64位機(jī)器

小試牛刀

我們先來個簡單的例子，看下 WaitGroup 是怎么使用的。示例中使用 Add(5) 表示我們有 5個子任務(wù)，然后起了 5個協(xié)程去完成任務(wù)，主協(xié)程使用 Wait() 方法等待子協(xié)程執(zhí)行完畢，輸出一共等待的時間。

func main() {
    var waitGroup sync.WaitGroup

    start := time.Now()
    waitGroup.Add(5)
    for i := 0; i < 5; i++ {
        go func() {
            defer waitGroup.Done()
            time.Sleep(time.Second)
            fmt.Println("done")
        }()
    }

    waitGroup.Wait()
    fmt.Println(time.Now().Sub(start).Seconds())
}

/*
done
done
done
done
done
1.000306089
*/

總覽

WaitGroup 一共有三個方法：

(wg *WaitGroup) Add(delta int)
(wg *WaitGroup) Done()
(wg *WaitGroup) Wait()

Add 方法用于設(shè)置 WaitGroup 的計數(shù)值，可以理解為子任務(wù)的數(shù)量
Done 方法用于將 WaitGroup 的計數(shù)值減一，可以理解為完成一個子任務(wù)
Wait 方法用于阻塞調(diào)用者，直到 WaitGroup 的計數(shù)值為0，即所有子任務(wù)都完成

正常來說，我們使用的時候，需要先確定子任務(wù)的數(shù)量，然后調(diào)用 Add() 方法傳入相應(yīng)的數(shù)量，在每個子任務(wù)的協(xié)程中，調(diào)用 Done()，需要等待的協(xié)程調(diào)用 Wait() 方法，狀態(tài)流轉(zhuǎn)如下圖：

底層實現(xiàn)

結(jié)構(gòu)體

type WaitGroup struct {
    noCopy noCopy // noCopy 字段標(biāo)識，由于 WaitGroup 不能復(fù)制，方便工具檢測

    state1 [3]uint32  // 12個字節(jié)，8個字節(jié)標(biāo)識 計數(shù)值和等待數(shù)量，4個字節(jié)用于標(biāo)識信號量
}

state1 是個復(fù)合字段，會拆分為兩部分： 64位（8個字節(jié)）的 statep 作為一個整體用于原子操作, 其中前面4個字節(jié)表示計數(shù)值，后面四個字節(jié)表示等待數(shù)量；剩余 32位（4個字節(jié)）semap 用于標(biāo)識信號量。

Go語言中對于64位的變量進(jìn)行原子操作，需要保證該變量是64位對齊的，也就是要保證這 8個字節(jié) 的首地址是 8 的整數(shù)倍。因此當(dāng) state1 的首地址是 8 的整數(shù)倍時，取前8個字節(jié)作為 statep ，后4個字節(jié)作為 semap；當(dāng) state1 的首地址不是 8 的整數(shù)倍時，取后8個字節(jié)作為 statep ，前4個字節(jié)作為 semap。

func (wg *WaitGroup) state() (statep *uint64, semap *uint32) {
    
    // 首地址是8的倍數(shù)時，前8個字節(jié)為 statep, 后四個字節(jié)為 semap
    if uintptr(unsafe.Pointer(&wg.state1))%8 == 0 {
        return (*uint64)(unsafe.Pointer(&wg.state1)), &wg.state1[2]
    } else { 
        
    // 后8個字節(jié)為 statep, 前四個字節(jié)為 semap    
        return (*uint64)(unsafe.Pointer(&wg.state1[1])), &wg.state1[0]
    }
}

Add

Add 方法用于添加一個計數(shù)值（負(fù)數(shù)相當(dāng)于減），當(dāng)計數(shù)值變?yōu)?后， Wait 方法阻塞的所有等待者都會被釋放
計數(shù)值變?yōu)樨?fù)數(shù)是非法操作，產(chǎn)生 panic
當(dāng)計數(shù)值為0時（初始狀態(tài)），Add 方法不能和 Wait 方法并發(fā)調(diào)用，需要保證 Add 方法在 Wait 方法之前調(diào)用，否則會 panic

func (wg *WaitGroup) Add(delta int) {

    // 拿到計數(shù)值等待者變量 statep 和 信號量 semap
    statep, semap := wg.state()

    // 計數(shù)值加上 delta: statep 的前四個字節(jié)是計數(shù)值，因此將 delta 前移 32位
    state := atomic.AddUint64(statep, uint64(delta)<<32)

    // 計數(shù)值
    v := int32(state >> 32)

    // 等待者數(shù)量
    w := uint32(state)

    // 如果加上 delta 之后，計數(shù)值變?yōu)樨?fù)數(shù)，不合法，panic
    if v < 0 {
        panic("sync: negative WaitGroup counter")
    }

    // delta > 0 && v == int32(delta) : 表示從 0 開始添加計數(shù)值
    // w!=0 ：表示已經(jīng)有了等待者
    // 說明在添加計數(shù)值的時候，同時添加了等待者，非法操作。添加等待者需要在添加計數(shù)值之后
    if w != 0 && delta > 0 && v == int32(delta) {
        panic("sync: WaitGroup misuse: Add called concurrently with Wait")
    }

    // v>0 : 計數(shù)值不等于0，不需要喚醒等待者，直接返回
    // w==0: 沒有等待者，不需要喚醒，直接返回
    if v > 0 || w == 0 {
        return
    }

    // 再次檢查數(shù)據(jù)是否一致
    if *statep != state {
        panic("sync: WaitGroup misuse: Add called concurrently with Wait")
    }

    // 到這里說明計數(shù)值為0，且等待者大于0，需要喚醒所有的等待者，并把系統(tǒng)置為初始狀態(tài)（0狀態(tài)）
  
  // 將計數(shù)值和等待者數(shù)量都置為0
    *statep = 0

    // 喚醒等待者
    for ; w != 0; w-- {
        runtime_Semrelease(semap, false, 0)
    }
}

Done

// 完成一個任務(wù)，將計數(shù)值減一，當(dāng)計數(shù)值減為0時，需要喚醒所有的等待者
func (wg *WaitGroup) Done() {
    wg.Add(-1)
}

Wait

// 調(diào)用 Wait 方法會被阻塞，直到 計數(shù)值 變?yōu)?
func (wg *WaitGroup) Wait() {

    // 獲取計數(shù)、等待數(shù)和信號量
    statep, semap := wg.state()

    for {
        state := atomic.LoadUint64(statep)

        // 計數(shù)值
        v := int32(state >> 32)

        // 等待者數(shù)量
        w := uint32(state)

        // 計數(shù)值數(shù)量為0，直接返回，無需等待
        if v == 0 {
            return
        }

        // 到這里說明計數(shù)值數(shù)量大于0
        // 增加等待者數(shù)量：這里會有競爭，比如多個 Wait 調(diào)用，或者在同時調(diào)用 Add 方法，增加不成功會繼續(xù) for 循環(huán)
        if atomic.CompareAndSwapUint64(statep, state, state+1) {
            // 增加成功后，阻塞在信號量這里，等待被喚醒
            runtime_Semacquire(semap)

            // 被喚醒的時候，應(yīng)該是0狀態(tài)。如果重用 WaitGroup，需要等 Wait 返回
            if *statep != 0 {
                panic("sync: WaitGroup is reused before previous Wait has returned")
            }
            return
        }
    }
}

易錯點(diǎn)

上面分析源碼可以看到幾個會產(chǎn)生 panic 的點(diǎn)，這也是我們使用 WaitGroup 需要注意的地方

1.計數(shù)值變?yōu)樨?fù)數(shù)

調(diào)用 Add 時參數(shù)值傳負(fù)數(shù)

func main() {
var wg sync.WaitGroup
wg.Add(1)
wg.Add(-1)
wg.Add(-1)
}

多次調(diào)用 Done 方法

func main() {
    var wg sync.WaitGroup

    wg.Add(1)
    go func() {
        fmt.Println("test")
        wg.Done()
        wg.Done()
    }()

    time.Sleep(time.Second)
    wg.Wait()
}

2.Add 和 Wait 并發(fā)調(diào)用

Add 和 Wait 并發(fā)調(diào)用，有可能達(dá)不到我們預(yù)期的效果，甚至 panic。如下示例中，我們想要等待 3 個子任務(wù)都執(zhí)行完后再執(zhí)行主任務(wù)，但實際情況可能是子任務(wù)還沒起來，主任務(wù)就繼續(xù)往下執(zhí)行了。

func doSomething(wg *sync.WaitGroup) {
    wg.Add(1)
    fmt.Println("do  something")
    defer wg.Done()
}

func main() {
    var wg sync.WaitGroup
    for i := 0; i < 3; i++ {
        go doSomething(&wg)
    }

    wg.Wait()
    fmt.Println("main")
}

//main
//do  something
//do  something

正確的使用方式，應(yīng)該是在調(diào)用 Wait 前先調(diào)用 Add

func doSomething(wg *sync.WaitGroup) {
    defer wg.Done()
    fmt.Println("do  something")
}

func main() {
    var wg sync.WaitGroup
    wg.Add(3)
    for i := 0; i < 3; i++ {
        go doSomething(&wg)
    }

    wg.Wait()
    fmt.Println("main")
}

//do  something
//do  something
//do  something
//main

3.沒有等 Wait 返回，就重用 WaitGroup

func main() {
    var wg sync.WaitGroup

    wg.Add(1)
    go func() {
        fmt.Println("do something")
        wg.Done()
        wg.Add(1)
    }()

    wg.Wait()
}

4.復(fù)制使用

我們知道 Go 語言中的參數(shù)傳遞，都是值傳遞，就會產(chǎn)生復(fù)制操作。因此在向函數(shù)傳遞 WaitGroup 時，使用指針進(jìn)行操作。

// 錯誤使用方式，沒有使用指針
func doSomething(wg sync.WaitGroup) {
    fmt.Println("do  something")
    defer wg.Done()
}

func main() {
    var wg sync.WaitGroup
    wg.Add(3)
    for i := 0; i < 3; i++ {
        // 這里沒使用指針，wg狀態(tài)一直不會改變，導(dǎo)致 Wait 一直阻塞
        go doSomething(wg)
    }

    wg.Wait()
    fmt.Println("main")
}