一個小例子，給你講透典型的 Go 併發操作

技术颜良發表於2024-09-08

一個小例子，給你講透典型的 Go 併發操作

原創訢亮程式設計師新亮

2024年09月08日 16:57 天津聽全文

程式設計師新亮

GitHub 9K+ Star | 技術交流分享

206篇原創內容

公眾號

如果你有一個任務可以分解成多個子任務進行處理，同時每個子任務沒有先後執行順序的限制，等到全部子任務執行完畢後，再進行下一步處理。這時每個子任務的執行可以併發處理，這種情景下適合使用 sync.WaitGroup。

雖然 sync.WaitGroup 使用起來比較簡單，但是一不留神很有可能踩到坑裡。

sync.WaitGroup 正確使用

比如，有一個任務需要執行 3 個子任務，那麼可以這樣寫：

func main() {
 var wg sync.WaitGroup

 wg.Add(3)

 go handlerTask1(&wg)
 go handlerTask2(&wg)
 go handlerTask3(&wg)

 wg.Wait()

 fmt.Println("全部任務執行完畢.")
}

func handlerTask1(wg *sync.WaitGroup) {
 defer wg.Done()
 fmt.Println("執行任務 1")
}

func handlerTask2(wg *sync.WaitGroup) {
 defer wg.Done()
 fmt.Println("執行任務 2")
}

func handlerTask3(wg *sync.WaitGroup) {
 defer wg.Done()
 fmt.Println("執行任務 3")
}

執行輸出：

執行任務 3
執行任務 1
執行任務 2
全部任務執行完畢.

sync.WaitGroup 閉坑指南

01

// 正確
go handlerTask1(&wg)

// 錯誤
go handlerTask1(wg)

執行子任務時，使用的 sync.WaitGroup 一定要是 wg 的引用型別！

02

注意不要將 wg.Add() 放在 go handlerTask1(&wg) 中！

例如：

// 錯誤
var wg sync.WaitGroup

go handlerTask1(&wg)

wg.Wait()

...

func handlerTask1(wg *sync.WaitGroup) {
 wg.Add(1)
 defer wg.Done()
 fmt.Println("執行任務 1")
}

注意 wg.Add() 一定要在 wg.Wait() 執行前執行！

小結

注意 wg.Add() 和 wg.Done() 的計數器保持一致！其實 wg.Done() 就是執行的 wg.Add(-1) 。

其實 sync.WaitGroup 使用場景比較侷限，僅適用於等待全部子任務執行完畢後，再進行下一步處理。如果需求是當第一個子任務執行失敗時，通知其他子任務停止執行，這時 sync.WaitGroup 是無法滿足的，需要使用到上下文（context）。

sync.WaitGroup + Context

在處理併發任務時，若需在任一子任務失敗時終止所有其他子任務，以下示例提供了一種實現方法。

package main

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

// handlerTask 處理單個任務
func handlerTask(ctx context.Context, taskId int, wg *sync.WaitGroup, cancel context.CancelFunc) {
 defer wg.Done()

 fmt.Printf("Request %d is processing...\n", taskId)

 // 模擬請求處理，如果 taskId 為1，則模擬失敗
 if taskId == 1 {
  fmt.Printf("Request %d failed\n", taskId)
  cancel() // 取消 context，通知其他請求停止
  return
 }

 // 監聽 context.Done() 通道
 select {
 case <-ctx.Done():
  fmt.Printf("Request %d is already cancelled\n", taskId)
  return
 default:
  // 繼續執行
  time.Sleep(3 * time.Second) // 模擬耗時操作
  fmt.Printf("Request %d succeeded\n", taskId)
 }
}

func main() {
 ctx, cancel := context.WithCancel(context.Background())
 var wg sync.WaitGroup

 wg.Add(3)
 go handlerTask(ctx, 1, &wg, cancel)
 go handlerTask(ctx, 2, &wg, cancel)
 go handlerTask(ctx, 3, &wg, cancel)

 // 等待所有任務完成或被取消
 go func() {
  wg.Wait()
  fmt.Println("All requests are finished or cancelled")
 }()

 // 給一些時間來處理，防止主 goroutine 過早退出
 time.Sleep(5 * time.Second)
}

解釋

1.任務 1 開始處理:

輸出 Request 1 is processing...。
檢查任務 ID，發現是 1，輸出 Request 1 failed。
呼叫 cancel()，取消上下文 ctx。

2.任務 2 和任務 3 開始處理:

輸出 Request 2 is processing... 和 Request 3 is processing...。
由於任務 1 已經呼叫 cancel()，上下文 ctx 已經被取消。
進入 select 語句，檢查 ctx.Done() 通道。
發現 ctx.Done() 通道已被關閉，輸出 Request 2 is already cancelled 和 Request 3 is already cancelled。

3.等待所有任務完成或被取消: