Go語言Context包原始碼學習

0前言

context包作為使用go進行server端開發的重要工具，其原始碼只有791行，不包含註釋的話預計在500行左右，非常值得我們去深入探討學習，於是在本篇筆記中我們一起來觀察原始碼的實現，知其然更要知其所以然。（當前使用go版本為1.22.2）

1核心資料結構

整體的介面實現和結構體embed圖

1.1Context介面

context介面定義了四個方法：

• Deadline方法返回context是否為timerctx，以及它的結束時間
• Done方法返回該ctx的done channel
• Err方法返回該ctx被取消的原因
• Value方法返回key對應的value

2emptyCtx

先來觀察原始碼

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
}

emptyctx實現了context介面中的所有方法，對於每個方法返回的都是空值，它沒有值、不能被取消以及沒有截止時間，它只作為一個空context的載體，相當於所有ctx的祖先。

如何建立一個context？

context包提供了Background()方法和TODO()方法，都用於建立一個空的context。

func Background() Context {
	return backgroundCtx{}
}
func TODO() Context {
	return todoCtx{}
}
type backgroundCtx struct{ emptyCtx }

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

type todoCtx struct{ emptyCtx }

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

它們返回的都為空的context，雖然方法名不同，但是效果是一樣的，那麼什麼時候使用background，什麼時候使用TODO呢，這是官方給出的註釋

“TODO 會返回一個非空的、為空的 [Context]。 程式碼應該在不清楚應該使用哪個 [Context] 或者 [Context] 尚未可用（因為周圍的函式尚未被擴充套件以接受 [Context] 引數）時使用 context.TODO。”

“background返回一個非空的、空的上下文物件。它不會被取消，沒有值，也沒有截止時間。它通常被主函式、初始化和測試使用，作為進入請求的頂級上下文。”

3cancelctx

先來觀察cancelctx結構體的實現

type cancelCtx struct {
	Context

	mu       sync.Mutex            // protects following fields
	done     atomic.Value          // of 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
	cause    error                 // set to non-nil by the first cancel call
}

• cancelCtx 內嵌了 Context 作為其父 context，根據go語言的特性，cancelCtx結構體就隱式實現了context介面中的所有方法，可以被當作一個介面來被呼叫，但是其方法還需要被具體的實現賦值才能進行呼叫。並且可以得知，cancelCtx的父類一定也是一個Context。

• mu是cancelCtx的內建鎖，用來協調併發場景下的資源獲取

• done的實際型別為chan struct{}，透過atomic包來實現併發安全，可以用於反應該ctx的生命週期情況，done是懶漢式建立的，只有第一次呼叫Done()方法時才會被建立，在下文的Done方法中會提到

• children用於關聯和子ctx的關係，當取消該ctx時，可以接連通知子ctx進行關閉，及時釋放資源。

• err用於返回ctx關閉的原因，呼叫的是context包定義的內建error

  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 }
  

• cause返回的是該ctx失效更底層的原因，例如導致DeadlineExceeded err的具體原因是“database connection timeout”

  // Example use:
  //
  //	ctx, cancel := context.WithCancelCause(parent)
  //	cancel(myError)
  //	ctx.Err() // returns context.Canceled
  //	context.Cause(ctx) // returns myError
  

3.1Done方法的實現

func (c *cancelCtx) Done() <-chan struct{} {
	d := c.done.Load() //獲取cancelCtx的done通道
	if d != nil { //假若已經建立過了done通道，則直接返回
		return d.(chan struct{}) 
	}
	c.mu.Lock() //上鎖
	defer c.mu.Unlock()
	d = c.done.Load() //Double check是否建立done通道，因為在上鎖前，可能其他goroutine呼叫了該ctx的Done方法。
	if d == nil { //如果仍然未建立
		d = make(chan struct{}) //建立該done channel
		c.done.Store(d) //儲存
	}
	return d.(chan struct{})
}

透過程式碼可以看見，ctx的done只有當被呼叫過Done方法時才會被建立，那麼為什麼這樣子設計呢？很容易想到主要目的就是為了節省了不必要的資源浪費，提高效率，在很多情況下建立context並不需要監聽done通道，只有在需要時才被建立，符合go語言的設計理念，只有需要的時候才引入。

3.2value方法的實現

func (c *cancelCtx) Value(key any) any {
	if key == &cancelCtxKey {
		return c
	}
	return value(c.Context, key)
}

若引數key與cancelCtxKey相符，則返回當前ctx本身

否則，就向父層 層層尋找

原始碼錶現的非常晦澀，為了具體知道這個Value方法是做什麼的，我們先來看對於cancelCtxKey的定義

// &cancelCtxKey is the key that a cancelCtx returns itself for.
var cancelCtxKey int

cancelCtxKey是一個私有的、唯一的識別符號，它用於返回cancelCtx它本身。

對於該cancelCtxKey具體使用場景，下面還會講到

3.3建立cancelCtx的WithCancel

//WithCancel返回一個帶有新的Done通道的父context的副本。只有當父ctx被關閉，或者返回的cancel方法被呼叫時，該ctx的Done通道才會被關閉。
func WithCancel(parent Context) (ctx Context, cancel CancelFunc) {
	c := withCancel(parent)
	return c, func() { c.cancel(true, Canceled, nil) }
}

該方法返回了一個cancelCtx，以及關閉它的cancel方法。

在1.20版本中，新增了一個WithCancelCause方法，該方法返回了一個cancelctx和它的CancelCauseFunc，我們也來看一下

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

在程式碼方面基本和withCancel方法一致，但是返回的CancelCauseFunc可以用於給使用者自定義ctx被取消的原因，例如

ctx, cancel := context.WithCancelCause(parent)
cancel(myError)
ctx.Err() // returns context.Canceled
context.Cause(ctx) // returns myError

接著來看withCancel

func withCancel(parent Context) *cancelCtx {
	if parent == nil { //若父ctx是nil，那麼不能建立
		panic("cannot create context from nil parent")
	}
	c := &cancelCtx{}
	c.propagateCancel(parent, c) //傳遞cancel給新建立的ctx
	return c
}

主要來看propagateCancel做了什麼：

//該方法主要用於建立父context和子context的聯絡，如果父context也是一個cancelCtx，它需要保證父context被取消時，子context也能跟著被取消。
func (c *cancelCtx) propagateCancel(parent Context, child canceler) {
	c.Context = parent //將父context內嵌入子context中

	done := parent.Done() //獲取父ctx的Done通道
	if done == nil { //如果通道不存在，那麼說明父context不是cancelCtx，不需要為它們兩個之間建立聯絡，因為父context永遠不會關閉。
		return // parent is never canceled
	}

	select {
	case <-done: //非堵塞地獲取父done通道的狀態，如果done通道以及被closed，那麼這裡會接受到一個零值，如果沒有被closed，會執行default後面的語句。
		//接受到零值，說明done通道以及被關閉了，父context已經被取消，此時應該立即呼叫子context的cancel方法，取消子context。
		child.cancel(false, parent.Err(), Cause(parent))
		return
	default:
	}
	//否則，執行以下程式碼
    //此處判斷父context是否是一個context包實現的cancelCtx，如何理解會在下文講述該方法時繼續說明
	if p, ok := parentCancelCtx(parent); ok {
		// 如果父context是一個cancelCtx，或者是從某個cancelCtx衍生出來的context
		p.mu.Lock() //加鎖
		if p.err != nil {
			// 如果存在err，說明已經被取消，此時也應該呼叫child的cancel方法取消子ctx
			child.cancel(false, p.err, p.cause)
		} else {
			if p.children == nil { //如果這是p的第一個子cancelCtx，需要初始化map的記憶體
				p.children = make(map[canceler]struct{})
			}
			p.children[child] = struct{}{}//繫結子ctx到父ctx的map中，用於當父ctx取消時，能通知所有的child跟著取消。
		}
		p.mu.Unlock()
		return
	}
	//如果父ctx不是cancelctx並且實現了AfterFuncer介面，即實現了AfterFunc方法（該方法會在ctx被取消後唯一一次呼叫），那麼就需要為父ctx再設定一個afterFunc方法，用於取消child並且傳遞err和cause
	if a, ok := parent.(afterFuncer); ok {
		// parent implements an AfterFunc method.
		c.mu.Lock()
		stop := a.AfterFunc(func() {
			child.cancel(false, parent.Err(), Cause(parent))
		})
		c.Context = stopCtx{
			Context: parent,
			stop:    stop, //這裡的stop方法可以用於當父ctx呼叫afterFunc的時候，取消父ctx對cancel函式的呼叫（看需求）
		}
		c.mu.Unlock()
		return
	}
	//下面的情況為增加一個goroutine，監聽父ctx自己實現的Done channel（使用者自定義的）
	goroutines.Add(1)
	go func() {
		select {
		case <-parent.Done()://如果監聽到父ctx自定義實現的Done channel關閉時，就關閉child
			child.cancel(false, parent.Err(), Cause(parent))
		case <-child.Done(): //如果child先關閉，那麼就立即釋放協程，避免協程洩露
		}
	}()
}

propagateCancel方法主要的作用是，保證了對於父ctx被取消時，為了能及時取消子ctx，避免不必要的資源浪費，建立父ctx和子ctx之間的聯絡。使用流程圖來表示該方法如下（省略了檢查afterFunc）

在該方法中，比較難以理解的地方是第二個if,"if p, ok := parentCancelCtx(parent); ok"究竟做了什麼，為此我們跟進原始碼檢視：

// parentCancelCtx 返回父級物件的底層 cancelCtx。
// 它透過查詢 parent.Value(&cancelCtxKey) 來找到最內層的 enclosing cancelCtx，然後檢查 parent.Done() 是否與該 cancelCtx 匹配。（如果不匹配，則該 cancelCtx 可能已被封裝在提供了不同 done 通道的自定義實現中，在這種情況下，我們不應該繞過它。）
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
}

三個if分別做了什麼：

• 第一個if：獲取parentCtx的done channel，並且檢視情況，若已經關閉或者父Ctx不可取消，此時返回false，回到propagate方法中最終會cancel掉child

• 第二個if：可以理解為透過Value方法找到離parentCtx“最近的”cancelCtx p，一般情況下，如果parentCtx就是一個CancelCtx，這時候就是parentCtx它本身。如果p不存在，也就是沒有可以取消的ctx，此時也會返回false。

• 第三個if：找到了p後，還讀取了p的done channel，這時候一般情況，pdone 當然會 == done，因此最終會返回p和true，那麼什麼時候會不相等呢？為什麼會不相等呢？為此，看下方的層次圖來理解這個if

在這個情況假設下，ParentCtx2是從ParentCtx1衍生出來的，ParentCtx1是一個標準的CancelCtx，而ParentCtx2是一個使用者自定義了ctx，它內層繼承了ParentCtx1並且自己實現了Done()方法，這時候程式碼中的“p”找到的就是ParentCtx1，而parent是ParentCtx2，此時的done和pdone就是兩個不同的channel了，這時候cihldctx應該監聽哪一個done channel呢？答案是監聽使用者自定義實現的Ctx的chennel，因為我們不應該繞過使用者實現的Done channel，這更加符合ctx到層次邏輯。假如這時候不去判斷pdone == done，直接返回的指標就是ParentCtx1的指標了。

3.4Cancel方法實現

接下來我們來看cancel方法是如何實現的。

//cancel 關閉 c.done，cancel 每一個c的children。如果removeFromParent為true，將會把c從parentCtx的child中移除。
func (c *cancelCtx) cancel(removeFromParent bool, err, cause error) {
	if err == nil { //必須存在err
		panic("context: internal error: missing cancel error")
	}
	if cause == nil { //沒有自定義設定cause，預設為err
		cause = err
	}
	c.mu.Lock() //上鎖
	if c.err != nil { //double check
		c.mu.Unlock()
		return // already canceled
	}
	c.err = err
	c.cause = cause
	d, _ := c.done.Load().(chan struct{})
	if d == nil { //done沒有被建立，直接儲存context包內以及建立了的關閉的channel，不需要再次建立
		c.done.Store(closedchan)
	} else {
		close(d) //關閉done
	}
	for child := range c.children { //關閉每一個子context
		// NOTE: acquiring the child's lock while holding parent's lock.
		child.cancel(false, err, cause)
	}
	c.children = nil
	c.mu.Unlock()

	if removeFromParent {
		removeChild(c.Context, c) //從父child中移除自己
	}
}

程式碼比較簡單，問題在於什麼時候removFromParent為true呢？為什麼為true？

回到withCancel方法中，我們可以看到返回的cancel方法中，此時removeFromParent就為true

return c, func() { c.cancel(true, Canceled, nil) }

當使用者主動呼叫cancel()時，就會將子ctx從父ctx中的child刪除。因為此時沒有必要再在父ctx中接受父或祖先的cancel通知。而當呼叫cancel函式內部，對child執行的cancel就為false，這是因為後面設定了c.children = nil，這時候是從父ctx的方向關閉了子ctx對其的連結。

4afterFuncCtx

afterFuncCtx是1.20版本後引入的新ctx，它的作用是當ctx被取消後，能執行一次自定義的F函式，一般用於回收資源等。

type afterFuncCtx struct {
	cancelCtx
	once sync.Once // either starts running f or stops f from running
	f    func()
}

可以看到afterFuncCtx embed了cancelCtx，在此基礎上新增了once和f。once保證了f只會被執行一次。接著我們來看如何實現一個afterFuncCtx

4.1AfterFunc

func AfterFunc(ctx Context, f func()) (stop func() bool) {
	a := &afterFuncCtx{ //建立區域性變數afterFuncCtx，記錄cancel時需要被執行的func
		f: f,
	}
	a.cancelCtx.propagateCancel(ctx, a)
	return func() bool { //返回stop方法，如果需要不執行func，則呼叫stop
		stopped := false
		a.once.Do(func() { //嘗試stop func函式，如果成功則stop會被設定為true
			stopped = true
		})
		if stopped { //stop後，a沒有存在的意義，進行取消。
			a.cancel(true, Canceled, nil)
		}
		return stopped //true為成功取消，false表示f已經被執行或者正在被執行。
	}
}

這是一個閉包實現，閉包是指在函式內部定義的函式（如這裡返回的 stop 函式），它會“捕獲”並儲存定義時可訪問的所有外部變數。在 AfterFunc 方法中，雖然 a 是區域性變數，但返回的 stop 方法引用了 a，形成了一個閉包，閉包會將 a 的記憶體保留在堆上，即使 AfterFunc 方法返回後，a 依然存在。所以當雖然沒有返回a，但是返回的stop方法任然能呼叫a，a的生命週期超出了afterFunc方法。

以下是一個示例

package main

import (
	"context"
	"fmt"
	"time"
)

func main() {
	// 建立一個 2 秒後超時的 context
	ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second)
	defer cancel()

	// 註冊 AfterFunc，context 完成時將呼叫清理操作
	stop := context.AfterFunc(ctx, func() {
		fmt.Println("清理操作正在執行...")
	})

	// 等待 3 秒
	time.Sleep(3 * time.Second)

	// 嘗試停止清理操作
	stopped := stop()
	if stopped {
		fmt.Println("成功停止清理操作")
	} else {
		fmt.Println("清理操作已經開始或已停止")
	}

	// 等待 context 超時
	<-ctx.Done()

	fmt.Println("程式結束")
}

4timerCtx

接下來來看timerCtx

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

	deadline time.Time
}

timer直接內嵌了cancelCtx以實現Done和Err，並且新添了timer和deadline欄位，deadline用於檢視ctx的截止時間，timer用於完成過時取消ctx。

4.1WithDeadline

先來看WithTimeout方法

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

WithTimeout接受一個存活時間來建立一個timerCtx，可以看到最終都是呼叫了WithDeadline

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

WithDeadline又呼叫了WithDeadlineCause，返回了timerCtx和一個CancelFunc。

4.2WithDeadlineCause

func WithDeadlineCause(parent Context, d time.Time, cause error) (Context, CancelFunc) {
	if parent == nil {
		panic("cannot create context from nil parent")
	}
	if cur, ok := parent.Deadline(); ok && cur.Before(d) { //如果parent的deadline更早，則直接返回parent的副本，不需要再建立timer
		// The current deadline is already sooner than the new one.
		return WithCancel(parent)
	}
	c := &timerCtx{
		deadline: d,
	}
	c.cancelCtx.propagateCancel(parent, c)
	dur := time.Until(d)
	if dur <= 0 {//檢測是否在建立過程中已經過了ddl
		c.cancel(true, DeadlineExceeded, cause) // deadline has already passed
		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
			c.cancel(true, DeadlineExceeded, cause)
		})
	}
	return c, func() { c.cancel(true, Canceled, nil) }//返回timerCtx和cancelfunc
}

比較好理解，所以接著往下看cancel方法

4.3.cancel

func (c *timerCtx) cancel(removeFromParent bool, err, cause error) {
	c.cancelCtx.cancel(false, err, cause)
	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()
}

這裡值得注意的地方只有第二行，它呼叫的是c.cancelCtx.cancel，取消的並非是父ctx，而是它本身，這裡的作用是為了區分cancel方法的實現，c.cancelCtx是父ctx的一個副本，並不是父ctx，所以真正的parent是c.cancelCtx.Context。

5valueCtx

type valueCtx struct {
	Context
	key, val any
}

valueCtx內嵌了Context介面所以擁有該介面的所有方法，以及新增了k-v pair。

5.1WithValue

// 提供的鍵必須可比較，並且不應是字串或任何其他內建型別，以避免在使用上下文時與其他包發生衝突。使用WithValue的使用者應為其鍵定義自己的型別。為了避免在將值賦給介面{}時分配記憶體，上下文鍵通常具有具體的型別struct{}。或者，應將匯出的上下文鍵變數的靜態型別設定為指標或介面。
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() { //必須保證key是可比較的
		panic("key is not comparable")
	}
	return &valueCtx{parent, key, val}
}

5.2value

func value(c Context, key any) any {
	for {
		switch ctx := c.(type) { //對context進行型別斷言
		case *valueCtx:
			if key == ctx.key { //如果key就是當前ctx的key，則直接返回val
				return ctx.val
			}
			c = ctx.Context //否則向父ctx查詢
		case *cancelCtx:
			if key == &cancelCtxKey {
				return c
			}
			c = ctx.Context
		case withoutCancelCtx:
			if key == &cancelCtxKey {
				// This implements Cause(ctx) == nil
				// when ctx is created using WithoutCancel.
				return nil
			}
			c = ctx.c
		case *timerCtx:
			if key == &cancelCtxKey {
				return &ctx.cancelCtx
			}
			c = ctx.Context
		case backgroundCtx, todoCtx:
			return nil //不存在key value
		default:
			return c.Value(key) //返回使用者實現的Value
		}
	}
}

可以看見，valueCtx查詢value的過程，類似於連結串列查詢，它是自底向上的，並且時間複雜度為O（N），它並不適用於存放大量的kv資料，原因有以下：

• 線性時間複雜度O（N），耗時太長
• 一個 valueCtx 例項只能存一個 kv 對，因此 n 個 kv 對會巢狀 n 個 valueCtx，造成空間浪費
• 不支援基於 k 的去重，相同 k 可能重複存在，並基於起點的不同，返回不同的 v. 由此得知，valueContext 的定位類似於請求頭，只適合存放少量作用域較大的全域性 meta 資料.

感謝觀看，參考部落格：

Golang context 實現原理 (qq.com)

深入Go：Context-騰訊雲開發者社群-騰訊雲 (tencent.com)

Go context的使用和原始碼分析_&cancelctxkey-CSDN部落格