Go語言 context包源碼學習
你必須非常努力，才能看起來毫不費力！
微信搜索公眾號[ 漫漫Coding路 ]，一起From Zero To Hero !
前言
日常 Go 開發中，Context 包是用的最多的一個了，幾乎所有函數的第一個參數都是 ctx，那麼我們為什麼要傳遞 Context 呢，Context 又有哪些用法，底層實現是如何呢？相信你也一定會有探索的慾望，那麼就跟着本篇文章，一起來學習吧！
需求一
開發中肯定會調用別的函數，比如 A 調用 B，在調用過程中經常會設置超時時間，比如超過2s 就不等待 B 的結果了，直接返回，那麼我們需要怎麼做呢？
// 睡眠5s，模擬長時間操作
func FuncB() (interface{}, error) {
time.Sleep(5 * time.Second)
return struct{}{}, nil
}
func FuncA() (interface{}, error) {
var res interface{}
var err error
ch := make(chan interface{})
// 調用FuncB()，將結果保存至 channel 中
go func() {
res, err = FuncB()
ch <- res
}()
// 設置一個2s的定時器
timer := time.NewTimer(2 * time.Second)
// 監測是定時器先結束，還是 FuncB 先返回結果
select {
// 超時，返回默認值
case <-timer.C:
return "default", err
// FuncB 先返回結果，關閉定時器，返回 FuncB 的結果
case r := <-ch:
if !timer.Stop() {
<-timer.C
}
return r, err
}
}
func main() {
res, err := FuncA()
fmt.Println(res, err)
}
上面我們的實現，可以實現超過等待時間後，A 不等待 B，但是 B 並沒有感受到取消信號，如果 B 是個計算密度型的函數，我們也希望B 感知到取消信號，及時取消計算並返回，減少資源浪費。
另一種情況，如果存在多層調用，比如A 調用 B、C，B 調用 D、E，C調用 E、F，在超過 A 的超時時間後，我們希望取消信號能夠一層層的傳遞下去，後續所有被調用到的函數都能感知到，及時返回。
需求二
在多層調用的時候，A->B->C->D，有些數據需要固定傳輸，比如 LogID，通過打印相同的 LogID，我們就能夠追溯某一次調用，方便問題的排查。如果每次都需要傳參的話，未免太麻煩了，我們可以使用 Context 來保存。通過設置一個固定的 Key，打印日誌時從中取出 value 作為 LogID。
const LogKey = "LogKey"
// 模擬一個日誌打印，每次從 Context 中取出 LogKey 對應的 Value 作為LogID
type Logger struct{}
func (logger *Logger) info(ctx context.Context, msg string) {
logId, ok := ctx.Value(LogKey).(string)
if !ok {
logId = uuid.New().String()
}
fmt.Println(logId + " " + msg)
}
var logger Logger
// 日誌打印 並 調用 FuncB
func FuncA(ctx context.Context) {
logger.info(ctx, "FuncA")
FuncB(ctx)
}
func FuncB(ctx context.Context) {
logger.info(ctx, "FuncB")
}
// 獲取初始化的，帶有 LogID 的 Context，一般在程序入口做
func getLogCtx(ctx context.Context) context.Context {
logId, ok := ctx.Value(LogKey).(string)
if ok {
return ctx
}
logId = uuid.NewString()
return context.WithValue(ctx, LogKey, logId)
}
func main() {
ctx = getLogCtx(context.Background())
FuncA(ctx)
}
這利用到了本篇文章講到的 valueCtx，繼續往下看，一起來學習 valueCtx 是怎麼實現的吧！
Context 接口
type Context interface {
Deadline() (deadline time.Time, ok bool)
Done() <-chan struct{}
Err() error
Value(key interface{}) interface{}
}
Context 接口比較簡單，定義了四個方法：
- Deadline() 方法返回兩個值，deadline 表示 Context 將會在什麼時間點取消，ok 表示是否設置了deadline。當 ok=false 時，表示沒有設置deadline，那麼此時 deadline 將會是個零值。多次調用這個方法返回同樣的結果。
- Done() 返回一個只讀的 channel，類型為 chan struct{}，如果當前的 Context 不支持取消，Done 返回 nil。我們知道，如果一個 channel 中沒有數據，讀取數據會阻塞；而如果channel被關閉，則可以讀取到數據，因此可以監聽 Done 返回的 channel，來獲取 Context 取消的信號。
- Err() 返回 Done 返回的 channel 被關閉的原因。當 channel 未被關閉時，Err() 返回 nil；channel 被關閉時則返回相應的值，比如 Canceled 、DeadlineExceeded。Err() 返回一個非 nil 值之後，後面再次調用會返回相同的值。
- Value() 返回 Context 保存的鍵值對中，key 對應的 value，如果 key 不存在則返回 nil。
Done() 是一個比較常用的方法，下面是一個比較經典的流式處理任務的示例：監聽 ctx.Done() 是否被關閉來判斷任務是否需要取消，需要取消則返回相應的原因；沒有取消則將計算的結果寫入到 out channel中。
func Stream(ctx context.Context, out chan<- Value) error {
for {
// 處理數據
v, err := DoSomething(ctx)
if err != nil {
return err
}
// ctx.Done() 讀取到數據，說明獲取到了任務取消的信號
select {
case <-ctx.Done():
return ctx.Err()
// 否則將結果輸出，繼續計算
case out <- v:
}
}
}
Value() 也是一個比較常用的方法，用於在上下文中傳遞一些數據。使用 context.WithValue() 方法存入 key 和 value，通過 Value() 方法則可以根據 key 拿到 value。
func main() {
ctx := context.Background()
c := context.WithValue(ctx, "key", "value")
v, ok := c.Value("key").(string)
fmt.Println(v, ok)
}
emptyCtx
Context 接口並不需要我們自己去手動實現，一般我們都是直接使用 context 包中提供的 Background() 方法和 TODO() 方法，來獲取最基礎的 Context。
var (
background = new(emptyCtx)
todo = new(emptyCtx)
)
func Background() Context {
return background
}
func TODO() Context {
return todo
}
Background() 方法一般用在 main 函數，或者程序的初始化方法中；在我們不知道使用哪個 Context，或者上文沒有傳遞 Context時，可以使用 TODO()。
Background() 和 TODO() 都是基於 emptyCtx 生成的，從名字可以看出來，emptyCtx 是一個空的Context，沒有 deadline、不能被取消、沒有鍵值對。
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
}
func (e *emptyCtx) String() string {
switch e {
case background:
return "context.Background"
case todo:
return "context.TODO"
}
return "unknown empty Context"
}
除了上面兩個最基本的 Context 外，context 包中提供了功能更加豐富的 Context，包括 valueCtx、cancelCtx、timerCtx，下面我們就挨個來看下。
valueCtx
使用示例
我們一般使用 context.WithValue() 方法向 Context 存入鍵值對，然後通過 Value() 方法根據 key 得到 value，此種功能的實現就依賴 valueCtx。
func main() {
ctx := context.Background()
c := context.WithValue(ctx, "myKey", "myValue")
v1 := c.Value("myKey")
fmt.Println(v1.(string))
v2 := c.Value("hello")
fmt.Println(v2) // nil
}
類型定義
valueCtx 結構體中嵌套了 Context，使用 key 、value 來保存鍵值對：
type valueCtx struct {
Context
key, val interface{}
}
WithValue
context包 對外暴露了 WithValue 方法，基於一個 parent context 來創建一個 valueCtx。從下面的源碼中可以看出，key 必須是可比較的！
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}
}
*valueCtx 實現了 Value()，可以根據 key 得到 value。這是一個向上遞歸尋找的過程，如果 key 不在當前 valueCtx 中，會繼續向上找 parent Context，直到找到最頂層的 Context，一般最頂層的是 emptyCtx，而 emtpyCtx.Value() 返回 nil。
func (c *valueCtx) Value(key interface{}) interface{} {
if c.key == key {
return c.val
}
return c.Context.Value(key)
}
cancelCtx
cancelCtx 是一個用於取消任務的 Context，任務通過監聽 Context 是否被取消，來決定是否繼續處理任務還是直接返回。
如下示例中，我們在 main 函數定義了一個 cancelCtx，並在 2s 後調用 cancel() 取消 Context，即我們希望 doSomething() 在 2s 內完成任務，否則就可以直接返回，不需要再繼續計算浪費資源了。
doSomething() 方法內部，我們使用 select 監聽任務是否完成，以及 Context 是否已經取消，哪個先到就執行哪個分支。方法模擬了一個 5s 的任務，main 函數等待時間是2s，因此沒有完成任務；如果main函數等待時間改為10s，則任務完成並會返回結果。
這只是一層調用，真實情況下可能會有多級調用，比如 doSomething 可能又會調用其他任務，一旦 parent Context 取消，後續的所有任務都應該取消。
func doSomething(ctx context.Context) (interface{}, error) {
res := make(chan interface{})
go func() {
fmt.Println("do something")
time.Sleep(time.Second * 5)
res <- "done"
}()
select {
case <-ctx.Done():
return nil, ctx.Err()
case value := <-res:
return value, nil
}
}
func main() {
ctx, cancel := context.WithCancel(context.Background())
go func() {
time.Sleep(time.Second * 2)
cancel()
}()
res, err := doSomething(ctx)
fmt.Println(res, err) // nil , context canceled
}
接下來就讓我們來研究下，cancelCtx 是如何實現取消的吧
類型定義
- canceler 接口包含 cancel() 和 Done() 方法，*cancelCtx 和 *timerCtx 均實現了這個接口。
- closedchan 是一個被關閉的channel，可以用於後面 Done() 返回
- canceled 是一個 err，用於 Context 被取消的原因
type canceler interface {
cancel(removeFromParent bool, err error)
Done() <-chan struct{}
}
// closedchan is a reusable closed channel.
var closedchan = make(chan struct{})
func init() {
close(closedchan)
}
var Canceled = errors.New("context canceled")
CancelFunc 是一個函數類型定義，是一個取消函數，有如下規範：
- CancelFunc 告訴一個任務停止工作
- CancelFunc 不會等待任務結束
- CancelFunc 支持並發調用
- 第一次調用後，後續的調用不會產生任何效果
type CancelFunc func()
&cancelCtxKey 是一個固定的key，用來返回 cancelCtx 自身
var cancelCtxKey int
cancelCtx
cancelCtx 是可以被取消的，它嵌套了 Context 接口，實現了 canceler 接口。cancelCtx 使用 children 字段保存同樣實現 canceler 接口的子節點，當 cancelCtx 被取消時，所有的子節點也會取消。
type cancelCtx struct {
Context
mu sync.Mutex // 保護如下字段，保證線程安全
done atomic.Value // 保存 channel，懶加載，調用 cancel 方法時會關閉這個 channel
children map[canceler]struct{} // 保存子節點，第一次調用 cancel 方法時會置為 nil
err error // 保存為什麼被取消，默認為nil，第一次調用 cancel 會賦值
}
*cancelCtx 的 Value() 方法 和 *valueCtx 的 Value() 方法類似，只不過加了個固定的key: &cancelCtxKey。當key 為 &cancelCtxKey 時返回自身
func (c *cancelCtx) Value(key interface{}) interface{} {
if key == &cancelCtxKey {
return c
}
return c.Context.Value(key)
}
*cancelCtx 的 done 字段是懶加載的，只有在調用 Done() 方法 或者 cancel() 時才會賦值。
func (c *cancelCtx) Done() <-chan struct{} {
d := c.done.Load()
// 如果已經有值了，直接返回
if d != nil {
return d.(chan struct{})
}
// 沒有值，加鎖賦值
c.mu.Lock()
defer c.mu.Unlock()
d = c.done.Load()
if d == nil {
d = make(chan struct{})
c.done.Store(d)
}
return d.(chan struct{})
}
Err 方法返回 cancelCtx 的 err 字段
func (c *cancelCtx) Err() error {
c.mu.Lock()
err := c.err
c.mu.Unlock()
return err
}
WithCancel
那麼我們如何新建一個 cancelCtx呢？context 包提供了 WithCancel() 方法，讓我們基於一個 Context 來創建一個 cancelCtx。WithCancel() 方法返回兩個字段，一個是基於傳入的 Context 生成的 cancelCtx，另一個是 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)
return &c, func() { c.cancel(true, Canceled) }
}
WithCancel 調用了兩個外部方法：newCancelCtx 、propagateCancel。newCancelCtx 比較簡單，根據傳入的 context，返回了一個 cancelCtx 結構體。
func newCancelCtx(parent Context) cancelCtx {
return cancelCtx{Context: parent}
}
propagateCancel 從名字可以看出，就是將 cancel 傳播。如果父Context支持取消，那麼我們需要建立一個通知機制，這樣父節點取消的時候，通知子節點也取消，層層傳播。
在 propagateCancel 中，如果 父Context 是 cancelCtx 類型且未取消，會將 子Context 掛在它下面，形成一個樹結構；其餘情況都不會掛載。
func propagateCancel(parent Context, child canceler) {
// 如果 parent 不支持取消，那麼就不支持取消傳播，直接返回
done := parent.Done()
if done == nil {
return
}
// 到這裡說明 done 不為 nil，parent 支持取消
select {
case <-done:
// 如果 parent 此時已經取消了，那麼直接告訴子節點也取消
child.cancel(false, parent.Err())
return
default:
}
// 到這裡說明此時 parent 還未取消
// 如果 parent 是未取消的 cancelCtx
if p, ok := parentCancelCtx(parent); ok {
// 加鎖，防止並發更新
p.mu.Lock()
// 再次判斷，因為有可能上一個獲得鎖的進行了取消操作。
// 如果 parent 已經取消了，那麼子節點也直接取消
if p.err != nil {
child.cancel(false, p.err)
} else {
// 把子Context 掛到父節點 parent cancelCtx 的 children字段下
// 之後 parent cancelCtx 取消時，能通知到所有的 子Context
if p.children == nil {
p.children = make(map[canceler]struct{})
}
p.children[child] = struct{}{}
}
p.mu.Unlock()
} else {
// parent 不是 cancelCtx 類型，可能是用戶自己實現的Context
atomic.AddInt32(&goroutines, +1)
// 啟動一個協程監聽，如果 parent 取消了，子 Context 也取消
go func() {
select {
case <-parent.Done():
child.cancel(false, parent.Err())
case <-child.Done():
}
}()
}
}
cancel 方法就是來取消 cancelCtx，主要的工作是：關閉c.done 中的channel，給 err 賦值，然後級聯取消所有 子Context。如果 removeFromParent 為 true，會從父節點中刪除以該節點為樹頂的樹。
cancel() 方法只負責自己管轄的範圍，即自己以及自己的子節點，然後根據配置判斷是否需要從父節點中移除自己為頂點的樹。如果子節點還有子節點，那麼由子節點負責處理，不用自己負責了。
propagateCancel() 中有三處調用了 cancel() 方法，傳入的 removeFromParent 都為 false，是因為當時根本沒有掛載，不需要移除。而 WithCancel 返回的 CancelFunc ，傳入的 removeFromParent 為 true，是因為調用 propagateCancel 有可能產生掛載，當產生掛載時，調用 cancel() 就需要移除了。
func (c *cancelCtx) cancel(removeFromParent bool, err error) {
// err 是指取消的原因，必傳，cancelCtx 中是 errors.New("context canceled")
if err == nil {
panic("context: internal error: missing cancel error")
}
// 涉及到保護字段值的修改，都需要加鎖
c.mu.Lock()
// 如果該Context已經取消過了，直接返回。多次調用cancel，不會產生額外效果
if c.err != nil {
c.mu.Unlock()
return
}
// 給 err 賦值，這裡 err 一定不為 nil
c.err = err
// close channel
d, _ := c.done.Load().(chan struct{})
// 因為c.done 是懶加載，有可能存在 nil 的情況
// 如果 c.done 中沒有值，直接賦值 closedchan；否則直接 close
if d == nil {
c.done.Store(closedchan)
} else {
close(d)
}
// 遍歷當前 cancelCtx 所有的子Context，讓子節點也 cancel
// 因為當前的Context 會主動把子Context移除，子Context 不用主動從parent中脫離
// 因此 child.cancel 傳入的 removeFromParent 為false
for child := range c.children {
child.cancel(false, err)
}
// 將 children 置空，相當於移除自己的所有子Context
c.children = nil
c.mu.Unlock()
// 如果當前 cancelCtx 需要從上層的 cancelCtx移除，調用removeChild方法
// c.Context 就是自己的父Context
if removeFromParent {
removeChild(c.Context, c)
}
}
從propagateCancel方法中可以看到，只有parent 屬於 cancelCtx 類型 ，才會將自己掛載。因此 removeChild 會再次判斷 parent 是否為 cancelCtx，和之前的邏輯保持一致。找到的話，再將自己移除，需要注意的是，移除會把自己及其自己下面的所有子節點都移除。
如果上一步 propagateCancel 方法將自己掛載到了 A 上，但是在調用 cancel() 時，A 已經取消過了，此時 parentCancelCtx() 會返回 false。不過這沒有關係，A 取消時已經將掛載的子節點移除了，當前的子節點不用將自己從 A 中移除了。
func removeChild(parent Context, child canceler) {
// parent 是否為未取消的 cancelCtx
p, ok := parentCancelCtx(parent)
if !ok {
return
}
// 獲取 parent cancelCtx 的鎖，修改保護字段 children
p.mu.Lock()
// 將自己從 parent cancelCtx 的 children 中刪除
if p.children != nil {
delete(p.children, child)
}
p.mu.Unlock()
}
parentCancelCtx 判斷 parent 是否為 未取消的 *cancelCtx。取消與否容易判斷，難判斷的是 parent 是否為 *cancelCtx，因為有可能其他結構體內嵌了 cancelCtx，比如 timerCtx，會通過比對 channel 來確定。
func parentCancelCtx(parent Context) (*cancelCtx, bool) {
// 如果 parent context 的 done 為 nil， 說明不支持 cancel，那麼就不可能是 cancelCtx
// 如果 parent context 的 done 為 closedchan， 說明 parent context 已經 cancel 了
done := parent.Done()
if done == closedchan || done == nil {
return nil, false
}
// 到這裡說明支持取消，且沒有被取消
// 如果 parent context 屬於原生的 *cancelCtx 或衍生類型，需要繼續進行後續判斷
// 如果 parent context 無法轉換到 *cancelCtx，則認為非 cancelCtx，返回 nil,fasle
p, ok := parent.Value(&cancelCtxKey).(*cancelCtx)
if !ok {
return nil, false
}
// 經過上面的判斷後，說明 parent context 可以被轉換為 *cancelCtx，這時存在多種情況:
// - parent context 就是 *cancelCtx
// - parent context 是標準庫中的 timerCtx
// - parent context 是個自己自定義包裝的 cancelCtx
//
// 針對這 3 種情況需要進行判斷，判斷方法就是:
// 判斷 parent context 通過 Done() 方法獲取的 done channel 與 Value 查找到的 context 的 done channel 是否一致
//
// 一致情況說明 parent context 為 cancelCtx 或 timerCtx 或 自定義的 cancelCtx 且未重寫 Done()，
// 這種情況下可以認為拿到了底層的 *cancelCtx
//
// 不一致情況說明 parent context 是一個自定義的 cancelCtx 且重寫了 Done() 方法，並且並未返回標準 *cancelCtx 的
// 的 done channel，這種情況需要單獨處理，故返回 nil, false
pdone, _ := p.done.Load().(chan struct{})
if pdone != done {
return nil, false
}
return p, true
}
timerCtx
簡介
timerCtx 嵌入了 cancelCtx，並新增了一個 timer 和 deadline 字段。timerCtx 的取消能力是復用 cancelCtx 的，只是在這個基礎上增加了定時取消而已。
在我們的使用過程中，有可能還沒到 deadline，任務就提前完成了，此時需要手動調用 CancelFunc。
func slowOperationWithTimeout(ctx context.Context) (Result, error) {
ctx, cancel := context.WithTimeout(ctx, 100*time.Millisecond)
defer cancel() // 如果未到截止時間，slowOperation就完成了，儘早調用 cancel() 釋放資源
return slowOperation(ctx)
}
類型定義
type timerCtx struct {
cancelCtx // 內嵌 cancelCtx
timer *time.Timer // 受 cancelCtx.mu 互斥鎖的保護
deadline time.Time // 截止時間
}
Deadline() 返回 deadline 字段的值
func (c *timerCtx) Deadline() (deadline time.Time, ok bool) {
return c.deadline, true
}
WithDeadline
WithDeadline 基於parent Context 和 時間點 d，返回了一個定時取消的 Context，以及一個 CancelFunc。返回的Context 有三種情況被取消：1. 到達了指定時間，就會主動取消；2. 手動調用了 CancelFunc；3. 父Context取消，導致該Context被取消。這三種情況哪種先到，就會首次觸發取消操作，後續的再次取消不會產生任何效果。
如果傳入 parent Context 的 deadline 比指定的時間 d 還要早，此時 d 就沒用處了，直接依賴 parent 取消傳播就可以了。
func WithDeadline(parent Context, d time.Time) (Context, CancelFunc) {
// 傳入的 parent 不能為 nil
if parent == nil {
panic("cannot create context from nil parent")
}
// parent 也有 deadline，並且比 d 還要早，直接依賴 parent 的取消傳播即可
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,
}
// 設置傳播，如果parent 屬於 cancelCtx，會掛載到 children 字段上
propagateCancel(parent, c)
// 距離截止時間 d 還有多久
dur := time.Until(d)
if dur <= 0 {
// 已經到了截止時間，直接取消，同時從 parent 中取消掛載
// 由於是超時，取消時的 err 是 DeadlineExceeded
c.cancel(true, DeadlineExceeded)
// 再返回 c 和 CancelFunc，已經取消掛載了，此時的 CancelFunc 不會從 parent 中取消掛載
// 後面再次調用 CancelFunc 不會產生任何效果了
// 主動取消的話，err 是 Canceled
return c, func() { c.cancel(false, Canceled) }
}
// 還沒有到截止時間，定義一個定時器，過了 dur 會自動取消
c.mu.Lock()
defer c.mu.Unlock()
if c.err == nil {
c.timer = time.AfterFunc(dur, func() {
// 由於是到了截止時間才取消，err 是 DeadlineExceeded
c.cancel(true, DeadlineExceeded)
})
}
// 返回 c 和 cancelFunc，主動取消的 err 是 Canceled
return c, func() { c.cancel(true, Canceled) }
}
接下來我們看下 cancel 方法，timerCtx 的 cancel 方法 就是調用內嵌 cancelCtx 的 cancel() 方法，默認是不從父節點移除
func (c *timerCtx) cancel(removeFromParent bool, err error) {
c.cancelCtx.cancel(false, err)
// 從父節點中移除
if removeFromParent {
removeChild(c.cancelCtx.Context, c)
}
// 把定時器停了，釋放資源
// 有可能還沒到deadline，手動觸發了 CancelFunc，此時把 timer 停了
c.mu.Lock()
if c.timer != nil {
c.timer.Stop()
c.timer = nil
}
c.mu.Unlock()
}
WithTimeout
WithTimeout 就是基於 WithDeadline，deadline 就是基於當前時間計算的
func WithTimeout(parent Context, timeout time.Duration) (Context, CancelFunc) {
return WithDeadline(parent, time.Now().Add(timeout))
}
總結
本篇文章，我們通過源碼+示例的方式，一起學習了 context 包相關的結構以及實現邏輯，包括如下內容
Context 接口：定義了一些接口方法和規範
emptyCtx：空的Context，Background() 和 TODO() 方法就是使用的 emptyCtx
valueCtx：用於保存鍵值對，查詢時是遞歸查詢，可以用於 LogID 這種全局 id 的保存
cancelCtx：可以取消的Context，用於取消信號的傳遞
timerCtx：定時取消的 cancelCtx
更多
個人博客: //lifelmy.github.io/
微信公眾號：漫漫Coding路
