Context 超時處理應該這麼做

1、超時處理有問題代碼

只通過 select 監聽超時信號，沒有處理正常邏輯處理。

package main

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

var wg sync.WaitGroup

func main() {
 // 定義一個3s超時context 所有衍生出來的goroutine必須在3s內完成 否則會cancel掉
 ctx, cancel := context.WithTimeout(context.Background(), 3*time.Second)
 defer cancel()

 go func() {
  wg.Add(2)
  go watch(ctx, 1)
  go watch(ctx, 2)
  wg.Wait()
 }()

// 這裏有問題，因爲不管是超時還是未超時都會進入<-ctx.Done() 也就是都會等待3s 這不是我們想要的
 select {
 case <-ctx.Done():
  fmt.Printf("watch %d %s\n", 0, ctx.Err())
 }

 fmt.Println("finished")
}

func watch(ctx context.Context, flag int) {
 defer wg.Done()

 func() {
  fmt.Printf("doing something flag:%d\n", flag)
  time.Sleep(10*time.Second)
  fmt.Println("finished flag:", flag)
 }()
}

2、如何解決？

很簡單，就是在所有 goroutine 結束之後發送一個 done 信號給主 goroutine，這樣 select 就會監聽到這個信號，從而安全退出。

// 增加done信號
 done := make(chan struct{})
 go func() {
  wg.Add(2)
  go watch(ctx, 1)
  go watch(ctx, 2)
  wg.Wait()
  done <- struct{}{} // 發送done信號
 }()

 select {
 case <-done: // 接受done信號
  fmt.Println("Done")
 case <-ctx.Done():
  fmt.Printf("watch %d %s\n", 0, ctx.Err())
 }

本文由 Readfog 進行 AMP 轉碼，版權歸原作者所有。
來源：https://mp.weixin.qq.com/s/fmvxsXgnLm2bS3qaqZ-JLA