深度解析go context實現原理及其源碼

時間 2021-02-27

原文原文鏈接

Context 基本使用方法

首先，咱們來看一下 Context 接口包含哪些方法，這些方法都是幹什麼用的。微信

包 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 方法時，都會返回和第一次調用相同的結果。ui

Done 方法返回一個 Channel 對象。在 Context 被取消時，此 Channel 會被 close，若是沒被取消，可能會返回 nil。後續的 Done 調用老是返回相同的結果。當 Done 被 close 的時候，你能夠經過 ctx.Err 獲取錯誤信息。Done 這個方法名其實起得並很差，由於名字太過籠統，不能明確反映 Done 被 close 的緣由，由於 cancel、timeout、deadline 均可能致使 Done 被 close，不過，目前尚未一個更合適的方法名稱。atom

關於 Done 方法，你必需要記住的知識點就是：若是 Done 沒有被 close，Err 方法返回 nil；若是 Done 被 close，Err 方法會返回 Done 被 close 的緣由。spa

Context使用場景

上下文信息傳遞（request-scoped），好比處理 http 請求、在請求處理鏈路上傳遞信息
控制子 goroutine 的運行
超時控制的方法調用
能夠取消的方法調用

valueCtx

valueCtx 是基於 parent Context 生成一個新的 Context，保存了一個key-value鍵值對。它主要用來傳遞上下文信息。線程

使用示例

ctx := context.Background()
ctx = context.WithValue(ctx, "key1", "0001")
ctx = context.WithValue(ctx, "key2", "0001")
ctx = context.WithValue(ctx, "key3", "0001")
ctx = context.WithValue(ctx, "key4", "0004")
fmt.Println(ctx.Value("key1")) // 0001

查找過程如圖所示：3d

結構體

type valueCtx struct {
   Context  // parent Context
   key, val interface{}  // key-value
}

func (c *valueCtx) Value(key interface{}) interface{} {
   // 若key值 等於 當前valueCtx存儲的key值 
   // 則取出其value並返回
   if c.key == key {
      return c.val
   }
   // 不然遞歸調用valueCtx中Value方法，獲取其parent Context中存儲的key-value
   return c.Context.Value(key)
}

經過觀察 valueCtx 結構體，它利用一個 Context 變量表示其父節點的 context ，這樣 valueCtx 也繼承了父節點的全部信息；而且它持有一個 key-value 鍵值對，說明它還能夠攜帶額外的信息。它還覆蓋了 Value 方法，優先從本身的存儲中檢查這個 key，不存在的話會從 parent 中繼續檢查。

WithValue

WithValue 就是向 context 中添加鍵值對：

func WithValue(parent Context, key, val interface{}) 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}
}

經過代碼能夠看出，向 context 中添加鍵值對並非在原 context 基礎上添加的，而是新建一個 valueCtx 子節點，將原 context 做爲父節點。以此類推，就會造成一個 context 鏈。在查找過程當中，若是當前 valueCtx 不存在key值，還會向 parent Context 去查找，若是 parent 仍是 valueCtx 的話，仍是遵循相同的原則：valueCtx 會嵌入 parent，因此仍是會查找 parent 的 Value 方法的。

cancleCtx

在咱們開發過程當中，咱們經常會遇到一些場景，須要主動取消長時間的任務或者停止任務，這個時候就可使用cancelCtx。經過調用cancel函數就可停止goroutine，進而去釋放所佔用的資源。

須要注意的是，不是隻有中途停止任務時才調用cancel函數，只要任務執行完畢後，就須要調用 cancel，這樣，這個 Context 才能釋放它的資源（通知它的 children 處理 cancel，從它的 parent 中把本身移除，甚至釋放相關的 goroutine）。

使用示例

func main() {
  // gen 在單獨的 goroutine 中生成整數 而後將它們發送到返回的管道
  gen := func(ctx context.Context) <-chan int {
     dst := make(chan int)
     n := 1
     go func() {
        for {
           select {
           case <-ctx.Done():
              return // returning not to leak the goroutine
           case dst <- n:
              n++
           }
        }
     }()
     return dst
  }
  ctx, cancel := context.WithCancel(context.Background())
  // 代碼完畢後調用cancel函數釋放goroutine所佔用的資源
  defer cancel() // cancel when we are finished consuming integers
  // 遍歷循環獲取管道中的值
  for n := range gen(ctx) {
     fmt.Println(n)
     if n == 5 {
        break
     }
  }
}

建立一個 gen函數，在gen函數中建立一個goroutine，專門用來生成整數，而後將他們發送到返回的管道。經過 context.WithCancel 建立可取消的 context ，最後遍歷循環獲取管道中值，當n的值爲5時，退出循環，結束進程。最後調用cancel函數釋放goroutine所佔用的資源。

結構體

type cancelCtx struct {
    Context
    mu       sync.Mutex            
    done     chan struct{}         
    children map[canceler]struct{} 
    err      error                 
}

cancelCtx和valueCtx相似，結構體中都有一個Context做爲其父節點；變量done表示關閉信號傳遞；變量children表示當前節點所擁有的子節點，err用於存儲錯誤信息表示任務結束的緣由。

接下來，看看cancelCtx實現的方法：

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

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
}

func (c *cancelCtx) cancel(removeFromParent bool, err error) {
   if err == nil {
      panic("context: internal error: missing cancel error")
   }
   c.mu.Lock()
   if c.err != nil {
      c.mu.Unlock()
      return // already canceled
   }
   c.err = err
   // 設置一個關閉的channel或者將done channel關閉，用以發送關閉信號
   if c.done == nil {
      c.done = closedchan
   } else {
      close(c.done)
   }
   // 遍歷循環將字節點context取消
   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()
   if removeFromParent {
      // 將當前context節點從父節點上移除
      removeChild(c.Context, c)
   }
}

cancelCtx結構體實現Done和cancel方法，Done方法實現了將done初始化。cancel方法用於將當前節點從父節點上移除以及移除當前節點下的全部子節點。

cancelCtx 被取消時，它的 Err 字段就是下面這個 Canceled 錯誤：

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

WithCancel

WithCancel函數用來建立一個可取消的context，即cancelCtx類型的context。

WithCancel函數返回值有兩個，一個爲parent 的副本Context，另外一個爲觸發取消操做的CancelFunc。

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

// newCancelCtx returns an initialized cancelCtx.
func newCancelCtx(parent Context) cancelCtx {
   // 將parent做爲父節點context生成一個新的子節點
   return cancelCtx{Context: parent}
}

func propagateCancel(parent Context, child canceler) {
   done := parent.Done()
   if done == nil {
      return // parent is never canceled
   }
   
   select {
   case <-done:
      // parent is already canceled
      child.cancel(false, parent.Err())
      return
   default:
   }
   
   // 獲取最近的類型爲cancelCtx的祖先節點
   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{})
         }
         // 將當前子節點加入最近cancelCtx祖先節點的children中
         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():
         }
      }()
   }
}

調用 WithCancel函數時，首先會調用 newCancelCtx函數建立一個以parent做爲父節點的context。而後調用propagateCancel函數，用來創建當前context節點與parent節點之間的關係。

在propagateCancel函數中，若是parent節點爲nil，說明parent以上的路徑沒有可取消的cancelCtx，則不須要處理。

不然經過parentCancelCtx函數過去當前節點最近的類型爲cancelCtx的祖先節點，首先須要判斷該祖先節點是否被取消，若已被取消就取消當前節點；不然將當前節點加入祖先節點的children列表中。

不然的話，則須要新起一個 goroutine，由它來監聽 parent 的 Done 是否已關閉。一旦parent.Done()返回的channel關閉，即context鏈中某個祖先節點context被取消，則將當前context也取消。

WithTimeout

WithTimeout 實際上是和 WithDeadline 同樣，只不過一個參數是超時時間，一個參數是截止時間。超時時間加上當前時間，其實就是截止時間，所以，WithTimeout 的實現是：

func WithTimeout(parent Context, timeout time.Duration) (Context, CancelFunc) { 
  // 當前時間+timeout就是deadline
  return WithDeadline(parent, time.Now().Add(timeout))
}

WithDeadline

WithDeadline 會返回一個 parent 的副本，而且設置了一個不晚於參數 d 的截止時間，類型爲 timerCtx（或者是 cancelCtx）。

若是它的截止時間晚於 parent 的截止時間，那麼就以 parent 的截止時間爲準，並返回一個類型爲 cancelCtx 的 Context，由於 parent 的截止時間到了，就會取消這個 cancelCtx。

若是當前時間已經超過了截止時間，就直接返回一個已經被 cancel 的 timerCtx。不然就會啓動一個定時器，到截止時間取消這個 timerCtx。

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

截止時間到了
cancel 函數被調用
parent 的 Done 被 close

使用示例

func main() {
  d := time.Now().Add(time.Second * 3)
  ctx, cancel := context.WithDeadline(context.Background(), d)
  defer cancel()
  select {
  case <-time.After(3 * time.Second):
     fmt.Println("overslept")
  case <-ctx.Done():
     fmt.Println(ctx.Err())
  }
}

WithDeadline

func WithDeadline(parent Context, d time.Time) (Context, CancelFunc) {
   if parent == nil {
      panic("cannot create context from nil parent")
   }
   // 若是parent的截止時間更早，直接返回一個cancelCtx便可
   if cur, ok := parent.Deadline(); ok && cur.Before(d) {
      return WithCancel(parent)
   }
   c := &timerCtx{
      cancelCtx: newCancelCtx(parent),
      deadline:  d,
   }
   // 創建新建context與可取消context祖先節點的取消關聯關係
   propagateCancel(parent, c)
   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的截止時間是否早於當前timerCtx，若爲true的話，直接返回一個cancelCtx便可。不然須要調用propagateCancel函數建議新建context與可取消context祖先節點的取消關聯關係，創建關聯關係以後，若當前時間已經超過截止時間後，直接cancel。不然的話，需設置一個定時器，到截止時間後取消。

總結

context主要用於父子任務之間的同步取消信號，本質上是一種協程調度的方式。另外在使用context時有兩點值得注意：上游任務僅僅使用context通知下游任務再也不須要，但不會直接干涉和中斷下游任務的執行，由下游任務自行決定後續的處理操做，也就是說context的取消操做是無侵入的；context是線程安全的，由於context自己是不可變的（immutable），所以能夠放心地在多個協程中傳遞使用。

到這裏，Context 的源碼已解讀完畢，但願對您有收穫，我們下期再見。