木鸟杂记

分布式系统，数据库，存储

Golang 笔记（二）：Context 源码剖析

发表于 更新于 分类于 ，

go-context-tree-construction.png

概述

Context 是 Go 中一个比较独特而常用的概念，用好了往往能事半功倍。但如果不知其然而滥用，则往往变成 “为赋新词强说愁”，轻则影响代码结构，重则埋下许多bug。

Golang 使用树形派生的方式构造 Context，通过在不同过程 [1] 中传递 deadline 和 cancel 信号，来管理处理某个任务所涉及到的一组 goroutine 的生命周期，防止 goroutine 泄露。并且可以通过附加在 Context 上的 Value 来传递/共享一些跨越整个请求间的数据。

Context 最常用来追踪 RPC/HTTP 等耗时的、跨进程的 IO 请求的生命周期，从而让外层调用者可以主动地或者自动地取消该请求，进而告诉子过程回收用到的所有 goroutine 和相关资源。

Context 本质上是一种在 API 间树形嵌套调用时传递信号的机制。本文将从接口、派生、源码分析、使用等几个方面来逐一解析 Context。

作者：木鸟杂记 https://www.qtmuniao.com/2020/07/12/go-context/, 转载请注明出处

Context 接口

Context 接口如下：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
// Context 用以在多 API 间传递 deadline、cancelation 信号和请求的键值对。
// Context 中的方法能够安全的被多个 goroutine 并发调用。
type Context interface {
    // Done 返回一个只读 channel，该 channel 在 Context 被取消或者超时时关闭
    Done() <-chan struct{}

    // Err 返回 Context 结束时的出错信息
    Err() error

    // 如果 Context 被设置了超时，Deadline 将会返回超时时限。
    Deadline() (deadline time.Time, ok bool)

    // Value 返回关联到相关 Key 上的值，或者 nil.
    Value(key interface{}) interface{}
}

上面是简略注释，接口详细信息可以访问 Context 的 godoc。

  • Done() 方法返回一个只读的 channel，当 Context 被主动取消或者超时自动取消时，该 Context 所有派生 Context 的 done channel 都被 close 。所有子过程通过该字段收到 close 信号后，应该立即中断执行、释放资源然后返回。

  • Err() 在上述 channel 被 close 前会返回 nil，在被 close 后会返回该 Context 被关闭的信息，error 类型，只有两种，被取消或者超时：

1
2
var Canceled = errors.New("context canceled")
var DeadlineExceeded error = deadlineExceededError{}

  • Deadline() 如果本 Context 被设置了时限，则该函数返回 ok=true 和对应的到期时间点。否则，返回 ok=false和 nil。

  • Value() 返回绑定在该 Context 链（我称为回溯链，下文会展开说明）上的给定的 Key 的值，如果没有，则返回 nil。注意，不要用于在函数中传参，其本意在于共享一些横跨整个 Context 生命周期范围的值。Key 可以是任何可比较的类型。为了防止 Key 冲突，最好将 Key 的类型定义为非导出类型，然后为其定义访问器。看一个通过 Context 共享用户信息的例子：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
package user

import "context"

// User 是要存于 Context 中的 Value 类型.
type User struct {...}

// key 定义为了非导出类型，以避免和其他 package 中的 key 冲突
type key int

// userKey 是 Context 中用来关联 user.User 的 key，是非导出变量
// 客户端需要用 user.NewContext 和 user.FromContext 构建包含
// user 的 Context 和从 Context 中提取相应 user 
var userKey key

// NewContext 返回一个带有用户值 u 的 Context.
func NewContext(ctx context.Context, u *User) context.Context {
  return context.WithValue(ctx, userKey, u)
}

// FromContext 从 Context 中提取 user，如果有的话.
func FromContext(ctx context.Context) (*User, bool) {
  u, ok := ctx.Value(userKey).(*User)
  return u, ok
}

Context 派生

Context 设计之妙在于可以从已有 Context 进行树形派生，以管理一组过程的生命周期。我们上面说了单个 Context 实例是不可变的，但可以通过 context 包提供的三种方法：WithCancel 、 WithTimeout 和 WithValue 来进行派生并附加一些属性（可取消、时限、键值），以构造一组树形组织的 Context。

go-context-tree.png

当根 Context 结束时，所有由其派生出的 Context 也会被一并取消。也就是说，父 Context 的生命周期涵盖所有子 Context 的生命周期。

context.Background() 通常用作根节点，它不会超时，不能被取消。

1
2
3
// Background 返回一个空 Context。它不能被取消，没有时限，没有附加键值。Background 通常用在
// main函数、init 函数、test 入口，作为某个耗时过程的根 Context。
func Background() Context

WithCancel 和 WithTimeout 可以从父 Context 进行派生，返回受限于父 Context 生命周期的新 Context。

通过 WithCancel 从 context.Background() 派生出的 Context 要注意在对应过程完结后及时 cancel，否则会造成 Context 泄露。

使用 WithTimeout 可以控制某个过程的处理时限。具体过程为，到点后， Context 发送信号到 Done Channel，子过程检测到 Context Done Channel [2] 中的信号，会立即退出。

1
2
3
4
5
6
7
8
9
10
11
// WithCancel 返回一份父 Context 的拷贝，和一个 cancel 函数。当父 Context 被关闭或者 
// 此 cancel 函数被调用时，该 Context 的 Done Channel 会立即被关闭.
func WithCancel(parent Context) (ctx Context, cancel CancelFunc)

// 调用 CancelFunc 取消对应 Context.
type CancelFunc func()

// WithTimeout 返回一份父 Context 的拷贝，和一个 cancel 函数。当父 Context 被关闭、
// cancel 函数被调用或者设定时限到达时，该 Context 的 Done Channel 会立即关闭。在 cancel 函数
// 被调用时，如果其内部 timer 仍在运行，将会被停掉。
func WithTimeout(parent Context, timeout time.Duration) (Context, CancelFunc)

WithValue 可以给 Context 附加上整个处理过程中的键值。

1
2
// WithValue 返回一个父 Context 的副本，并且附加上给定的键值对.
func WithValue(parent Context, key interface{}, val interface{}) Context

Context 源码解析

Go context 使用嵌入类，以类似继承的方式组织几个 Context 类：emptyCtx、valueCtx、 cancelCtx、timerCtx。

go-context-implementation.png

形象的来说，通过嵌入的方式，Go 对树形组织的 Context 体系中的每个 Context 节点都构造了一个指向父亲实例”指针”。从另一个角度来说，这是一种经典代码组织模式——组合模式，每一层只增量 or 覆盖实现自己所关注的功能，然后通过路由调用来复用已有的实现。

空实现 emptyCtx

emptyCtx 实现了一个空的 Context，所有接口函数都是空实现。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
type emptyCtx int

func (*emptyCtx) Deadline() (deadline time.Time, ok bool) {
  return
}
func (*emptyCtx) Done() <-chan struct{} {
  return nil // 返回 nil，从语法上说是空实现，从语义上说是该 Context 永远不会被关闭。
}
//... 其他的省略，类似都是满足语法要求的空函数体

func (e *emptyCtx) String() string {
  switch e {
  case background:
    return "context.Background"
  case todo:
    return "context.TODO"
  }
  return "unknown empty Context"
}

context.Background() 和 context.TODO() 返回的都是 emptyCtx 的实例。但其语义略有不同。前者做为 Context 树的根节点，后者通常在不知道用啥时用。

1
2
3
4
5
6
7
8
9
10
11
12
var (
  background = new(emptyCtx)
  todo       = new(emptyCtx)
)

func Background() Context {
  return background
}

func TODO() Context {
  return todo
}

附加单键值 valueCtx

valueCtx 嵌入了一个 Context 接口以进行 Context 派生，并且附加了一个 KV 对。从 context.WithValue 函数可以看出，每附加一个键值对，都得套上一层新的 valueCtx。在使用 Value(key interface) 接口访问某 Key 时，会沿着 Context 树回溯链不断向上遍历所有 Context 直到 emptyCtx：

  1. 如果遇到 valueCtx 实例，则比较其 key 和给定 key 是否相等
  2. 如果遇到其他 Context 实例，就直接向上转发。但这里有个特例，为了获取给定 Context 所有祖先节点中最近的cancelCtx，go 用了一个特殊的 key：cancelCtxKey，遇到该 key 时，cancelCtx 会返回自身。这个在 cancelCtx 实现中会提到。

对于其他的接口调用（Done, Err, Deadline），会路由到嵌入的 Context 上去。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
type valueCtx struct {
  Context // 嵌入，指向父 Context
  key, val interface{}
}

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

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} // 附加上 kv，并引用父 Context
}

可取消的 cancelCtx

context 包中核心实现在 cancelCtx 中，包括构造树形结构、进行级联取消。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
type cancelCtx struct {
  Context

  mu       sync.Mutex            // 保证下面三个字段的互斥访问
  done     chan struct{}         // 惰式初始化，被第一个 cancel() 调用所关闭
  children map[canceler]struct{} // 被第一个 cancel() 调用置 nil
  err      error                 // 被第一个 cancel() 调用置非 nil
}

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

func (c *cancelCtx) Done() <-chan struct{} {
  c.mu.Lock()
  if c.done == nil {
    c.done = make(chan struct{})
  }
  d := c.done
  c.mu.Unlock()
  return d
}

Value() 函数的实现有点意思，遇到特殊 key：cancelCtxKey 时，会返回自身。这个其实是复用了 Value 函数的回溯逻辑，从而在 Context 树回溯链中遍历时，可以找到给定 Context 的第一个祖先 cancelCtx 实例。

children 保存的是子树中所有路径向下走的第一个可以 cancel 的 Context (实现了 canceler 接口，比如 cancelCtx 或 timerCtx 节点)，可以参考后面的图来形象理解。

下面将逐一详细说明。

回溯链

回溯链是各个 context 包在实现时利用 go 语言嵌入（embedding）的特性来构造的，主要用于：

  1. Value() 函数被调用时沿着回溯链向上查找匹配的键值对。
  2. 复用 Value() 的逻辑查找最近 cancelCtx 祖先，以构造 Context 树。

在 valueCtx、cancelCtx、timerCtx 中只有 cancelCtx 直接（valueCtx 和 timerCtx 都是通过嵌入实现，调用该方法会直接转发到 cancelCtx 或者 emptyCtx ）实现了非空 Done() 方法，因此 done := parent.Done() 会返回第一个祖先 cancelCtx 中的 done channel。但如果 Context 树中有第三方实现的 Context 接口的实例时，parent.Done() 就有可能返回其他 channel。

因此，如果 p.done != done ，说明在回溯链中遇到的第一个实现非空 Done() Context 是第三方 Context ，而非 cancelCtx。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
// parentCancelCtx 返回 parent 的第一个祖先 cancelCtx 节点
func parentCancelCtx(parent Context) (*cancelCtx, bool) {
  done := parent.Done() // 调用回溯链中第一个实现了 Done() 的实例(第三方Context类/cancelCtx)
  if done == closedchan || done == nil {
    return nil, false
  }
  p, ok := parent.Value(&cancelCtxKey).(*cancelCtx) // 回溯链中第一个 cancelCtx 实例
  if !ok {
    return nil, false
  }
  p.mu.Lock()
  ok = p.done == done
  p.mu.Unlock()
  if !ok { // 说明回溯链中第一个实现 Done() 的实例不是 cancelCtx 的实例
    return nil, false
  }
  return p, true
}

树构建

Context 树的构建是在调用 context.WithCancel() 调用时通过 propagateCancel 进行的。

1
2
3
4
5
func WithCancel(parent Context) (ctx Context, cancel CancelFunc) {
  c := newCancelCtx(parent)
  propagateCancel(parent, &c)
  return &c, func() { c.cancel(true, Canceled) }
}

Context 树，本质上可以细化为 canceler （*cancelCtx 和 *timerCtx）树，因为在级联取消时只需找到子树中所有的 canceler ，因此在实现时只需在树中保存所有 canceler 的关系即可（跳过 valueCtx），简单且高效。

1
2
3
4
5
6
// A canceler is a context type that can be canceled directly. The
// implementations are *cancelCtx and *timerCtx.
type canceler interface {
  cancel(removeFromParent bool, err error)
  Done() <-chan struct{}
}

具体实现为，沿着回溯链找到第一个实现了 Done() 方法的实例，

  1. 如果为 canceler 的实例，则其必有 children 字段，并且实现了 cancel 方法（canceler），将该 context 放进 children 数组即可。此后，父 cancelCtx 在 cancel 时会递归遍历所有 children，逐一 cancel。
  2. 如果为非 canceler 的第三方 Context 实例，则我们不知其内部实现，因此只能为每个新加的子 Context 启动一个守护 goroutine，当 父 Context 取消时，取消该 Context。

需要注意的是，由于 Context 可能会被多个 goroutine 并行访问，因此在更改类字段时，需要再一次检查父节点是否已经被取消，若父 Context 被取消，则立即取消子 Context 并退出。 

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
func propagateCancel(parent Context, child canceler) {
  done := parent.Done()
  if done == nil {
    return // 父节点不可取消
  }

  select {
  case <-done:
    // 父节点已经取消
    child.cancel(false, parent.Err())
    return
  default:
  }

  if p, ok := parentCancelCtx(parent); ok { // 找到一个 cancelCtx 实例
    p.mu.Lock()
    if p.err != nil {
      // 父节点已经被取消
      child.cancel(false, p.err)
    } else {
      if p.children == nil {
        p.children = make(map[canceler]struct{}) // 惰式创建
      }
      p.children[child] = struct{}{}
    }
    p.mu.Unlock()
  } else {                                // 找到一个非 cancelCtx 实例
    atomic.AddInt32(&goroutines, +1)
    go func() {
      select {
      case <-parent.Done():
        child.cancel(false, parent.Err())
      case <-child.Done(): 
      }
    }()
  }
}

下面用一张图来解释下回溯链和树组织， C0 是 emptyCtx，通常由 context.Background() 得来，作为 Context 树的根节点。C1~C4 依次通过嵌入的方式从各自父节点派生而来。图中的虚线是由嵌入（embedded）而构成的回溯链，实线是由 cancelCtx children 数组而保存的父子关系。

parentCancelCtx(C2) 和 parentCancelCtx(C4) 都为 C1，则 C1 的 children 数组中保存的为 C2 和 C4。构建了这两层关系后，就可以沿着回溯链向上查询 Value 值，包括找到第一个祖先 cancelCtx；也可以沿着 children 关系往下进行级联取消。

go-context-tree-construction.png

当然，图中所有 Context 都是针对 go 包中的系统 Context，没有画出有第三方 Context 的情况。而实际代码由于增加了对第三方 Context 的处理逻辑，稍微难懂一些。区分系统 Context 实现和用户自定义 Context 的关键点在于是否实现了 canceler 接口。

第三方 Context 实现了此接口就可以进行树形组织，并且在上游 cancelCtx 取消时，递归调用 children 的 cancel 进行级联取消。否则只能通过为每个第三方 Context 启动一个 goroutine 来监听上游取消事件，以对第三方 Context 进行取消了。

级联取消

下面是级联取消中的关键函数 cancelCtx.cancel 的实现。在本 cancelCtx 取消时，需要级联取消以该 cancelCtx 为根节点的 Context 树中的所有 Context，并将根 cancelCtx 从其从父节点中摘除，以让 GC 回收该 cancelCtx 子树所有节点的资源。

cancelCtx.cancel 是非导出函数，不能在 context 包外调用，因此持有 Context 的内层过程不能自己取消自己，须由返回的 CancelFunc （简单的包裹了cancelCtx.cancel ）来取消，其句柄一般为外层过程所持有。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
func (c *cancelCtx) cancel(removeFromParent bool, err error) {
  if err == nil { // 需要给定取消的理由，Canceled or DeadlineExceeded
    panic("context: internal error: missing cancel error")
  }
  
  c.mu.Lock()
  if c.err != nil {
    c.mu.Unlock()
    return // 已经被其他 goroutine 取消
  }
  
  // 记下错误，并关闭 done
  c.err = err
  if c.done == nil {
    c.done = closedchan
  } else {
    close(c.done)
  }
  
  // 级联取消
  for child := range c.children {
    // NOTE: 持有父 Context 的同时获取了子 Context 的锁
    child.cancel(false, err)
  }
  c.children = nil
  c.mu.Unlock()

  // 子树根需要摘除，子树中其他节点则不再需要
  if removeFromParent {
    removeChild(c.Context, c)
  }
}

timerCtx

timerCtx 在嵌入 cancelCtx 的基础上增加了一个计时器 timer，根据用户设置的时限，到点取消。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
type timerCtx struct {
  cancelCtx
  timer *time.Timer // Under cancelCtx.mu

  deadline time.Time
}

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

func (c *timerCtx) cancel(removeFromParent bool, err error) {
  // 级联取消子树中所有 Context
  c.cancelCtx.cancel(false, err)
  
  if removeFromParent {
    // 单独调用以摘除此节点，因为是摘除 c，而非 c.cancelCtx
    removeChild(c.cancelCtx.Context, c)
  }
  
  // 关闭计时器
  c.mu.Lock()
  if c.timer != nil {
    c.timer.Stop()
    c.timer = nil
  }
  c.mu.Unlock()
}

设置超时取消是在 context.WithDeadline() 中完成的。如果祖先节点时限早于本节点，只需返回一个 cancelCtx 即可，因为祖先节点到点后在级联取消时会将其取消。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
func WithDeadline(parent Context, d time.Time) (Context, CancelFunc) {
  if cur, ok := parent.Deadline(); ok && cur.Before(d) {
    // 祖先节点的时限更早
    return WithCancel(parent)
  }
  
  c := &timerCtx{             
    cancelCtx: newCancelCtx(parent), // 使用一个新的 cancelCtx 实现部分 cancel 功能
    deadline:  d,
  }
  propagateCancel(parent, c) // 构建 Context 取消树，注意传入的是 c 而非 c.cancelCtx
  dur := time.Until(d)       // 测试时限是否设的太近以至于已经结束了
  if dur <= 0 {
    c.cancel(true, DeadlineExceeded)
    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) }
}

Context 使用

使用了 Context 的子过程须保证在 Context 被关闭时及时退出并释放资源。也就是说，使用 Context 需要遵循上述原则才能保证级联取消时释放资源的效果。因此，Context 本质上是一种树形分发信号的机制，可以用 Context 树追踪过程调用树，当外层过程取消时，使用 Context 级联通知所有被调用过程。

以下是一个典型子过程的检查 Context 以确定是否需要退出的代码片段：

1
2
3
4
5
6
7
8
9
for ; ; time.Sleep(time.Second) {
    select {
    case <-context.Done():
      return
    default:
    }
  
    // 一些耗时操作
}

可以看出，Context 接口本身并没有 Cancel 方法，这和 Done() 返回的 channel 是只读的是一个道理：Context 关闭信号的发送方和接收方通常不在一个函数中。比如，当父 goroutine 启动了一些子 goroutine 来干活时，只能是父 goroutine 来关闭 done channel，子 goroutine 来检测 channel 的关闭信号。即不能在子 goroutine 中 取消父 goroutine 中传递过来的 Context。

Context 注意

Context 有一些使用实践需要遵循：

  1. Context 通常作为函数中第一个参数
  2. 不要在 struct 中存储 Context，每个函数都要显式的传递 Context。不过实践中可以根据 struct 的生命周期来灵活组合。
  3. 不要使用 nil Context，尽管语法上允许。不知道使用什么值合适时，可以使用 context.TODO() 。
  4. Context value 是为了在请求生命周期中共享数据，而非作为函数中传递额外参数的方法。因为这是一种隐式的语义，极易造成 bug；要想传额外参数，还是要在函数中显式声明。
  5. Context 是 immutable 的，因此是线程安全的，可以在多个 goroutine 中传递并使用同一个 Context。

注

[1] 文中的过程，指的是计算密集型或者 IO 密集型的耗时函数，或者 goroutine。

[2] Context 的 Done Channel，指的是 context.Done() 返回的 channel。它是 Context 内的关键数据结构，作为沟通不同过程的的渠道。需要结束时，父过程向该 channel 发送信号，子过程读取该 channel 信号后做扫尾工作并且退出。

参考

  1. go doc context：https://golang.org/pkg/context/
  2. code review conmments: https://github.com/golang/go/wiki/CodeReviewComments#contexts
  3. go blog context：https://blog.golang.org/context
  4. go context 源码：https://golang.org/src/context/context.go?s=8419:8483#L222
  5. go 语言设计与实现： https://draveness.me/golang/docs/part3-runtime/ch06-concurrency/golang-context/

我是青藤木鸟，一个喜欢摄影、专注大规模数据系统的程序员，欢迎关注我的公众号：“木鸟杂记”，有更多的分布式系统、存储和数据库相关的文章，欢迎关注。 关注公众号后，回复“资料”可以获取我总结一份分布式数据库学习资料。 回复“优惠券”可以获取我的大规模数据系统付费专栏《系统日知录》的八折优惠券。

我们还有相关的分布式系统和数据库的群，可以添加我的微信号：qtmuniao，我拉你入群。加我时记得备注：“分布式系统群”。 另外，如果你不想加群，还有一个分布式系统和数据库的论坛（点这里），欢迎来玩耍。

wx-distributed-system-s.jpg

不妨一读