【Go進階—併發程式設計】Context

與昊發表於2022-03-10

Context 是 Go 應用開發常用的併發控制技術,它與 WaitGroup 最大的不同點是 Context 對於派生 goroutine 有更強的控制力,它可以控制多級的 goroutine。

儘管有很多的爭議,但是在很多場景下使用 Context 都很方便,所以現在它已經在 Go 生態圈中傳播開來了,包括很多的 Web 應用框架,都切換成了標準庫的 Context。標準庫中的 database/sql、os/exec、net、net/http 等包中都使用到了 Context。而且,如果遇到了下面的一些場景,也可以考慮使用 Context:

  • 上下文資訊傳遞 ,比如處理 http 請求、在請求處理鏈路上傳遞資訊;
  • 控制子 goroutine 的執行;
  • 超時控制的方法呼叫;
  • 可以取消的方法呼叫。

實現原理

介面定義

包 context 定義了 Context 介面,Context 的具體實現包括 4 個方法,分別是 Deadline、Done、Err 和 Value,如下所示:

type Context interface {
    Deadline() (deadline time.Time, ok bool)
    Done() <-chan struct{}
    Err() error
    Value(key interface{}) interface{}
}

Deadline 方法會返回這個 Context 被取消的截止日期。如果沒有設定截止日期,ok 的值是 false。後續每次呼叫這個物件的 Deadline 方法時,都會返回和第一次呼叫相同的結果。

Done 方法返回一個 Channel 物件,基本上都會在 select 語句中使用。在 Context 被取消時,此 Channel 會被 close,如果沒被取消,可能會返回 nil。當 Done 被 close 的時候,可以通過 ctx.Err 獲取錯誤資訊。

關於 Err 方法,你必須要記住的知識點就是:如果 Done 沒有被 close,Err 方法返回 nil;如果 Done 被 close,Err 方法會返回 Done 被 close 的原因。

Value 返回此 ctx 中和指定的 key 相關聯的 value。

Context 中實現了 2 個常用的生成頂層 Context 的方法:

  • context.Background():返回一個非 nil 的、空的 Context,沒有任何值,不會被 cancel,不會超時,沒有截止日期。一般用在主函式、初始化、測試以及建立根 Context 的時候。
  • context.TODO():返回一個非 nil 的、空的 Context,沒有任何值,不會被 cancel,不會超時,沒有截止日期。當你不清楚是否該用 Context,或者目前還不知道要傳遞一些什麼上下文資訊的時候,就可以使用這個方法。

事實上,它們兩個底層的實現是一模一樣的。絕大多數情況下可以直接使用 context.Background。

在使用 Context 的時候,有一些約定俗成的規則:

  1. 一般函式使用 Context 的時候,會把這個引數放在第一個引數的位置。
  2. 從來不把 nil 當做 Context 型別的引數值,可以使用 context.Background() 建立一個空的上下文物件,也不要使用 nil。
  3. Context 只用來臨時做函式之間的上下文透傳,不能持久化 Context 或者把 Context 長久儲存。把 Context 持久化到資料庫、本地檔案或者全域性變數、快取中都是錯誤的用法。
  4. key 的型別不推薦字串型別或者其它內建型別,否則容易在包之間使用 Context 時候產生衝突。使用 WithValue 時,key 的型別應該是自己定義的型別。
  5. 常常使用 struct{} 作為底層型別定義 key 的型別。對於 exported key 的靜態型別,常常是介面或者指標。這樣可以儘量減少記憶體分配。

context 包中實現 Context 介面的 struct,除了用於 context.Background() 的 emptyCtx 外,還有 cancelCtx、timerCtx 和 valueCtx 三種。

cancelCtx

type cancelCtx struct {
    Context

    mu       sync.Mutex            // 互斥鎖
    done     atomic.Value          // 呼叫 cancel 時會關閉的 channel
    children map[canceler]struct{} // 記錄了由此 Context 派生的所有 child,此 Context 被 cancel 時會同時 cancel 所有 child
    err      error                 // 錯誤資訊
}
WithCancel

cancelCtx 是通過 WithCancel 方法生成的。我們常常在一些需要主動取消長時間的任務時,建立這種型別的 Context,然後把這個 Context 傳給長時間執行任務的 goroutine。當需要中止任務時,我們就可以 cancel 這個 Context,這樣長時間執行任務的 goroutine,就可以通過檢查這個 Context,知道 Context 已經被取消了。

WithCancel 返回值中的第二個值是一個 cancel 函式。記住,不是隻有你想中途放棄,才去呼叫 cancel,只要你的任務正常完成了,就需要呼叫 cancel,這樣,這個 Context 才能釋放它的資源(通知它的 children 處理 cancel,從它的 parent 中把自己移除,甚至釋放相關的 goroutine)。

看下核心原始碼:

func WithCancel(parent Context) (ctx Context, cancel CancelFunc) {
    if parent == nil {
        panic("cannot create context from nil parent")
    }
    c := newCancelCtx(parent)
    propagateCancel(parent, &c)
    return &c, func() { c.cancel(true, Canceled) }
}

func newCancelCtx(parent Context) cancelCtx {
    return cancelCtx{Context: parent}
}

func propagateCancel(parent Context, child canceler) {
    done := parent.Done()

    if p, ok := parentCancelCtx(parent); ok {
        p.mu.Lock()
        if p.err != nil {
            // parent 已經取消了,直接取消子 Context
            child.cancel(false, p.err)
        } else {
            // 將 child 新增到 parent 的 children 切片
            if p.children == nil {
                p.children = make(map[canceler]struct{})
            }
            p.children[child] = struct{}{}
        }
        p.mu.Unlock()
    } else {
        atomic.AddInt32(&goroutines, +1)
        // 沒有 parent 可以“掛載”,啟動一個 goroutine 監聽 parent 的 cancel,同時 cancel 自身
        go func() {
            select {
            case <-parent.Done():
                child.cancel(false, parent.Err())
            case <-child.Done():
            }
        }()
    }
}

程式碼中呼叫的 propagateCancel 方法會順著 parent 路徑往上找,直到找到一個 cancelCtx,或者為 nil。如果不為空,就把自己加入到這個 cancelCtx 的 child,以便這個 cancelCtx 被取消的時候通知自己。如果為空,會新起一個 goroutine,由它來監聽 parent 的 Done 是否已關閉。

當這個 cancelCtx 的 cancel 函式被呼叫的時候,或者 parent 的 Done 被 close 的時候,這個 cancelCtx 的 Done 才會被 close。

cancel 是向下傳遞的,如果一個 WithCancel 生成的 Context 被 cancel 時,如果它的子 Context(也有可能是孫,或者更低,依賴子的型別)也是 cancelCtx 型別的,就會被 cancel。

cancel

cancel 方法的作用是 close 自己及其後代的 done 通道,達到通知取消的目的。WithCancel 方法的第二個返回值 cancel 就是本函式。來看一下主要程式碼實現:

func (c *cancelCtx) cancel(removeFromParent bool, err error) {
    c.mu.Lock()
    // 設定 cancel 的原因
    c.err = err 
    // 關閉自身的 done 通道
    d, _ := c.done.Load().(chan struct{})
    if d == nil {
        c.done.Store(closedchan)
    } else {
        close(d)
    }
    // 遍歷所有 children,逐個呼叫 cancel 方法
    for child := range c.children {
        child.cancel(false, err)
    }
    c.children = nil
    c.mu.Unlock()

    // 正常情況下,需要將自身從 parent 的 children 切片中刪除
    if removeFromParent {
        removeChild(c.Context, c)
    }
}

timerCtx

type timerCtx struct {
    cancelCtx
    timer *time.Timer 

    deadline time.Time
}

timerCtx 在 cancelCtx 基礎上增加了 deadline 用於標示自動 cancel 的最終時間,而 timer 就是一個觸發自動 cancel 的定時器。timerCtx 可以由 WithDeadline 和 WithTimeout 生成, WithTimeout 實際呼叫了 WithDeadline,二者實現原理一致,只不過使用語境不一樣:WithDeadline 是指定最後期限,WithTimeout 是指定最長存活時間。

WithDeadline

來看一下 WithDeadline 方法的實現:

func WithDeadline(parent Context, d time.Time) (Context, CancelFunc) {
    // 如果parent的截止時間更早,直接返回一個cancelCtx即可
    if cur, ok := parent.Deadline(); ok && cur.Before(d) {
        return WithCancel(parent)
    }
    c := &timerCtx{
        cancelCtx: newCancelCtx(parent),
        deadline:  d,
    }
    propagateCancel(parent, c) // 同cancelCtx的處理邏輯
    dur := time.Until(d)
    if dur <= 0 { //當前時間已經超過了截止時間,直接cancel
        c.cancel(true, DeadlineExceeded)
        return c, func() { c.cancel(false, Canceled) }
    }
    c.mu.Lock()
    defer c.mu.Unlock()
    if c.err == nil {
        // 設定一個定時器,到截止時間後取消
        c.timer = time.AfterFunc(dur, func() {
            c.cancel(true, DeadlineExceeded)
        })
    }
    return c, func() { c.cancel(true, Canceled) }
}

WithDeadline 會返回一個 parent 的副本,並且設定了一個不晚於引數 d 的截止時間,型別為 timerCtx(或者是 cancelCtx)。

如果它的截止時間晚於 parent 的截止時間,那麼就以 parent 的截止時間為準,並返回一個型別為 cancelCtx 的 Context,因為 parent 的截止時間到了,就會取消這個 cancelCtx。如果當前時間已經超過了截止時間,就直接返回一個已經被 cancel 的 timerCtx。否則就會啟動一個定時器,到截止時間取消這個 timerCtx。

綜合起來,timerCtx 的 Done 被 Close 掉,主要是由下面的某個事件觸發的:

  • 截止時間到了;
  • cancel 函式被呼叫;
  • parent 的 Done 被 close。

和 cancelCtx 一樣,WithDeadline(WithTimeout)返回的 cancel 一定要呼叫,並且要儘可能早地被呼叫,這樣才能儘早釋放資源,不要單純地依賴截止時間被動取消。

valueCtx

type valueCtx struct {
    Context
    key, val interface{}
}

valueCtx 只是在 Context 基礎上增加了一個 key-value 對,用於在各級協程間傳遞一些資料。

WithValue 基於 parent Context 生成一個新的 valueCtx,儲存了一個 key-value 鍵值對。valueCtx 覆蓋了 Value 方法,優先從自己的儲存中檢查這個 key,不存在的話會從 parent 中繼續檢查。

相關文章