Go 裡的超時控制

crossoverJie發表於2021-10-28
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 便會退出，也就達到了超時的目的。

output:

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

相關文章