Golang中Context包使用場景和示例詳解
本文結合示例代碼講解一下context包的幾種使用場景。
控制子協程退出
context包提供了一種機制，可以在多個goroutine之間進行通信和控制。使用Context包能夠有效地控制程序的并發性，提高程序的健壯性和性能。
Golang是沒有辦法讓其他goroutine退出的，goroutine只能自己退出。之所以說context包可以控制子協程退出意思是子協程可以接收到主協程發出的退出信號，然后自己退出。看如下示例代碼：
package main
import (
"context"
"errors"
"sync"
)
func request(ctx context.Context, url string) error {
result := make(chan int)
err := make(chan error)
go func() {
// 假如isSuccess是請求返回的結果，成功則通過result傳遞成功信息，錯誤通過error傳遞錯誤信息
isSuccess := true
if isSuccess {
result <- 1
} else {
err <- errors.New("some error happen")
}
}()
select {
case <-ctx.Done():
// 其他請求失敗
return ctx.Err()
case e := <-err:
// 本次請求失敗，返回錯誤信息
return e
case <-result:
// 本此請求成功，不返回錯誤信息
return nil
}
}
func main() {
ctx, cancel := context.WithCancel(context.Background())
// 調用接口a
err := request(ctx, "https://xxx.com/a")
if err != nil {
return
}
wg := sync.WaitGroup{}
// 調用接口b
wg.Add(1)
go func() {
defer wg.Done()
err := request(ctx, "https://xxx.com/b")
if err != nil {
cancel()
}
}()
// 調用接口c
wg.Add(1)
go func() {
defer wg.Done()
err := request(ctx, "https://xxx.com/c")
if err != nil {
cancel()
}
}()
wg.Wait()
}首先調用context.WithCancel方法構造了一個Context和返回了一個cancel函數，其他goroutine調用的方法都傳入了這個Context作為第一個參數，當主goroutine想要告訴所有goroutine需要退出的時候，通過調用cancel函數把退出的信息告訴所有的goroutine。所有goroutine通過監聽ctx.Done返回的channel得到退出信號然后退出。
超時控制
例如查詢數據庫、調用RPC服務、調用HTTP接口等場景，這些操作都是阻塞的，如果一直不返回數據的話，會影響產品的用戶體驗。針對這些情況的解決方式通常是設置一個超時間，超過后自動取消操作。
使用context包中的WithDeadline和WithTimeout方法可以實現這個解決方法。先看如下例子：
package main
import (
"context"
"fmt"
"time"
)
func main() {
ctx, cancel := context.WithTimeout(context.Background(), 50*time.Millisecond)
defer cancel()
select {
case <-time.After(1 * time.Second):
fmt.Println("overslept")
case <-ctx.Done():
fmt.Println(ctx.Err()) // 輸出 "context deadline exceeded"
}
}因為設置的超時時間是50毫秒，所以select會進入第二個case，會輸出“context deadline exceeded”。
Golang的net/http包發起http請求的時候是實現了超時控制的，看如下代碼：
package main
import (
"context"
"io"
"log"
"net/http"
"time"
)
func main() {
req, err := http.NewRequest(http.MethodGet, "https://www.baidu.com", nil)
if err != nil {
log.Fatal(err)
}
// 構造一個超時間為50毫秒的Context
ctx, cancel := context.WithTimeout(context.Background(), 50*time.Millisecond)
defer cancel()
req = req.WithContext(ctx)
c := &http.Client{}
res, err := c.Do(req)
if err != nil {
log.Fatal(err)
}
defer res.Body.Close()
out, err := io.ReadAll(res.Body)
if err != nil {
log.Fatal(err)
}
log.Println(string(out))
}執行后會輸出“context deadline exceeded”，如果將context.WithTimeout方法的timeout參數調大一些，就可以看到正常的返回數據。
上下文傳遞數據
這個作用在鏈路追蹤中非常重要，鏈路追蹤需要將traceID層層往下傳遞，在服務間傳遞。
type traceIdKey struct{}{}
// 定義固定的Key
var TraceIdKey = traceIdKey{}
func ServeHTTP(w http.ResponseWriter, req *http.Request){
// 首先從請求中獲取到traceID
traceId := getTraceIdFromRequest(req)
// 將Key存入Context中
ctx := context.WithValue(req.Context(), TraceIdKey, traceId)
// 設置超時時間
ctx = context.WithTimeout(ctx, time.Second)
// 攜帶traceId發起rpc請求
repResp := RequestRPC(ctx, ...)
// 攜帶traceId查詢DB
dbResp := RequestDB(ctx, ...)
// ...
}
func RequestRPC(ctx context.Context, ...) interface{} {
// 獲取traceid，在調用rpc時記錄日志
traceId, _ := ctx.Value(TraceIdKey)
// 發起請求
// ...
return
}接收到請求后，通過req獲取到traceId并記錄到Context中，在調用其他RPC服務和查詢DB時，傳入構造的Context。在后續代碼中，可以通過Context拿到存入的traceId。
