Golang Context源码解析

        自从接触Go语言以来，Context就一直在用，不过一直弄不清楚这东西究竟有啥用，平常都是CV代码，最多就是稀里糊涂用个Background()。不过最近在开发Golang反向代理的时候遇到一个传值问题，因为调用的是底层的库不支持传递变量，所以甚是困惑，后来在stackoverflow搜索了一番，发现已有类似问题，给出的答案很简单，就是通过Context的WithValue传值，而且底层库自身的参数也包含Context，所以困惑的问题一下就解决了。回过头来想一下，如此简单明显的问题都不会，说明对Golang了解的过于片面，这可不是一个优秀工程师应有的表现。然后就看了下Context包的源码。

        Context包其实写的挺简单的，代码量也不多，很适合阅读。看完之后给我感觉就是Context就像一个共享内存，在不同的goroutine之间共享一些信息，比如值、或者信号等等，设计比较简洁，代码也比较有参考价值。

        下面直接看代码：

package context

import (
	"errors"
	"internal/reflectlite"
	"sync"
	"sync/atomic"
	"time"
)

/// Context 是一个接口，包含四个方法Deadline(), Done(), Err(),Value()，Context主要用于传递信号、值等。
type Context interface {	
	/// 如果有设置当前工作完成的截止时间的话，则返回该时间，否则返回空，timerCtx 内部实现就有个变量保存了截止时间，另外需要保证该接口调用的幂等性。
	Deadline() (deadline time.Time, ok bool)

	/// Done()返回一个用于通知工作结束的chan,该通道一般结合select使用，如果该context没有取消函数，则返回nil
	/// 当调用cancel的时候，该通道会被关闭，则所有监听该通道的子协程都会收到closed信号
	Done() <-chan struct{}

	/// 返回Done()被关闭的原因，未被closed的时候返回nil
	Err() error

	/// 返回context中该key对应的值，当前context找不到则从派生路径上找，都找不到返回nil
	Value(key interface{}) interface{}
}


/// 默认的context被取消错误
var Canceled = errors.New("context canceled")


/// 超时取消错误
var DeadlineExceeded error = deadlineExceededError{}

type deadlineExceededError struct{}

func (deadlineExceededError) Error() string   { return "context deadline exceeded" }
func (deadlineExceededError) Timeout() bool   { return true }
func (deadlineExceededError) Temporary() bool { return true }


/// emptyCtx 顾名思义，就是一个空的context，它实现了Context，不过啥都没做，相当于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
}

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

var (
	/// Context包定义了两个私有指针变量background和todo，这两者都是emptyCtx
	background = new(emptyCtx)
	todo       = new(emptyCtx)
)


/// 返回私有变量background，主要用于作为顶层的context，类似于c++的基类
func Background() Context {
	return background
}


/// 返回私有变量todo，这东西和Background()唯一的区别就是命名不同
func TODO() Context {
	return todo
}


/// 自定义取消函数类型 当调用该类型函数的时候需要停止当前work,并无需等待。另外该类型的实现可能同时被多个协程调用，需要考虑并发的问题，另外首次调用之后，后续的调用直接返回即可。
/// 这个不难理解，如果一个context已经取消了，再次取消应该直接返回
type CancelFunc func()

/// WithCancel 返回一个cancelCtx 和一个取消函数，cancelCtx应该是除了context包中Context之外最重要的结构了，其它的WithDeadline、WithTimeout实际上都是基于WithCancel扩展的
/// 实现上就是初始化一个cancelCtx，然后和parent context绑定，绑定的目的是当parent取消的时候，能够及时取消子context，具体实现看下面代码
func WithCancel(parent Context) (ctx Context, cancel CancelFunc) {
    /// 不允许从nil context派生
	if parent == nil { 
		panic("cannot create context from nil parent")
	}
    /// 这一步就是初始化cancelCtx
	c := newCancelCtx(parent) 
	
    /// 绑定到父context上
    propagateCancel(parent, &c) 

    /// 返回一个cancel函数，这个函数就是调用cancel取消自己，同时第一个参数表明需要将ctx从parent context的孩子队列中移除，移除就是解除和parent context的绑定关系
	return &c, func() { c.cancel(true, Canceled) } 
}


/// 初始化一个cancelCtx
func newCancelCtx(parent Context) cancelCtx {
	return cancelCtx{Context: parent}
}


/// 保存go协程创建的数量，测试用
var goroutines int32


/// propagate 这个单词英文意思是繁衍，增殖; 使遗传; 扩散; 使蔓延，在这里指绑定父context和子context，绑定的目的是使得父context的取消能够影响到子context，否则父context取消不能及时通知子context
func propagateCancel(parent Context, child canceler) {
	done := parent.Done() 
	if done == nil {
		/// 如果父context 没有cancel channel，说明父context不会执行cancel，所以这里不用绑定了
		return // parent is never canceled
	}

	select {
    /// 如果父context已经取消了，则将子context也取消掉，然后就可以返回了
	case <-done:
		// parent is already canceled 
		child.cancel(false, parent.Err())
		return
	default:
	}

	/// 这一步是获取派生路径上的cancelCtx，如果有的话，则绑定到父context的children队列
	if p, ok := parentCancelCtx(parent); ok { 
		p.mu.Lock() /// 加把锁保护下，防止并发
		if p.err != nil { /// 再次检查父context是否取消，如果取消了则相应的也应该将子context取消
			// parent has already been canceled
			child.cancel(false, p.err)
		} else {
			// 如果父context没有取消则将子context加入到父context的children队列里面，这样当父context取消的时候，它能够取消所有子context
			if p.children == nil { /// 初始化队列
				p.children = make(map[canceler]struct{})
			}
			p.children[child] = struct{}{} /// 添加到队列
		}
		p.mu.Unlock()
	} else {
		/// 走到这里有两种情形，一种父context是cancelCtx，不过已经取消了，另一种是父context不是cancelCtx,这种情况有可能是用户自定义的cancelCtx，这种情况下就不用加到父contex队列里面去了，因为不知道父context有没有队列这个属性
		atomic.AddInt32(&goroutines, +1)
		go func() {
			select {
			case <-parent.Done(): // 如果父context先取消，则接着取消子context
				child.cancel(false, parent.Err())
			case <-child.Done():  // 如果子context先取消 则直接返回就OK了
			}
		}()
	}
}

// 顾名思义，这个专门给cancelCtx用的key
var cancelCtxKey int 

// 获取派生路径上父cancelCtx
func parentCancelCtx(parent Context) (*cancelCtx, bool) {
	done := parent.Done()

	// 如果父context已取消或者没有父context的Done channel为空，则直接返回
	if done == closedchan || done == nil {
		return nil, false
	}
	/// 获取父辈派生路径上最近的cancelCtx, 这里使用了cancelCtxKey, 只有cancelCtx类型才实现了它
	/// 可以思考下这里为啥不用 p, ok := parent.(*cancelCtx)???，我觉得这样写也可以，只不过不符合原本的context设计规范，就像上面的propagateCancel，无非是加入到childen列表还是另起一个协程监听的区别，如果理解不正确，欢迎指正
	p, ok := parent.Value(&cancelCtxKey).(*cancelCtx)
	if !ok {
		return nil, false
	}
	p.mu.Lock()
	// 判断和父context是否一样，不一样说明是父节点的父节点
	ok = p.done == done
	p.mu.Unlock()
	if !ok {
		return nil, false
	}
	// 走到这里说明找到了派生路径上最近的一个cancelCtx类型
	return p, true
}


// 从父context中移除孩子节点自己，其实就是解除和parent的绑定
func removeChild(parent Context, child canceler) {
	/// 如果父节点不是cancelCtx类型，说明父context没有childen队列，因此无需删除直接返回
	p, ok := parentCancelCtx(parent)
	if !ok {
		return
	}
	p.mu.Lock()
	/// 存在父节点，从父节点中删除自己
	if p.children != nil {
		delete(p.children, child)
	}
	p.mu.Unlock()
}


/// canceler 接口，包含两个方法，cancelCtx实现了它
type canceler interface {
	cancel(removeFromParent bool, err error) // cancel函数，第一个参数表明是否需要从父context的孩子队列中移除，一般情况下如果child context主动取消，这时候就需要从父context中移除自身，使得父context取消的时候无需再取消子context

    // 返回Done chan
	Done() <-chan struct{}
}

/// 顾名思义，创建一个取消的channel,啥时候取消呢，包初始化的时候
var closedchan = make(chan struct{})

/// 初始化操作直接将私有变量closedchan 关掉, 话说所有cancelCtx返回的cancel后，Done()返回同一个channel
func init() {
	close(closedchan)
}


/// cancelCtx 实现很好理解，首先需要一个变量保存父context，其次需要一个变量保存子context列表，另外还有个channel用于通知，对于这些变量，还需要加把锁保护下
type cancelCtx struct {
	Context

    /// 加把锁 并发保护下
	mu       sync.Mutex      
    /// 懒人模式，需要的时候才创建	
	done     chan struct{}          

    /// 子context列表 
	children map[canceler]struct{} // set to nil by the first cancel call 

    /// cancel的原因
	err      error                 
}

/// cancelCtx Value()实现很巧妙，可以在自己的代码里借鉴下
func (c *cancelCtx) Value(key interface{}) interface{} {
	if key == &cancelCtxKey { // 这里用到一个私有变量，只有cancelCtx类型才会返回自己
		return c
	}
	return c.Context.Value(key) 
}

/// 返回cancelCtx Done通道
func (c *cancelCtx) Done() <-chan struct{} {
	c.mu.Lock()
	if c.done == nil { // 没有初始化则初始化，返回channel,调用cancel后才会关闭
		c.done = make(chan struct{})
	}
	d := c.done
	c.mu.Unlock()
	return d
}

/// 设置cancel err
func (c *cancelCtx) Err() error { 
	c.mu.Lock()
	err := c.err
	c.mu.Unlock()
	return err
}

type stringer interface {
	String() string
}

func contextName(c Context) string {
	if s, ok := c.(stringer); ok {
		return s.String()
	}
	return reflectlite.TypeOf(c).String()
}

func (c *cancelCtx) String() string {
	return contextName(c.Context) + ".WithCancel"
}


/// cancelCtx的cancel方法 主要逻辑就是检查自身是否关闭，没关闭则关闭一下，另外看是否需要将自身从父context移除
func (c *cancelCtx) cancel(removeFromParent bool, err error) {
    // 调用cancel方法必须传递一个err进来
	if err == nil {
		panic("context: internal error: missing cancel error")
	}
	c.mu.Lock()       // 先抢占下锁，防止并发调用
	if c.err != nil { // 看看是否被其它协程取消了，如果是的话，直接返回
		c.mu.Unlock()
		return // already canceled
	}
	// 取消的时候 err赋值
	c.err = err
    
    // 这里为nil说明还没有人监听channel,这里直接返回默认的close channel, 这样其它子协程即便监听也会立刻返回
	if c.done == nil { 
		c.done = closedchan
	} else { // 如果不为nil，说明有人监听了，这里直接关闭
		close(c.done)
	}

	for child := range c.children {		
		/// 将自身衍生的所有子context都取消掉，传递false，表示不需要从父context中移除自己，因为父context自己做了这步操作
		child.cancel(false, err)
	}

	// 清空子context
	c.children = nil
	c.mu.Unlock()

	// 是否从父context中移除，通常cancel自己的话会传递true，标明将自己从父context中移除
	if removeFromParent {
		removeChild(c.Context, c)
	}
}


/// WithDeadline返回一个timerCtx，本质上也是一个cancelCtx，相较于cancelCtx必须主动调用cancel之外，timerCtx还多了一种cancel调用方法，即设定一个定时器，到了指定时间自动取消
func WithDeadline(parent Context, d time.Time) (Context, CancelFunc) {
    /// parent context 不允许为nil
	if parent == nil {
		panic("cannot create context from nil parent")
	}

	/// 如果父context的截止时间比当前设定的要早，则直接返回一个cancelCtx将自身和parent 绑定下，这是因为父context一定先取消，这时候只需要和父context绑定等ta调cancel取消自己就好了
	if cur, ok := parent.Deadline(); ok && cur.Before(d) {
		// The current deadline is already sooner than the new one.
		return WithCancel(parent)
	}

	/// 如果父context取消的时间比自身还要晚，那就不能等待父context取消自己，需要自己主动起一个定时器取消，所以这里新建一个timerCtx
	c := &timerCtx{
		cancelCtx: newCancelCtx(parent),
		deadline:  d,
	}

	/// 和cancelCtx一样，绑定到父context上
	propagateCancel(parent, c)

	dur := time.Until(d)
    /// 如果已经到了截止时间，则直接取消就好了
	if dur <= 0 {
		c.cancel(true, DeadlineExceeded) 
        /// 这里传递false是因为上一步传了true
		return c, func() { c.cancel(false, Canceled) } 
	}

	c.mu.Lock()
	defer c.mu.Unlock()
    /// 再次检查是否被取消，当父context取消的时候这里err不为nil
	if c.err == nil { 
		c.timer = time.AfterFunc(dur, func() { /// 起一个定时器，超时则调用cancel 函数
			c.cancel(true, DeadlineExceeded)
		})
	}
	return c, func() { c.cancel(true, Canceled) }
}


/// 内部封装了一个cancelCtx，另外需要一个定时器和一个时间变量保存相关配置
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) String() string {
	return contextName(c.cancelCtx.Context) + ".WithDeadline(" +
		c.deadline.String() + " [" +
		time.Until(c.deadline).String() + "])"
}

/// timerCtx 和cancelCtx 取消唯一不同的地方在于timerCtx还需要关闭定时器
func (c *timerCtx) cancel(removeFromParent bool, err error) {

    /// 先通知子context关闭
	c.cancelCtx.cancel(false, err)
    
    /// 自己主动调用情况下 则需要将自身从父context中移除
	if removeFromParent { 
		removeChild(c.cancelCtx.Context, c)
	}

	c.mu.Lock()
	if c.timer != nil { /// 关闭定时器
		c.timer.Stop()
		c.timer = nil
	}
	c.mu.Unlock()
}

/// WithTimeout 这个实现其实和WithDeadline一样，区别在于Deadline可以指定任意时间，WithTimeout则是在当前时间加上timeout
func WithTimeout(parent Context, timeout time.Duration) (Context, CancelFunc) {
	return WithDeadline(parent, time.Now().Add(timeout))
}


/// WithValue返回一个存储K,V的valueCtx, 这里需要key是可比较的，并且不建议使用内置的string等类型作为key，因为那样会导致冲突，另外
/// 不应该把WithValue用来传递普通参数。通常情况下，WithValue 用于传递一些跨API的数据，比如用于链路跟踪的traceID等等
/// valueCtx实现也很简单，主要就是一个K，V加上一个parent context
func WithValue(parent Context, key, val interface{}) Context {    
    /// 同样不允许parent为nil
	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}
}

/// valueCtx实现更简单，就包含一个parent ctx和一对K,V，因为只有一对，所以也不需要锁保护
type valueCtx struct {
	Context
	key, val interface{}
}

// stringify tries a bit to stringify v, without using fmt, since we don't
// want context depending on the unicode tables. This is only used by
// *valueCtx.String().
func stringify(v interface{}) string {
	switch s := v.(type) {
	case stringer:
		return s.String()
	case string:
		return s
	}
	return "<not Stringer>"
}

func (c *valueCtx) String() string {
	return contextName(c.Context) + ".WithValue(type " +
		reflectlite.TypeOf(c.key).String() +
		", val " + stringify(c.val) + ")"
}

/// 查找key对应的value，找不到则从派生路径上去找
func (c *valueCtx) Value(key interface{}) interface{} {
	if c.key == key {
		return c.val
	}
	return c.Context.Value(key)
}

参考文档：

1. Go 语言并发编程与 Context | Go 语言设计与实现