Context

context 包主要用于在 Goroutine 之间传递上下文信息，包括：取消信号、超时时间、截止时间等。

优点在于在多个 Goroutine 组成的树中同步取消信号以减少对资源的消耗和占用，并且它是并发安全的。

为什么

设置场景：

当请求被取消或是处理时间太长，这有可能是使用者关闭了浏览器或是已经超过了请求方规定的超时时间，请求方直接放弃了这次请求结果。这时，所有正在为这个请求工作的 goroutine 需要快速退出，系统就可以回收相关的资源。

当在业务的高峰期，某个下游服务的响应变慢，而当前系统的请求又没有超时控制，或者超时时间设置地过大，那么等待下游服务返回数据的协程就会越来越多。

协程也是要消耗系统资源的，后果就是协程数激增，内存占用飙涨，甚至导致服务不可用。更严重的会导致雪崩效应，整个服务对外表现为不可用。

context 包就是为了解决上面所说的这些问题而开发的。

举个例子：

func TestContextDemo(t *testing.T) {
	ctx, cancel := context.WithTimeout(context.Background(), 1*time.Second)
	defer cancel()

	// 情况一：当请求只需要 500 ms
	go handle(ctx, 500*time.Millisecond)
	// 情况二：当请求需要 1500 ms
	//go handle(ctx, 1500*time.Millisecond)
	select {
	case <-ctx.Done():
		fmt.Println("main", ctx.Err())
	}
	time.Sleep(1 * time.Second)
}
func handle(ctx context.Context, duration time.Duration) {
	select {
	case <-ctx.Done():
		fmt.Println("handle", ctx.Err())
	case <-time.After(duration):
		fmt.Println("process request with", duration)
	}
}

当请求需要 500 ms时候会打印

=== RUN   TestContextDemo
process request with 500ms
main context deadline exceeded
--- PASS: TestContextDemo (2.02s)

同时你可以看到这三行打印的顺序。

当请求需要 1500 ms时候会打印

=== RUN   TestContextDemo
main context deadline exceeded
handle context deadline exceeded
--- PASS: TestContextDemo (2.03s)

怎么用

用代码来举例吧！

func TestContext(t *testing.T) {
	// 是链路的起点
	ctx := context.Background()
	// 仅在不确定应该使用哪种上下文时使用
	//ctx := context.TODO()
	// 设置键值对
	ctx = context.WithValue(ctx, "key", "value")
	// 取消信号
	ctx, cancel := context.WithCancel(ctx)
	// 一般用完就会进行取消
	//defer cancel()
	go func() {
		time.Sleep(time.Second * 2)
		cancel()
	}()
	// 如果你不进行关闭的话，会一直等待，就不会接着打印。
	<-ctx.Done()
	t.Log("context:", ctx.Err())

	// 设置截止时间
	context.WithDeadline(ctx, time.Now().Add(time.Second))
	// 设置超时时间
	context.WithTimeout(ctx, time.Second)
}

主要常用的是这几个API：

• WithCancel：手动取消上下文。
• WithDeadline：设置取消的截止时间
• WithTimeout：设置取消的超时时间
• WithValue：共享数据

一般常用于这几个场景：

• 传递共享数据：传递请求对应用户的认证令牌以及用于进行分布式追踪的请求 ID等
• 取消 goroutine：等待服务端退出信号、取消HTTP请求的执行、防止 goroutine 泄漏等

所以然

context.Context是接口，定义了四个方法：

1. Deadline — 返回 context.Context 被取消的时间，也就是截止日期；
2. Done — 返回一个 Channel，这个 Channel 会在当前工作完成或者上下文被取消后关闭，多次调用 Done 方法会返回同一个 Channel；
3. Err— 返回context.Context结束的原因，它只会在Done方法对应的 Channel 关闭时返回非空的值；
   1. 如果 context.Context 被取消，会返回 Canceled 错误；
   2. 如果 context.Context 超时，会返回 DeadlineExceeded 错误；
4. Value — 从 context.Context 中获取键对应的值；父节点是无法获取子节点的设置的键值对。

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

第二个比较重要接口：

type canceler interface {
	cancel(removeFromParent bool, err error)
	Done() <-chan struct{}
}

主要实现与取消的相关的方法。

*cancelCtx 和 *timerCtx是这两个结构体的指针实现了 canceler 接口。

Background

context.Background() 这个方法就是返回一个私有结构体 backgroundCtx，TODO方法也是返回特定的结构体todoCtx。

这两个结构体主要通过嵌入方式，包含了emptyCtx 类型与方法。

type backgroundCtx struct{ emptyCtx }

type todoCtx struct{ emptyCtx }


type emptyCtx struct{}

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 any) any {
	return nil
}

emptyCtx就是一个很简单上下文类型，没有参数没有值。

WithValue

context.WithValue这个方法也相对比较简单，返回的是一个valueCtx结构体类型。

func WithValue(parent Context, key, val any) Context {
	//parent key 的校验
	return &valueCtx{parent, key, val}
}

type valueCtx struct {
	Context
	key, val any
}

valueCtx嵌入了 Context 类型，就会继承所有方法，但是它重新实现了Value方法，所以除了 Value 之外的 Err、Deadline 等方法代理到父上下文parent中。

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

func value(c Context, key any) any {
	for {
		switch ctx := c.(type) {
		case *valueCtx:
			if key == ctx.key {
				return ctx.val
			}
			c = ctx.Context
		case *cancelCtx:
			if key == &cancelCtxKey {
				return c
			}
			c = ctx.Context
		case withoutCancelCtx:
			if key == &cancelCtxKey {
				// This implements Cause(ctx) == nil
				// when ctx is created using WithoutCancel.
				return nil
			}
			c = ctx.c
		case *timerCtx:
			if key == &cancelCtxKey {
				return &ctx.cancelCtx
			}
			c = ctx.Context
		case backgroundCtx, todoCtx:
			return nil
		default:
			return c.Value(key)
		}
	}
}

首先会在context.valueCtx内部进行匹配，如果匹配不成功，将会从父上下文context进行查找，需要判断Context是否是valueCtx、cancelCtx、timerCtx以及emptyCtx等；根据不同的类型做不同处理。

• cancelCtx和timerCtx先进行cancelCtxKey判断；

• emptyCtx直接返回nil；

• valueCtx则不断向上去找parent级上下文对应的值；

最后直至找到或者返回nil。

cancelCtxKey，这是一个context包中的私有变量，当对cancelCtx调用Value方法并用这个key作为参数时，返回cancelCtx本身；

// &cancelCtxKey is the key that a cancelCtx returns itself for.
var cancelCtxKey int

func (c *cancelCtx) Value(key any) any {
	if key == &cancelCtxKey {
		return c
	}
	return value(c.Context, key)
}

WithCancel

WithCancel会返回一个新的子context实例并用于取消该上下文的函数。当执行返回的取消函数，当前上下文以及它的子上下文都会被取消，所有的 Goroutine 都会同步收到这一取消信号。

func WithCancel(parent Context) (ctx Context, cancel CancelFunc) {
	c := withCancel(parent)
	return c, func() { c.cancel(true, Canceled, nil) }
}

func withCancel(parent Context) *cancelCtx {
	if parent == nil {
		panic("cannot create context from nil parent")
	}
	c := &cancelCtx{}
	c.propagateCancel(parent, c)
	return c
}

func (c *cancelCtx) propagateCancel(parent Context, child canceler) {
	c.Context = parent
	// 父节点是emptyCtx类型
	done := parent.Done()
	if done == nil {
		return // parent is never canceled
	}
	// 如果父上下文不是emptyCtx类型，使用select来判断一下父上下文的done channel是不是已经被关闭掉了
    // 关闭则调用child的cancel方法
    // select其实会阻塞，但这里给了一个default方法，所以如果父上下文的done channel没有被关闭则继续之心后续代码
    // 这里相当于利用了select的阻塞性来做if-else判断
	select {
	case <-done:
		// parent is already canceled
		child.cancel(false, parent.Err(), Cause(parent))
		return
	default:
	}
	// parentCancelCtx目的在于寻找父上下文中最底层的cancelCtx，因为像timerCtx等结构体会内嵌cancelCtx
	if p, ok := parentCancelCtx(parent); ok {
		// parent is a *cancelCtx, or derives from one.
        // 设置内层 cancelCtx 与 child 之间的关联关系
		p.mu.Lock()
		if p.err != nil {
			// parent has already been canceled
			child.cancel(false, p.err, p.cause)
		} else {
            // 父上下文未被取消，创建子节点，并挂靠上去
			if p.children == nil {
				p.children = make(map[canceler]struct{})
			}
			p.children[child] = struct{}{}
		}
		p.mu.Unlock()
		return
	}
	// 判断是否存在钩子函数
	if a, ok := parent.(afterFuncer); ok {
		// parent implements an AfterFunc method.
		c.mu.Lock()
        // 添加一个取消子上下文的钩子函数，返回钩子函数的 【取消注册 AfterFunc 的停止函数】。
		stop := a.AfterFunc(func() {
			child.cancel(false, parent.Err(), Cause(parent))
		}) 
        // 创建 stopCtx 作为 cancel父级上下文。
		c.Context = stopCtx{
			Context: parent,
			stop:    stop,
		}
		c.mu.Unlock()
		return
	}

	goroutines.Add(1)
	go func() {
		select {
            //当父节点取消，则取消子节点。
		case <-parent.Done():
			child.cancel(false, parent.Err(), Cause(parent))
            //子节点自己取消了，那就退出 select。
		case <-child.Done():
		}
	}()
}

propagateCancel作用是在 parent 和 child 之间同步取消和结束的信号，保证在 parent 被取消时，child 也会收到对应的信号，不会出现状态不一致的情况。

func parentCancelCtx(parent Context) (*cancelCtx, bool) {
	done := parent.Done()
	if done == closedchan || done == nil {
		return nil, false
	}
	p, ok := parent.Value(&cancelCtxKey).(*cancelCtx)
	if !ok {
		return nil, false
	}
	pdone, _ := p.done.Load().(chan struct{})
	if pdone != done {
		return nil, false
	}
	return p, true
}

parentCancelCtx核心代码是通过cancelCtxKey获取父上下文中的内嵌cancelCtx，供后续设置关联关系。其余的就是考虑到的额外情况：

1. 如果parent的done已经是closedchan或者是nil，就没必要获取内层的cancelCtx来建立层级关系，直接用parent本身与child做好关联cancel即可。
2. 对于内嵌cancelCtx可能加了一些自定义方法，比如复写了Done或者cancel，那么这个情况下，用户自己负责处理。

对于 cancelCtx 的 cancel 方法，也是核心代码。

func (c *cancelCtx) cancel(removeFromParent bool, err, cause error) {
    // 必须传错误
	if err == nil {
		panic("context: internal error: missing cancel error")
	}
	if cause == nil {
		cause = err
	}
    // 后续需要更新，进行抢锁
	c.mu.Lock()
    // 因为在抢锁过程中，可能c.err已经被某个协程修改了，需要对c.err判断是否还是nil，才能进行更新
	if c.err != nil {
		c.mu.Unlock()
		return // already canceled
	}
    // 赋值
	c.err = err
	c.cause = cause
    // 读取done的值
	d, _ := c.done.Load().(chan struct{})
	if d == nil {
		c.done.Store(closedchan)
	} else {
		close(d)
	}
    // 遍历c的children调用他们的cancel；
	for child := range c.children {
		// NOTE: acquiring the child's lock while holding parent's lock.
		child.cancel(false, err, cause)
	}
	c.children = nil
	c.mu.Unlock()
	// 从父节点中移除自己 
	if removeFromParent {
		removeChild(c.Context, c)
	}
}

总体来看，cancel() 方法的功能就是关闭 channel：c.done；递归地取消它的所有子节点；从父节点从删除自己。

核心是通过关闭channel，将取消的信号传递给所有的子节点，propagateCancel起了一个协程，来获取取消信号。

在调用子节点 cancel 方法的时候，传入的第一个参数 removeFromParent 是 false，可是在WithCancel函数中调用cancel方法传入的是 true。

func removeChild(parent Context, child canceler) {
    // 当父上下文是 stopCtx 需要取消注册的AfterFunc函数
	if s, ok := parent.(stopCtx); ok {
		s.stop()
		return
	}
    // 找到 CancelCtx
	p, ok := parentCancelCtx(parent)
	if !ok {
		return
	}
    // 加锁，保证原子性
	p.mu.Lock()
    // 删除子节点
	if p.children != nil {
		delete(p.children, child)
	}
	p.mu.Unlock()
}

问：什么时候会传 true 呢？

答：创建一个可取消的 context 节点时，调用返回的 cancelFunc 时，会将这个 context 从它的父节点里删除，所以context 节点的子节点也都被取消了。

其实 stopCtx是主要用作保存取消注册AfterFunc 的停止函数，是cancelCtx的父级上下文。

WithTimeout\WithDeadline

这两个函数都是由timerCtx创建的，timerCtx内嵌了cancelCtx，在基础上携带了一个计时器和截止时间。

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

	deadline time.Time
}

先来看看WithDeadline函数。

func WithDeadline(parent Context, d time.Time) (Context, CancelFunc) {
	return WithDeadlineCause(parent, d, nil)
}


func WithDeadlineCause(parent Context, d time.Time, cause error) (Context, CancelFunc) {
	if parent == nil {
		panic("cannot create context from nil parent")
	}
    // 如果parent的deadline小于当前时间，直接创建cancelCtx，里面会调用propagateCancel方法
	if cur, ok := parent.Deadline(); ok && cur.Before(d) {
		// The current deadline is already sooner than the new one.
		return WithCancel(parent)
	}
    // 创建timerCtx，
	c := &timerCtx{
		deadline: d,
	}
     // 设置层级取消关联
	c.cancelCtx.propagateCancel(parent, c)
	dur := time.Until(d)
    // 如果已经超时直接取消
	if dur <= 0 {
		c.cancel(true, DeadlineExceeded, cause) // deadline has already passed
		return c, func() { c.cancel(false, Canceled, nil) }
	}
	c.mu.Lock()
	defer c.mu.Unlock()
    // 如果没有超时并且没有被调用过cancel，那么设置timer，超时则调用cancel方法；      
	if c.err == nil {
		c.timer = time.AfterFunc(dur, func() {
			c.cancel(true, DeadlineExceeded, cause)
		})
	}
	return c, func() { c.cancel(true, Canceled, nil) }
}

func (c *timerCtx) cancel(removeFromParent bool, err, cause error) {
    // 直接调用 cancelCtx 的取消方法
	c.cancelCtx.cancel(false, err, cause)
	if removeFromParent {
		// Remove this timerCtx from its parent cancelCtx's children.
		removeChild(c.cancelCtx.Context, c)
	}
	c.mu.Lock()
    // 关掉定时器，这样，在deadline 到来时，不会再次取消
	if c.timer != nil {
		c.timer.Stop()
		c.timer = nil
	}
	c.mu.Unlock()
}

WithTimeout函数就更加简单了。

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

参考资料

• www.cnblogs.com/qcrao-2018/…
• 【Go语言设计与实现】draveness.me/golang/docs…
• www.cnblogs.com/dojo-lzz/p/…