Go 组件 Context 源码学习笔记
为什么要使用 Context?
Go 强大而简洁的并发能力
要论 go 最津津乐道的功能莫过于 go 强大而简洁的并发能力。
func main(){ go func(){ fmt.Println("Hello World") }() }
通过简单的 go func(){},go 可以快速生成新的协程并运行。
想象一个没有 context 的世界
有并发的地方就有江湖。每个编程语言都有各自的并发编程方式,也有不同的并发控制方法,比如 java 通过 join()来做主子线程同步。go 里面常用于协程间通信和管理的有 channel 和 sync 包。比如 channel 可以通知协程做特定操作(退出,阻塞等),sync 可以加锁和同步。
假如我要实现一个可以同时关闭所有协程的程序,可以这样实现:
closed := make(chan struct{}) for i := 0; i < 2; i++ { // do something go func(i int) { select { case <-closed: fmt.Printf("%d Closed\n", i) } }(i) } // 发送指令关闭所有协程 close(closed) time.Sleep(1 * time.Second)
因为 go 的协程不支持直接从外部退出,不像 C++和 Java 有个线程 ID 可以操作。所以只能通过协程自己退出的方式。一般来说通过 channel 来控制是最方便的。
如果我想加点功能,比如到时间后退出,只要给 channel 增加关闭条件即可:
closed := make(chan struct{}) for i := 0; i < 2; i++ { go func(i int) { // do something select { case <-closed: fmt.Printf("%d Timeout\n", i) } }(i) } // 加个时间条件 ta := time.After(5 * time.Second) select { case <-ta: close(closed) } time.Sleep(1 * time.Second)
用 Context 精简代码
上面的代码已经够简单了,但是还是显得有些复杂。比如每次都要在协程内部增加对 channel 的判断,也要在外部设置关闭条件。试想一下,如果程序要限制的是总时长,而不是单个操作的时长,这样每个操作要限制多少时间也是个难题。
这个时候就轮到 Context 登场了。Context 顾名思义是协程的上下文,主要用于跟踪协程的状态,可以做一些简单的协程控制,也能记录一些协程信息。
下面试着用 Context 改造下前面的例子。
// 空的父context pctx := context.TODO() // 子context(携带有超时信息),cancel函数(可以主动触发取消) //ctx, cancel := context.WithTimeout(pctx, 5*time.Second) ctx, _ := context.WithTimeout(pctx, 5*time.Second) for i := 0; i < 2; i++ { go func(i int) { // do something // 大部分工具库内置了对ctx的判断,下面的部分几乎可以省略 select { case <-ctx.Done(): fmt.Printf("%d Done\n", i) } }(i) } // 调用cancel会直接关闭ctx.Done()返回的管道,不用等到超时 //cancel() time.Sleep(6 * time.Second)
通过 Context 可以进一步简化控制代码,且更为友好的是,大多数 go 库,如 http、各种 db driver、grpc 等都内置了对 ctx.Done()的判断,我们只需要将 ctx 传入即可。
Context 基础用法
接下来介绍 Context 的基础用法,最为重要的就是 3 个基础能力,取消、超时、附加值。
新建一个 Context
ctx := context.TODO() ctx := context.Background()
这两个方法返回的内容是一样的,都是返回一个空的 Context,这个 Context 一般用来做父 Context。
WithCancel
// 函数声明 func WithCancel(parent Context) (ctx Context, cancel CancelFunc) // 用法:返回一个子Context和主动取消函数 ctx, cancel := context.WithCancel(parentCtx)
这个函数相当重要,会根据传入的 Context 生成一个子 Context 和一个取消函数。当父 Context 有相关取消操作,或者直接调用 cancel 函数的话,子 Context 就会被取消。
举个日常业务中常用的例子:
// 一般操作比较耗时或者涉及远程调用等,都会在输入参数里带上一个ctx,这也是公司代码规范里提倡的 func Do(ctx context.Context, ...) { ctx, cancel := context.WithCancel(parentCtx) // 实现某些业务逻辑 // 当遇到某种条件,比如程序出错,就取消掉子Context,这样子Context绑定的协程也可以跟着退出 if err != nil { cancel() } }
WithTimeout
// 函数声明 func WithTimeout(parent Context, timeout time.Duration) (Context, CancelFunc) // 用法:返回一个子Context(会在一段时间后自动取消),主动取消函数 ctx := context.WithTimeout(parentCtx, 5*time.Second)
这个函数在日常工作中使用得非常多,简单来说就是给 Context 附加一个超时控制,当超时 ctx.Done()返回的 channel 就能读取到值,协程可以通过这个方式来判断执行时间是否满足要求。
举个日常业务中常用的例子:
// 一般操作比较耗时或者涉及远程调用等,都会在输入参数里带上一个ctx,这也是公司代码规范里提倡的 func Do(ctx context.Context, ...) { ctx, cancel := context.WithTimeout(parentCtx) // 实现某些业务逻辑 for { select { // 轮询检测是否已经超时 case <-ctx.Done(): return // 有时也会附加一些错误判断 case <-errCh: cancel() default: } } }
现在大部分 go 库都实现了超时判断逻辑,我们只需要传入 ctx 就好。
WithDeadline
// 函数声明 func WithDeadline(parent Context, d time.Time) (Context, CancelFunc) // 用法:返回一个子Context(会在指定的时间自动取消),主动取消函数 ctx, cancel := context.WithDeadline(parentCtx, time.Now().Add(5*time.Second))
这个函数感觉用得比较少,和 WithTimeout 相比的话就是使用的是截止时间。
WithValue
// 函数声明 func WithValue(parent Context, key, val interface{}) Context // 用法: 传入父Context和(key, value),相当于存一个kv ctx := context.WithValue(parentCtx, "name", 123) // 用法:将key对应的值取出 v := ctx.Value("name")
这个函数常用来保存一些链路追踪信息,比如 API 服务里会有来保存一些来源 ip、请求参数等。
因为这个方法实在是太常用了,比如 grpc-go 里的 metadata 就使用这个方法将结构体存储在 ctx 里。
func NewOutgoingContext(ctx context.Context, md MD) context.Context { return context.WithValue(ctx, mdOutgoingKey{}, rawMD{md: md}) }
Context 源码实现
Context 是一个接口
虽然我们平时写代码时直接 context.Context 拿来就用,但实际上 context.Context 是一个接口,源码里是有多种不同的实现的,借此实现不同的功能。
type Context interface { // 返回这个ctx预期的结束时间 Deadline() (deadline time.Time, ok bool) // 返回一个channel,当执行结束或者取消时被close,我们平时可以用这个来判断ctx绑定的协程是否该退出。实现里用的懒汉模式,所以一开始可能会返回nil Done() <-chan struct{} // 如果未完成,返回nil。已完成源码里目前就两种错误,已被取消或者已超时 Err() error // 返回ctx绑定的key对应的value值 Value(key interface{}) interface{} }
Context 们是一棵树
context 整体是一个树形结构,不同的 ctx 间可能是兄弟节点或者是父子节点的关系。
同时由于 Context 接口有多种不同的实现,所以树的节点可能也是多种不同的 ctx 实现。总的来说我觉得 Context 的特点是:
- 树形结构,每次调用 WithCancel, WithValue, WithTimeout, WithDeadline 实际是为当前节点在追加子节点;
- 继承性,某个节点被取消,其对应的子树也会全部被取消;
- 多样性,节点存在不同的实现,故每个节点会附带不同的功能。
Context 的果子们
在源码里实际只有 4 种实现,要弄懂 context 的源码其实把这 4 种对应的实现学习一下就行,他们分别是:
- emptyCtx:一个空的 ctx,一般用于做根节点;
- cancelCtx:核心,用来处理取消相关的操作;
- timerCtx:用来处理超时相关操作;
- valueCtx:附加值的实现方法。
现在先简单对这几个实现有个概念,后面会对其中核心关键的部分讲解下。
Context 类图
从类图中可以看出,源码里有 4 种结构和 3 种接口,相对于其他 go 库源码来说是比较简单的。
核心的接口是 Context,里面包含了最常用的判断是否处理完成的 Done()方法 。其他所有结构都通过 ① 实现方法或 ② 组合的方式来实现该接口。
核心的结构是 cancelCtx,被 timerCtx 包含。cancelCtx 和 timerCtx 可以说代表了 Context 库最核心的取消和超时相关的实现,也最为复杂些。
Context 源码
因为篇幅关系,不会把每一行源码都拎出来,会挑比较重点的方法讲下。由于平时我们使用都是通过几个固定的方法入口,所以会围绕这几个方法讲下。
emptyCtx
对外体现
var ( background = new(emptyCtx) todo = new(emptyCtx) ) func Background() Context { return background } func TODO() Context { return todo }
TODO(),Background()其实都是返回一个 emptyCtx。
实现
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" }
这个结构非常简单,都是返回 nil。emptyCtx 主要用于新建一个独立的树。比方说,我想在协程里做些异步操作,但是又想脱离主协程的 ctx 控制如使用独立的超时限制,就可以使用这种方式。但是在整个 go 程序里只有 todo 和 background 两个大根节点,所以TODO()和Background()其实是新建第二层级的子树。
func demo(ctx context.Context){ nctx := context.TODO() nctx := context.WithTimeout(nctx, 5*time.Second) ... }
valueCtx
对外体现
// 设置key, value值 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} } // 根据key读取value func (c *valueCtx) Value(key interface{}) interface{} { if c.key == key { return c.val } return c.Context.Value(key) }
通过公共方法设置值,再通过 valueCtx 的内部方法获取值。后面再仔细讲下 Value 的实现方式。
实现
type valueCtx struct { Context key, val interface{} } // 根据key读取value func (c *valueCtx) Value(key interface{}) interface{} { // 每个ctx只绑定一个key,匹配则返回。否则向上追溯到匹配为止 if c.key == key { return c.val } return c.Context.Value(key) }
从实现上可以看出,每当我们往 ctx 里调 WithValue 塞值时,都会生成一个新的子节点。调用的次数多了,生成的子树就很庞大。
若当前节点的 key 和传入的 key 不匹配会沿着继承关系向上递归查找。递归到根就变成 nil,表示当前 key 在该子树序列里没存。
cancelCtx
介绍完上面两种比较简单的结构后,终于要来到复杂的 cancelCtx。cancelCtx 和 timerCtx 关联性很强,基本上弄懂一个,另外一个也差不多了。
对外方法
func WithCancel(parent Context) (ctx Context, cancel CancelFunc) { // 新建一个cancelCtx c := newCancelCtx(parent) // 将父节点的取消函数和子节点关联,做到父节点取消,子节点也跟着取消 propagateCancel(parent, &c) // 返回当前节点和主动取消函数(调用会将自身从父节点移除,并返回一个已取消错误) return &c, func() { c.cancel(true, Canceled) } }
对外的方法里包含的几个方法都是重点的方法,后面主要讲下
结构
type cancelCtx struct { Context mu sync.Mutex // protects following fields done 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 }
- done:用于判断是否完成
- cancel:存子取消节点
- err:取消时的错误,超时或主动取消
type canceler interface { cancel(removeFromParent bool, err error) Done() <-chan struct{} }
这个接口约定了可以取消的 context,比如 cancelCtx 和 timerCtx 是可以取消的,emptyCtx 和 valueCtx 是不可以取消的。
初始化
// newCancelCtx returns an initialized cancelCtx. func newCancelCtx(parent Context) cancelCtx { return cancelCtx{Context: parent} }
初始化就是将父节点设置了一下,其他不设置。
cancelCtx 的取消实现
// cancel closes c.done, cancels each of c's children, and, if // removeFromParent is true, removes c from its parent's children. func (c *cancelCtx) cancel(removeFromParent bool, err error) { // 取消无论是通过父节点还是自身主动取消,err都不为空 if err == nil { panic("context: internal error: missing cancel error") } c.mu.Lock() if c.err != nil { // c.err 不为空表示已经被取消过,比如父节点取消时子节点可能已经主动调用过取消函数 c.mu.Unlock() return // already canceled } c.err = err if c.done == nil { // closedchan 是一个已经关闭的channel,要特殊处理是因为c.done是懒加载的方式。只有调用c.Done()时才会实际创建 c.done = closedchan } else { close(c.done) } // 递归取消子节点 for child := range c.children { // NOTE: acquiring the child's lock while holding parent's lock. child.cancel(false, err) } c.children = nil c.mu.Unlock() // 从父节点中移除当前节点 if removeFromParent { removeChild(c.Context, c) } }
整个过程可以总结为
- 前置判断,看是否为异常情况
- 关闭 c.done,这样外部调用 cancelCtx.Done()就会有返回结果
- 递归调用子节点的 cancel 方法
- 视情况从父节点中移除子节点
这里child.cancel(false, err)不从父节点移除子节点是因为当前节点操作已取过锁,移除操作会再取锁造成冲突,故先全部 cancel 后再将 children 置为 nil 一次性移除。
propagateCancel 绑定父子节点的取消关系
// propagateCancel arranges for child to be canceled when parent is. func propagateCancel(parent Context, child canceler) { done := parent.Done() if done == nil { // 若当前节点追溯到根没有cancelCtx或者timerCtx的话,表示当前节点的祖先没有可以取消的结构,后面的父子绑定的操作就可以不用做了,可参考下图 return // parent is never canceled } select { case <-done: // 父节点已取消就直接取消子节点,无需移除是因为父子关系还没加到parent.children // parent is already canceled child.cancel(false, parent.Err()) return default: } // 获取最近的可取消的祖先 if p, ok := parentCancelCtx(parent); ok { p.mu.Lock() if p.err != nil { // 和前面一样,如果祖先节点已经取消过了,后面就没必要绑定,直接取消就好 // parent has already been canceled child.cancel(false, p.err) } else { // 绑定父子关系 if p.children == nil { p.children = make(map[canceler]struct{}) } p.children[child] = struct{}{} } p.mu.Unlock() } else { // 当ctx是开发者自定义的并继承context.Context接口会进入这个分支,另起一个协程来监听取消动作,因为开发者自定义的习惯可能和源码中用c.done和c.err的判断方式有所不同 atomic.AddInt32(&goroutines, +1) go func() { select { case <-parent.Done(): child.cancel(false, parent.Err()) case <-child.Done(): } }() } }
① 当祖先继承链里没有 cancelCtx 或 timerCtx 等实现时,Done()方法总是返回 nil,可以作为前置判断
② parentCancelCtx 取的是可以取消的最近祖先节点
总结
总结一下,cancelCtx 的作用其实就两个:
- 绑定父子节点,同步取消信号,父节点取消子节点也跟着取消
- 提供主动取消函数
timerCtx
结构体
type timerCtx struct { cancelCtx timer *time.Timer // Under cancelCtx.mu. deadline time.Time }
相比 cancelCtx 多了一个计时器和截止时间
取消方法
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() }
取消方法就是直接调用 cancelCtx 的取消外加计时器停止
对外方法
func WithDeadline(parent Context, d time.Time) (Context, CancelFunc) { if cur, ok := parent.Deadline(); ok && cur.Before(d) { // 传入的截止时间在父节点截止时间之后,则父节点取消时会同步取消当前子节点,不需要额外再设置计费器了,可以当普通的cancelCtx对待。 // 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 passed 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 的主要功能就是用于控制协程退出和附加链路信息。核心实现的结构体有 4 个,最复杂的是 cancelCtx,最常用的是 cancelCtx 和 valueCtx。整体呈树状结构,父子节点间同步取消信号。