聊一聊 golang 中的 Context 的實作

編程語言中的 Context

Context 的直接翻譯是背景關系或環境，在編程語言中，翻譯成運行環境更合適，

比如一段程式，在執行之初，我們可以設定一個環境引數：最大運行時間，一旦超過這個時間，程式也應該隨之終止，

在 golang 中， Context 被用來在各個 goroutine 之間傳遞取消信號、超時時間、截止時間、key-value等環境引數，

golang 中的 Context 的實作

golang中的Context包很小，除去注釋，只有200多行，非常適合通過原始碼閱讀來了解它的設計思路，

注：本文中的golang 均指 go 1.14

介面 Context 的定義

golang 中 Context 是一個介面型別，具體定義如下：

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

Deadline()

Deadline() 回傳的是當前 Context 生命周期的截止時間，

Done()

Done() 回傳的是一個只讀的 channel，如果能從這個 channel 中讀到任何值，則表明context的生命周期結束，

Err()

這個比較簡單，就是回傳例外，

Value(key interface{})

Value(key interface{}) 回傳的是 Context 存盤的 key 對應的 value，如果在當前的 Context 中沒有找到，就會從父 Context 中尋找，一直尋找到最后一層，

4種基本的context型別

emptyCtx 定義及實作

type emptyCtx int

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 interface{}) interface{} {
	return nil
}

看 emptyCtx 很輕松，因為它什么都沒做，僅僅是實作了 Context 這個介面，在 context 包中，有一個全域變數 background，值為 new(emptyCtx)，它的作用就是做個跟 Context，其它型別的 Context 都是在 background 的基礎上擴展功能，

cancelCtx 定義及實作

先看下 cancelCtx 的定義和創建，

// 定義
type cancelCtx struct {
	Context

	mu       sync.Mutex            // protects following fields
	done     chan struct{}         // created lazily, closed by first cancel call
	children map[canceler]struct{} // set to nil by the first cancel call
	err      error                 // set to non-nil by the first cancel call
}

// 創建
func WithCancel(parent Context) (ctx Context, cancel CancelFunc) {
	c := newCancelCtx(parent)
	propagateCancel(parent, &c)
	return &c, func() { c.cancel(true, Canceled) }
}

// newCancelCtx returns an initialized cancelCtx.
func newCancelCtx(parent Context) cancelCtx {
	return cancelCtx{Context: parent}
}

總體來說，cancelCtx 的創建就是把父 Context 復制到 cancelCtx 的成員 Context 上，然后把父 Context 的一些信號廣播到子 Context 上，最后回傳了 cancelCtx 的參考，以及一個 cancelFunc，

我們看一下 cancel 實作的細節：

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
	if c.done == nil {
		c.done = closedchan
	} else {
		close(c.done)
	}
	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 {
		removeChild(c.Context, c)
	}
}

cancel 有兩個引數，一個是 removeFromParent，表示當前的取消操作是否需要把自己從父 Context 中移除，第二個引數就是執行取消操作需要回傳的錯誤提示，

根據 cancel 的流程，如果 c.done 是 nil (父 Context 是 emptyCtx 的情況)，就賦值 closedchan，（ closedchan 是一個被關閉的channel）；如果不是nil，就直接關閉，然后遞回關閉子 Context，

這里注意一下，關閉子 Context 的時候，removeFromParent 引數傳值是 false，這是因為當前 Context 在關閉的時候，把 child 置成了 nil，所以子 Context 就不用再執行一次從父 Context 移除自身的操作了，

最后，我們重點說一說 propagateCancel 函式，

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

	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 {
		atomic.AddInt32(&goroutines, +1)
		go func() {
			select {
			case <-parent.Done():
				child.cancel(false, parent.Err())
			case <-child.Done():
			}
		}()
	}
}

從函式名 propagateCancel 大概能看出看出來這個函式的功能，即 “傳播取消（信號）”，回想一下，父 Context 是如何判斷有沒有收到取消信號的？是根據它的私有成員 ctx.done 來判斷的，那子 Context 如何能接收到這個信號呢？這就是函式 propagateCancel 干的事情，把 ctx.done 賦值給子 Context 的私有成員 done，子 Context 就可以獲取到取消的信號，

propagateCancel 的實際處理要更為復雜一些，首先是判斷判斷父 Context 有沒有被 cancel 掉？如果已經 cancel 掉，那么直接 cancel 掉當前的子 Context；如果沒有的話，就會斷言父 Context 是否是emptyCtx 型別，如果是，就通過父 Context 的成員 children 把子 Context 掛在父 Context 下面；如果不是，就啟一個協程監聽父 Context 信號，

解釋一下為什么會 斷言父 Context 是否是emptyCtx 型別 ？想象一下，如果是你來寫這段邏輯，會怎么寫？最簡單的方法就是每個子 Context 啟一個協程，監聽取消信號，這種方式能確實能實作取消信號廣播的功能，但缺點就是如果子 Context 過多，協程就會很多，一直占用系統資源；而如果父 Context 的型別是 cancelCtx，那么它就能通過成員 children 遞回的取消子 Context，一邊是 n 個協程監聽取消信號，一遍是一個協程就能遞回取消所有子 Context，哪種方式消耗資源少，一目了然，

timerCtx 定義及實作

先看以下 timerCtx 的定義和創建：

type timerCtx struct {
	cancelCtx
	timer *time.Timer // Under cancelCtx.mu.

	deadline time.Time
}

func WithDeadline(parent Context, d time.Time) (Context, CancelFunc) {
	if cur, ok := parent.Deadline(); ok && cur.Before(d) {
		// The current deadline is already sooner than the new one.
		return WithCancel(parent)
	}
	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)
		})
	}
	return c, func() { c.cancel(true, Canceled) }
}

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

有了前面的 cancelCtx 的基礎后，看 timerCtx 會清晰很多，timerCtx 的結構簡單一些，timeCtx 有三個成員，第一個是 cancelCtx，這意味這 timerCtx 的取消的操作其實是通過 cancelCtx 實作的；第二個成員是 timer，這是一個定時器，干的事情就是到 deadline 的時候，執行 cancel 操作；第三個成員就是 deadline，

當然，除了等定時器到期自動執行 cancel 操作，也可以主動執行：

func (c *timerCtx) cancel(removeFromParent bool, err error) {
	c.cancelCtx.cancel(false, err)
	if removeFromParent {
		// Remove this timerCtx from its parent cancelCtx's children.
		removeChild(c.cancelCtx.Context, c)
	}
	c.mu.Lock()
	if c.timer != nil {
		c.timer.Stop()
		c.timer = nil
	}
	c.mu.Unlock()
}

如果主動執行 cancel 操作，除了會遞回取消子 Context，還是終止定時器，

valueCtx 的定義和創建

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

func WithValue(parent Context, key, val interface{}) Context {
	if key == nil {
		panic("nil key")
	}
	if !reflectlite.TypeOf(key).Comparable() {
		panic("key is not comparable")
	}
	return &valueCtx{parent, key, val}
}

func (c *valueCtx) Value(key interface{}) interface{} {
	if c.key == key {
		return c.val
	}
	return c.Context.Value(key)
}

valueCtx 也很簡單，一個 Context 型別的成員，還有兩個都是 interface{} 型別的成員 key，value，

從 valueCtx 的創建能看到，如果想給 Context 存盤一個鍵值對，只能通過 WithValue 函式創建，且每個 Context 只能存盤一對，取值的方式是遞回尋找父 Context 存盤的鍵值對，所以一個 Context 相當于存盤了全部父節點的鍵值對，

另外可以看到，valueCtx 的成員是 Context 型別，不是 cancelCtx 型別，這一點需要注意，所以不同的業務場景需要選擇不同的 Context，

golang 中 Context 的使用

golang 中 Context 的使用套路是在最開始的時候，創建一個 root Context，這個 root Context 就是 emptyCtx 的一個實體，

var (
	background = new(emptyCtx)
)

func Background() Context {
	return background
}

接著是根據各個場景，創建不同型別的 Context，

此外，官方博客也給出了 Context 使用的一些建議：

1. 不能在其它型別的結構下放 Context 型別的成員，
2. Context 型別應該作為函式的第一個引數使用，簡寫是 ctx
3. 不要用 nil 來代替本該傳入的 Context，實在不行可以先傳 context.Todo() (和 background 類似)，
4. 不要把函式內部的引數添加到 ctx 中，ctx 中應該存一些貫穿始終的資料，
5. Context 是并發安全的，所以不用擔心多個執行緒同時使用，

結尾

golang 的 Context 就講到這里，由于篇幅原因，總覺得還有不少地方沒有講清楚，下回有機會結合業務場景講一下 Context 的具體使用，

參考

1. 深度解密Go語言之context
2. 由淺入深聊聊Golang的context