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 是一個接口，用來描述一個程序的上下文。接口中提供了四個抽象的方法，定義以下：

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

• 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 提供了兩個方法作初始化：

func Background() Context{}
func TODO() Context {}

上面方法均會返回空的 Context，可是 Background 通常是全部 Context 的基礎，全部 Context 的源頭都應該是它。TODO 方法通常用於當傳入的方法不肯定是哪一種類型的 Context 時，爲了不 Context 的參數爲nil而初始化的 Context。

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

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) {}

上面三種方法比較相似，均會基於 parent Context 生成一個子 ctx，以及一個 Cancel 方法。若是調用了cancel 方法，ctx 以及基於 ctx 構造的子 context 都會被取消。不一樣點在於 WithCancel 必須要手動調用 cancel 方法，WithDeadline
能夠設置一個時間點，WithTimeout 是設置調用的持續時間，到指定時間後，會調用 cancel 作取消操做。

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

func WithValue(parent Context, key, val interface{}) Context {}

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.

使用方法

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

超時查詢的例子

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

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

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

在查詢時，須要先從pool中獲取一個db的連接，代碼大概以下：

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

// 阻塞從req中獲取連接，若是超時，直接返回
select {
case <-ctx.Done():
  // 獲取連接超時了，直接返回錯誤
  // do something
  return nil, ctx.Err()
case ret, ok := <-req:
  // 拿到連接，校驗並返回
  return ret.conn, ret.err
}

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

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

select {
default:
// 校驗是否已經超時，若是超時直接返回
case <-ctx.Done():
  return nil, ctx.Err()
}
// 若是尚未超時，調用驅動作查詢
return queryer.Query(query, dargs)

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

能夠看出，在派生出的帶有超時取消功能的 Context 時，內部方法在作異步操做（好比獲取連接，查詢等）時會先查看是否已經
Done了，若是Done，說明請求已超時，直接返回錯誤；不然繼續等待，或者作下一步工做。這裏也能夠看出，要作到超時控制，須要不斷判斷 Done() 是否已關閉。

鏈路追蹤的例子

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

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

// 建議把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
}

上述代碼中，當拿到請求後，咱們經過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實例，代碼以下。

type cancelCtx struct {
  Context

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

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

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

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

// 把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 <-parent.Done():
        child.cancel(false, parent.Err())
      case <-child.Done():
      }
    }()
  }
}

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

func (c *cancelCtx) Done() <-chan struct{} {
  c.mu.Lock()
  if c.done == nil {
    c.done = make(chan struct{})
  }
  d := c.done
  c.mu.Unlock()
  return d
}

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

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)
  }
}

這裏能夠看到，當執行cancel時，除了會關閉當前的cancel外，還作了兩件事，① 全部的child 都調用cancel方法，② 因爲該上下文已經關閉，須要從父上下文中移除當前的上下文。

定時取消功能的上下文

WithDeadline, WithTimeout 提供了實現定時功能的 Context 方法，返回一個timerCtx結構體。WithDeadline 是給定了執行截至時間，WithTimeout 是倒計時時間，WithTImeout 是基於WithDeadline實現的，所以咱們僅看其中的WithDeadline
便可。WithDeadline 內部實現是基於cancelCtx 的。相對於 cancelCtx 增長了一個計時器，並記錄了 Deadline 時間點。下面是timerCtx 結構體：

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

WithDeadline 的實現：

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) }
}

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

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

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()
}

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

傳遞值的上下文

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

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

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

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}
}

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

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

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 安全。