Go 里的超時控制也很簡單
前言
日常開發中我們大概率會遇到超時控制的場景,比如一個批量耗時任務、網絡請求等;一個良好的超時控制可以有效的避免一些問題(比如 goroutine 泄露、資源不釋放等)。
Timer
在 go 中實現超時控制的方法非常簡單,首先第一種方案是 Time.After(d Duration):
- func main() {
- fmt.Println(time.Now())
- x := <-time.After(3 * time.Second)
- fmt.Println(x)
- }
output:
- 2021-10-27 23:06:04.304596 +0800 CST m=+0.000085653
- 2021-10-27 23:06:07.306311 +0800 CST m=+3.001711390
time.After() 會返回一個 Channel,該 Channel 會在延時 d 段時間后寫入數據。
有了這個特性就可以實現一些異步控制超時的場景:
- func main() {
- ch := make(chan struct{}, 1)
- go func() {
- fmt.Println("do something...")
- time.Sleep(4*time.Second)
- ch<- struct{}{}
- }()
- select {
- case <-ch:
- fmt.Println("done")
- case <-time.After(3*time.Second):
- fmt.Println("timeout")
- }
- }
這里假設有一個 goroutine 在跑一個耗時任務,利用 select 有一個 channel 獲取到數據便退出的特性,當 goroutine 沒有在有限時間內完成任務時,主 goroutine 便會退出,也就達到了超時的目的。
- do something...
- timeout
timer.After 取消,同時 Channel 發出消息,也可以關閉通道等通知方式。
注意 Channel 最好是有大小,防止阻塞 goroutine ,導致泄露。
Context
第二種方案是利用 context,go 的 context 功能強大;
利用 context.WithTimeout() 方法會返回一個具有超時功能的上下文。
- ch := make(chan string)
- timeout, cancel := context.WithTimeout(context.Background(), 3*time.Second)
- defer cancel()
- go func() {
- time.Sleep(time.Second * 4)
- ch <- "done"
- }()
- select {
- case res := <-ch:
- fmt.Println(res)
- case <-timeout.Done():
- fmt.Println("timout", timeout.Err())
- }
同樣的用法,context 的 Done() 函數會返回一個 channel,該 channel 會在當前工作完成或者是上下文取消生效。
- timout context deadline exceeded
通過 timeout.Err() 也能知道當前 context 關閉的原因。
goroutine 傳遞 context
使用 context 還有一個好處是,可以利用其天然在多個 goroutine 中傳遞的特性,讓所有傳遞了該 context 的 goroutine 同時接收到取消通知,這點在多 go 中應用非常廣泛。
- func main() {
- total := 12
- var num int32
- log.Println("begin")
- ctx, cancelFunc := context.WithTimeout(context.Background(), 3*time.Second)
- for i := 0; i < total; i++ {
- go func() {
- //time.Sleep(3 * time.Second)
- atomic.AddInt32(&num, 1)
- if atomic.LoadInt32(&num) == 10 {
- cancelFunc()
- }
- }()
- }
- for i := 0; i < 5; i++ {
- go func() {
- select {
- case <-ctx.Done():
- log.Println("ctx1 done", ctx.Err())
- }
- for i := 0; i < 2; i++ {
- go func() {
- select {
- case <-ctx.Done():
- log.Println("ctx2 done", ctx.Err())
- }
- }()
- }
- }()
- }
- time.Sleep(time.Second*5)
- log.Println("end", ctx.Err())
- fmt.Printf("執行完畢 %v", num)
- }
在以上例子中,無論 goroutine 嵌套了多少層,都是可以在 context 取消時獲得消息(當然前提是 context 得傳遞走)
某些特殊情況需要提前取消 context 時,也可以手動調用 cancelFunc() 函數。
Gin 中的案例
Gin 提供的 Shutdown(ctx) 函數也充分使用了 context。
- ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
- defer cancel()
- if err := srv.Shutdown(ctx); err != nil {
- log.Fatal("Server Shutdown:", err)
- }
- log.Println("Server exiting")
比如以上代碼便是超時等待 10s 進行 Gin 的資源釋放,實現的原理也和上文的例子相同。
總結
因為寫 go 的時間不長,所以自己寫了一個練手的項目:一個接口壓力測試工具。
其中一個很常見的需求就是壓測 N 秒后退出,這里正好就應用到了相關知識點,同樣是初學 go 的小伙伴可以參考。
https://github.com/crossoverJie/ptg/blob/d0781fcb5551281cf6d90a86b70130149e1525a6/duration.go#L41