go Context 设计与实现

版本：go 1.19

在前一篇文章中我们讨论了 go Context 的一些常见使用方式，今天我们再来从源码的角度深入了解一下 Context 的设计与实现。

Context 的源码数量不多，去掉注释大概只有两三百行，但是包含的信息量巨大（所以本文也比较长），而且设计得非常巧妙，值得读一读。

然后，下面的 图解 propagateCancel 这一小节的几个图描述了 Context 的工作机制，如果不想看代码，可以直接拉到下面。

再了解一下 chan

在开始本文之前，先来了解一下 Context 实现的关键：chan，对于 chan（再准确一点，我们这里讨论的其实是只读 chan），我们需要清楚以下几点：

• // 创建一个 chan，类型是 struct{} ch := make(chan struct{}) go func() { select { // 这个 case 会在 chan 关闭或者收到值的时候执行， // 在这里的情况是关闭了 chan。 case v, ok := // 输出 "chan ch is closed." fmt.Println("chan ch is closed.") } // 关闭 chan 之后得到的是 ch 的零值，也就是一个空结构体实例 fmt.Println(v) // {} } }() // 关闭 chan，所有从 chan 读取的操作都会立即返回。 // 关闭 chan 之后， // 返回一个 channel，当 context 被取消或者到了 deadline 的时候， // 这个 channel 会被 close，从而 } // 在 channel Done 返回的 channel 关闭后，返回 context 取消原因。 Err() error // 返回 context 是否会被取消以及自动取消时间（即 deadline） // ok 为 true，表明设置了 deadline，第一个返回值就是设置的 deadline // ok 为 false，表示没有设置 deadline，第一个返回值没意义。 Deadline() (deadline time.Time, ok bool) // 获取 key 对应的 value Value(key interface{}) interface{} } ctx, _ := context.WithTimeout(context.Background(), time.Second) // 输出 ctx 的 deadline，具体时间为 1 秒之后 spew.Dump(ctx.Deadline()) ctx1 := context.Background() // ctx1 的超时时间是一个零值 spew.Dump(ctx1.Deadline()) // 输出： // (time.Time) 2022-11-19 11:45:38.702281 +0800 CST m=+1.000233039 // (bool) true // (time.Time) 0001-01-01 00:00:00 +0000 UTC // (bool) false } cancel(removeFromParent bool, err error) Done() } } return } func (*emptyCtx) Done() } { return nil } func (*emptyCtx) Err() error { return nil } func (*emptyCtx) Value(key any) any { return nil } func (e *emptyCtx) String() string { switch e { case background: return "context.Background" case todo: return "context.TODO" } return "unknown empty Context" } return background } func TODO() Context { return todo } // cancelCtx 也实现了 `Context` 接口 Context // mu 用以保护后面的 done、children、err 字段 mu sync.Mutex // 是一个 chan struct{}，懒汉式创建， // 在第一次 cancel 的时候被关闭 done atomic.Value // 记录所有可以取消的子 Context // 在第一次 cancel 的时候会被设置为 nil。 children map[canceler]struct{} // 在第一次 cancel 的时候会被设置为非 nil 的值 err error } } { // 如果 done 这个 chan 已经初始化了，就直接返回。 d := c.done.Load() if d != nil { return d.(chan struct{}) } // 如果 done 还没初始化，则会进行初始化。 // 也就是上面说的 "懒汉式" 的创建方式，只有在需要的时候才会初始化。 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{}) } // 使用 mu 保证并发安全，本质是 return c.err c.mu.Lock() err := c.err c.mu.Unlock() return err } // 必须传递一个 err 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() 的时候会返回这个原因 c.err = err // 关闭 done 这个通道，通知其他协程 d, _ := c.done.Load().(chan struct{}) if d == nil { c.done.Store(closedchan) } else { close(d) } // 遍历它的所有子结点，并对其子结点进行取消操作 for child := range c.children { child.cancel(false, err) } // 将子结点置空 c.children = nil c.mu.Unlock() // 从父结点中移除自己 if removeFromParent { removeChild(c.Context, c) } } // 必须从其他 Context 派生 if parent == nil { panic("cannot create context from nil parent") } c := newCancelCtx(parent) // 将 c 挂靠到 parent 的 children 属性中， // 从而在 parent 取消的时候，可以感知得到。 propagateCancel(parent, &c) // 具体实现后面有详细说明 return &c, func() { c.cancel(true, Canceled) } } // 创建一个 cancelCtx 实例 func newCancelCtx(parent Context) cancelCtx { return cancelCtx{Context: parent} } // 判断 parent 的 done 是否已经关闭或者并没有 done chan。 // 如果是，则返回 nil 和 false done := parent.Done() // done 为 nil 表示 parent 或者到根 Context 这条路径上并没有 *cancelCtx（只有 valueCtx 或 emptyCtx）。 if done == closedchan || done == nil { return nil, false } // 判断 parent 是否是一个 cancelCtx // 如果不是，则返回 nil 和 false // 讲道理，parent.Value(&cancelCtxKey) 的返回值只有两种情况： // emptyCtx（找到根节点也没找到）或者 *cancelCtx（找到了） // （parent.Value 实现细节见下面的 value 那一小节） p, ok := parent.Value(&cancelCtxKey).(*cancelCtx) if !ok { return nil, false } // 执行到这里的时候：p 是一个 *cancelCtx // 判断 parent.Done() 和 p.done 是否相等： // 不等则意味着 *cancelCtx 已经被包装在自定义实现中了，这个时候，我们不应该绕过它。 // 详细请参考：go issue 28728（google） pdone, _ := p.done.Load().(chan struct{}) if pdone != done { return nil, false } return p, true } // 从父结点移除自己（从 parent 移除 child） func removeChild(parent Context, child canceler) { p, ok := parentCancelCtx(parent) if !ok { return } // 从父结点的 children 中移除 child p.mu.Lock() if p.children != nil { delete(p.children, child) } p.mu.Unlock() } // 如果 Context 树上完全不存在 cancelCtx，则直接返回 done := parent.Done() if done == nil { return // parent is never canceled } // 如果 parent 已经取消，则直接取消 child select { case // 找到了，但是已经取消了，则取消 child p.mu.Lock() if p.err != nil { // parent has already been canceled child.cancel(false, p.err) } else { // 找到了，尚未取消。 // 将 child 写入到 p 的 children 属性中。 // p.children 是懒汉式创建的。 if p.children == nil { p.children = make(map[canceler]struct{}) } p.children[child] = struct{}{} } p.mu.Unlock() } else { // 执行到这里的原因是： // 用户自定义了 Done 通道（跟 parent 不是同一个 done）， // 所以不能以父节点路径上的 done 来决定 child 是否取消， // 需要通过启动新协程的方式来监听 Done 通道，从而可以正常取消 parent 的孩子节点。 atomic.AddInt32(&goroutines, +1) go func() { select { case Context ch } } func (a *A) Done() } { return a.ch } func TestCancel(t *testing.T) { // 创建一个 cancel context ctx, cancel0 := WithCancel(TODO()) // 创建一个 A 实例 // 这个实例可以内嵌了 Context，所以可以当作 Context 使用， // 但是我们覆盖了 Context 本身的 Done 方法。 ch := make(}) a := A{ctx, ch} // 因为我们覆盖了 Done 方法，所以 go 底层会认为是开发者想要 // 自行控制协程取消，所以在 WithCancel 的时候并不会把 ctx1 // 挂载到 a 的 children 属性下，这样一来， // go 底层只能再启动一个协程来监听 a 的 Done chan， // 从而可以在 a cancel 的时候可以正常通知到 ctx1。 ctx1, cancel := WithCancel(&a) go func() { time.Sleep(time.Millisecond * 10) cancel() }() // ctx2 会写入到 ctx1 的 children 属性中， // 这样就不需要启动新的协程来监测 ctx1 的 done。 ctx2, cancel2 := WithCancel(ctx1) time.Sleep(time.Millisecond * 20) } cancelCtx timer *time.Timer // Under cancelCtx.mu. deadline time.Time } // 不能基于 nil 创建一个 timerCtx if parent == nil { panic("cannot create context from nil parent") } // 如果当前设置的超时时间比 parent 设置的超时时间更长， // 那么不用 timerCtx 开启定时器了，因为 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, } // 将刚创建的 timerCtx 挂载到父 Context 中 propagateCancel(parent, c) // 判断还有多久到达 deadline dur := time.Until(d) // 如果 deadline 已经过去了，那么直接执行 timerCtx 的 cancel 逻辑, // 同时移除跟父节点的关联。（创建了还没来得及启动定时器就到期了） if dur c.cancel(true, DeadlineExceeded) // deadline has already passed // cancel 不再需要从父节点移除自身，上一行已经移除了 return c, func() { c.cancel(false, Canceled) } } // 启动一个定时器，在到达 deadline 的时候执行 cancel 操作。 c.mu.Lock() defer c.mu.Unlock() if c.err == nil { c.timer = time.AfterFunc(dur, func() { c.cancel(true, DeadlineExceeded) }) } // 除了定时器之外，返回一个 CancelFunc 给用户提供自行取消的方式。 return c, func() { c.cancel(true, Canceled) } } return WithDeadline(parent, time.Now().Add(timeout)) } Context key, val any } // parent 不能为 nil if parent == nil { panic("cannot create context from nil parent") } // key 不能为 nil if key == nil { panic("nil key") } // key 必须是可以比较的， // 因为在获取值的时候需要进行 key 的比较。 if !reflectlite.TypeOf(key).Comparable() { panic("key is not comparable") } // 返回 valueCtx return &valueCtx{parent, key, val} } Context key, val any } for { switch ctx := c.(type) { case *valueCtx: if key == ctx.key { return ctx.val } c = ctx.Context case *cancelCtx: if key == &cancelCtxKey { return c } c = ctx.Context case *timerCtx: if key == &cancelCtxKey { return &ctx.cancelCtx } c = ctx.Context case *emptyCtx: return nil default: return c.Value(key) } } }