Golang 教學系列 - 何謂 WaitGroup? 等待 Goroutine 的好幫手！

發表於 2020-12-18 更新於 2024-04-28 分類於 Golang Disqus：文章字數： 3.3k 所需閱讀時間 ≈ 6 分鐘

接續上一篇文章：Golang 教學系列 - 何謂 Goroutine，上次為了避免 main func 執行結束導致 goroutine 被關閉，使用 time.Sleep 的方式，但是這樣的方式並不彈性，所以今天介紹另外一種等待 goroutine 的方式，那就是 WaitGroup。

如果想看更清楚的講解可以同步看我的影片來參考這篇文章唷：Golang 教學系列 - 何謂 WaitGroup? 等待 Goroutine 的好幫手！| 肯尼攻城獅

使用 WaitGroup 來找質數的例子

直接上程式碼：

package main

import (
	"fmt"
	"sync"
	"time"
)

func main() {
	wg := new(sync.WaitGroup)
	num := 300000
	wg.Add(num)
	start := time.Now().Unix()
	for i := 1; i <= num; i++ {
		go findPrimes(i, wg)
	}
	wg.Wait()
	end := time.Now().Unix()
	fmt.Println(end - start, "seconds")
}

func findPrimes(num int, wg *sync.WaitGroup) {
	defer wg.Done()
	if num == 1 {
		return
	} else if num == 2 {
		fmt.Println(num)
	} else {
		for i := 2; i < num; i++ {
			if num % i == 0 {
				return
			}
		}
		fmt.Println(num)
	}
}

使用 sync package 裡面的 WaitGroup Struct，宣告方式可以透過 new 或是 &sync.WaitGroup{} 也可以。
然後 WaitGroup 有提供一個 Add func，這個意思是想要等待 Goroutine 的數量，這邊因為是會起 30000 個 Goroutine 所以就 wg.Add(num)。
而使用 go keyword 的 func 要記得將 wg 當作參數傳進去，這是因為每當有一個 Goroutine 做完工作之後就要呼叫 wg.Done() 通知 WaitGroup 可以減少等待一個 Goroutine，事實上，在 Done () 的實作就是將前面 Add 30000 的數目減一的意思是一樣的。

通常呼叫 Done () 的方式會使用 defer，來確保在結束 func 前要記得執行。
最後在 main func 裡面加入了 wg.Wait() 這個其實執行之後就會將 main 來進行 block 住，不繼續往下執行程式碼，直到所有的 Goroutine 都已經結束，Wait () 這個函式才會跳出來結束。
所以這樣的好處就是非常的彈性，我們不需要透過 time.Sleep 來手動控制要等待幾秒。

WaitGroup 用法介紹

WaitGroup Struct 講解

這邊滿推薦直接點進去 Golang 原始碼來看註解，會更了解：

// A WaitGroup waits for a collection of goroutines to finish.
// The main goroutine calls Add to set the number of
// goroutines to wait for. Then each of the goroutines
// runs and calls Done when finished. At the same time,
// Wait can be used to block until all goroutines have finished.
//
// A WaitGroup must not be copied after first use.
type WaitGroup struct {
	noCopy noCopy

	// 64-bit value: high 32 bits are counter, low 32 bits are waiter count.
	// 64-bit atomic operations require 64-bit alignment, but 32-bit
	// compilers do not ensure it. So we allocate 12 bytes and then use
	// the aligned 8 bytes in them as state, and the other 4 as storage
	// for the sema.
	state1 [3]uint32
}

從上面註解可以得知：

WaitGroup 是用來等待一群 Goroutine 的功用
透過 Add 來加入 Goroutine，而每一個 Goroutine 做完之後都需要呼叫 Done
同時使用 Wait 可以來進行 block，直到所有的 Goroutine 都已經結束。

宣告 WaitGroup 的方式

直接上程式碼：

package main

import "sync"

func main() {
	var wg1 sync.WaitGroup
	wg2 := new(sync.WaitGroup)
	wg3 := &sync.WaitGroup{}
	var wg4 = &sync.WaitGroup{}
}

WaitGroup 要透過 Pointer 來做操作
因為看原始碼可以知道每一個 Struct 內的 func 都是透過 WaitGroup 來進行操作的：
1
2
3
func (wg *WaitGroup) Done() {
wg.Add(-1)
}
但是以上四種宣告方式都可以，在使用第一種的時候因為 Golang 語法糖的關係，在使用的時候會自動幫你轉成用 Pointer 來做操作
之所以要有 Pointer 是因為這樣才會指到同一個 Struct，不然會造成 call by value 的方式

來看第一個例子

直接上程式碼：

package main

import (
	"fmt"
	"math/rand"
	"sync"
	"time"
)

func main() {
	wg := &sync.WaitGroup{}
	jobs := 100
	wg.Add(jobs)
	for i := 1; i <= 100; i++ {
		go doTask(wg)
	}
	wg.Wait()
	fmt.Println("jobs all done!")
}

func doTask(wg *sync.WaitGroup) {
	defer func() {
		fmt.Println("one task done")
		wg.Done()
	}()
	number := rand.Intn(10) + 1
	time.Sleep(time.Duration(number) * time.Second)
}

這程式碼的示範是：

宣告一個 WaitGroup
定義假設有 100 工作，每一個工作由一個 Goroutine 來負責，所以這邊透過一個 for loop 來啟動 100 Goroutine 來執行 doTask()
然後記得要先將 100 的 Goroutine 加入 WaitGroup 在啟動 Goroutine。
最後使用 Wait 來等待所有 Goroutine 結束工作
在 doTask 裡面記得要用 defer wg.Done() 的操作，用 defer 也是為了避免忘記做這件事情，因為有時候 func 可能會很長而忘記使用。
這邊 doTask 因為為了示範，所以透過隨機產生 number 然後透過 time.Sleep 來 sleep，讓這 100Goroutine 不會太快結束，然後透過在 defer 來 Println("one task done") 來看出這個 Goroutine 結束了。

總結

這篇文章主要講解如何使用 WaitGroup 來等待 Goroutine，也講了一些小例子，下篇文章要講的是在使用 WaitGroup 的時候會遇到一些 error 的情況，還有可能會碰到一些的坑。

最後最後！請聽我一言！

如果你還沒有註冊 Like Coin，你可以透過我的邀請註冊連結來免費註冊，註冊完後就可以在文章最下方幫我按下 Like 按鈕，而 Like 最多可以點五次，如何一來你不用付出任何一塊錢，就能給我寫這篇文章最大的回饋！