Golang context 原理分析

1. 说在前面
2. 场景分析
- 2.1 链式传递
- 2.2 主动取消
- 2.3 任务超时
- 2.4 数据存储
3. 源码解读
- 3.1 一个核心数据结构
  - 3.1.1 Context
- 3.2 四种具体实现
  - 3.2.1 emptyCtx
  - 3.2.2 cancelCtx
  - 3.2.3 timerCtx
  - 3.2.4 valueCtx
- 3.3 六个核心方法
  - 3.3.1 Background() && TODO()
  - 3.3.2 WithCancel()
  - 3.3.3 WithDeadline()
  - 3.3.4 WithTimeout()
  - 3.3.5 WithValue()
4. 一些思考
- 思考1：emptyCtx 为什么不是 struct{}类型？
- 思考2：backgound 和 todo 有什么区别？
- 思考3：cancelCtx 怎么保证父亲 👨 取消的同时取消儿子 👦？
- 思考4：valueCtx 可以用于数据存储吗？

1. 说在前面

context 是 golang 中的经典工具，主要在异步场景中用于实现并发协调控制以及对 goroutine 的生命周期管理。除此之外，context 还兼有一定的数据存储能力。本文旨在剖析 context 的核心工作原理。

本文使用到的 Go 版本为 1.18，源码位置 src/context/context.go

2. 场景分析

2.1 链式传递

在 Go 中可以认为协程的组织是一种链式传递，每一个子协程的创建都是基于父协程，但是父协程对子协程的控制则是通过 context 实现；同样的，每一个 context 也都是基于父 context 创建，最终形成链式结构，根 context 就是 emptyCtx。

2.2 主动取消

取消场景，是父协程任务取消的时候，将子协程一并取消。

在下面这个案例中，子协程的任务需要 2s 才能执行完，但是父协程 1s 后执行报错主动 cancel 任务的执行，cancel() 方法会通知子协程一并取消任务的执行。

package mainimport ("context""errors""fmt""time"
)func testCtx(ctx context.Context) {cancelCtx, cancel := context.WithCancel(ctx)defer cancel()go func(ctx context.Context) {cancelCtx, cancel := context.WithCancel(ctx)defer cancel()// do somethingselect {case <-time.After(2 * time.Second):// 假设任务需要 2s 完成fmt.Println("work done")case <-cancelCtx.Done():fmt.Println("work canceled")return}// do something}(cancelCtx)// 模拟执行过程中父协程报错取消任务<-time.After(1 * time.Second)err := errors.New("fake err")if err != nil {cancel()}
}func main() {ctx := context.Background()testCtx(ctx)<-time.After(3 * time.Second)
}

2.3 任务超时

超时场景，是父协程任务超时的时候会触发取消流程，需将子协程一并取消。

在下面这个案例中，子协程的任务需要 2s 才能执行完，但是父协程 1s 后任务超时，开始执行取消任务的流程，通知子协程一并取消任务的执行。

package mainimport ("context""fmt""time"
)func testCtx(ctx context.Context) {timerCtx, cancel := context.WithTimeout(ctx, 1*time.Second)defer cancel()go func(ctx context.Context) {timerCtx, cancel := context.WithTimeout(ctx, 2*time.Second)defer cancel()// do somethingselect {case <-time.After(2 * time.Second):// 假设任务需要 2s 完成fmt.Println("work done")case <-timerCtx.Done():fmt.Println("work canceled")return}// do something}(timerCtx)// 模拟等待子协程退出（偷懒）<-time.After(2 * time.Second)
}func main() {ctx := context.Background()testCtx(ctx)<-time.After(3 * time.Second)
}

2.4 数据存储

context 可以用于数据存储，通常是用于存储一些元数据。

package mainimport ("context""fmt""time"
)func testCtx(ctx context.Context) {valueCtx := context.WithValue(ctx, "name", "father")fmt.Println(valueCtx.Value("name"))go func(ctx context.Context) {valueCtx := context.WithValue(ctx, "name", "child")fmt.Println(valueCtx.Value("name"))}(valueCtx)
}func main() {ctx := context.Background()testCtx(ctx)<-time.After(3 * time.Second)
}

3. 源码解读

3.1 一个核心数据结构

3.1.1 Context

首先 Context 本质是官方提供的一个 interface，实现了该 interface 定义的都被称之为 context。

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

Deadline() 返回 ctx 生命终了时间，如果没有则 ok 为 false
Done() 返回一个 channel 用于判断 ctx 是否已经结束
Err() 用于 ctx 结束后获取错误信息
Value() 获取 ctx 中存入的键值对

3.2 四种具体实现

3.2.1 emptyCtx

官方实现的一个空 ctx 版本，默认都是返回空值，通常是作为所有 ctx 的根。

type emptyCtx intfunc (*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
}

3.2.2 cancelCtx

具有取消功能，父 ctx 取消的时候将所有子 ctx 一并取消。

类型定义

type cancelCtx struct {Contextmu       sync.Mutex            // 保护临界资源done     atomic.Value          // chan struct{} 类型，用于判断 ctx 是否取消，第一次调用 cancel() 后 closechildren map[canceler]struct{} // 存储所有的子 ctxerr      error                 // ctx 取消后记录错误信息
}

Done() 方法实现

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

如果 done 不为空则直接返回
加锁🔒
如果 done 为空则创建，这里也说说明一点：done 是懒加载的，第一次调用 Done() 方法才会创建 done
返回 done
解锁🔒

Value() 方法实现

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

如果传入的 key 是 cancelCtxKey 则返回自身，这里制定了一种私有协议（外部无法访问 cancelCtxKey），用于后面判断一个父 ctx 是否是一个 cancelCtx 类型
否则调用 value 方法找到 key 对应的 value

Err() 方法实现

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

加锁🔒保护，返回 err

Deadline() 方法实现

未实现，继承父 Context

3.2.3 timerCtx

具有定时取消的功能，因为是继承自 cancelCtx 所以同样具有主动取消功能

类型定义

type timerCtx struct {cancelCtxtimer *time.Timer  // timerCtx 内部维护的一个定时器deadline time.Time // ctx 的终了时间
}

Done() 方法实现

未实现，继承父 cancelCtx

Value() 方法实现

未实现，继承父 cancelCtx

Err() 方法实现

未实现，继承父 cancelCtx

Deadline() 方法实现

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

返回 dealline

3.2.4 valueCtx

具有存储数据的功能，通常是一些元数据信息

类型定义

type valueCtx struct {Contextkey, val any // 记录数据
}

Done() 方法实现

未实现，继承父 Context

Value() 方法实现

func (c *valueCtx) Value(key any) any {if c.key == key {return c.val}return value(c.Context, key)
}func value(c Context, key any) any {for {switch ctx := c.(type) {case *valueCtx:if key == ctx.key {return ctx.val}c = ctx.Contextcase *cancelCtx:if key == &cancelCtxKey {return c}c = ctx.Contextcase *timerCtx:if key == &cancelCtxKey {return &ctx.cancelCtx}c = ctx.Contextcase *emptyCtx:return nildefault:return c.Value(key)}}
}

判断传入的 key 是否等于当前 ctx 的 k，如果相等则返回
否则就从父 ctx 中找，一直到 emptyCtx 找不到就返回 nil

Err() 方法实现

未实现，继承父 Context

Deadline() 方法实现

未实现，继承父 Context

3.3 六个核心方法

3.3.1 Background() && TODO()

用于获取 emptyCtx，本质上没有区别，仅仅是语义上的区别。

var (background = new(emptyCtx)todo       = new(emptyCtx)
)func Background() Context {return background
}func TODO() Context {return todo
}

3.3.2 WithCancel()

函数说明：传入一个父 ctx 返回一个子 ctx 和一个 cancel 函数

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) }
}func newCancelCtx(parent Context) cancelCtx {return cancelCtx{Context: parent}
}

如果传入的 parent 是空的，panic
基于 parent 创建 cancelCtx
propagateCancel 用于传递 cancel 的特性，用于保证父 ctx 取消的时候，子 ctx 也取消
返回创建的 cancelCtx 和一个闭包函数，闭包函数调用 cancel 取消创建的 ctx

func propagateCancel(parent Context, child canceler) {done := parent.Done()if done == nil {return // parent is never canceled}select {case <-done:// parent is already canceledchild.cancel(false, parent.Err())returndefault:}if p, ok := parentCancelCtx(parent); ok {p.mu.Lock()if p.err != nil {// parent has already been canceledchild.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():}}()}
}

parent 是一个永远不会被取消的 ctx，直接 return
如果 parent 已经被取消了，把当前新创建的 cancelCtx 也取消，并 return
如果 parent 是一个 cancelCtx 类型则将新创建的 cancelCtx 加入到 parent 的 children 中
如果 parent 不是 cancelCtx 类型，则启动一个 goroutine 来保证 parent 取消的时候当前 cancelCtx 也被取消掉

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 = errd, _ := c.done.Load().(chan struct{})if d == nil {c.done.Store(closedchan)} else {close(d)}for child := range c.children {// NOTE: acquiring the child's lock while holding parent's lock.child.cancel(false, err)}c.children = nilc.mu.Unlock()if removeFromParent {removeChild(c.Context, c)}
}

前置判断
将 done chan 关闭
遍历 children，将所有的子 ctx 都一并取消掉
如果 parent 是 cancelCtx 类型，需要将当前 ctx 从 parent 的 children 中删除

3.3.3 WithDeadline()

函数说明：传入一个父 ctx，和一个终了时间，返回一个子 ctx 和一个 cancel 函数；timerCtx 继承自 cancelCtx，拥有 cancelCtx 的一切特性

func WithDeadline(parent Context, d time.Time) (Context, CancelFunc) {if parent == nil {panic("cannot create context from nil parent")}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 passedreturn 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) }
}

如果 parent 也是一个 timerCtx 并且传入的终了时间在 parent 的终了时间之后，那么新创建的 ctx 就没必要拥有定时特性，使用 WithCancel 构造一个 cancelCtx 返回即可
否则创建一个 timerCtx
propagateCancel 传递 cancel 的特性
如果 deadline 时间已经过了，直接 cancel 然后 return
创建一个定时任务，定时结束触发 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()
}

调用 parent 的 cancel 函数
如果 parent 是 cancelCtx 类型，需要将当前 ctx 从 parent 的 children 中删除
定时器 stop

3.3.4 WithTimeout()

仅仅对 WithDeadline 进行了简单封装

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

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

4. 一些思考

思考1：emptyCtx 为什么不是 struct{}类型？

struct{} 作为一个空类型并不占用底层存储空间，所以它的多个不同对象有可能会使用相同的地址，无法区分出 background 和 todo 对象。

思考2：backgound 和 todo 有什么区别？

本质没有区别，都是 emptyCtx，更多的是语义上的区别，background 通常作为所有 ctx 链的最顶层。

思考3：cancelCtx 怎么保证父亲 👨 取消的同时取消儿子 👦？

机制 1：cancelCtx 有一个 children 字段记录了所有的子节点，当父节点被取消的时候会给所有子节点来一刀 🔪，依次传递最终将所有子、孙子、孙孙子都刀 🔪 了。父亲取消的时候也会通知爷爷，让爷爷从 children 中删除父亲。

机制 2：如果父亲不是一个 cancelCtx 类型，则不会有 children 属性怎么办？当使用 WithCancel()创建的时候，发现父亲不是 cancelCtx 就会启动一个守护协程判断父亲是否 Done()，如果父亲 over 了，就会干掉儿子并退出；否则儿子先挂了，也会退出。

思考4：valueCtx 可以用于数据存储吗？

valueCtx 不适合视为存储介质，存放大量的 kv 数据，它的定位类似于请求头，只适合存放少量作用域较大的全局 meta 数据：一个 valueCtx 实例只能存一个 kv 对，因此 n 个 kv 对会嵌套 n 个 valueCtx，造成空间浪费；
基于 k 寻找 v 的过程是线性的，时间复杂度 O(N)；不支持基于 k 的去重，相同 k 可能重复存在，并基于起点的不同，返回不同的 v。