Golang Context 原理與實戰

本文讓我們一起來學習 golang Context 的使用和標準庫中的Context的實現。

golang context 包 一開始只是 Google 內部使用的一個 Golang 包，在 Golang 1.7的版本中正式被引入標準庫。下面開始學習。

簡單介紹

在學習 context 包之前，先看幾種日常開發中經常會碰到的業務場景：

1. 業務需要對訪問的數據庫，RPC ，或API接口，為了防止這些依賴導致我們的服務超時，需要針對性的做超時控制。
2. 為了詳細瞭解服務性能，記錄詳細的調用鏈Log。

上面兩種場景在web中是比較常見的，context 包就是為了方便我們應對此類場景而使用的。

接下來, 我們首先學習 context 包有哪些方法供我們使用；接著舉一些例子，使用 context 包應用在我們上述場景中去解決我們遇到的問題；最後從源碼角度學習 context 內部實現，瞭解 context 的實現原理。

Context 包

Context 定義

context 包中實現了多種 Context 對象。Context 是一個接口，用來描述一個程序的上下文。接口中提供了四個抽象的方法，定義如下：

<code>type Context interface {
  Deadline() (deadline time.Time, ok bool)
  Done() /<code>

• Deadline() 返回的是上下文的截至時間，如果沒有設定，ok 為 false
• Done() 當執行的上下文被取消後，Done返回的chan就會被close。如果這個上下文不會被取消，返回nil
• Err() 有幾種情況:如果Done() 返回 chan 沒有關閉，返回nil如果Done() 返回的chan 關閉了， Err 返回一個非nil的值，解釋為什麼會Done()如果Canceled，返回 "Canceled"如果超過了 Deadline，返回 "DeadlineEsceeded"
• Value(key) 返回上下文中 key 對應的 value 值

Context 構造

為了使用 Context，我們需要了解 Context 是怎麼構造的。

Context 提供了兩個方法做初始化：

<code>func Background() Context{}
func TODO() Context {}/<code>

上面方法均會返回空的 Context，但是 Background 一般是所有 Context 的基礎，所有 Context 的源頭都應該是它。TODO 方法一般用於當傳入的方法不確定是哪種類型的 Context 時，為了避免 Context 的參數為nil而初始化的 Context。

其他的 Context 都是基於已經構造好的 Context 來實現的。一個 Context 可以派生多個子 context。基於 Context 派生新Context 的方法如下：

<code>func WithCancel(parent Context) (ctx Context, cancel CancelFunc){}
func WithDeadline(parent Context, d time.Time) (Context, CancelFunc) {}
func WithTimeout(parent Context, timeout time.Duration) (Context, CancelFunc) {}/<code>

上面三種方法比較類似，均會基於 parent Context 生成一個子 ctx，以及一個 Cancel 方法。如果調用了cancel 方法，ctx 以及基於 ctx 構造的子 context 都會被取消。不同點在於 WithCancel 必需要手動調用 cancel 方法，WithDeadline

可以設置一個時間點，WithTimeout 是設置調用的持續時間，到指定時間後，會調用 cancel 做取消操作。

除了上面的構造方式，還有一類是用來創建傳遞 traceId， token 等重要數據的 Context。

<code>func WithValue(parent Context, key, val interface{}) Context {}/<code>

withValue 會構造一個新的context，新的context 會包含一對 Key-Value 數據，可以通過Context.Value(Key) 獲取存在 ctx 中的 Value 值。

通過上面的理解可以直到，Context 是一個樹狀結構，一個 Context 可以派生出多個不一樣的Context。我們大概可以畫一個如下的樹狀圖：

一個background，衍生出一個帶有traceId的valueCtx，然後valueCtx衍生出一個帶有cancelCtx

的context。最終在一些db查詢，http查詢，rpc沙遜等異步調用中體現。如果出現超時，直接把這些異步調用取消，減少消耗的資源，我們也可以在調用時，通過Value 方法拿到traceId，並記錄下對應請求的數據。

當然，除了上面的幾種 Context 外，我們也可以基於上述的 Context 接口實現新的Context.

使用方法

下面我們舉幾個例子，學習上面講到的方法。

超時查詢的例子

在做數據庫查詢時，需要對數據的查詢做超時控制，例如：

<code>ctx = context.WithTimeout(context.Background(), time.Second)
rows, err := pool.QueryContext(ctx, "select * from products where id = ?", 100)/<code>

上面的代碼基於 Background 派生出一個帶有超時取消功能的ctx，傳入帶有context查詢的方法中，如果超過1s未返回結果，則取消本次的查詢。使用起來非常方便。為了瞭解查詢內部是如何做到超時取消的，我們看看DB內部是如何使用傳入的ctx的。

在查詢時，需要先從pool中獲取一個db的鏈接，代碼大概如下：

<code>// src/database/sql/sql.go
// func (db *DB) conn(ctx context.Context, strategy connReuseStrategy) *driverConn, error)

// 阻塞從req中獲取鏈接，如果超時，直接返回
select {
case /<code>

req 也是一個chan，是等待鏈接返回的chan，如果Done() 返回的chan 關閉後，則不再關心req的返回了，我們的查詢就超時了。

在做SQL Prepare、SQL Query 等操作時，也會有類似方法：

<code>select {
default:
// 校驗是否已經超時，如果超時直接返回
case /<code>

上面在做查詢時，首先判斷是否已經超時了，如果超時，則直接返回錯誤，否則才進行查詢。

可以看出，在派生出的帶有超時取消功能的 Context 時，內部方法在做異步操作（比如獲取鏈接，查詢等）時會先查看是否已經

Done了，如果Done，說明請求已超時，直接返回錯誤；否則繼續等待，或者做下一步工作。這裡也可以看出，要做到超時控制，需要不斷判斷 Done() 是否已關閉。

鏈路追蹤的例子

在做鏈路追蹤時，Context 也是非常重要的。（所謂鏈路追蹤，是說可以追蹤某一個請求所依賴的模塊，比如db，redis，rpc下游，接口下游等服務，從這些依賴服務中找到請求中的時間消耗）

下面舉一個鏈路追蹤的例子：

<code>// 建議把key 類型不導出，防止被覆蓋
type traceIdKey struct{}{}

// 定義固定的Key
var TraceIdKey = traceIdKey{}

func ServeHTTP(w http.ResponseWriter, req *http.Request){
  // 首先從請求中拿到traceId
  // 可以把traceId 放在header裡，也可以放在body中
  // 還可以自己建立一個 （如果自己是請求源頭的話）
  traceId := getTraceIdFromRequest(req)

  // Key 存入 ctx 中
  ctx := context.WithValue(req.Context(), TraceIdKey, traceId)

  // 設置接口1s 超時
  ctx = context.WithTimeout(ctx, time.Second)

  // query RPC 時可以攜帶 traceId
  repResp := RequestRPC(ctx, ...)

  // query DB 時可以攜帶 traceId
  dbResp := RequestDB(ctx, ...)

  // ...
}

func RequestRPC(ctx context.Context, ...) interface{} {
    // 獲取traceid，在調用rpc時記錄日誌
    traceId, _ := ctx.Value(TraceIdKey)
    // request

    // do log
    return
}/<code>

上述代碼中，當拿到請求後，我們通過req 獲取traceId， 並記錄在ctx中，在調用RPC，DB等時，傳入我們構造的ctx，在後續代碼中，我們可以通過ctx拿到我們存入的traceId，使用traceId 記錄請求的日誌，方便後續做問題定位。

當然，一般情況下，context 不會單純的僅僅是用於 traceId 的記錄，或者超時的控制。很有可能二者兼有之。

如何實現

知其然也需知其所以然。想要充分利用好 Context，我們還需要學習 Context 的實現。下面我們一起學習不同的 Context 是如何實現 Context 接口的，

空上下文

Background(), Empty() 均會返回一個空的 Context emptyCtx。emptyCtx 對象在方法 Deadline(), Done(), Err(), Value(interface{}) 中均會返回nil，String() 方法會返回對應的字符串。這個實現比較簡單，我們這裡暫時不討論。

有取消功能的上下文

WithCancel 構造的context 是一個cancelCtx實例，代碼如下。

<code>type cancelCtx struct {
  Context

  // 互斥鎖，保證context協程安全
  mu       sync.Mutex
  // cancel 的時候，close 這個chan
  done     chan struct{}
  // 派生的context
  children map[canceler]struct{}
  err      error
}/<code>

WithCancel 方法首先會基於 parent 構建一個新的 Context，代碼如下：

<code>func WithCancel(parent Context) (ctx Context, cancel CancelFunc) {
  c := newCancelCtx(parent)  // 新的上下文
  propagateCancel(parent, &c) // 掛到parent 上
  return &c, func() { c.cancel(true, Canceled) }
}/<code>

其中，propagateCancel 方法會判斷 parent 是否已經取消，如果取消，則直接調用方法取消；如果沒有取消，會在parent的children 追加一個child。這裡就可以看出，context 樹狀結構的實現。 下面是propateCancel 的實現：

<code>// 把child 掛在到parent 下
func propagateCancel(parent Context, child canceler) {
  // 如果parent 為空，則直接返回
  if parent.Done() == nil {
    return // parent is never canceled
  }
  
  // 獲取parent類型
  if p, ok := parentCancelCtx(parent); ok {
    p.mu.Lock()
    if p.err != nil {
      // parent has already been canceled
      child.cancel(false, p.err)
    } else {
      if p.children == nil {
        p.children = make(map[canceler]struct{})
      }
      p.children[child] = struct{}{}
    }
    p.mu.Unlock()
  } else {
    // 啟動goroutine，等待parent/child Done
    go func() {
      select {
      case /<code>

Done() 實現比較簡單，就是返回一個chan，等待chan 關閉。可以看出 Done 操作是在調用時才會構造 chan done，done 變量是延時初始化的。

<code>func (c *cancelCtx) Done() /<code>

在手動取消 Context 時，會調用 cancelCtx 的 cancel 方法，代碼如下：

<code>func (c *cancelCtx) cancel(removeFromParent bool, err error) {
  // 一些判斷,關閉 ctx.done chan
  // ...
  if c.done == nil {
    c.done = closedchan
  } else {
    close(c.done)
  }

  // 廣播到所有的child，需要cancel goroutine 了
  for child := range c.children {
    // NOTE: acquiring the child's lock while holding parent's lock.
    child.cancel(false, err)
  }
  c.children = nil
  c.mu.Unlock()

  // 然後從父context 中，刪除當前的context
  if removeFromParent {
    removeChild(c.Context, c)
  }
}/<code>

這裡可以看到，當執行cancel時，除了會關閉當前的cancel外，還做了兩件事，① 所有的child 都調用cancel方法，② 由於該上下文已經關閉，需要從父上下文中移除當前的上下文。

定時取消功能的上下文

WithDeadline, WithTimeout 提供了實現定時功能的 Context 方法，返回一個timerCtx結構體。WithDeadline 是給定了執行截至時間，WithTimeout 是倒計時時間，WithTImeout 是基於WithDeadline實現的，因此我們僅看其中的WithDeadline

即可。WithDeadline 內部實現是基於cancelCtx 的。相對於 cancelCtx 增加了一個計時器，並記錄了 Deadline 時間點。下面是timerCtx 結構體：

<code>type timerCtx struct {
  cancelCtx
  // 計時器
  timer *time.Timer
  // 截止時間
  deadline time.Time
}/<code>

WithDeadline 的實現：

<code>func WithDeadline(parent Context, d time.Time) (Context, CancelFunc) {
  // 若父上下文結束時間早於child，
  // 則child直接掛載在parent上下文下即可
  if cur, ok := parent.Deadline(); ok && cur.Before(d) {
    return WithCancel(parent)
  }

  // 創建個timerCtx, 設置deadline
  c := &timerCtx{
    cancelCtx: newCancelCtx(parent),
    deadline:  d,
  }

  // 將context掛在parent 之下
  propagateCancel(parent, c)

  // 計算倒計時時間
  dur := time.Until(d)
  if dur <= 0 {
    c.cancel(true, DeadlineExceeded) // deadline has already passed
    return c, func() { c.cancel(false, Canceled) }
  }
  c.mu.Lock()
  defer c.mu.Unlock()
  if c.err == nil {
    // 設定一個計時器，到時調用cancel
    c.timer = time.AfterFunc(dur, func() {
      c.cancel(true, DeadlineExceeded)
    })
  }
  return c, func() { c.cancel(true, Canceled) }
}/<code>

構造方法中，將新的context 掛在到parent下，並創建了倒計時器定期觸發cancel。

timerCtx 的cancel 操作，和cancelCtx 的cancel 操作是非常類似的。在cancelCtx 的基礎上，做了關閉定時器的操作

<code>func (c *timerCtx) cancel(removeFromParent bool, err error) {
  // 調用cancelCtx 的cancel 方法 關閉chan，並通知子context。
  c.cancelCtx.cancel(false, err)
  // 從parent 中移除
  if removeFromParent {
    removeChild(c.cancelCtx.Context, c)
  }
  c.mu.Lock()
  // 關掉定時器
  if c.timer != nil {
    c.timer.Stop()
    c.timer = nil
  }
  c.mu.Unlock()
}/<code>

timeCtx 的 Done 操作直接複用了cancelCtx 的 Done 操作，直接關閉 chan done 成員。

傳遞值的上下文

WithValue 構造的上下文與上面幾種有區別，其構造的context 原型如下：

<code>type valueCtx struct {
  // 保留了父節點的context
  Context
  key, val interface{}
}/<code>

每個context 包含了一個Key-Value組合。valueCtx 保留了父節點的Context，但沒有像cancelCtx 一樣保留子節點的Context. 下面是valueCtx的構造方法：

<code>func WithValue(parent Context, key, val interface{}) Context {
  if key == nil {
    panic("nil key")
  }
  // key 必須是課比較的，不然無法獲取Value
  if !reflect.TypeOf(key).Comparable() {
    panic("key is not comparable")
  }
  return &valueCtx{parent, key, val}
}/<code>

直接將Key-Value賦值給struct 即可完成構造。下面是獲取Value 的方法：

<code>func (c *valueCtx) Value(key interface{}) interface{} {
  if c.key == key {
    return c.val
  }
  // 從父context 中獲取
  return c.Context.Value(key)
}/<code>

Value 的獲取是採用鏈式獲取的方法。如果當前 Context 中找不到，則從父Context中獲取。如果我們希望一個context 多放幾條數據時，可以保存一個map 數據到 context 中。這裡不建議多次構造context來存放數據。畢竟取數據的成本也是比較高的。

注意事項

最後，在使用中應該注意如下幾點：

• context.Background 用在請求進來的時候，所有其他context 來源於它。
• 在傳入的conttext 不確定使用的是那種類型的時候，傳入TODO context （不應該傳入一個nil 的context)
• context.Value 不應該傳入可選的參數，應該是每個請求都一定會自帶的一些數據。（比如說traceId，授權token 之類的）。在Value 使用時，建議把Key 定義為全局const 變量，並且key 的類型不可導出，防止數據存在衝突。
• context goroutines 安全。