Go 併發控制：context 源碼解讀

context 是 Go 語言的特色設計之一，主要作用有兩個：控制鏈路和安全傳值，並且 context 是併發安全的。context 在 Go 1.17 版本被引入，經過數年的迭代，在設計和用法上已經趨於穩定，本文以最新的 Go 1.23.0 版本源碼爲基礎，帶你深入理解 context 的設計和實現。

context 設計

context 被設計爲一個接口，名爲Context。爲了支持不同特性，這個接口有多種結構體實現。而每個結構體又提供了一個或多個 exported 函數（大寫字母開頭的公開函數）作爲構造函數來實例化 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 any) any
}

可以看到Context 只有 4 個方法，可謂大道至簡。

其中Canceled 和DeadlineExceeded 兩個錯誤定義如下：

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

var DeadlineExceeded error = deadlineExceededError{}

type deadlineExceededError struct{}

func (deadlineExceededError) Error() string   { return "context deadline exceeded" }
func (deadlineExceededError) Timeout() bool   { return true }
func (deadlineExceededError) Temporary() bool { return true }

這裏採用了典型的 Sentinel error 用法。並且從deadlineExceededError 實現的方法來看，其鼓勵行爲斷言而非類型斷言。

context 實現

接下來我們對Context 接口的具體實現進行逐一講解。

emptyCtx

emptyCtx 是最基礎的 context 實現，定義如下：

type emptyCtx struct{}

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
}

它確實 “基礎”，也確實 “empty”，所有實現都爲空，沒有代碼邏輯，僅是一個 context 架子。

backgroundCtx 和 todoCtx

backgroundCtx 定義如下：

type backgroundCtx struct{ emptyCtx }

func (backgroundCtx) String() string {
 return "context.Background"
}

它內嵌了emptyCtx，也僅比emptyCtx 多實現了一個String() 方法。

todoCtx 實現同理：

type todoCtx struct{ emptyCtx }

func (todoCtx) String() string {
 return "context.TODO"
}
Background() 和 TODO()

我們在使用 context 時，往往會用context.Background() 或context.TODO() 來定義了最頂層 context，這兩個方法實現如下：

func Background() Context {
 return backgroundCtx{}
}

func TODO() Context {
 return todoCtx{}
}

沒錯，最常用的 context 代碼實現就是這麼簡單，它們是整個 context 鏈路的基礎。

cancelCtx

cancelCtx 結構體定義如下：

type cancelCtx struct {
 Context // “繼承”的父 Context

 mu       sync.Mutex            // 持有鎖保護下面這些字段
 done     atomic.Value          // 值爲 chan struct{} 類型，會被懶惰創建，在第一次調用取消函數 cancel() 時被關閉，表示 Context 已被取消
 children map[canceler]struct{} // 所有可以被取消的子 Context 集合，它們在第一次調用取消函數 cancel() 時被級聯取消，然後置爲 nil
 err      error                 // 取消原因，在第一次調用取消函數 cancel() 時被設置值
 cause    error                 // 取消根因，在第一次調用取消函數 cancel() 時被設置值
}

cancelCtx 直接內嵌了Context 接口，也就是說，它支持任意其他類型的 context 實現作爲父上下文（parent context）。

前文說過，context 是併發安全的，所以cancelCtx 內部持有一把互斥鎖，保證安全的操作結構體屬性。

done 屬性爲atomic.Value 類型，是爲了支持原子操作，使用它可以減少互斥鎖的使用頻率，稍後你將在Done() 方法中看到。它的值是chan struct{} 類型。

children 屬性是一個集合，記錄了當前 context 的所有子上下文（child context）。這樣，父子 context 就產生了鏈路關係，以此爲基礎實現父 context 取消時，級聯的取消所有子 context。

err 和cause 分別記錄了 context 被取消的原因和根因，err 是 context 包內部產生的，cause 則是我們在使用WithXxxCause() 方法構造 context 對象時傳入的。

這裏涉及的canceler 定義如下：

type canceler interface {
 cancel(removeFromParent bool, err, cause error) // 取消函數
 Done() <-chan struct{}                          // 通過返回的 channel 能夠知道是否被取消
}

它是一個接口，表示一個可以被取消的對象。也就是說，在 context 包中設計的支持取消的 context 類型都需要提供這兩個方法。父 context 取消時會調用子 context 的cancel() 方法進行級聯取消；並且有取消功能的 context 必須要實現Done() 方法，這樣使用者才能通過監聽 done channel 知道這個 context 是否被取消。

cancelCtx 的Done() 方法實現如下：

func (c *cancelCtx) Done() <-chan struct{} {
 // 使用 double-check 來提升性能
 d := c.done.Load() // 原子操作，比互斥鎖更加輕量
 if d != nil {      // 如果存在 channel 直接返回
  return d.(chan struct{})
 }
 c.mu.Lock() // 如果不存在 channel，則要先加鎖，然後創建 channel 並返回
 defer c.mu.Unlock()
 d = c.done.Load()
 if d == nil { // 爲保證併發安全，再做一次檢查
  d = make(chan struct{})
  c.done.Store(d)
 }
 return d.(chan struct{})
}

這裏使用了double-check 來提升程序的性能，這也是done 屬性爲什麼被設計成atomic.Value 類型的原因。首先使用c.done.Load() 來判斷標識 context 是否取消的chan struct{} 是否存在，存在則直接返回，不存在纔會加鎖創建。

cancelCtx 的cancel() 方法實現如下：

func (c *cancelCtx) cancel(removeFromParent bool, err, cause error) {
 if err == nil {
  panic("context: internal error: missing cancel error")
 }
 if cause == nil { // 如果沒有設置根因，取 err
  cause = err
 }
 c.mu.Lock()
 if c.err != nil { // 如果 err 不爲空，說明已經被取消，直接返回
  c.mu.Unlock()
  return
 }

 // NOTE: 只有第一次調用 cancel 纔會執行之後的代碼

 // 記錄錯誤和根因
 c.err = err
 c.cause = cause
 d, _ := c.done.Load().(chan struct{})
 if d == nil { // 如果 done 爲空，直接設置一個已關閉的 channel
  c.done.Store(closedchan)
 } else { // 如果 done 有值，將其關閉
  close(d)
 }
 // 級聯取消所有子 Context
 for child := range c.children {
  // NOTE: 獲取子 Context 的鎖，同時持有父 Context 的鎖
  child.cancel(false, err, cause)
 }
 c.children = nil // 清空子 Context 集合，因爲已經完成了 Context 樹整個鏈路的取消操作
 c.mu.Unlock()

 if removeFromParent { // 從父 Context 的 children 集合中移除當前 Context
  removeChild(c.Context, c)
 }
}

這個方法用來取消cancelCtx，它接收 3 個參數，removeFromParent 表示是否要從父 context 的children 屬性集合中移除當前的cancelCtx；err 和cause 則分別表示取消的錯誤原因和根因。

在第 9 行，因爲使用了c.err != nil 來判斷err 是否爲空，如果不爲空，說明 context 已經被取消，直接返回。所以，多次調用cancel() 方法效果相同。

當第一次調用cancel() 方法時會記錄err 和cause。接着判斷 done channel 是否存在，不存在就直接設置爲一個已經關閉的 channel 對象closedchan；如果存在則調用close(d) 將其關閉。

接着，會遍歷c.children 屬性對當前cancelCtx 的所有子 context 進行級聯取消，即依次調用它們的cancel() 方法。然後清空children 集合。

最終根據參數removeFromParent 的值決定是否要從父 context 的children 屬性集合中移除cancelCtx。

這裏涉及的closedchan 定義如下：

// closedchan 表示一個已關閉的 channel
var closedchan = make(chan struct{})

// 導入 context 包時直接關閉 closedchan
func init() {
 close(closedchan)
}

在 context 包被導入時就直接關閉了。

removeChild() 函數的具體實現如下：

func removeChild(parent Context, child canceler) {
 if s, ok := parent.(stopCtx); ok {
  s.stop()
  return
 }
 p, ok := parentCancelCtx(parent)
 if !ok {
  return
 }
 p.mu.Lock()
 if p.children != nil {
  delete(p.children, child)
 }
 p.mu.Unlock()
}

首先判斷父 context 是否爲stopCtx 類型，如果是，則調用其s.stop() 方法。關於stopCtx 類型暫時不必深究，後文中講解*cancelCtx.propagateCancel() 方法時我會更詳細的解釋。

接着調用parentCancelCtx() 函數向上查找父 context 或其鏈路中是否存在*cancelCtx 對象，如果不存在，直接返回；如果存在，從其children 屬性集合中移除當前 context。

parentCancelCtx() 函數實現如下：

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
}

如果父 context 的Done() 方法返回closedchan，說明已經被取消了；如果返回nil，則說明父 context 永遠不會被取消。這兩種情況，都不必繼續向上查找*cancelCtx 對象了，直接返回false 表示未找到。

接下來使用&cancelCtxKey 作爲key 從父 context 中查找value，並且斷言查找到的對象是否爲*cancelCtx 類型，如果!ok 說明未找到，返回false；否則，說明找到的*cancelCtx。

然後對*cancelCtx 進行進一步的檢查，確保返回的*cancelCtx 的 done channel 與父 context 的 done channel 是匹配的，如果不匹配，說明*cancelCtx 已經被包裝在一個自定義實現中，爲了避免影響自定義 context 實現，這種情況下返回false 表示未找到；如果匹配，才返回*cancelCtx 對象和true 表示找到了。

cancelCtx 還實現了Context 接口的Value() 和Err() 方法：

func (c *cancelCtx) Value(key any) any {
 // 使用 &cancelCtxKey 標記需要返回自身
 // 這是一個未導出的（unexported）類型，所以僅作爲 context 包內部實現的一個“協議”，對用戶不可見
 if key == &cancelCtxKey {
  return c
 }
 // 接着向上遍歷父 Context 鏈路，查詢 key
 return value(c.Context, key)
}

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

Err() 方法的實現非常簡單，沒什麼好說的。

cancelCtx 實現了Value() 方法，這是爲了實現一個特殊的 “內部協議”。這個方法裏有一個特殊的判斷if key == &cancelCtxKey，如果成立，則不去查找給定key 所對應的value；如果不成立才調用value() 函數繼續進行查找。

cancelCtxKey 就是一個普通的變量：

var cancelCtxKey int

上面介紹的parentCancelCtx() 函數中，之所以能夠使用parent.Value(&cancelCtxKey).(*cancelCtx) 獲取到*cancelCtx 對象，就是通過在Value() 方法中這個特殊的 “協議” 來實現的。

從Value() 方法的實現來看，只要調用*cancelCtx.Value() 方法時傳入&cancelCtxKey 作爲查找的key，就返回*cancelCtx 對象本身。

注意：&cancelCtxKey 是一個unexported類型的指針變量，所以外部無法使用，只作爲 “內部協議”。

這個設計有點奇技淫巧的意思，不過卻很有用。

這裏涉及的value() 函數我們暫且不繼續深究，後文再來講解。

此外，cancelCtx 也實現了自己的String() 方法：

type stringer interface {
 String() string
}

func contextName(c Context) string {
 if s, ok := c.(stringer); ok {
  return s.String()
 }
 return reflectlite.TypeOf(c).String()
}

func (c *cancelCtx) String() string {
 return contextName(c.Context) + ".WithCancel"
}
WithCancel() 和 WithCancelCause()

看完了cancelCtx 的實現，接下來看下我們如何構造一個cancelCtx。

context 包提供了兩種構造cancelCtx 的方法，分別是WithCancel() 和WithCancelCause()。

WithCancel() 實現如下：

type CancelFunc func()

func WithCancel(parent Context) (ctx Context, cancel CancelFunc) {
 c := withCancel(parent)
 return c, func() { c.cancel(true, Canceled, nil) }
}

WithCancel() 根據給定的父 context 構造一個新的具有取消功能的cancelCtx 並返回，其核心邏輯是代理給withCancel() 函數去實現的。

WithCancelCause() 實現如下：

type CancelCauseFunc func(cause error)

func WithCancelCause(parent Context) (ctx Context, cancel CancelCauseFunc) {
 c := withCancel(parent)
 return c, func(cause error) { c.cancel(true, Canceled, cause) }
}

WithCancelCause() 與WithCancel() 類似，但返回CancelCauseFunc 而不是CancelFunc。可以發現二者的唯一區別就是返回的函數是否支持設置 context 被取消的根因cause。

那麼接下來就看看withCancel() 函數是如何實現的：

func withCancel(parent Context) *cancelCtx {
 if parent == nil {
  panic("cannot create context from nil parent")
 }
 c := &cancelCtx{}            // 帶取消功能的 Context
 c.propagateCancel(parent, c) // 將新構造的 Context 向上傳播掛載到父 Context 的 children 屬性中，這樣當父 Context 取消時子 Context 對象 c 也會級聯取消
 return c
}

這個函數邏輯並不多，這裏構造了一個cancelCtx 並返回，核心邏輯都交給了propagateCancel() 方法。

propagateCancel() 方法實現如下：

func (c *cancelCtx) propagateCancel(parent Context, child canceler) {
 c.Context = parent // “繼承”父 Context，這裏可以是任何實現了 Context 接口的類型

 // NOTE: 父 Context 沒有實現取消功能
 done := parent.Done()
 if done == nil { // 如果父 Context 的 Done() 方法返回 nil，說明父 Context 沒有取消的功能，那麼無需傳播子 Context 的 cancel 功能到父 Context
  return
 }

 // NOTE: 父 Context 已經被取消
 select {
 case <-done: // 直接取消子 Context，且取消原因設置爲父 Context 的取消原因
  child.cancel(false, parent.Err(), Cause(parent))
  return
 default:
 }

 // NOTE: 父 Context 還未取消
 if p, ok := parentCancelCtx(parent); ok { // 如果父 Context 是 *cancelCtx 或者從 *cancelCtx 派生而來
  p.mu.Lock()
  if p.err != nil {
   // 如果父 Context 的 err 屬性有值，說明已經被取消，直接取消子 Context
   child.cancel(false, p.err, p.cause)
  } else {
   if p.children == nil { // 延遲創建父 Context 的 children 屬性
    p.children = make(map[canceler]struct{})
   }
   p.children[child] = struct{}{} // 將 child 加入到這個 *cancelCtx 的 children 集合中
  }
  p.mu.Unlock()
  return
 }

 // NOTE: 父 Context 實現了 afterFuncer 接口
 if a, ok := parent.(afterFuncer); ok { // 測試文件 afterfunc_test.go 中 *afterFuncCtx 實現了 afterFuncer 接口
  c.mu.Lock()
  stop := a.AfterFunc(func() { // 註冊子 Context 取消功能到父 Context，當父 Context 取消時，能級聯取消子 Context
   child.cancel(false, parent.Err(), Cause(parent))
  })
  c.Context = stopCtx{ // 將當前 *cancelCtx 的直接父 Context 設置爲 stopCtx
   Context: parent, // stopCtx 的父 Context 設置爲 parent
   stop:    stop,
  }
  c.mu.Unlock()
  return
 }

 // NOTE: 父 Context 不是已知類型，但實現了取消功能
 goroutines.Add(1) // 記錄下開啓了幾個 goroutine，用於測試代碼
 go func() {       // 開起一個 goroutine，監聽父 Context 是否被取消，如果取消則級聯取消子 Context
  select {
  case <-parent.Done(): // 父 Context 被取消
   child.cancel(false, parent.Err(), Cause(parent))
  case <-child.Done(): // 自己被取消
  }
 }()
}

propagateCancel() 方法將cancelCtx 對象向上傳播掛載到父 context 的children 屬性集合中，這樣當父 context 被取消時，子 context 也會被級聯取消。這個方法邏輯稍微有點多，也是 context 包中最複雜的方法了，拿下它，後面的代碼就都很簡單了。

首先將parent 參數記錄到cancelCtx.Context 屬性中，作爲父 context。接下來會對父 context 做各種判斷，以此來決定如何處理子 context。

第 5 行通過parent.Done() 拿到父 context 的 done channel，如果值爲nil，則說明父 context 沒有取消功能，所以不必傳播子 context 的取消功能到父 context。

第 11 行使用select...case... 來監聽<-done 是否被關閉，如果已關閉，則說明父 context 已經被取消，那麼直接調用child.cancel() 取消子 context。因爲 context 的取消功能是從上到下級聯取消，所以父 context 被取消，那麼子 context 也一定要取消。

如果父 context 尚未取消，則在第 19 行判斷父 context 是否爲*cancelCtx 或者從*cancelCtx 派生而來。如果是，則判斷父 context 的err 屬性是否有值，有值則說明父 context 已經被取消，那麼直接取消子 context；否則將子 context 加入到這個*cancelCtx 類型的父 context 的children 屬性集合中。

如果父 context 不是*cancelCtx 類型，在第 35 行判斷父 context 是否實現了afterFuncer 接口。如果實現了，則新建一個stopCtx 作爲當前*cancelCtx 的父 context。

最終，如果之前對父 context 的判斷都不成立，則開啓一個新的 goroutine 來監聽父 context 和子 context 的取消信號。如果父 context 被取消，則級聯取消子 context；如果子 context 被取消，則直接退出 goroutine。

至此propagateCancel() 方法的主要邏輯就梳理完了。

不過，在當前的 context 包實現中，其實在第 35 行判斷父 context 是否實現了afterFuncer 接口的 case 永遠不會發生。afterFuncer 接口定義如下：

type afterFuncer interface {
 AfterFunc(func()) func() bool
}

在 Go 1.23.0 版本 context 包的源碼中，並沒有一個 context 實現了afterFuncer 接口。所以stopCtx 也並沒有被真正使用。所以我纔在前文講解removeChild() 函數時說stopCtx 類型不必深究。

不過我們還是簡單看一下stopCtx 的定義：

type stopCtx struct {
 Context
 stop func() bool
}

它同樣嵌入了Context 接口，stop 方法用於註銷AfterFunc。

NOTE:

其實 afterFuncer 接口在 context/afterfunc_test.go 文件中有一個 afterFuncContext 類型是實現了的，只不過是測試代碼，所以我們還是無法使用。

我在 issues/61672 中找到了一些關於 afterFuncer 的討論，在我看來這是一個爲了填早期設計的坑而定義的，如果能重來，大概率 Context 不會被設計成接口，而是結構體。

此外，這裏還用到了Cause() 函數從parent 中提取根因，Cause() 函數實現如下：

func Cause(c Context) error {
 if cc, ok := c.Value(&cancelCtxKey).(*cancelCtx); ok {
  cc.mu.Lock()
  defer cc.mu.Unlock()
  return cc.cause
 }
 return c.Err()
}

這裏同樣使用特殊的key``&cancelCtxKey 來查找 context 鏈路中的*cancelCtx，如果找到，則返回*cancelCtx.cause，否則將 context 的錯誤原因作爲根因。

針對cancelCtx 類型的源碼講解就到這裏，可以說cancelCtx 是最複雜的 context 實現了，後文中要講解的timerCtx 和afterFuncCtx 都是基於它實現的。

timerCtx

timerCtx 結構體定義如下：

type timerCtx struct {
 cancelCtx             // “繼承”了 cancelCtx
 timer     *time.Timer // Under cancelCtx.mu.

 deadline time.Time
}

timerCtx 內部嵌入了cancelCtx 以 “繼承”Done() 和Err() 方法。並且它還關聯了一個定時器timer 和截止時間deadline，以此來實現在截止時間到期時，自動取消 context。

timerCtx 實現的方法如下：

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

func (c *timerCtx) String() string {
 return contextName(c.cancelCtx.Context) + ".WithDeadline(" +
  c.deadline.String() + " [" +
  time.Until(c.deadline).String() + "])"
}

func (c *timerCtx) cancel(removeFromParent bool, err, cause error) {
 c.cancelCtx.cancel(false, err, cause)
 if removeFromParent {
  // 將此 *timerCtx 從其父 *cancelCtx 的 children 集合中刪除
  removeChild(c.cancelCtx.Context, c)
 }
 c.mu.Lock()
 if c.timer != nil {
  c.timer.Stop()
  c.timer = nil
 }
 c.mu.Unlock()
}

你是否還記得我們在講解Context 接口時提到，Deadline() 方法返回的ok 值爲false 時說明 context 沒有設置截止時間。這裏返回true 則說明timerCtx 支持設置截止時間。

timerCtx 也實現了自己的String() 方法。其實所有 context 實現都有自己的String() 方法。

timerCtx 的並沒有直接使用cancelCtx 的取消方法，而是自己也實現了cancel() 方法。內部調用的removeChild() 函數我們在前文講解cancelCtx 時已經見過了。這裏唯一需要注意的一點是，如果timer 屬性不爲nil 則調用timer.Stop() 將其停止，並將屬性值置爲nil，以此讓timer 對象儘早被 GC 回收。

WithDeadline() 和 WithTimeoutCause()

我們先來看timerCtx 的第一個構造函數WithDeadline()：

func WithDeadline(parent Context, d time.Time) (Context, CancelFunc) {
 return WithDeadlineCause(parent, d, nil)
}

WithDeadline() 直接將邏輯代理給了WithDeadlineCause() 來處理，WithDeadlineCause() 實現如下：

func WithDeadlineCause(parent Context, d time.Time, cause error) (Context, CancelFunc) {
 if parent == nil {
  panic("cannot create context from nil parent")
 }
 // 如果父 Context 的截止時間已經比傳入的 d 更早，直接返回一個 *cancelCtx（無需構造 *timerCtx 等待定時器判斷截止時間到了才取消 Context）
 if cur, ok := parent.Deadline(); ok && cur.Before(d) {
  return WithCancel(parent)
 }
 c := &timerCtx{ // 構造一個帶有定時器和截止時間功能的 Context
  deadline: d,
 }
 // 這裏使用 cancelCtx 結構體默認值，初始化 timerCtx 時沒有顯式初始化 cancelCtx 字段
 c.cancelCtx.propagateCancel(parent, c) // 向父 Context 傳播 cancel 功能，這樣當父 Context 取消時當前 Context 也會被級聯取消
 dur := time.Until(d)
 if dur <= 0 { // 截止日期已過，直接取消
  c.cancel(true, DeadlineExceeded, cause)
  return c, func() { c.cancel(false, Canceled, nil) }
 }
 c.mu.Lock()
 defer c.mu.Unlock()
 if c.err == nil {
  c.timer = time.AfterFunc(dur, func() { // 等待截止時間到期，自動調用 cancel 取消 Context
   c.cancel(true, DeadlineExceeded, cause)
  })
 }
 return c, func() { c.cancel(true, Canceled, nil) }
}

可以發現WithDeadline(parent, d) 等價於WithDeadlineCause(parent, d, nil)。

WithDeadlineCause() 實現代碼不多，首先對parent 是否爲nil 做了檢查。接着檢查父 context 的截止時間是否比傳入的d 更早，如果是，則直接創建一個*cancelCtx 並返回，無需創建*timerCtx。這是因爲 context 具有級聯取消的能力，既然父 context 的截止時間更早，則父 context 一定先於子 context 取消，所以子 context 會被級聯取消，這就沒必要再大費周章的構造*timerCtx 來定時取消子 context 了。

如果上述條件不成立，則構造一個帶有定時器和截止時間功能的*timerCtx。並且，同樣需要調用cancelCtx.propagateCancel() 向上傳播取消功能。

接着判斷是否已到截止時間，如果到了，則直接取消 context。否則使用time.AfterFunc() 來實現延遲取消 context。

WithTimeout() 和 WithTimeoutCause()

WithTimeout() 和WithTimeoutCause() 兩個方法同樣用於構造timerCtx，其實現如下：

func WithTimeout(parent Context, timeout time.Duration) (Context, CancelFunc) {
 return WithDeadline(parent, time.Now().Add(timeout))
}

func WithTimeoutCause(parent Context, timeout time.Duration, cause error) (Context, CancelFunc) {
 return WithDeadlineCause(parent, time.Now().Add(timeout), cause)
}

WithTimeout() 內部調用了WithDeadline()，而WithTimeoutCause() 內部則調用了WithDeadlineCause()。

WithDeadline() 和WithDeadlineCause() 接收一個絕對時間d time.Time。

WithTimeout() 和WithTimeoutCause() 接收一個相對時間timeout time.Duration，並在內部將其轉換爲絕對時間。

所以timerCtx 類型的構造函數有 4 個。

withoutCancelCtx

withoutCancelCtx 故名思義，是沒有取消功能的 context，它可以打斷 context 控制鏈路中級聯取消的能力。

withoutCancelCtx 結構體定義非常簡單，只有一個屬性c 用來保存父 context：

type withoutCancelCtx struct {
 c Context
}

withoutCancelCtx 實現方法如下：

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

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

func (withoutCancelCtx) Err() error {
 return nil
}

func (c withoutCancelCtx) Value(key any) any {
 return value(c, key)
}

func (c withoutCancelCtx) String() string {
 return contextName(c.c) + ".WithoutCancel"
}

withoutCancelCtx 雖然沒有取消功能，但實現了Value 方法，可以根據key 查詢value。這樣才能保證整個 context 鏈路中傳值的能力不被中斷。

WithoutCancel()

不僅withoutCancelCtx 結構體設計簡單，它的構造函數WithoutCancel() 同樣非常簡單：

func WithoutCancel(parent Context) Context {
 if parent == nil {
  panic("cannot create context from nil parent")
 }
 return withoutCancelCtx{parent}
}

這裏只對父 context 是否爲nil 做了檢查，然後就直接返回實例化的withoutCancelCtx 對象了。

valueCtx

我們前面介紹的 context 從設計上來說都是爲了實現控制鏈路的，與其他 context 不同，valueCtx 用於實現在 context 鏈路中進行安全傳值。

valueCtx 實現如下：

type valueCtx struct {
 Context
 key, val any // 存儲的鍵值對，注意一個 Context 僅能保存一對 key/value，這樣就能實現併發讀的安全，copy-on-write
}

func (c *valueCtx) Value(key any) any {
 if c.key == key { // 在自己的鍵值對中查找
  return c.val
 }
 return value(c.Context, key) // 沿着父 Context 向上查找
}

valueCtx 結構體內部嵌入了Context 接口，這樣可以直接複用父 context 實現的方法。key 和value 字段則用於存儲鍵值對。可以發現，一個valueCtx 對象只能存儲一對key/value。

在用戶調用Value() 方法查找給定key 關聯的value 時，首先判斷是否在當前 context 中，如果不在，則交給value() 函數來處理。

在介紹*cancelCtx.Value() 方法時，我們並沒有深入講解value() 函數，那麼現在是時候看下value() 函數是如何實現的了：

func value(c Context, key any) any {
 for {
  switch ctx := c.(type) { // 斷言 Context 類型
  case *valueCtx: // 表示一個用於安全傳遞數據的 Context
   if key == ctx.key { // 與當前 Context 的 key 匹配，直接返回對應的值 val
    return ctx.val
   }
   c = ctx.Context // key 不匹配，繼續向上遍歷父 Context
  case *cancelCtx: // 表示一個帶有取消功能的 Context
   if key == &cancelCtxKey { // 檢查 key 是否等於 &cancelCtxKey（這是一個指向 *cancelCtx 的特殊鍵），如果匹配，就返回自身（即 c 對象）
    return c
   }
   c = ctx.Context // key 不匹配，繼續向上遍歷父 Context
  case withoutCancelCtx: // 表示一個不帶取消功能的 Context（使用 WithoutCancel() 創建出來的 Context 類型）
   if key == &cancelCtxKey { // 檢查 key 是否等於 &cancelCtxKey，如果匹配，說明要查找的是取消信號的特殊鍵，就返回 nil，因爲這種 Context 沒有取消信號
    return nil
   }
   c = ctx.c // 如果 key 不匹配，則繼續向上遍歷父 Context
  case *timerCtx: // 表示一個帶有定時器的 Context
   if key == &cancelCtxKey { // 檢查 key 是否等於 &cancelCtxKey，如果匹配，返回其包裝的 *cancelCtx
    return &ctx.cancelCtx
   }
   c = ctx.Context // key 不匹配，繼續向上遍歷父 Context
  case backgroundCtx, todoCtx: // 這兩個類型是無值的 Context（通常這是 Context 樹的根），所以直接返回 nil
   return nil
  default: // 如果沒有匹配任何已知的 Context 類型，則調用 Context 的 Value 方法去查找 key 對應的值
   return c.Value(key)
  }
 }
}

這裏代碼看似複雜，實際上邏輯非常簡單。啓用一個for 無限循環，沿着傳進來的 context 對象c 的父路徑，循環查找匹配的key，直到找到目標value 或走到鏈路根節點返回nil。

for 循環中，首先會斷言當前 context 對象c 的類型，如果是*valueCtx，判斷key 是否匹配，匹配則直接返回ctx.val，不匹配則將父 context 取出賦值給c，進行下一輪循環；如果是*cancelCtx 或*timerCtx，判斷key 是否匹配&cancelCtxKey 這個特殊值，匹配則根據我們前文講過的 “內部協議” 返回當前*cancelCtx，否則將父 context 取出賦值給c，進行下一輪循環；如果是withoutCancelCtx，當key 匹配&cancelCtxKey 時返回nil，因爲這個 context 的實現不支持取消功能，key 不匹配同樣將父 context 取出賦值給c，進行下一輪循環；如果是backgroundCtx 或todoCtx，則說明已經遍歷到 context 鏈路的頂點，所以直接返回nil，表示未查找到；如果所有已知類型都沒匹配，則調用其Value() 方法繼續查找。

所以，從源碼中我們能夠看出，context 根據給定的key 查找value 時，是自下而上查找的。

此外，valueCtx 同樣實現了自己的String() 方法：

func stringify(v any) string {
 switch s := v.(type) {
 case stringer: // 實現了 String() 方法，就返回 String() 內容
  return s.String()
 case string: // 字符串類型就返回字符串內容
  return s
 case nil: // nil 返回字符串格式
  return "<nil>"
 }
 // 其他類型會返回對象類型名的字符串格式，而不是對象值的字符串形式
 return reflectlite.TypeOf(v).String()
}

// 代碼示例：context.WithValue(context.Background(), "a", 1)
// 輸出示例：context.Background.WithValue(a, int)
func (c *valueCtx) String() string {
 // 取父 Context 的 string 形式 + .WithValue(k, v)
 return contextName(c.Context) + ".WithValue(" +
  stringify(c.key) + ", " +
  stringify(c.val) + ")"
}
WithValue()

valueCtx 構造函數WithValue() 實現如下：

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

這裏對parent 和key 都做了檢查，注意key 一定是可比較類型。

可以發現，valueCtx 並沒有使用互斥鎖，這是因爲每次新增key/value 時，都會新建一個新的valueCtx，並將parent 賦值給valueCtx。這種 copy-on-write 的思想，保證絕不修改現有的 context 對象，那麼程序中併發讀取值時就不會產生 data race，同時也能保證併發安全。

afterFuncCtx

我們最後還未介紹的 context 就僅剩一個afterFuncCtx 類型了，其實現如下：

type afterFuncCtx struct {
 cancelCtx           // “繼承”了 cancelCtx
 once      sync.Once // 要麼用來開始執行 f，要麼用來阻止 f 被執行
 f         func()
}

與timerCtx 一樣，afterFuncCtx 內部也嵌入了cancelCtx。此外它還有兩個屬性once 和f，once 保證一個操作僅執行一次，要麼用來開始執行f，要麼用來阻止f 被執行，f 是一個延遲函數，在構造函數AfterFunc() 中被傳入賦值。

afterFuncCtx 實現了自己的cancel() 方法：

func (a *afterFuncCtx) cancel(removeFromParent bool, err, cause error) {
 a.cancelCtx.cancel(false, err, cause) // 取消 cancelCtx
 if removeFromParent {
  removeChild(a.Context, a) // 將當前 *afterFuncCtx 從 cancelCtx 的父 Context 的 children 屬性中移除
 }
 a.once.Do(func() { // 確保僅執行一次
  go a.f() // 開啓新的 goroutine 執行 f，如果在調用 a.cancel() 之前 stop 函數被調用，stop 函數中的 a.once.Do 優先被執行，則此處就不會執行
 })
}

afterFuncCtx 在取消時，首先會取消父cancelCtx。然後根據參數removeFromParent 決定是否從父 context 的children 屬性中移除。最後使用once.Do() 確保f 函數僅執行一次。

AfterFunc()

afterFuncCtx 構造函數AfterFunc() 實現如下：

func AfterFunc(ctx Context, f func()) (stop func() bool) {
 a := &afterFuncCtx{
  f: f,
 }
 // 調用 cancelCtx 的向上傳播方法，將 a 的取消功能掛載到父 ctx 的 children 屬性中，實現級聯取消
 a.cancelCtx.propagateCancel(ctx, a)
 return func() bool { // 返回一個停止函數，用於阻止 f 被執行
  stopped := false
  a.once.Do(func() { // 確保僅執行一次
   stopped = true // 如果此處被執行，則 a.cancel 方法內部的 a.once.Do 就不會重複執行，即阻止 f 被執行
  })
  if stopped { // 第一次調用，取消 Context
   a.cancel(true, Canceled, nil)
  }
  return stopped
 }
}

與其他 context 構造函數不同，AfterFunc() 並不會返回構造的afterFuncCtx 對象，而是返回一個stop() 函數。其實AfterFunc() 的功能是爲 context 註冊一個延遲函數，當 context 被取消時，開啓新的 goroutine 異步執行f()。而stop() 函數的作用則是用來阻止f() 被執行。

因爲stop() 函數和cancel() 方法內部使用的a.once.Do() 是同一個，所以二者只能有一個會被執行。可以總結stop() 函數和cancel() 方法執行邏輯如下：

至此，context 包的源碼就全部解讀完成了。

總結

context 包在 Go 1.7 版本被引入，核心功能是控制鏈路和安全傳值，且併發安全。

context 被設計爲一個Context 接口和多個實現了此接口的結構體。一切 context 鏈路都會從一個空的emptyCtx 開始，由context.Background() 或context.TODO() 來定義了最頂層 context，接着使用WithXxx() 方法在原有的 context 基礎上附加新的功能，形成 context 鏈路。

context 鏈路最終可能發展成一個樹形結構，不過你要清楚，控制鏈路是從上到下的，父 context 取消，則會及聯的取消所有帶有取消功能的子孫 context；但通過給定key 查找value 則是自下而上的，而這就會導致從不同的起點出發，查找 context 中相同key 對應的value 可能不同。

我畫了一張 context 樹形結構圖：

在這幅圖中，從控制鏈路的角度出發，如果我們取消 context 3️⃣，則 context 7️⃣ 會被級聯取消，因爲 6️⃣ 不支持取消，控制鏈路會被打斷，所以 9️⃣ 不會被取消；如果取消 context 7️⃣，則 context 3️⃣ 不會被取消，因爲控制鏈路是從上到下的。

從安全傳值的角度出發，根據給定key 查找value，假如 context 2️⃣ 中存儲的是key: value2，context 8️⃣ 中存儲的是key: value8，那麼從 context 2️⃣ 4️⃣ 5️⃣ 中看到的就是key: value2；從 context 8️⃣ 🔟 中看到的則是key: value8。

我用代碼構造了這幅圖中的 context 樹，放在了這裏 https://github.com/jianghushinian/blog-go-example/blob/main/context/main.go，你可以點擊進去跟着代碼來實驗一下。也可以將代碼 clone 到本地，進行修改，嘗試執行和分析結果，以此來加深你對 context 的理解。

本文示例源碼我都放在了 GitHub 中 https://github.com/jianghushinian/blog-go-example/tree/main/context，歡迎點擊查看。

希望此文能對你有所啓發。

聯繫我

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

猜你喜歡