context-Context Code Review

概述

context.Context 類型表示上下文信息，作用是在 API 通信、進程通信、函數調用之間傳遞超時時間、取消信號和其他數據值。

規則

使用 context.Context 時應遵循以下規則，以保證不同包之間的接口一致性，並啓用靜態分析工具來檢查上下文的鏈路傳播和可能產生的數據競態。

不要將 context.Context 嵌入結構體類型，應該明確地將其作爲函數參數傳遞，並且應該是第一個參數，典型的命名爲 ctx
不要傳遞一個值爲 nil 的 context.Context, 即使函數內部並不使用這個參數。如果確定不了合適的 context.Context 類型，可以傳遞 context.TODO 類型
僅在 API 訪問、進程通信之間使用 context.Value, 而不是將其作爲可選參數傳遞給函數
context.Context 可以在多個 goroutine 之間同時使用，同一個 context.Context 可以傳遞給不同的 goroutine 裏面的函數

4 種類型

emptyCtx: 所有 Context 類型的基類型，使用 context.TODO(), context.Background 創建
valueCtx: 可以在 Context 中傳遞 1 個 K-V 鍵值對，它爲 key 字段實現 Value 方法，並將調用委託給嵌入的 Context
cancelCtx: 取消執行，當執行被取消時，子節點中實現 canceler 接口的 Context 同樣被取消
timerCtx: 在 cancelCtx 基礎上封裝一層，支持基於超時和定時器的 cancelCtx

內部實現

我們來探究一下 context.Context 的內部實現，文件路徑爲 $GOROOT/src/context/context.go，筆者的 Go 版本爲 go1.19 linux/amd64。

Context 接口

Context 接口用於在多個 API 之間傳遞超時時間、取消信號和其他數據值，接口的方法可以被多個 goroutine 併發調用。

type Context interface {
 // Deadline 返回 Context 執行截止時間
 // 如果沒有設置截止時間, 返回 false
 // 連續多次調用 Deadline 會返回相同的結果
 Deadline() (deadline time.Time, ok bool)

 // Done 返回一個只讀 channel, 該 channel 在 Context 超時或被取消時關閉
 // 如果 context 永遠不會被取消，Done 返回 nil
 // 連續多次調用 Done 會返回相同的結果
 // Done channel 的關閉可能在 cancel 函數返回後異步執行

 // 當 cancel 調用時，WithCancel 執行 Done 關閉
  // 當 deadline 到期，WithDeadline 執行 Done 關閉
 // 當 timeout 超時， WithTimeout 執行 Done 關閉

 // Done 是在 select 語句中使用的

 //  示例
 //  DoSomething 函數生成 values 併發送至 out，直到 DoSomething 返回一個錯誤或 ctx.Done 已經關閉
 //  func Stream(ctx context.Context, out chan<- Value) error {
 //   for {
 //    v, err := DoSomething(ctx)
 //    if err != nil {
 //     return err
 //    }
 //    select {
 //    case <-ctx.Done():
 //     return ctx.Err()
 //    case out <- v:
 //    }
 //   }
 //  }
 //
 // See https://blog.golang.org/pipelines for more examples of how to use
 // a Done channel for cancellation.
 Done() <-chan struct{}

 // 如果 Done 還未關閉，返回 nil
 // 如果 Done 已經關閉，根據下述情況返回具體的錯誤
 //    如果 context 被取消，返回 Canceled
 //    如果 deadline 到期，返回 DeadlineExceeded
 // 如果返回一個非 nil 錯誤, 那麼連續多次調用 Err 會返回同樣的錯誤
 Err() error

 // Value 返回 Context 關聯的 key 對應的值，如果沒有值關聯，返回 nil
 // 連續多次使用相同的 key 參數，調用 Value, 會返回相同的結果

  // key 標識 Context 的特定值
 // 希望在 Context 中存儲值的函數，通常會在全局變量中分配一個鍵
 // 然後將該鍵作爲參數傳入 context.WithValue 和 Context.Value
 // key 可以是任何支持 equality 的類型，應該將 key 定義爲未導出類型，以避免衝突

 // 示例
 //  // 包 user 定義了存儲在 Context 中的 User 數據類型
 //  package user
 //
 //  import "context"
 //
 //  // User 是存儲在 Context 中的值的類型
 //  type User struct {...}
 //
 //  // key 是此包中定義的鍵的未導出類型
 //  // 這樣可以防止與其他包中定義的 key 發生衝突
 //  type key int
 //
 //  // userKey 是存儲在 Context 中的 user.User 值的 key, 是未導出的
 //  // 調用方使用 user.NewContext 和 user.FromContext, 而不是直接使用 key
 //  var userKey key
 //
 //  // NewContext 返回一個包含 User 類型值的 Context
 //  func NewContext(ctx context.Context, u *User) context.Context {
 //   return context.WithValue(ctx, userKey, u)
 //  }
 //
 //  // FromContext 返回一個存儲在 Context 中的 User 類型值
 //  func FromContext(ctx context.Context) (*User, bool) {
 //   u, ok := ctx.Value(userKey).(*User)
 //   return u, ok
 //  }
 Value(key any) any
}

錯誤類型

Canceled 錯誤類型表示 Context 被取消執行，DeadlineExceeded 錯誤類型表示 Context 超時。

var Canceled = errors.New("context canceled")

var DeadlineExceeded error = deadlineExceededError{}

type deadlineExceededError struct{}

func (deadlineExceededError) Error() string  { return "context deadline exceeded" }

emptyCtx 類型

emptyCtx 類型沒有傳遞值，沒有超時時間，也永遠不會被調用，該類型實現了 Context 接口，但都是空實現。

需要注意的是，emptyCtx 不能定義爲 struct{} 類型 (爲了節省內存)，因爲該類型的變量必須有不同的地址 (空 struct{} 地址是唯一的)。

type emptyCtx int

func (*emptyCtx) Deadline() (deadline time.Time, ok bool) {
    return
}

func (*emptyCtx) Done() <-chan struct{} {
    return nil
}

func (*emptyCtx) Err() error {
    return nil
}

func (*emptyCtx) Value(key any) any {
    return nil
}

Background, TODO 類型

Background, TODO 本質上都是 emptyCtx，只是定義了不同的別名用以區分，兩者的語義分別如下:

Background: 一般作爲 Context 樹的根結點
TODO: 一般在無法確定應該使用哪種 Context 類型時，使用 TODO

var (
    background = new(emptyCtx)
    todo       = new(emptyCtx)
)

func (e *emptyCtx) String() string {
 switch e {
 case background:
  return "context.Background"
 case todo:
  return "context.TODO"
 }
 return "unknown empty Context"
}

// 語義：作爲 Context 樹的根結點
// Background 返回一個非 nil, 沒有 values，沒有 deadline，也永遠不會被調用的 empty Context
// 典型使用場景有主函數，初始化，測試，以及作爲頂級 Context 傳入請求
func Background() Context {
 return background
}

// 語義：不知道應該使用哪種 Context 類型時，使用 TODO
// TODO 返回一個非 nil 的 empty Context
// 當不清楚應該使用哪種 Context 類型，或者 Context 類型還不可用時（函數的調用棧存在不接收處理 Context 參數的情況），使用 context.TODO
func TODO() Context {
 return todo
}

canceler 接口

// canceler 是一個可以被直接取消的 Context 類型
// *cancelCtx 和 *timerCtx 實現了這個接口
type canceler interface {
 cancel(removeFromParent bool, err error)
 Done() <-chan struct{}
}

Cancel 類型

// CancelFunc 通知一個操作放棄其執行（取消對應 Context）
// CancelFunc 不會等待執行停止
// CancelFunc 可以被多個 goroutine 併發調用
// CancelFunc 只有第一次調用有效，後續調用什麼也不執行
type CancelFunc func()

// cancelCtx 實現取消操作，當它被取消時，同時取消所有實現了 canceler 接口的子節點
type cancelCtx struct {
 Context

 mu       sync.Mutex            // 保證下面三個字段的互斥訪問
 done     atomic.Value          // 懶惰式初始化，被第一個 cancel() 調用關閉
 children map[canceler]struct{} // 被第一個 cancel() 重置爲 nil
 err      error                 // 被第一個 cancel() 重置爲 non-nil
}

// 複用了 Value 函數的回溯邏輯
// 從而在 Context 樹回溯鏈中遍歷時，可以找到給定 Context 的第一個祖先 cancelCtx 實例
func (c *cancelCtx) Value(key any) any {
    // 特殊的 key: cancelCtxKey
 if key == &cancelCtxKey {
        return c
    }
    return value(c.Context, key)
}

func (c *cancelCtx) Done() <-chan struct{} {
 d := c.done.Load()
 if d != nil {
  return d.(chan struct{})
 }

 c.mu.Lock()
 defer c.mu.Unlock()
  // double check
 d = c.done.Load()
 if d == nil {
  // 惰式初始化
  d = make(chan struct{})
  c.done.Store(d)
 }
 return d.(chan struct{})
}

func (c *cancelCtx) Err() error {
 c.mu.Lock()
 err := c.err
 c.mu.Unlock()
 return err
}

// WithCancel 返回包裝後的 Context 和用於取消該 Context 的函數 cancel
// 當 cancel 函數被調用，或者父 Context 被關閉時，返回的 Context 的 Done channel 會關閉

// 在 Context 涉及的上下文操作完成後，應該立即調用 cancel 函數，釋放與 Context 關聯的資源
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) }
}

// newCancelCtx 返回一個初始化後的 cancelCtx
func newCancelCtx(parent Context) cancelCtx {
 return cancelCtx{Context: parent}
}

cancelCtx.cancel 方法

// 小技巧
// closedchan 表示一個可重用的已關閉 channel
var closedchan = make(chan struct{})

func init() {
    // 初始化時關閉
    close(closedchan)
}

// cancel 關閉 c.done channel, 取消 c 的所有子節點
// 如果 removeFromParent 爲 true, 將 c 從其父節點中刪除
func (c *cancelCtx) cancel(removeFromParent bool, err error) {
 if err == nil {
  // 需要給出取消的理由: Canceled or DeadlineExceeded
  panic("context: internal error: missing cancel error")
 }

 c.mu.Lock()
 if c.err != nil {
  // 已經被其他 goroutine 取消
  c.mu.Unlock()
  return
 }

 // 記錄錯誤
 c.err = err
 d, _ := c.done.Load().(chan struct{})
 if d == nil {
  // 惰式創建
  c.done.Store(closedchan)
 } else {
  // 關閉 c.done
  close(d)
 }

 // 級聯取消
 for child := range c.children {
  // NOTE: 持有父節點鎖的同時獲取子節點的鎖
  child.cancel(false, err)
 }
 c.children = nil
 c.mu.Unlock()

 // 將 c 從其父節點中刪除
 if removeFromParent {
  removeChild(c.Context, c)
 }
}

Timer 類型

// timerCtx 類型包含一個定時器和一個超時時間
// 通過嵌入 cancelCtx 複用 Done() 方法和 Err() 方法
// 通過停止定時器然後委託給 cancelCtx.cancel() 方法
type timerCtx struct {
 cancelCtx
 timer *time.Timer

 deadline time.Time
}

func (c *timerCtx) Deadline() (deadline time.Time, ok bool) {
    return c.deadline, true
}

func (c *timerCtx) String() string {
 ...
}

// 在 Context 涉及的上下文操作完成後，應該立即調用 cancel 函數，釋放與 Context 關聯的資源
func WithTimeout(parent Context, timeout time.Duration) (Context, CancelFunc) {
    return WithDeadline(parent, time.Now().Add(timeout))
}

timerCtx.cancel 方法

func (c *timerCtx) cancel(removeFromParent bool, err error) {
 // 級聯取消子樹中所有 context
 c.cancelCtx.cancel(false, err)
 if removeFromParent {
  // 從父節點中刪除當前節點
  removeChild(c.cancelCtx.Context, c)
 }

 c.mu.Lock()
 if c.timer != nil {
  // 關閉定時器 (避免 timer 泄漏)
  c.timer.Stop()
  c.timer = nil
 }
 c.mu.Unlock()
}

Value 類型

// valueCtx 包含 1 個 K-V 鍵值對，它爲 key 字段實現 Value 方法，並將調用委託給嵌入的 Context
type valueCtx struct {
    Context
    key, val any
}

func (c *valueCtx) Value(key any) any {
    if c.key == key {
        return c.val
    }
 return value(c.Context, key)
}

// WithValue 將給定鍵值對包裝進 Context 並返回
// key 必須具有可比性，不應該是 string 或其他內置類型，以避免調用方之間產生衝突，最好使用自定義類型作爲 key
func WithValue(parent Context, key, val any) Context {
 if parent == nil {
  panic("cannot create context from nil parent")
 }
 if key == nil {
  panic("nil key")
 }
 if !reflectlite.TypeOf(key).Comparable() {
  panic("key is not comparable")
 }
 return &valueCtx{parent, key, val}
}

parentCancelCtx 方法

// parentCancelCtx 返回參數的第一個祖先 *cancelCtx
// 它通過查找 parent.Value(&cancelCtxKey), 來找到最內層的（回溯鏈中第一個） *cancelCtx
// 然後檢查 parent.Done() 與 *cancelCtx 是否匹配
// 如果不匹配，說明 *cancelCtx 被包裝在一個實現了 context.Context 接口的自定義 Context 中
func parentCancelCtx(parent Context) (*cancelCtx, bool) {
 done := parent.Done()
 if done == closedchan || done == nil {
  return nil, false
 }
 p, ok := parent.Value(&cancelCtxKey).(*cancelCtx)
 if !ok {
  return nil, false
 }
 pdone, _ := p.done.Load().(chan struct{})
 if pdone != done {
  return nil, false
 }
 return p, true
}

propagateCancel 方法

// propagateCancel 實現節點編排: 當參數 (父節點) 取消時，同時取消其子節點
func propagateCancel(parent Context, child canceler) {
 done := parent.Done()
 if done == nil {
  // 不可取消 (父節點是 emptyCtx 類型，或者自定義實現的 Context 類型)
  return
 }

 select {
 case <-done:
  // 已經取消
  child.cancel(false, parent.Err())
  return
 default:
 }

 if p, ok := parentCancelCtx(parent); ok {
  // 找到一個 cancelCtx 實例
  p.mu.Lock()
  if p.err != nil {
   // parent 已經取消
   child.cancel(false, p.err)
  } else {
   if p.children == nil {
      // 惰式創建
    p.children = make(map[canceler]struct{})
   }
   p.children[child] = struct{}{}
  }
  p.mu.Unlock()
 } else {
  // 找到一個非 cancelCtx 實例
  // 增加 goroutines 計數
  atomic.AddInt32(&goroutines, +1)

  go func() {
   select {
   case <-parent.Done():
    child.cancel(false, parent.Err())
   case <-child.Done():
   }
  }()
 }
}

WithDeadline 方法

// WithDeadline 返回包裝後的 Context 和用於取消該 Context 的函數 cancel
// 包裝之後的 Context 的截止時間調整爲參數 d
// 如果 parent 的截止時間已經早於參數 d, WithDeadline(parent, d) 在語義上就等於 parent
// 下面幾種情況，不論哪種情況先發生，返回的 Context 的 Done channel 都會被關閉
//      1. 當返回的函數 cancel 被調用
//      2. 截止時間已到
//      3. 父節點的 Done channel 被關閉

// 在 Context 涉及的上下文操作完成後，應該立即調用 cancel 函數，釋放與 Context 關聯的資源
func WithDeadline(parent Context, d time.Time) (Context, CancelFunc) {
 if parent == nil {
  panic("cannot create context from nil parent")
 }

 if cur, ok := parent.Deadline(); ok && cur.Before(d) {
  // parent 的截止時間更早
  return WithCancel(parent)
 }

// 包裝一個新的 cancelCtx 實現部分接口
 c := &timerCtx{
  cancelCtx: newCancelCtx(parent),
  deadline:  d,
 }

 // 構建 Context 取消樹，傳入的是 c 而非 c.cancelCtx
 propagateCancel(parent, c)

 dur := time.Until(d)
 if dur <= 0 {
  // 截止時間設置的太接近，已經超時
  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) }
}

小結

從應用層面來說，context.Context 主要用來在多個 goroutine 之間傳遞截止時間、取消信號，尤其是在多個層級函數調用棧中，所以我們可以看到大多數第三方庫提供的 API 中，第一個參數都是 context.Context, 此外，雖然 context.Context 也可以傳遞值，但實際中很少使用這種方式，典型的場景比如傳遞 Request-ID, Token, Auth-Key 等。

從內部實現來說，context.Context 主要是基於 樹形結構 + 互斥鎖，這兩者讀者應該都比較熟悉，這裏就不做贅述了。這裏列舉一些內部實現中值得我們學習的代碼技巧，比如 將具體錯誤公開爲一個變量 的 DeadlineExceeded, 提供友好 API 的空實現 background, todo, 通過通道來表示一種變化標識符 的 closedchan。

擴展閱讀

關於使用 Context.WithValue() 傳遞值的問題，在《Concurrency in Go》一書中，作者提到了幾個觀點:

值可以跨 API 訪問
值應該是不可變的
值應該是簡單的數據類型
值應該是具體數據，而非某個類型或方法
不同的值不會影響到後續的鏈路邏輯

該書作者還給出了幾種常見的 傳遞值，並與上面的 5 點要求進行了對比:

SdyH9K

Reference

golang context[1]
Concurrency in Go[2]

鏈接

[1] golang context: https://blog.golang.org/context

[2] Concurrency in Go: https://book.douban.com/subject/26994591/

本文由 Readfog 進行 AMP 轉碼，版權歸原作者所有。
來源：https://mp.weixin.qq.com/s/C2L-reec4DejgzgUoZzP0A

概述

規則