Go Context 源码阅读

[wewin]

读源码一是为了学习好的代码风格，二是为了对 Go 这门语言能有深入的了解，能成长为一名合格的 Gopher

Context

Context 翻译成中文就是 '上下文' 的意思，准确的说它是 goroutine 的上下文, Go 1.7 开始引入 context，我看大代码是 Go 1.14.6 context 代码位于 go 源码的 src/context 文件中，这个包代码很少并且包含大量的注释，很方便我们阅读源代码

Context 的作用是用来传递 goroutine 之间的上下文信息，包括 取消信号, 超时信号，截止时间，请求信息(session, cookie)，控制一批 goroutine 的生命周期. 在 Go 中我们往往使用 channel + select 的方式来控制协成的生命周期。但是对于复杂的场景，比如 Go 中通常一个协程会衍生出很多子协程， 分别处理不同的事情，这些携程往往具有相同生命周期，具有通用的变量，如果 goroutine 的层级较深使用 channel + select 不太方便，这个时候就可以使用 context.

Context 的底层原理实现

在 context 包中定义了 Context 这种 interface 类型

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

这个 interface 包含 Deadline, Done, Err, Value 四个方法 Deadline: 返回 context 是否会被取消，以及取消的时间 Done: 是在 context 被取消或者 deadline 后返回一个被关闭的 channel Err: 在 channel 关闭后，返回 context 取消原因 Value: 用来获取 key 对应的 value

同时在这个包中有一个 emptyCtx，它实现了 Context 接口

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 的一个最小实现，方法很简单，要么直接返回，要么直接返回 nil 。

在 Go 中初始化一个 context，我们经常使用 context.Background() 或者 context.TODO(), 从如下源码中可以看出这两个方法实际上返回的就是一个 emptyCtx

var (
	background = new(emptyCtx)
	todo       = new(emptyCtx)
)

func Background() Context {
	return background
}

func TODO() Context {
	return todo
}

context.Background() 和 context.TODO() 看着除了方法名不一样，其他都是一样的。 在使用中我们需要区分两种 context 的使用场景，context.Background() 通常是用在 main 、测试， 或者最高层的 context (相当于根 context) 的初始化 context 的时候，而 TODO context 则是当我们不清楚使用什么 context 的时候使用

除了 context.Background() 和 context.TODO() 初始化 context 的方法，context 包还为我们提供了如下四个生成 context 的函数

咱们先看看函数签名如下：

WithCancel(parent Context) (ctx Context, cancel CancelFunc)
WithDeadline(parent Context, d time.Time) (Context, CancelFunc)
WithTimeout(parent Context, timeout time.Duration) (Context, CancelFunc)
WithValue(parent Context, key, val interface{}) Context

WithCancel

WithCancel 用于生成一个可取消的 context

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

通常我们代码里如下使用

ctx := context.Background()
cancel, ctx := context.WithCancel(ctx)

它接受一个 context (父 context)，返回一个 context 和一个可取消该 context 的方法,

WithCancel 首先调用了 newCancelCtx 私有方法，生成了一个 cancelCtx 结构体，然后在调用 propagateCancel 方法将 new context 挂载到父 context 上，我们着重看下 newCancelCtx 和 propagateCancel 方法

// newCancelCtx returns an initialized cancelCtx.
func newCancelCtx(parent Context) cancelCtx {
	return cancelCtx{Context: parent}
}

// 新的结构体，包含一个 Context，和 4 个私有属性
type cancelCtx struct {
	// 指向的是父 context
	// context 链类似一个链表，但是 Context 指向是父 context
    Context
    
    // 互斥锁，保证字段的安全
    mu       sync.Mutex            // protects following fields 
    // 在 context 取消后首先关闭该 chan
    done     chan struct{}         // created lazily, closed by first cancel call
    // 从此 context 衍生出的子 context 挂载在这里
    children map[canceler]struct{} // set to nil by the first cancel call
    // cancel 的原因
    err      error                 // set to non-nil by the first cancel call
}

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

newCancelCtx 代码比较少，他一个父 context，返回一个 cancelCtx 类型的 context, 父 context 被赋值到了 cancelCtx 的 Context 字段 需要注意的是 newCancelCtx 返回的是一个 cancelCtx 类型，该 cancelCtx 实现了 canceler interface 的 cancel 和 Done 方法 cancel 方法用来取消这个 context 以及这个 context children map 上的子 context Done 返回的怎是一个关闭的 channel, 用来表示该 context 是否 cancel

cancel 方法源码：

func (c *cancelCtx) cancel(removeFromParent bool, err error) {
	if err == nil {
		panic("context: internal error: missing cancel error")
	}
	c.mu.Lock()
	// 该 context 已经 Done
	if c.err != nil {
		c.mu.Unlock()
		return // already canceled
	}
	c.err = err
	// 关闭 channel
	if c.done == nil {
		c.done = closedchan
	} else {
		close(c.done)
	}
	// 循环取消该 context 的子 context
	for child := range c.children {
		// NOTE: acquiring the child's lock while holding parent's lock.
		child.cancel(false, err)
	}
	// 取消子 context 后，将 children 字段置空
	c.children = nil
	c.mu.Unlock()

	// 如果指定了 removeFromParent = true
	// 则需要将该 context 从其父 context 的 children map 字段中删除
	if removeFromParent {
		removeChild(c.Context, c)
	}
}

// Done 方法则返回该 context 的 done channel
// context 关闭后，外部接受到 channel 的 close 信号
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
}

cancel 方法实现特别简单，通过 context 结构体中的 done chan struct{} 这个字段实现的 调用 cancel 方法，本质上就是对该 context 的 done channel 字段执行 close 操作 此外，如果入参 removeFromParent = ture, 会将此 context 从他的父 context 的 children map 上删除

下面是 propagateCancel 方法源码，propagateCancel 方法特别重要 该方法就是将 newCancelCtx 方法生成的新的 cancelCtx 挂在亲父 context 的 children map 上 在执行挂载的时候，如果父 context 已经取消，就将此 context 也取消掉

func propagateCancel(parent Context, child canceler) {
	done := parent.Done()
	// step1: 如果 parent 是不可 cancel 的
	// 此时直接返回，没有将 child context 挂到 children map 上的必要
	// 因为即便挂载上去，因为父 context 不能取消，子 context 也更无法通过父 context 来取消
	if done == nil {
		return // parent is never canceled
	}

	select {
	case <-done:
		// step2: 父 context 被取消，所以需要将此子 context 也取消
		child.cancel(false, parent.Err())
		return
	default:
	}

	// step3: 获取当前 context 的父 context
	// parentCancelCtx 是用来获取父 cancelCtx
	if p, ok := parentCancelCtx(parent); ok {
		p.mu.Lock()
		if p.err != nil { // 父 context 是 canceled （ context 如果取消的，err 字段一定不为空）
			// step3.1: 如果父 context 是已取消的，就需要将子 context 也取消了
			child.cancel(false, p.err)
		} else {
			// step3.2: 父 context 没有取消，将此子 context 挂到父 context 的 children map 字段
			if p.children == nil {
				p.children = make(map[canceler]struct{})
			}
			p.children[child] = struct{}{}
		}
		p.mu.Unlock()
	} else { 
		// step4: 这个 else 分支是比较难以理解的地方
		// 可以理解为，在并发模型下，如果其他 goroutine 将 parent 的 context 改成了一个 cancelCtx
		// 那么没有这个分支，会出现 parent done 的时候 child 不知道 parent done 信息
		// 导致 child context 无法 cancel, child context 控制的相关的 goroutine 就无法结束，出现内存泄漏
		atomic.AddInt32(&goroutines, +1)
		go func() {
			select {
			case <-parent.Done(): // 监听父 context 的 cancel 信息
				child.cancel(false, parent.Err())
			case <-child.Done(): // 如果 child 自身 cancel 就退出 select，避免当前这个 goroutine 内存泄漏
			}
		}()
	}
}

propagateCancel 方法是 context 中算是最复杂的一个方法了，它实现的功能是很简单的， 就是将 newCancelCtx 方法生成的新的 cancelCtx 挂在亲父 context 的 children map 上，不过过程中有很多细节处理，只有耐性阅读源码才能准确的理解这些处理上的细节

WithCancel 方法最后一句 return &c, func() { c.cancel(true, Canceled) }返回当前这个 cancelCtx 的 cancel 方法，作为该 context 外部控制该 context 取消的方法

通过 WithCancel 方法的分析，我们知道了 WithCancel 就是接受 context 参数，该参数作为 parent context 生成一个可以取消的 context， 并且会判断 parent context，如果他是一个未没有取消的 cancelCtx 类型的 context，就将当前新生成的 context 挂到 parent context 的 children map 上， 而 context 的是否取取消是通过 context 的 done 字段实现的，该字段是 chan struce{} 类型，取消一个 context 本质是将该 context 的 done 字段的的 channel 关闭

如下可见，cancel context 的 Done 方法，可以看出其返回的就是一个 close 的 channel

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
}

WithDeadline

看完 WithCancel 后，我们接着看 WithDeadline 底层实现，有了上面 WithCancel 的学习，看 WithDeadline 就比较容易了

func WithDeadline(parent Context, d time.Time) (Context, CancelFunc) {
	// 父 context 已经过期，或者父 context 的 deadline 是早于当前 context 的过期时间的
	// 就调用 WithCancel 创建一个 cancel context
	if cur, ok := parent.Deadline(); ok && cur.Before(d) {
		// The current deadline is already sooner than the new one.
		return WithCancel(parent)
	}
	// 创建一个 timerCtx，timerCtx 实现了 Context 接口
	// timerCtx 结构体包含了 (继承了) cancelCtx
	// 此外还有一个 *time.Timer 类型的 timer 字段和一个 time.Time 类型的 deadline 字段
	// timer 字段存储用来执行 deadline 的定时任务
	// deadline 是结束时间
	c := &timerCtx{
		// 创建一个 cancelCtx 的 context
		cancelCtx: newCancelCtx(parent),
		// context 的取消时间
		deadline:  d,
	}
	
	// 这里和 WithCancel 中一样，将 context 挂载到父 context children map 上
	propagateCancel(parent, c)
	
	// 结束时间已过, 取消当前 context
	dur := time.Until(d)
	if dur <= 0 {
		// 这里的 cancel 是 timerCtx 中实现的 cancel 方法
		// 而不是从 cancelCtx 中继承的 cancel，详见下文
		c.cancel(true, DeadlineExceeded) // deadline has already passed
		return c, func() { c.cancel(false, Canceled) }
	}
	c.mu.Lock()
	defer c.mu.Unlock()
	
	// 这里是 deadline context 实现的核心
	// 创建一个定时任务，到达时间指定的结束时间(d)的时候，执行此任务，将这个 context 取消掉
	if c.err == nil {
		c.timer = time.AfterFunc(dur, func() {
			c.cancel(true, DeadlineExceeded)
		})
	}
	return c, func() { c.cancel(true, Canceled) }
}

func (c *timerCtx) cancel(removeFromParent bool, err error) {
    // timerCtx 的 cancel 首先调用了 cancelCtx 的 cancel 方法，将此 context 关闭
    // 并将 children map 字段上挂的子 context 取消
    c.cancelCtx.cancel(false, err)
    
    // 如果指定了 removeFromParent = true
    // 将从父 context 的 children map 上将当前 context 移除
    if removeFromParent {
        removeChild(c.cancelCtx.Context, c)
    }
    
    c.mu.Lock()
    // 终止 timer 字段上挂载的定时任务
    // 因为上面已经主动 cancel 了，所以需要停止当前 context 上的取消任务了
    if c.timer != nil {
        c.timer.Stop()
        c.timer = nil
    }   
    c.mu.Unlock()
}

从 WithDeadline 源码中我们会发现 with deadline context 的实现很简单 底层的创建和 WithCancel 基本一致，不同点是 WithDeadline 创建出来的 context 提多了一个 deadline 和 timer 字段 deadline 字段用来记录该 context 的结束时间 timer 上面挂了一个定时任务，负责到了指定的 deadline 时执行 cancel 方法，取消当前 context

WithTimeout

看完 WithDeadline 的实现，相信大家能想到 WithTimeout 这种 context 的实现了， 将 WithTimeout 的相对时间，转为一个绝对时间就是，WithTimeout 就变成了一个 WithDeadline，context 源码包中也是这么实现的，大家可以自行阅读源码

WithValue

现在看看 WithValue 的实现, 从 WithValue 源码中，我们可以看到，WithValue 返回了一个 valueCtx，该 valueCtx 是实现了 Context 接口

func WithValue(parent Context, key, val interface{}) Context {
	// key 不能是 nil
	if key == nil {
		panic("nil key")
	}
	// key 必须是可比较的，因为在通过 key 来取 value 的时候，需要对比 key 是否相等
	// 可以看下文中的 valueCtx 的 Value 方法的实现，方法返回的时候会对 key 进行比较
	if !reflectlite.TypeOf(key).Comparable() {
		panic("key is not comparable")
	}
	// 创建一个 valueCtx 类型的 context
	return &valueCtx{parent, key, val}
}

// valueCtx 继承了 Context, 新增了 key, value 两个字段 // WithValue(parent Context, key, val interface{}) Context // 这里显而易见，key 、value 两个字段就是用来存储 WithValue 调用的时候传递进来的 key 和 value

type valueCtx struct {
	Context
	key, val interface{}
}

我们接着看 WithValue context 获取 value 的方法 这个方法实现很简单，需要注意的是获取值操作是递归获取的 先从 context 本身取值，如果取不到会沿着父链向上获取，最终会找到 emptyCtx, 此时返回值是 nil

func (c *valueCtx) Value(key interface{}) interface{} {
	if c.key == key {
		return c.val
	}
	// 递归获取 key 对应的 value
	return c.Context.Value(key)
}

context 的使用建议

以下建议来自 context 这个包里的详细注释，我做了一个简单翻译，我们在使用中应当遵循这些使用建议

• 不应当将 context 作为结构体的字段，应该将 context 作为函数参数往下传递
• context 作为参数使用的时候，应该作为第一个形参，如下

func DoSomething(ctx context.Context, arg Arg) error {
	// ... use ctx ...
}

• 不能传递 nil Context，如果不知道使用什么 context 的时候使用 context.TODO 即可
• 使用 context Values 传递值的时候，不应当将函数参数作为 value 传递，context 应当用来 传递通用性的变量，如请求相关的 accept-language，cookie，session 等
• context 是线程安全的

总结

Go 从 1.7 引入了 context，主要用于在 goroutine 之间传递取消信号、截止时间控制、超时时间控制以及一些通用型变量传递， 我读的源码是 Go 1.14.6 版本的，Go context 包代码特别简短，有大量的注释（注释比代码多，哈哈哈），很适合学习大家可以去读一读源码