快捷導(dǎo)航

Go語言context上下文管理的使用

更新時(shí)間：2022年03月09日 11:27:24 作者：golandscape

本文主要介紹了Go語言context上下文管理的使用，文中通過示例代碼介紹的非常詳細(xì)，具有一定的參考價(jià)值，感興趣的小伙伴們可以參考一下

context 有什么作用

context 主要用來在goroutine 之間傳遞上下文信息，包括：取消信號(hào)、超時(shí)時(shí)間、截止時(shí)間、k-v 等。

Go 常用來寫后臺(tái)服務(wù)，通常只需要幾行代碼，就可以搭建一個(gè) http server。

在 Go 的 server 里，通常每來一個(gè)請(qǐng)求都會(huì)啟動(dòng)若干個(gè) goroutine 同時(shí)工作：有些去數(shù)據(jù)庫拿數(shù)據(jù)，有些調(diào)用下游接口獲取相關(guān)數(shù)據(jù)……

這些 goroutine 需要共享這個(gè)請(qǐng)求的基本數(shù)據(jù)，例如登陸的 token，處理請(qǐng)求的最大超時(shí)時(shí)間（如果超過此值再返回?cái)?shù)據(jù)，請(qǐng)求方因?yàn)槌瑫r(shí)接收不到）等等。當(dāng)請(qǐng)求被取消或是處理時(shí)間太長，這有可能是使用者關(guān)閉了瀏覽器或是已經(jīng)超過了請(qǐng)求方規(guī)定的超時(shí)時(shí)間，請(qǐng)求方直接放棄了這次請(qǐng)求結(jié)果。這時(shí)，所有正在為這個(gè)請(qǐng)求工作的 goroutine 需要快速退出，因?yàn)樗鼈兊?ldquo;工作成果”不再被需要了。在相關(guān)聯(lián)的 goroutine 都退出后，系統(tǒng)就可以回收相關(guān)的資源。

在Go 里，我們不能直接殺死協(xié)程，協(xié)程的關(guān)閉一般會(huì)用 channel+select 方式來控制。但是在某些場(chǎng)景下，例如處理一個(gè)請(qǐng)求衍生了很多協(xié)程，這些協(xié)程之間是相互關(guān)聯(lián)的：需要共享一些全局變量、有共同的 deadline 等，而且可以同時(shí)被關(guān)閉。再用 channel+select 就會(huì)比較麻煩，這時(shí)就可以通過 context 來實(shí)現(xiàn)。

一句話：context 用來解決 goroutine 之間退出通知、元數(shù)據(jù)傳遞的功能。

context 使用起來非常方便。源碼里對(duì)外提供了一個(gè)創(chuàng)建根節(jié)點(diǎn) context 的函數(shù)：

func Background() Context

background 是一個(gè)空的 context，它不能被取消，沒有值，也沒有超時(shí)時(shí)間。有了根節(jié)點(diǎn) context，又提供了四個(gè)函數(shù)創(chuàng)建子節(jié)點(diǎn) context：

func WithCancel(parent Context) (ctx Context, cancel CancelFunc)
func WithDeadline(parent Context, deadline time.Time) (Context, CancelFunc)
func WithTimeout(parent Context, timeout time.Duration) (Context, CancelFunc)
func WithValue(parent Context, key, val interface{}) Context

context 會(huì)在函數(shù)傳遞間傳遞。只需要在適當(dāng)?shù)臅r(shí)間調(diào)用 cancel 函數(shù)向 goroutines 發(fā)出取消信號(hào)或者調(diào)用 Value 函數(shù)取出 context 中的值。

不要將 Context 塞到結(jié)構(gòu)體里。直接將 Context 類型作為函數(shù)的第一參數(shù)，而且一般都命名為ctx。
不要向函數(shù)傳入一個(gè) nil 的 context，如果你實(shí)在不知道傳什么，標(biāo)準(zhǔn)庫給你準(zhǔn)備好了一個(gè) context：todo。
不要把本應(yīng)該作為函數(shù)參數(shù)的類型塞到 context 中，context 存儲(chǔ)的應(yīng)該是一些共同的數(shù)據(jù)。例如：登陸的 session、cookie 等。
同一個(gè) context 可能會(huì)被傳遞到多個(gè) goroutine，別擔(dān)心，context 是并發(fā)安全的。

傳遞共享的數(shù)據(jù)

對(duì)于 Web 服務(wù)端開發(fā)，往往希望將一個(gè)請(qǐng)求處理的整個(gè)過程串起來，這就非常依賴于 Thread Local（對(duì)于 Go 可理解為單個(gè)協(xié)程所獨(dú)有）的變量，而在 Go 語言中并沒有這個(gè)概念，因此需要在函數(shù)調(diào)用的時(shí)候傳遞 context。

package main

import (
    "context"
    "fmt"
)
func main() {
    ctx := context.Background()
    process(ctx)
    ctx = context.WithValue(ctx, "traceId", "qcrao-2019")
    process(ctx)
}
func process(ctx context.Context) {
    traceId, ok := ctx.Value("traceId").(string)
    if ok {
        fmt.Printf("process over. trace_id=%s\n", traceId)
    } else {
        fmt.Printf("process over. no trace_id\n")
    }
}

運(yùn)行結(jié)果：

process over. no trace_id
process over. trace_id=qcrao-2019

第一次調(diào)用 process 函數(shù)時(shí)，ctx 是一個(gè)空的 context，自然取不出來 traceId。第二次，通過 WithValue 函數(shù)創(chuàng)建了一個(gè) context，并賦上了 traceId 這個(gè) key，自然就能取出來傳入的 value 值。

取消 goroutine

我們先來設(shè)想一個(gè)場(chǎng)景：打開外賣的訂單頁，地圖上顯示外賣小哥的位置，而且是每秒更新 1 次。app 端向后臺(tái)發(fā)起 websocket 連接（現(xiàn)實(shí)中可能是輪詢）請(qǐng)求后，后臺(tái)啟動(dòng)一個(gè)協(xié)程，每隔 1 秒計(jì)算 1 次小哥的位置，并發(fā)送給端。如果用戶退出此頁面，則后臺(tái)需要“取消”此過程，退出 goroutine，系統(tǒng)回收資源。

func Perform() {
    for {
        calculatePos()
        sendResult()
        time.Sleep(time.Second)
    }
}

如果需要實(shí)現(xiàn)“取消”功能，并且在不了解 context 功能的前提下，可能會(huì)這樣做：給函數(shù)增加一個(gè)指針型的 bool 變量，在 for 語句的開始處判斷 bool 變量是發(fā)由 true 變?yōu)?false，如果改變，則退出循環(huán)。

上面給出的簡(jiǎn)單做法，可以實(shí)現(xiàn)想要的效果，沒有問題，但是并不優(yōu)雅，并且一旦協(xié)程數(shù)量多了之后，并且各種嵌套，就會(huì)很麻煩。優(yōu)雅的做法，自然就要用到 context。

func Perform(ctx context.Context) {
    for {
        calculatePos()
        sendResult()
        select {
        case <-ctx.Done():
            // 被取消，直接返回
            return
        case <-time.After(time.Second):
            // block 1 秒鐘 
        }
    }
}

主流程可能是這樣的：

ctx, cancel := context.WithTimeout(context.Background(), time.Hour)
go Perform(ctx)
// ……
// app 端返回頁面，調(diào)用cancel 函數(shù)
cancel()

注意一個(gè)細(xì)節(jié)，WithTimeout 函數(shù)返回的 context 和 cancelFun 是分開的。context 本身并沒有取消函數(shù)，這樣做的原因是取消函數(shù)只能由外層函數(shù)調(diào)用，防止子節(jié)點(diǎn) context 調(diào)用取消函數(shù)，從而嚴(yán)格控制信息的流向：由父節(jié)點(diǎn) context 流向子節(jié)點(diǎn) context。

防止 goroutine 泄漏

前面那個(gè)例子里，goroutine 還是會(huì)執(zhí)行完，最后返回，可能多浪費(fèi)一些系統(tǒng)資源。這里改編一個(gè) “如果不用 context 取消，goroutine 就會(huì)泄漏的例子”

func gen() <-chan int {
    ch := make(chan int)
    go func() {
        var n int
        for {
            ch <- n
            n++
            time.Sleep(time.Second)
        }
    }()
    return ch
}

這是一個(gè)可以生成無限整數(shù)的協(xié)程，但如果我只需要它產(chǎn)生的前 5 個(gè)數(shù)，那么就會(huì)發(fā)生 goroutine 泄漏：

func main() {
    for n := range gen() {
        fmt.Println(n)
        if n == 5 {
            break
        }
    }
    // ……
}

當(dāng) n == 5 的時(shí)候，直接 break 掉。那么 gen 函數(shù)的協(xié)程就會(huì)執(zhí)行無限循環(huán)，永遠(yuǎn)不會(huì)停下來。發(fā)生了 goroutine 泄漏。

用 context 改進(jìn)這個(gè)例子：

func gen(ctx context.Context) <-chan int {
    ch := make(chan int)
    go func() {
        var n int
        for {
            select {
            case <-ctx.Done():
                return
            case ch <- n:
                n++
                time.Sleep(time.Second)
            }
        }
    }()
    return ch
}
func main() {
    ctx, cancel := context.WithCancel(context.Background())
    defer cancel() // 避免其他地方忘記 cancel，且重復(fù)調(diào)用不影響
    for n := range gen(ctx) {
        fmt.Println(n)
        if n == 5 {
            cancel()
            break
        }
    }
    // ……
}

增加一個(gè) context，在 break 前調(diào)用 cancel 函數(shù)，取消 goroutine。gen 函數(shù)在接收到取消信號(hào)后，直接退出，系統(tǒng)回收資源。