Context 源碼，再度重相逢

各位讀者朋友們大家好，我是隨波逐流的薯條。深秋了，前幾天氣溫驟降，北京的人和狗都不願意出門，趴在窩裏凍的打寒顫。我的書房裏沒裝空調，暖氣要十一月中旬纔來，每次想學習都得下很大的決心，所以這篇文章發出來時比預期又晚了幾天~

最近我的心也是冰冰的，我目前做在線數據開發，如果大家幹過這活肯定知道，數據開發最重要的是數據口徑，開發前一定得對清楚... 本人作爲在這上面踩了很多坑的人這幾天接需求時又掉進去了，一個需求在已經上線的情況下，不同來源的數據做 diff 總是對不上，一查就是口徑不對，來來回回改了根據口徑改了一遍邏輯，搞得我 tm 的真想和提供口徑的人打一架，md。

有點扯遠了，言歸正傳，這篇文章接着上次的  Context 這三個應用場景，你知嗎 繼續看看 context 源碼，讀者可能覺得【Context 源碼，再度重相逢】這個標題比較奇怪。起這個題目是因爲在下 讀 context 源碼時找了一些資料，最好的中文資料應是【碼農桃花源】qcrao 在 19 年寫過的一篇關於 context 解析的文章，所以我在猶豫要不要寫我的這篇，說實話代碼都看完了不寫出來吹吹牛逼總覺得有點虧。好在 rao 老闆分析 context 源碼基於的 Go 版本是 1.9.2，如今 Go 已經 1.17 了，context 的源碼也有很多更新，於是不才就來一篇基於 1.17.2 的 context 源碼分析，不多說了，發車！

源碼分析

context 的核心作用是存儲鍵值對和取消機制。存儲鍵值對比較簡單，取消機制比較複雜，先來看一下 Context 抽象出來的接口：

type Context interface {
  // 如果是timerCtx或者自定義的ctx實現了此方法，返回截止時間和true，否則返回false
 Deadline() (deadline time.Time, ok bool)
  // 這裏監聽取消信號
  Done() <-chan struct{}
  // ctx取消時，返回對應錯誤，有context canceled和context deadline exceeded
  Err() error 
  // 返回key的val
 Value(key interface{}) interface{}
}

ctx 存儲鍵值對

鍵值對 ctx 比較簡單，先來看一下它的邏輯：要新建一個存儲鍵值對的 ctx，需要調用WithValue，它返回一個valueCtx地址對象。valueCtx 結構體內部很簡單，有個 Context 接口和 k-v 對：

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

valueCtx實現了Value方法，邏輯也很簡單：

func (c *valueCtx) Value(key interface{}) interface{} {
  // key相同則返回key
  if c.key == key {
  return c.val
 }
  // 否則從父節點中調用Value方法繼續尋找key
 return c.Context.Value(key)
}

寫一段代碼看一下從valueCtx中查找某個 key 的過程：

func main() {
 ctx := context.Background()

 ctx1 := context.WithValue(ctx, "name", "uutc")
 ctx2 := context.WithValue(ctx1, "age", "18")
 ctx3 := context.WithValue(ctx2, "traceID", "89asd7yu9asghd")

 fmt.Println(ctx3.Value("name"))
}

valueCtx是個鏈表模型，當我們從 ctx3 中查找 name 這個 key, 最終要走到 ctx1 中才能返回對應的 value，如圖所示：

雖然鏈表的查找效率是 O(n) 的，但貌似一個請求裏面也不會有 1000 個 ctx，個人認爲 ctx 鏈表的查找效率可以接受，讀者有興趣可以給 go 團隊提個 pr，把鏈表改成紅黑樹試試，嘿嘿~

ctx 的取消機制

context 的取消機制我個人認爲可以分成兩種：第一種是普通取消，需要取消 ctx 時調用 cancel 函數。第二個是根據時間取消，用戶可以定義一個過期 time 或一個 deadline，到這個時間時自動取消。

普通取消

現在假裝沒看源碼，聯想一下如果我們自己實現。該如何寫取消。 建立 ctx 時, 是在 parent 的基礎上 copy 一份，然後添加自己的屬性, 不同協程可能持有不同的 ctx, 若想在請求層面做協程取消，就需要廣播機制，比如在下圖中：

img

若我們要取消 ctx2，應分爲向上取消和向下取消兩部分，向下取消要把當前節點的子節點都幹掉，在這裏需要 tx4、ctx5 都取消掉；而向上取消需要把他在父節點中刪除，在這裏需要把自己 (ctx2) 從父節點 ctx 的子節點列表中刪除；

取消這個動作本身並沒有神奇的地方。ctx 創建一個 channel，然後協程通過 select 去監聽這個 channel，沒數據時處於阻塞狀態，當調用取消函數，函數內部執行 close(chan) 操作, select 監聽到關閉信號執行 return，達到取消協程的目的，寫個 demo：

func main() {
 done := make(chan struct{})

 go func() {
  close(done)
 }()

 select {
 case <-done:
  println("exit!")
  return
 }
}

下面來看 go 源碼是怎麼實現的取消，首先抽象出了一個canceler接口，這個接口裏最重要的就是 cancel 方法，調用這個方法可以發送取消信號，有兩個結構體實現了這個接口，分別是 *cancelCtx(普通取消) 和 *timerCtx(時間取消)。

type canceler interface {
 cancel(removeFromParent bool, err error)
 Done() <-chan struct{}
}

cancelCtx 對應前文說的普通取消機制，它是 context 取消機制的基石，也是源碼中比較難理解的地方，先來看一下它的結構體：

type cancelCtx struct {
 Context

 mu       sync.Mutex            // context號稱併發安全的基石
 done     atomic.Value          // 用於接收ctx的取消信號，這個數據的類型做過優化，之前是 chan struct 類型
 children map[canceler]struct{} // 儲存此節點的實現取消接口的子節點，在根節點取消時，遍歷它給子節點發送取消信息
 err      error                 // 調用取消函數時會賦值這個變量
}

若我們要生成一個可取消的 ctx，需要調用 WithCancel 函數，這個函數的內部邏輯是：

func WithCancel(parent Context) (ctx Context, cancel CancelFunc) {
 if parent == nil {
  panic("cannot create context from nil parent")
 }
 c := newCancelCtx(parent)    // 基於父節點，new一個CancelCtx對象
 propagateCancel(parent, &c)  // 掛載c到parent上
 return &c, func() { c.cancel(true, Canceled) } // 返回子ctx，和返回函數
}

這裏邏輯比較重的地方是 propagateCancel 函數和 cancel 方法，propagateCancel 函數主要工作是把子節點掛載父節點上，下面來看看它的源碼：

func propagateCancel(parent Context, child canceler) {
 done := parent.Done()
  // 判斷父節點的done是否爲nil,若爲nil則爲不可取消的ctx, 直接返回
 if done == nil {
  return
 }
  // 看能否從done裏面讀到數據，若能說明父節點已取消，取消子節點，返回即可，不能的話繼續流轉到後續邏輯
 select {
 case <-done:
  child.cancel(false, parent.Err())
  return
 default:
 }

  // 調用parentCancelCtx函數，看是否能找到ctx上層最接近的可取消的父節點
 if p, ok := parentCancelCtx(parent); ok {
    //這裏是可以找到的情況
  p.mu.Lock()
  // 父節點有err，說明已經取消，直接取消子節點
    if p.err != nil {
   child.cancel(false, p.err)
  } else {
      // 把本節點掛載到父節點的children map中
   if p.children == nil {
    p.children = make(map[canceler]struct{})
   }
   p.children[child] = struct{}{}
  }
  p.mu.Unlock()
 } else {
    // 若沒有可取消的父節點掛載
  atomic.AddInt32(&goroutines, +1)
  // 新起一個協程
    go func() {
   select {
      // 監聽到父節點取消時，取消子節點
   case <-parent.Done():
    child.cancel(false, parent.Err())
      // 監聽到子節點取消時，什麼都不做，退出協程
   case <-child.Done():
   }
  }()
 }
}

我看這段源碼時產生了兩個問題：

  1. 函數內部的 parentCancelCtx 函數中有個 else 分支，什麼條件下會走到這裏

  2. 調用 cancel 方法需要傳遞 bool 值，何時傳 true，何時傳 false

經過一番研究，大概解決了這倆問題，下面依次做解答。

什麼條件下會走到 parentCancelCtx 函數的 else 分支

首先看下 parentCancelCtx 函數的邏輯。parentCancelCtx 函數用來查找 ctx 最近的一個可取消的父節點，這個函數也經過了優化，原代碼是：

func parentCancelCtx(parent Context) (*cancelCtx, bool) {
 for {
  switch c := parent.(type) {
  case *cancelCtx:
   return c, true
  case *timerCtx:
   return &c.cancelCtx, true
  case *valueCtx:
   parent = c.Context
  default:
   return nil, false
  }
 }
}

這段代碼比較簡單，起了一個 for 循環，遇到*cancelCtx和*timerCtx類型就返回，遇到*valueCtx則繼續向上查找 parent，直到找到了節點或者找不到爲止。

最新版本的代碼放棄粗暴的使用 for{} 循環尋找父節點，而是用 parent.Value 方法查到父節點是否能找到特定的 key，由於 Value 是遞歸的，所以這裏表面上看不出來循環的足跡：

func parentCancelCtx(parent Context) (*cancelCtx, bool) {
 done := parent.Done()
 if done == closedchan || done == nil {
  return nil, false
 }
 // Value是遞歸向上查找，直到找到有*cancelCtxKey 的ctx爲止
 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
}

知道了這個前提，我們繼續研究**什麼條件下會走到 parentCancelCtx 函數的 else 分支。**我自己實現了一個 Context，代碼如下

type ContextCancel struct {
 context.Context
}

func (*ContextCancel) Done() <-chan struct{} {
 ch := make(chan struct{}, 100)
 return ch
}

當調用這段代碼時，即可走到 else 分支，寫個 demo：

func main() {
 ctx := context.Background()
 ctx1, _ := context.WithCancel(ctx)
 ctx2 := context.WithValue(ctx1, "hello", "world")
 ctx3 := ContextCancel{ctx2}
 ctx4, _ := context.WithCancel(&ctx3)  // 這裏可以走到else分支

 println(ctx4)
}

與源碼中 CancelCtx 不同的是，我這裏的 Done 方法只是簡單返回，並沒有把 done 的值存到 Context 中去。所以在執行 parentCancelCtx 時，這裏會判斷失敗，返回 false：

pdone, _ := p.done.Load().(chan struct{})
if pdone != done {
 return nil, false
}

通過parent.Value(&cancelCtxKey).(*cancelCtx)雖然找到了 cancelCtx，但是在 load Done 方法值的時候卻鎩羽而歸，parentCancelCtx 這裏判斷失敗，最終返回 nil 和 false，最終走到 else 分支。所以這個 else 分支主要是預防用戶自己實現了一個定製的 Ctx 中，隨意提供了一個 Done chan 的情況的，由於找不到可取消的父節點的，只好新起一個協程做監聽。

調用 cancel 方法需要傳遞 bool 值，何時傳 true，何時傳 false

要明白這個問題，先來看一下 * cancelCtx 類型的 cancel 方法實現：

func (c *cancelCtx) cancel(removeFromParent bool, err error) {
 // 取消時必須傳入err，不然panic
  if err == nil {
  panic("context: internal error: missing cancel error")
 }
 c.mu.Lock()
  // 已經出錯了，說明已取消，直接返回
 if c.err != nil {
  c.mu.Unlock()
  return
 }
  // 用戶傳進來的err賦給c.err
 c.err = err
 d, _ := c.done.Load().(chan struct{})
 if d == nil {
    // 這裏其實和關閉chan差不多，因爲後續會用closedchan作判斷
  c.done.Store(closedchan)
 } else {
    // 關閉chan
  close(d)
 }
  // 這裏是向下取消，依次取消此節點所有的子節點
 for child := range c.children {
  child.cancel(false, err)
 }
  // 清空子節點
 c.children = nil
 c.mu.Unlock()
  // 這裏是向上取消，取消此節點和父節點的聯繫
 if removeFromParent {
  removeChild(c.Context, c)
 }
}

removeChild 函數的邏輯比較簡單，核心就是調用 delete 方法，在父節點的子節點中清空自己。

func removeChild(parent Context, child canceler) {
   p, ok := parentCancelCtx(parent)
   if !ok {
      return
   }
   p.mu.Lock()
   if p.children != nil {
      delete(p.children, child) // 這裏只是刪除一個
   }
   p.mu.Unlock()
}

看完這倆函數的邏輯後，這個問題也可以回答。當父節點調用 cancel 函數時傳遞 true， 其他情況傳遞 false。

true 用來向上刪除，核心邏輯是調用removeChild函數里面的的：

if p.children != nil {
  delete(p.children, child) // 這裏只是刪除一個
}

而 false 就是用來非向上刪除，只需要執行完 cancel 方法這段代碼即可：

for child := range c.children {
  child.cancel(false, err) // 這裏把子節點都幹掉
}

看到這裏，ctx 的普通取消機制基本差不多了，下面來看一下基於時間的取消機制。

時間取消

時間取消 ctx 可傳入兩種時間，第一種是傳入超時時間戳；第二種是傳入 ctx 持續時間，比如 2s 之後 ctx 取消，持續時間在實現上是在 time.Now 的基礎上加了個 timeout 湊個超時時間戳，本質上都是調用的WithDeadline函數。

WithDeadline 函數內部 new 了一個timerCtx，先來看一下這個結構體的內容：

type timerCtx struct {
 cancelCtx
 timer *time.Timer  // 一個統一的計時器，後續通過 time.AfterFunc 使用
 deadline time.Time // 過期時間戳
}

可以看到 timerCtx 內嵌了 cancelCtx，實際上在超時取消這件事上，timerCtx 更多負責的是超時相關的邏輯，而取消主要調用的 cancelCtx 的 cancel 方法。先來看一下WithDeadline函數的邏輯，看如何返回一個時間 Ctx：

func WithDeadline(parent Context, d time.Time) (Context, CancelFunc) {
 // 若父節點爲nil，panic
 if parent == nil {
  panic("cannot create context from nil parent")
 }
 // 如果parent有超時時間、且過期時間早於參數d，那parent取消時，child 一定需要取消，直接通過WithCancel走起
 if cur, ok := parent.Deadline(); ok && cur.Before(d) {
  // The current deadline is already sooner than the new one.
  return WithCancel(parent)
 }
 // 構造一個timerCtx, 主要傳入一個過期時間
 c := &timerCtx{
  cancelCtx: newCancelCtx(parent),
  deadline:  d,
 }
 // 把這個節點掛載到父節點上
 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 {
  // 否則等到過期時間時，執行取消操作
  c.timer = time.AfterFunc(dur, func() {
   c.cancel(true, DeadlineExceeded)
  })
 }
 // 返回一個ctx和一個取消函數
 return c, func() { c.cancel(true, Canceled) }
}

看完源碼可以知道，除了基於時間的取消，當調用 CancelFunc 時，也能取消超時 ctx。再來看一下*timerCtx的 cancel 方法的源碼：

func (c *timerCtx) cancel(removeFromParent bool, err error) {
 // 調用cancel的cancel取消掉它下游的ctx
 c.cancelCtx.cancel(false, err)
 // 取消掉它上游的ctx的連接
 if removeFromParent {
  removeChild(c.cancelCtx.Context, c)
 }
 // 把timer停掉
 c.mu.Lock()
 if c.timer != nil {
  c.timer.Stop()
  c.timer = nil
 }
 c.mu.Unlock()
}

至此，context 源碼分析的差不多了，關於這塊還有個挺常見的問題，Context.TODO 和 Context.Backend 的區別。從代碼上看，他倆沒啥區別，都只是emptyCtx的對象而已，emptyCtx源碼很簡單，這裏不再贅述。

源碼賞析

寫這篇文章時，我在想看源碼的好處什麼。個人認爲有兩點，第一可以從源碼角度看到一個概念的全部細節，第二個是可以學習大牛寫代碼的思路。實際上 context 的代碼也有個迭代過程，下面列舉一些閱讀源碼時學習到的點：

if >= 2 用 switch 替換

來自String() string方法。

atomic.Value 替換 chan struct 減少鎖使用

來自cancelCtx源碼，用 atomic.Value 類型替換了 chan struct{}。

加鎖前，先獲取值避免加鎖

這點在 go 源碼中隨處可見，簡單列舉幾處：

看源碼時就感覺這個 select 有點突兀，一看果然爲優化效率後加的~

String 邏輯賞析

之前的源碼只是粗暴的使用 Sprintf 函數，後來自己搞了個 stringer 接口，用反射去打印 Context。

一個 Bug

截止日期已經過了，cancel 已經執行過了，沒必要在返回取消函數中再從父 ctx 中取消自己了. 感覺 removeFromParent 有點沒抽象好，這不，作者自己都掉坑裏去了。

總結

個人感覺 context 代碼挺值得一看的：struct 裏面嵌套 interface，struct 並不對外暴露，而是提供多個Withxxx方法新建對象；interface 對外暴露，用戶可以根據需要構建自己的 Context；timerCtx struct 裏面嵌套 CancelCtx struct 以此來複用 cancel 的邏輯等等。關於這塊讀者有什麼心得可以在留言區評論，我會挑選三個最有價值的評論，每個評論發 20 塊紅包給你加個餐~

最後，給自己打個廣告

歡迎加入 隨波逐流的薯條 微信羣。

薯條目前有草帽羣、木葉羣、琦玉羣，羣交流內容不限於技術、投資、趣聞分享等話題。歡迎感興趣的同學入羣交流。

入羣請加薯條的個人微信：709834997。並備註：加入薯條微信羣。

歡迎關注我的公衆號~

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

猜你喜歡