Go 语言的 Context 源码分析
11月 24, 2020
研究 Context 的源码，有助于对结合运用 interface 和 struct 的理解，以及对其他三方框架对 Context 接口的重新实现也能有一定认识，在实际开发中更是非常有帮助。Context 的源码非常短小，加上大概一半的注释整个文件也才 500 行，500 行就可以和面试官扯一个小时，ROI 巨大，确定不研究下么。
必备知识 #
WithTimeout 使用 #
func main() {
ctx, cancel := context.WithTimeout(context.Background(), time.Second*2)
defer cancel()
work(ctx)
}
func work(ctx context.Context) {
done := make(chan bool)
go func(ctx context.Context) {
println("工作中...")
time.Sleep(time.Second * 3)
done <- true
}(ctx)
select {
case <-ctx.Done():
fmt.Println("超时了", ctx.Err())
case <-done:
fmt.Println("执行完成")
}
}
//工作中...
//超时了 context deadline exceeded
通过控制 WithTimeout 的时间，可以设置在多久后收到 ctx.Done() 的信号，调用 cancel 也能达到同样的效果。
WithValue 使用 #
package main
import (
"context"
"fmt"
)
func main() {
ctx := context.WithValue(context.Background(), "hello", "world")
ctx1 := context.WithValue(ctx, "foo", "bar")
fmt.Println(ctx.Value("hello"))
fmt.Println(ctx1.Value("hello"))
fmt.Println(ctx1.Value("foo"))
fmt.Println(ctx.Value("foo"))
}
func Done() <-chan struct{} {
return make(chan struct{})
}
// output:
// world
// world
// bar
// <nil>
如上面例子，WithValue 函数可以在上下文中定义变量并返回创建的新子 ctx，通过 ctx.Value() 可以获取到你之前定义过的值，需要注意的是，代码中 ctx 并不能获取到 ctx1 中的 value，而 ctx1 可以获取到 ctx 中的 value，说明父子关系的 context 之间“子对父”是有隔离的。
close 使用 #
这里需要再额外加一个 close 队列的例子，先看懂，下面的源码会用到。
package main
import (
"context"
"fmt"
"time"
)
func main() {
ctx, _ := context.WithTimeout(context.Background(), time.Second*3)
done := make(chan struct{})
go func(c context.Context) {
time.Sleep(time.Second)
close(done)
}(ctx)
select {
case <-ctx.Done():
fmt.Println("timeout")
case <-done:
fmt.Println("work done!")
}
}
//work done!
这里要注意的是，close 一个队列后，select 也能收到信号并进入到了相应的 case，可以说这是 cancel 方法的核心原理了。
结构体嵌套接口 #
interface 可以嵌套在 struct 里，这个知识点我是在看源码的时候才发现的，以前真不知道。源码中的 context.Context 是一个 interface，它定义了四个方法，通过上面的例子多少已经知道了一些使用方法，重点不是 Context，而是 cancelCtx 的实现如下：
type Context interface {
Deadline() (deadline time.Time, ok bool)
Done() <-chan struct{}
Err() error
Value(key interface{}) interface{}
}
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
}
这里的 cancelCtx 嵌套了 Context 接口，那么 cancelCtx 的实例就可以用所有实现了 Context 接口的其他结构体来初始化了，源码里这么用的：
func newCancelCtx(parent Context) cancelCtx {
return cancelCtx{Context: parent}
}
这样 newCancelCtx 返回的新的 cancelCtx 实例就初始化好了：不需要再实现这些接口里定义的方法，就自动获得了调用接口方法的能力。
var cancelCtxKey int
func (c *cancelCtx) Value(key interface{}) interface{} {
if key == &cancelCtxKey {
return c
}
return c.Context.Value(key)
}
这里先有个理解，下面的源码可以细看。
结构体嵌套结构体 #
说完接口嵌套在结构体里，这里就再拓展说一下结构体嵌套在结构体里的用法。在 Go 中结构体 A 嵌套另一个结构体 B，这个平日经常会遇到，通过嵌套，可以扩展结构体，而且减少重复代码。
package main
import "fmt"
type A struct {
Name string
}
type B struct {
A
Age int
}
func main() {
b := B{
A: A{
Name: "A_name",
},
Age: 18,
}
fmt.Println(b.Name)
}
// A_name
字段提升 #
在上面代码中，可以通过 b.Name 直接访问到 A 的 Name 属性，但是却不能在初始化 b 的时候直接定义，否则编译器会抛出错误。
源码分析 #
通过上面的例子，已经知道怎么使用了，现在再来分析源码。看 context 的源码时和 channel 一对比就发现了一点好处，就是没有关键字带来的编译转换的复杂度，也就是说如果要实现更复杂功能的 context，你行你也能上。
Background 和 TODO #
var (
background = new(emptyCtx)
todo = new(emptyCtx)
)
func Background() Context {
return background
}
func TODO() Context {
return todo
}
先看 👆Backgroud 和 TODO 的定义，他们都是通过 new 初始化了一个 emptyCtx，看起来 Background 和 TODO 类型没有什么区别，根据官方文档的描述，TODO 可以用在还不知道未来如何应用哪种上下文的情况下使用。
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
}
可见 emptyCtx 实现了 Context 接口，但是都是空方法，没有任何功能。这种 ctx 通常作为起始的 context 向下传递。
WithCancel #
WithCancel 应该是 Context 里最核心的实现了，后面几个类型的 Context 都是以此为基础进行拓展的，下面开始详细看下它的源码。
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) }
}
// newCancelCtx returns an initialized cancelCtx.
func newCancelCtx(parent Context) cancelCtx {
return cancelCtx{Context: parent}
}
👆 WithCancel 函数创建并返回了一个 cancelCtx 类型的上下文和一个 cancel 闭包。
cancelCtx #
// A cancelCtx can be canceled. When canceled, it also cancels any children
// that implement canceler.
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
}
type canceler interface {
cancel(removeFromParent bool, err error)
Done() <-chan struct{}
}
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
}
func (c *cancelCtx) Err() error {
c.mu.Lock()
err := c.err
c.mu.Unlock()
return err
}
cancelCtx 包含了一个 Context 的 interface，和多出来的几个字段：
mu互斥锁。done一个struct{}类型的队列，用于接收cancel方法发送的取消信号。children是一个map，key 为canceler的接口（定义了cancel方法和Done()方法），而 value 是一个空 struct，实际上并没有使用。err调用cancel时附带的错误信息。
propagateCancel #
func propagateCancel(parent Context, child canceler) {
done := parent.Done()
if done == nil {
return // parent is never canceled
}
select {
case <-done:
// 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 {
atomic.AddInt32(&goroutines, +1)
go func() {
select {
case <-parent.Done():
child.cancel(false, parent.Err())
case <-child.Done():
}
}()
}
}
propagateCancel 函数也是非常核心的函数之一，在上下文初始化时通过 children 绑定了父子关系，并在 parent 和 child 之间同步取消和结束的信号，保证在 parent 被取消时，child 也会收到对应的信号，不会出现状态不一致的情况。
cancel 方法 #
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 = err
if c.done == nil {
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)
}
}
👆 cancel 会递归遍历自己的 children 的 cancel 方法，并使用 close 关闭自己 done 队列。
WithDeadline #
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 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) }
}
type timerCtx struct {
cancelCtx
timer *time.Timer // Under cancelCtx.mu.
deadline time.Time
}
timerCtx #
timerCtx 继承了 cancelCtx 相关的变量和方法，同时又比 cancelCtx 多了 timer 和 deadline，c.timer 通过 time.AfterFunc 设置了一个定时器，在超过这个时间之后就会调用 cancel 方法。
WithTimeout #
func WithTimeout(parent Context, timeout time.Duration) (Context, CancelFunc) {
return WithDeadline(parent, time.Now().Add(timeout))
}
WithTimeout 只是对 WithDeadline 的参数的封装，后面就都一样了。
WithValue #
func WithValue(parent Context, key, val interface{}) 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}
}
type valueCtx struct {
Context
key, val interface{}
}
valueCtx #
WithValue 使用 valueCtx 从 Context 类型的 parent 创建新的上下文，除了多出 key 和 value 字段，没有新的内容了，值得多说一句，没看源码之前我还以为一个 context 里的 value 是一个 map，可以设置多个 key，看了源码才知道只能设置一个 key 值，那些可以设置多个 key 值的是三方库的实现。
func (c *valueCtx) Value(key interface{}) interface{} {
if c.key == key {
return c.val
}
return c.Context.Value(key)
}
小结 #
这里几乎按照源码顺序，每一段代码都做了分析，基本上可以从底层了解了 Context 的原理。一个 Context 可以无限生成子 ctx 使用，并可以被根 ctx 取消，事实上，这可能就是 context 被设计出来的意义，不管是从 interface 和 struct 运用的语法上，还是 Context 的实现上，相信你肯定也有收获。
参考 #
https://draveness.me/golang/docs/part3-runtime/ch06-concurrency/golang-context/