Golang 源码 Context

• 简介
• 示例
  • WithValue
  • withCancel
  • WithTimeout
• 源码
  • valueCtx
  • cancelCtx
  • timerCtx
• 参考链接

简介

Context 译为“上下文”，是 Golang 中常用的并发控制技术，与 WaitGroup 最大的不同点是 Context 对于派生的 goroutine 有更强的控制力，可以控制多级的 goroutine。WaitGroup 面向的是水平 goroutine，而 Context 是面向垂直（派生）goroutine。

使用规则

• 不要把 Context 放在结构体中，要以参数的方式传递
• 把 Context 作为第一个参数传递给入口和出口请求链路的每一个函数，放在第一位
• 如果需要传递 Context，不要传递 nil，否则 trace 追踪会断了链接
• 并发安全，可以在多个 goroutine 中传递
• 可以把 Context 对象传递给任意个数的 goroutine，对它执行 取消 操作，则所有 goroutine 都会收到取消信号。

Context 主要有两个功能（两个没有关系的功能）

• 在 Context 链上共享数据
• 控制 goroutine 退出

Context 仅仅是一个接口，根据具体实现方式的不同衍生了不同类型的 Context

• cancelCtx 实现了 Context 接口，通过 WithCancel() 创建 cancelCtx 实例
• timerCtx 实现了 Context 接口，通过 WithDeadline() 和 WithTimeout() 创建 timerCtx 实例
• valueCtx 实现了 Context 接口，通过 WithValue() 创建valueCtx实例

三种 Context 可以相互为父节点，从而可以组合得到不同的应用形式

示例

WithValue

// https://play.golang.org/p/undefined
func main() {

	c := make(chan string)
	ctx := context.WithValue(context.Background(), "hello", "world")

	go func(ctx context.Context) {
		fmt.Println("wait 50ms")
		time.Sleep(50 * time.Millisecond)

		c <- ctx.Value("hello").(string)

	}(ctx)

	fmt.Println("result:", <-c)
}

withCancel

// https://play.golang.org/p/BIgFWP9W1I7
func main() {

	c := make(chan struct{})
	ctx, cancel := context.WithCancel(context.Background())

	go func(cc context.Context) {
		t1 := time.Now()

		for {
			select {
			case <-cc.Done():
				fmt.Println("stop:", time.Since(t1))
				close(c)
				time.Sleep(2 * time.Second) // 等待关闭 c
			default:
				time.Sleep(10 * time.Millisecond)
				fmt.Println("running")
			}
		}

	}(ctx)

	time.Sleep(20 * time.Millisecond)
	cancel()

	<-c
}

WithTimeout

// https://play.golang.org/p/Ef_hflYPNIQ
func main() {

	c := make(chan struct{})
	ctx, cancel := context.WithTimeout(context.Background(), 50*time.Millisecond)
	defer cancel()

	go func(cc context.Context) {
		t1 := time.Now()

		for {
			select {
			case <-cc.Done():
				fmt.Println("stop:", time.Since(t1))
				close(c)
				time.Sleep(2 * time.Second) // 等待关闭 c
			default:
				time.Sleep(10 * time.Millisecond)
				fmt.Println("running")
			}
		}

	}(ctx)

	<-c
}

源码

以 Golang: tag 1.12 为例

Context 其实是一个 interface，只要实现了这个 interface 则都是 Context

// https://github.com/golang/go/blob/06472b99cdf59f00049f3cd8c9e05ba283cb2c56/src/context/context.go#L57
type Context interface {
	Deadline() (deadline time.Time, ok bool)
	Done() <-chan struct{}
	Err() error
	Value(key interface{}) interface{}
}

• Deadline 返回当前 Context 完成的截止日期，如果没有设置则 ok == false
• Done 返回一个 Channel，Context 完成则返回的 Channel 也是关闭的，如果永远不取消此 Context 则有可能返回 nil，连续调用返回相同的值
• Err 返回当前 Context 结束的原因，如果 Done 没有关闭返回 nil，如果 Done 关闭则返回具体原因
  • 如果被取消则返回 Canceled
  • 如果超时则返回 DeadlineExceeded
• Value Context 不仅仅控制 goroutine，还用于在树状 goroutine 传递信息，Value 就是根据 key 查询 map 中的 value

Context 中的 Background() 和 TODO() 为基础 Context，主要被其他四个 function 当做父 Context，通过源码知晓，Background() 和 TODO() 实现方法是一致的，但是普遍使用 Background() 作为其他 function 的父 Context

func Background() Context
func TODO() Context

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

func Background() Context {
	return background
}

func TODO() Context {
	return todo
}

Context 定义了 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 }

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

除了 emptyCtx 还有 cancelCtx、timerCtx、valueCtx，分别用于实现下面四种 function，基于父 Context 派生出具有复杂功能的 Context

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

valueCtx

valueCtx 只是在 parent Context 的基础上增加了一个 Key-value 对，用于在各级协程间传递一些数据

由于 valueCtx 即不需要 cancel 也不需要 deadline，那么只需要实现 Value() 接口即可

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

func (c *valueCtx) String() string {
	return fmt.Sprintf("%v.WithValue(%#v, %#v)", c.Context, c.key, c.val)
}

• Value() 实现比较简单，每个 Context 只会存储一对 key-value，在查询时，会优先查询当前 Context 的 key，如果在当前 Context 没有查找到 key 则会向父 Context（上一级 Context），可以通过 Context 查询到父 Context 的 value

func (c *valueCtx) Value(key interface{}) interface{} {
	if c.key == key {
		return c.val
	}
	return c.Context.Value(key)
}

• WithValue() 实现也很简单，会判断 key 是否是可以进行等值比较，然后返回构造成功的 valueCtx

func WithValue(parent Context, key, val interface{}) Context {
	if key == nil {
		panic("nil key")
	}
	if !reflect.TypeOf(key).Comparable() {
		panic("key is not comparable")
	}
	return &valueCtx{parent, key, val}
}

cancelCtx

cancelCtx 记录了由此 Context 派生所有的子 Context，此 Context 被 Cancel 时会把所有子 child 都 Cancel

cancelCtx 与 deadline 和 value 无关，所有只需要实现 Done() 和 Err() 方法即可

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() 只需要返回一个 channel 即可，对于 cancelCtx 只需要返回成员变量 done 即可，由于 cancelCtx 没有初始化函数，所有需要考虑 cancelCtx 中的 done 初始化问题

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
}

• Err() 接口和 Done 接口类似，只需要返回 cancelCtx 中的 err 即可

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

• cancel() 方法是用来取消自己的子 Context，如果 removeFromParent 参数为 true 则从它的父节点删除它

// https://github.com/golang/go/blob/06472b99cdf59f00049f3cd8c9e05ba283cb2c56/src/context/context.go#L344
func (c *cancelCtx) cancel(removeFromParent bool, err error) {

	// 读取关闭原因
	c.err = err

	// 关闭 c.done 通知子 Context
	close(c.done)

	// 遍历子 Context，逐个调用 cancel 方法
	for child := range c.children {
		// NOTE: acquiring the child's lock while holding parent's lock.
		child.cancel(false, err)
	}

	// 将自己从父 Context 删除
	if removeFromParent {
		removeChild(c.Context, c)
	}
}

• WithCancel() 通过 cancelCtx 返回一个实现 Context 接口，主要实现了，初始化一个 cancelCtx 实例，将 cancelCtx 添加到父 Context 中（如果父节点也可以被 cancel 的话）

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

func newCancelCtx(parent Context) cancelCtx {
	return cancelCtx{Context: parent}
}

• propagateCancel() 主要作用是在父 Context 和 子 Context 之间同步取消信号，保证父 Context 被取消的时，子 Context 会收到对应的信号，不会发生状态不一致的问题

1. 当父 Context Done 为空说明不会触发取消事件，直接返回
2. 在 Context 继承链上有节点支持取消上下文时，判断父 Context 是否已经触发了取消信号
   • 已经取消了，则当前 child 会立刻取消
   • 还没有取消，则把当前 Context 加到父 Context 的 child 列表中
3. 如果所有 Context 都不支持 cancel 则会启动一个协程监听父 Context 和子 Context，当父 Context 结束时结束子 Context

func propagateCancel(parent Context, child canceler) {
	if parent.Done() == nil {
		return // parent is never canceled
	}
	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 {
		go func() {
			select {
			case <-parent.Done():
				child.cancel(false, parent.Err())
			case <-child.Done():
			}
		}()
	}
}

timerCtx

timerCtx 在 cancelCtx 的基础上增加了 deadline 用于标识自动 cancel 的最终时间，timer 是一个触发自动 cancel 的定时器。timerCtx 需要在 cancelCtx 的基础上实现 Deadline() 和 cancel() 方法

由此衍生出了 WithDeadline() 和 WithTimeout()

• deadline: 指定最后期限，比如 Context 将于 2019.02.31 11:11:11 时结束
• timeout: 指定多久时间后解除，比如 Context 将于 30s 后结束

type timerCtx struct {
	cancelCtx
	timer *time.Timer // Under cancelCtx.mu.

	deadline time.Time
}

• Deadline() 只需要返回 Context 具体的 deadline 即可

func (c *timerCtx) Deadline() (deadline time.Time, ok bool) {
	return c.deadline, true
}

• cancel() 继承了 cancelCtx 的 cancel 方法，只需要把 timer 关闭即可

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()
}

• WithTimeout() 只是调用了 WithDeadline()

func WithTimeout(parent Context, timeout time.Duration) (Context, CancelFunc) {
	return WithDeadline(parent, time.Now().Add(timeout))
}

• WithDeadline() 初始化 timerCtx，并且启动定时器，在指定时间取消 Context，不同cancel 会有不同的错误原因

1. deadline 没有到来之前手动 cancel，则报 Canceled（具体为: context canceled）
2. 如果是 deadline 自动关闭则报 DeadlineExceeded（具体为: context deadline exceeded）

func WithDeadline(parent Context, d time.Time) (Context, CancelFunc) {

	// 目前的截止日期早于新(d time.Timer)的截止日期。
	if cur, ok := parent.Deadline(); ok && cur.Before(d) {
		// The current deadline is already sooner than the new one.
		return WithCancel(parent)
	}

	// 初始化 cancelCtx，并把节点添加到父 Context 中
	c := &timerCtx{
		cancelCtx: newCancelCtx(parent),
		deadline:  d,
	}
	propagateCancel(parent, c)

	// 当前时间超过了截止时间则调用 cancel 方法取消信号
	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) }
}

参考链接

• Go 专家编程