1. 前言

Context 是在 go1.7 版本引入的一个包，context 主要用来在 goroutine 之间传递上下文信息，包括：取消信号、超时时间、截止时间、k-v 等。之后的 go 版本 context package 也进行了许多更新，截至写本文时最新go1.21 版本加入了 AfterFunc API（源码在此），并且实现了 context merge 功能（在此之前也有第三方的 package 实现了该功能）。如果你之前并没有了解过 context，可以先阅读 go 官方 blog 中的文章 Go Concurrency Patterns: Context 来了解 context 的实际应用。

2. Context 能够做什么？

想象你写了一个 http server，server 需要接受来自 client 的请求。一般一个请求都会需要若干个 goroutine 同时工作。比如有的需要去查询数据库，有需要去调用下游接口的服务等等，这样就形成了下图所示的 goroutine之间的层级关系。

context.1.png

在这些 goroutine 需要共享这个请求的一些数据，比如用于鉴权的 token 和请求的最大超时时间。假设 client 因为超时关闭了连接，如果此时 server 还在处理该请求就会造成 server 资源的浪费，所以就需要一种机制能够将这种超时信号传递给和该请求相关的 goroutines。

context package 就是为了解决上述问题而设计的：在一组相关的 goroutines 之间传递共享的值，取消信号，超时时间等。

通过 context 在 goroutines 之间传递信息可由下图所示：

context.2.png

3. 源码分析

因为 AfterFunc API 是在 go1.21 （源码在此）中才引入的，所以本文先介绍 go1.20（源码在此）版本的 context package，后面介绍 AfterFunc API 时，再来介绍新引入的内容。

3.1 context interface

// A Context carries a deadline, a cancellation signal, and other values across
// API boundaries.
//
// Context's methods may be called by multiple goroutines simultaneously.
type Context interface {
	Deadline() (deadline time.Time, ok bool)
	Done() <-chan struct{}
        Err() error
	Value(key any) any
}

在 context interface 中声明了四个方法，分别为 Deadline，Done，Err，Value。

当 context 有截止时间（下文介绍的 timerCtx 和 deadlineCtx 属于这一类），返回的 ok == true，并且 deadline 的值为截止时间。如果 context 没有截止时间（下文介绍的 cancelCtx 和 valueCtx 属于这一类），ok == false，表示 deadline 没有被设置。
Done 返回的 channel c，只用于 close(c)，当某个 v := <-ctx.Done() 能够读取到值。则说明该 ctx 被关闭或者其 parent ctx 被关闭。
ctx.Done() 返回的 channel c 未被关闭，Err 方法返回 nil，如果通道 c 已关闭，Err 方法返回非 nil 的错误，可能的原因有如下：如果 context 已被取消，则返回 Canceled 错误。如果 context 的截止时间已过，则返回 DeadlineExceeded 错误。在 Err 返回非 nil 错误之后，后续对 Err 的调用将返回相同的错误。
Value 方法返回 ctx 中指定 key 的 value，如果该 key 没有关联的 value，则返回 nil。对于相同的 key，多次调用 Value 方法会返回相同的 value。

Canceled 和 DeadlineExceeded 定义如下：

var Canceled = errors.New("context canceled")

// DeadlineExceeded is the error returned by Context.Err when the context's
// deadline passes.
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 }

3.2 empty context

Context package 中定义了 empty context，该 context 主要的使用的场景是：当你不知道使用哪种 context 时，可以使用 emptyCtx。其中对外暴露了两个接口来使用 emptyCtx：context.Background() 和 context.TODO()。

emptyCtx 定义如下：

emptyCtx context 永远不会被取消，没有任何值，也没有截止时间。它不是 struct{} 类型（其底层类型为 int），因为该类型的变量必须具有不同的地址。

// An emptyCtx is never canceled, has no values, and has no deadline. It is not
// struct{}, since vars of this type must have distinct addresses.
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 any) any {
	return nil
}

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

context.Background() 和 context.TODO() 定义如下：

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

// Background returns a non-nil, empty Context. It is never canceled, has no
// values, and has no deadline. It is typically used by the main function,
// initialization, and tests, and as the top-level Context for incoming
// requests.
func Background() Context {
	return background
}

// TODO returns a non-nil, empty Context. Code should use context.TODO when
// it's unclear which Context to use or it is not yet available (because the
// surrounding function has not yet been extended to accept a Context
// parameter).
func TODO() Context {
	return todo
}

从源码可以看出，Background() 和 TODO() 函数返回的类型 *emptyCtx。

3.3 cancel context

Context package 中定义了 cancelCtx 用来传递取消信号，后面介绍的 timerCtx 也是是基于 cancelCtx 此实现的。

3.3.1 canceler interface 定义

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

3.3.2 cancelCtx struct 定义

// 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     atomic.Value          // of 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
	cause    error                 // set to non-nil by the first cancel call
}

Context 表示 parent context。
mu 用于多个 goroutines 互斥访问其它 filed，也就是说明 cancelCtx 是多线程安全的。
done 表示用来取消的 channel
children 表示 child context，如果该 context 被取消，则 child context 中能被取消的 context 都会被取消。
err 当第一次调用 cancel 函数时，设置为非 nil 值。
cause 当第一次调用 cancel 函数时，设置为非 nil 值。

3.3.3 创建 cancelCtx

使用对外开放的接口 WithCancel 创建新的 cancelCtx。

// A CancelFunc tells an operation to abandon its work.
// A CancelFunc does not wait for the work to stop.
// A CancelFunc may be called by multiple goroutines simultaneously.
// After the first call, subsequent calls to a CancelFunc do nothing.

type CancelFunc func()

func WithCancel(parent Context) (ctx Context, cancel CancelFunc) {
        //调用 withCancel 创建 cancelCtx
	c := withCancel(parent)
	return c, func() { c.cancel(true, Canceled, nil) }
}

其中我们可以看到返回的 CancelFunc 可以取消 context。如果你想主动取消 concelCtx。可以这么做：

ctx, cancel := WithCancel(context.Background)
 //主动调用 cancel func 取消 ctx。
cancel()

withCancel 代码如下：

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

newCancelCtx 定义如下：返回新的新的 *cancelCtx。

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

再介绍 progateCancel 函数之前，我们先来看一个创建 cancelCtx 的另一个函数 WithCancelCause。

// A CancelCauseFunc behaves like a CancelFunc but additionally sets the cancellation cause.
// This cause can be retrieved by calling Cause on the canceled Context or on
// any of its derived Contexts.
//
// If the context has already been canceled, CancelCauseFunc does not set the cause.
// For example, if childContext is derived from parentContext:
//   - if parentContext is canceled with cause1 before childContext is canceled with cause2,
//     then Cause(parentContext) == Cause(childContext) == cause1
//   - if childContext is canceled with cause2 before parentContext is canceled with cause1,
//     then Cause(parentContext) == cause1 and Cause(childContext) == cause2

type CancelCauseFunc func(cause error)

func WithCancelCause(parent Context) (ctx Context, cancel CancelCauseFunc) {
	c := withCancel(parent)
	return c, func(cause error) { c.cancel(true, Canceled, cause) }
}

可以看出 WithCancelCause 函数返回的 CancelCauseFunc 可以传入 cause error。

ctx, cancel := WithCancelCause(context.Background)
 //主动调用 cancel func 取消 ctx。
err := errors.New("timeout")
cancel(err)

3.3.4 Cause 函数

Cause 函数返回一个非 nil 的错误，解释了为什么 c 被取消。

c 或者 parent 其中之一的第一次取消设置了 cause。

如果该取消是通过调用 CancelCauseFunc(err) 发生的，则 Cause 返回 err。

否则，Cause(c) 返回与 c.Err() 相同的值。

如果 c 尚未被取消，则 Cause 返回 nil。

// Cause returns a non-nil error explaining why c was canceled.
// The first cancellation of c or one of its parents sets the cause.
// If that cancellation happened via a call to CancelCauseFunc(err),
// then Cause returns err.
// Otherwise Cause(c) returns the same value as c.Err().
// Cause returns nil if c has not been canceled yet.

func Cause(c Context) error {
	if cc, ok := c.Value(&cancelCtxKey).(*cancelCtx); ok {
		cc.mu.Lock()
		defer cc.mu.Unlock()
		return cc.cause
	}
	return nil
}

3.3.5 propagateCancel 函数

propagateCancel 函数用于建立 parent context 和 child context 之间的 取消信号传递的层级关系。

// propagateCancel arranges for child to be canceled when parent is.
func propagateCancel(parent Context, child canceler) {
	done := parent.Done()
   //如果 parent 不能够被取消，直接返回。比如 parent = context.Background() 
	if done == nil {
		return // parent is never canceled
	}
        
   //parent 是否被取消，如果已经被取消了，则直接取消 child。
	select {
	case <-done:
		// parent is already canceled
		child.cancel(false, parent.Err(), Cause(parent))
		return
	default:
	}
        
   //找到真正的 cancel context，可能 parent context 的某个 embedded filed 才是个 cancelCtx。
	if p, ok := parentCancelCtx(parent); ok {
		p.mu.Lock()
           //如果 parent 已经被取消，则直接取消 child。
		if p.err != nil {
			// parent has already been canceled
			child.cancel(false, p.err, p.cause)
		} else {
                   //如果 p.children 属性为 nil，则新建 p.children
			if p.children == nil {
				p.children = make(map[canceler]struct{})
			}
                  在 p.children 中 添加 child:struct{}{} 键值对，构建父子层级关系。
			p.children[child] = struct{}{}
		}
		p.mu.Unlock()
	} else {
        // goroutines.Add(1) 仅用作测试
		goroutines.Add(1)
        // 如果没有找到可取消的父 context。新启动一个 goroutine 监控父节点或子节点取消信号
		go func() {
			select {
			case <-parent.Done():
				child.cancel(false, parent.Err(), Cause(parent))
			case <-child.Done():
			}
		}()
	}
}

变量 goroutines 仅用作测试

// goroutines counts the number of goroutines ever created; for testing.
var goroutines atomic.Int32

这个函数的作用就是构建 parent 和 children 之间的层级关系，当 parent cancel 时，则 child 也能够收到取消的信号。

3.3.6 parentCancelCtx 函数

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

// parentCancelCtx returns the underlying *cancelCtx for parent.
// It does this by looking up parent.Value(&cancelCtxKey) to find
// the innermost enclosing *cancelCtx and then checking whether
// parent.Done() matches that *cancelCtx. (If not, the *cancelCtx
// has been wrapped in a custom implementation providing a
// different done channel, in which case we should not bypass it.)
func parentCancelCtx(parent Context) (*cancelCtx, bool) {
	done := parent.Done()
	if done == closedchan || done == nil {
		return nil, false
	}
  //如果 parent 是 cancelCtx 类型，则下面调用 cancelCtx 的 Value 方法。
	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 函数返回 parent 的底层 *cancelCtx。它通过查找 parent.Value(&cancelCtxKey) 来找到最内部的 *cancelCtx，然后检查 parent.Done() 是否与该*cancelCtx 匹配。（如果不匹配，则 *cancelCtx 已被包装在提供不同 done channel 的自定义实现中，在这种情况下，我们不应绕过它。）

这里需要注意的两点是：

实现了 context 接口的 parent 其实际类型不一定是 *cancelCtx，比如 timerContext，timerContext struct（后面会有介绍）有一个 field 是 cancelCtx。所以需要 parent.Value(&cancelCtxKey) 找到其最内部的 cancelCtx。
如果有一个 parent 是自定义的 customContext struct，实现了 context 接口，且其有一个 field 类型是 cancelCtx，并且 customContext 自己实现 Done() 方法，这就是为什么需要如下代码进行判断：

pdone, _ := p.done.Load().(chan struct{})
	if pdone != done {
		return nil, false
	}

concelCtx 的 Value 方法定义如下：

func (c *cancelCtx) Value(key any) any {
// key == &cancelCtxKey 则直接返回 c。
	if key == &cancelCtxKey {
		return c
	}
// 否则调用 value 函数。
	return value(c.Context, key)
}

value 的函数定义如下：

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 *timerCtx:
			if key == &cancelCtxKey {
				return ctx.cancelCtx
			}
			c = ctx.Context
		case *emptyCtx:
			return nil
		default:
			return c.Value(key)
		}
	}
}

可以看出 value 是一个 逐级向上 查找对应 key 的 value 的过程。

3.3.7 cancel 函数

cancel 函数需要将 cancel 信号传递给它的所有 children，定义如下：

// cancel closes c.done, cancels each of c's children, and, if
// removeFromParent is true, removes c from its parent's children.
// cancel sets c.cause to cause if this is the first time c is canceled.

func (c *cancelCtx) cancel(removeFromParent bool, err, cause error) {
   //如果 err == nil，直接 panic
	if err == nil {
		panic("context: internal error: missing cancel error")
	}
   //如果没有给 cause 赋值，cause = err
	if cause == nil {
		cause = err
	}
	c.mu.Lock()
   //如果已经被取消，则直接 return
	if c.err != nil {
		c.mu.Unlock()
		return // already canceled
	}
   //对 c.err 和 c.cause 进行赋值
	c.err = err
	c.cause = cause
   //如果 done 的值为 nil，则赋予一个 closechan。在 init 函数中，该 closedchan 已经被 close。
   //若不为 nil，直接 close。
	d, _ := c.done.Load().(chan struct{})
	if d == nil {
		c.done.Store(closedchan)
	} else {
		close(d)
	}
   //遍历 c.children，取消 child。
	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()
  //需要将自身从 parent 中的 children 移除，则进行移除。
	if removeFromParent {
		removeChild(c.Context, c)
	}
}

closedchan 是一个可以重用的变量，并且在 context package 的 init 函数中被关闭。

// closedchan is a reusable closed channel.
var closedchan = make(chan struct{})

func init() {
	close(closedchan)
}

另外需要注意的是，当 WithCancel 函数返回的 cancel 函数中，调用 cancelCtx 的 cancel 函数，第一个参数 removeFromParent 的值为 true。

func() { c.cancel(true, Canceled, nil) }

而在 cancelCtx 的 cancel 中，将 children 从 parent 中移除，调用的 cancel 函数的第一个参数 removeFromParent 的值为 false。

child.cancel(false, err, cause)

当 context1 将取消信号传递给 context 2时（如下图所示），调用 cancel 函数，context1 需要将自身从其 parent 的 children 中移除，所以需要将 removeFromParent 的值设置 true。在 cancel 函数中，调用 children （context 2）的 cancel 函数时，removeFromParent 参数可以设置为 false，因为此时 context2 已经从 context 1 的 children 中移除了，所以没有必要再将其从 context1 中的 children 中移除第二遍。

context.3.png

3.3.8 removeChild 函数

removeChild 函数将 child 从 parent 中移除，消除它们之间的层级关系，代码如下：

// removeChild removes a context from its parent.

func removeChild(parent Context, child canceler) {
	p, ok := parentCancelCtx(parent)
	if !ok {
		return
	}
	p.mu.Lock()
//如果 p.children 不为空，则删除 child。
	if p.children != nil {
		delete(p.children, child)
	}
	p.mu.Unlock()
}

3.3.9 其它函数

cancelCtx 的 Done() 函数采用了惰性初始化的方式，当需要时再创建。代码如下：

func (c *cancelCtx) Done() <-chan struct{} {
	d := c.done.Load()
	if d != nil {
		return d.(chan struct{})
	}
	c.mu.Lock()
	defer c.mu.Unlock()
	d = c.done.Load()
	if d == nil {
		d = make(chan struct{})
		c.done.Store(d)
	}
	return d.(chan struct{})
}

cancelCtx 的 Err() 函数如下，返回 err 的值。代码如下：

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

cancelCtx 的 String() 函数代码如下：

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"
}

3.4 timerCtx

3.4.1 timerCtx struct 定义

timerCtx 的定义如下：

// A timerCtx carries a timer and a deadline. It embeds a cancelCtx to
// implement Done and Err. It implements cancel by stopping its timer then
// delegating to cancelCtx.cancel.

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

	deadline time.Time
}

*cancelCtx cancelCtx 指针。
timer 定时器，当触发定时器，就执行 cancelCtx 的 cancel 函数。
deadline 表示定时器触发时间。

3.4.2 创建 timerCtx

使用开放的接口 WithDeadline 可以创建 timerCtx。代码如下：

// WithDeadline returns a copy of the parent context with the deadline adjusted
// to be no later than d. If the parent's deadline is already earlier than d,
// WithDeadline(parent, d) is semantically equivalent to parent. The returned
// context's Done channel is closed when the deadline expires, when the returned
// cancel function is called, or when the parent context's Done channel is
// closed, whichever happens first.
//
// Canceling this context releases resources associated with it, so code should
// call cancel as soon as the operations running in this Context complete.

func WithDeadline(parent Context, d time.Time) (Context, CancelFunc) {
   //parent 为 nil，直接 panic。
	if parent == nil {
		panic("cannot create context from nil parent")
	}
   //和 parent 的deadline 进行比较，如果 parent 有 deadline，且 parent deadline 早于传入的参数 d。直接返回一个以 parent 的 cancelCtx。
	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{
		cancelCtx: newCancelCtx(parent),
		deadline:  d,
	}
   //构建 cancel 信号层级
	propagateCancel(parent, c)
   //计算到 deadline 的持续时间
	dur := time.Until(d)
   //如果持续时间 <= 0。则直接取消创建的 timerCtx，并且返回 c 和 cancelFunc。
	if dur <= 0 {
		c.cancel(true, DeadlineExceeded, nil) // deadline has already passed
		return c, func() { c.cancel(false, Canceled, nil) }
	}
	c.mu.Lock()
	defer c.mu.Unlock()
   // 如果 c.err == nil，创建定时器任务，当定时器被触发，调用 timerCtx 的 cancel 函数，传播取消信号。
	if c.err == nil {
		c.timer = time.AfterFunc(dur, func() {
			c.cancel(true, DeadlineExceeded, nil)
		})
	}
	return c, func() { c.cancel(true, Canceled, nil) }
}

可以看出创建 timerCtx 的过程，判断了许多条件。如下：

parent 的 deadline 是否早于自己，如果是，则直接创建 cancleCtx。
参数 d 与当前时间比较，如果剩余 dur <= 0，则直接调用新创建的 timerCtx 的 cancel 函数。

并且 timerCancel 取消传递的信号的传递也是借助于 cancelCtx 实现的。

3.4.3 timerCtx cancel 函数

timerCtx 的 cancel 函数的过程如下：

先调用 field cancelCtx 的 cancel 函数。
将自身从 parent 中移除
停止定时器。

func (c *timerCtx) cancel(removeFromParent bool, err, cause error) {
	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()
	if c.timer != nil {
		c.timer.Stop()
		c.timer = nil
	}
	c.mu.Unlock()
}

3.4.4 其它函数

timerCtx 的 Deadline 函数定义如下：

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

timerCtx 的 String 函数定义如下：

func (c *timerCtx) String() string {
	return contextName(c.cancelCtx.Context) + ".WithDeadline(" +
		c.deadline.String() + " [" +
		time.Until(c.deadline).String() + "])"
}

3.5 Withtimeout 接口

从代码可以看出 Withtimeout 函数创建的仍然是一个 timerCtx，只不过 time.Now 加上 time.Duration 生成一个 deadline 参数，调用 WithDeadline 函数。

// WithTimeout returns WithDeadline(parent, time.Now().Add(timeout)).
//
// Canceling this context releases resources associated with it, so code should
// call cancel as soon as the operations running in this Context complete:
//
//	func slowOperationWithTimeout(ctx context.Context) (Result, error) {
//		ctx, cancel := context.WithTimeout(ctx, 100*time.Millisecond)
//		defer cancel()  // releases resources if slowOperation completes before timeout elapses
//		return slowOperation(ctx)
//	}

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

3.6 valueCtx

3.6.1 valueCtx struct 定义

valueCtx 包括三个部分：

Context
key
val

valueCtx 包含一个键值对和一个 Context。

// A valueCtx carries a key-value pair. It implements Value for that key and
// delegates all other calls to the embedded Context.
type valueCtx struct {
	Context
	key, val any
}

3.6.2 Value 函数

valueCtx 的 Value 首先判断 key 是否与自身的 c.key 相等，如果相等，直接返回 c.val，不相等，则调用 value 函数（3.3.6 小节中有相关的代码）。

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

valueCtx 查询 key 的过程可由下图表示：

valueCtx.drawio.png

因为查找方向是往上走的，所以，父节点没法获取子节点存储的值，子节点却可以获取父节点的值。

3.6.3 其它函数

valueCtx 的 String 函数，代码如下：

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

stringify 函数，代码如下：

// 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 any) string {
	switch s := v.(type) {
	case stringer:
		return s.String()
	case string:
		return s
	}
	return "<not Stringer>"
}

4. context package go1.21 版本

在第三节源码分析中已经介绍了 go1.20 版本的 context package，最新的 go1.21 版本中引入了 AfterFunc API 以及相关的内容。

4.1 backgroundCtx and todoCtx

type backgroundCtx struct{ emptyCtx }

func (backgroundCtx) String() string {
	return "context.Background"
}

type todoCtx struct{ emptyCtx }

func (todoCtx) String() string {
	return "context.TODO"
}

// Background returns a non-nil, empty [Context]. It is never canceled, has no
// values, and has no deadline. It is typically used by the main function,
// initialization, and tests, and as the top-level Context for incoming
// requests.

func Background() Context {
	return backgroundCtx{}
}

// TODO returns a non-nil, empty [Context]. Code should use context.TODO when
// it's unclear which Context to use or it is not yet available (because the
// surrounding function has not yet been extended to accept a Context
// parameter).

func TODO() Context {
	return todoCtx{}
}

从上面可以看出，定义了两个新的结构体 backgroundCtx 和 todoCxt，且均只包含一个字段：emptyCtx。 Background() 和 TODO() 函数分别返回默认初始化的结构体 backgroundCtx 和 todoCtx。

4.2 WithoutCancel

新增加了一个 withoutCancleCtx 结构体类型，定义如下：

type withoutCancelCtx struct {
	c Context
}

利用 WithoutCancel 创建新的 withoutCancelCtx。

func WithoutCancel(parent Context) Context {
	if parent == nil {
		panic("cannot create context from nil parent")
	}
	return withoutCancelCtx{parent}
}

WithoutCancel 返回一个 parent 的副本，当 parent 被取消时，不会被取消。
返回的 withoutCancelCtx 没有 Deadline 或 Err，并且其 Done channel 为nil。
在返回的 withoutCancelCtx 上调用 Cause 会返回nil。

func (withoutCancelCtx) Deadline() (deadline time.Time, ok bool) {
	return
}

func (withoutCancelCtx) Done() <-chan struct{} {
	return nil
}

func (withoutCancelCtx) Err() error {
	return nil
}

func (c withoutCancelCtx) Value(key any) any {
	return value(c, key)
}

func (c withoutCancelCtx) String() string {
	return contextName(c.c) + ".WithoutCancel"
}

4.3 value 函数

因为新增了 backgroundCxt，todoCtx 和 withoutCancelCtx 结构体类型，value 函数的代码如下：

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

4.3 AfterFunc API

在 go proposal 中 context: add Merge #36503，大佬们在讨论是否需要加入 Merge 这个API，函数签名如下面所示：

func Merge(parent1, parent2 context.Context) (context.Context, context.CancelFunc)

Merge 将 context 2 合并到 context 1，并且返回一个新的 result context。

当 parentt1 done 或者 parent 2 done 或者它们的 cancel called。则 result context done。
result context 的 deadline = min(parent1.Deadline, parent2.Deadline)
result context 具有 context parent1 和 parent2 中的 k-v（通过value 函数调用获得），并且如果 parent1 和 parent2 具有相同的 key，则优先使用 parent1 中的值。

经过大佬们一系列讨论，最终 Damien Neil 提了个 proposal context: add AfterFunc #57928，如果想详细了解过程，可以阅读此链接中的内容。

4.3.1 afterFuncCtx

定义了一个新的结构体 afterFuncCtx，其中包含三个字段：

cancelCtx 用于取消
once sync.Once 用于调用函数 f 或者终止 f 的运行
f 函数

type afterFuncCtx struct {
	cancelCtx
	once sync.Once // either starts running f or stops f from running
	f    func()
}

afterFuncCtx 的 cancel 函数：

func (a *afterFuncCtx) cancel(removeFromParent bool, err, cause error) {
  //取消 cancelCtx
	a.cancelCtx.cancel(false, err, cause)
	if removeFromParent {
           //从 parent context 中 移除 a。
		removeChild(a.Context, a)
	}
  //执行 f 函数
	a.once.Do(func() {
		go a.f()
	})
}

afterFuncCtx 被取消时，会调用 f 函数，并且只会调用一次。

4.3.2 AfterFunc 函数

AfterFunc 函数定义如下：

// AfterFunc arranges to call f in its own goroutine after ctx is done
// (cancelled or timed out).
// If ctx is already done, AfterFunc calls f immediately in its own goroutine.
//
// Multiple calls to AfterFunc on a context operate independently;
// one does not replace another.
//
// Calling the returned stop function stops the association of ctx with f.
// It returns true if the call stopped f from being run.
// If stop returns false,
// either the context is done and f has been started in its own goroutine;
// or f was already stopped.
// The stop function does not wait for f to complete before returning.
// If the caller needs to know whether f is completed,
// it must coordinate with f explicitly.
//
// If ctx has a "AfterFunc(func()) func() bool" method,
// AfterFunc will use it to schedule the call.

func AfterFunc(ctx Context, f func()) (stop func() bool) {
	a := &afterFuncCtx{
		f: f,
	}
	a.cancelCtx.propagateCancel(ctx, a)
	return func() bool {
		stopped := false
		a.once.Do(func() {
			stopped = true
		})
		if stopped {
			a.cancel(true, Canceled, nil)
		}
		return stopped
	}
}

AfterFunc 传入两个参数 ctx 和 f。

其中 f 用于创建 afterFuncCtx a，利用 a.cancelCtx 的 propagateCancel 来构建 ctx 和 a 的层级关系。

在返回的 stop 函数中，如果 a.once.Do 能够被调用（也就是 a 并未被取消，如果取消，在 afterFuncCtx 中已经调用了 a.once.Do 函数，则在 stop 中就不能调用 a.once.Do），再调用 a.cancel(true, Canceled, nil) 函数（保证 a.cancel 只能调用一次），返回 true。否则，返回false。

一个使用 AfterFunc 的例子

package main

import (
	"context"
	"fmt"
	"sync"
	"time"
)

func main() {
	waitOnCond := func(ctx context.Context, cond *sync.Cond) error {
		stopf := context.AfterFunc(ctx, cond.Broadcast)
		defer stopf()
		cond.Wait()
		return ctx.Err()
	}

	ctx, cancel := context.WithTimeout(context.Background(), 1*time.Millisecond)
	defer cancel()

	var mu sync.Mutex
	cond := sync.NewCond(&mu)

	mu.Lock()
	err := waitOnCond(ctx, cond)
	fmt.Println(err)

}

调用 waitOnCond，当 ctx 的 deadline 已到，就会传递取消信号到 afterFuncCtx （AfterFunc 中创建的，并且与 timerCtx 构成了层级关系）。然后 afterFuncCtx 调用自身的 cancel 函数。

这样就实现了当传递给 AfterFunc 的第一个参数 ctx 取消时，就可以调用传递给 AfterFunc 的第二个参数 f。调用如下代码：

a.once.Do(func() {
		go a.f()
	})

可由如下示意图所示：

afterFunc.drawio.png

AfterFunc 也可用于 merge context。示列代码如下：

package main

import (
	"context"
	"errors"
	"fmt"
)

// mergeCancel returns a context that contains the values of ctx,
// and which is canceled when either ctx or cancelCtx is canceled.
func mergeContext(ctx, cancelCtx context.Context) (context.Context, context.CancelFunc) {
    mergedCtx, cancel := context.WithCancelCause(ctx)
    stop := context.AfterFunc(cancelCtx, func() {
        cancel(context.Cause(cancelCtx))
    })
    return mergedCtx, func() {
        stop()
	     cancel(context.Canceled)
    }
}

func main() {
	ctx1, cancel1 := context.WithCancelCause(context.Background())
	defer cancel1(errors.New("ctx1 canceled"))

	ctx2, cancel2 := context.WithCancelCause(context.Background())

	mergedCtx, mergedCancel := mergeContext(ctx1, ctx2)
	defer mergedCancel()

	cancel2(errors.New("ctx2 canceled"))
	<-mergedCtx.Done()
	fmt.Println(context.Cause(mergedCtx))
}

Output:

ctx2 canceled 上述代码将 ctx1 和 ctx2 合并到 mergedCtx。可以先看一下 ctx 的示意图：

merged.drawio.png

从上图可以看出，当 ctx1 或者 ctx2 取消时，能够将 cancel 信号传递到 mergedCtx。

调用 ctx1 的 cancel1 func 取消，是因为在 mergeContext 函数，使用 context.WithCancelCause(ctx) 函数，构建了 ctx1 和 mergedCtx 之间的层级关系。
调用 ctx2 的 cancel2 func 取消，是因为在 mergeContext 中调用了 context.AfterFunc 函数，构建了 ctx2 和 mergedCtx 之间的层级关系。
如果调用 mergedCancel 函数，在 stop 函数，会将 mergedCtx 与 ctx2 通过 context.AfterFunc 建立的层级关系消除掉，调用 context.WithCancelCause 函数返回的 cancel 函数，将 ctx1 与 mergedCtx 之间的层级关系消除掉。

这样使用 AfterFunc 函数可以完成 merge context 的作用。

4.3.3 propagateCancel

cancelCtx 的 propagateCancel 函数中，也进行了更新。

// propagateCancel arranges for child to be canceled when parent is.
// It sets the parent context of cancelCtx.

func (c *cancelCtx) propagateCancel(parent Context, child canceler) {
	c.Context = parent

	done := parent.Done()
	if done == nil {
		return // parent is never canceled
	}

	select {
	case <-done:
		// parent is already canceled
		child.cancel(false, parent.Err(), Cause(parent))
		return
	default:
	}

	if p, ok := parentCancelCtx(parent); ok {
		// parent is a *cancelCtx, or derives from one.
		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()
		stop := a.AfterFunc(func() {
			child.cancel(false, parent.Err(), Cause(parent))
		})
		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))
		case <-child.Done():
		}
	}()
}

我们首先来看下 stopCtx：

其中包含两个字段 Context 和 stop 函数。

// A stopCtx is used as the parent context of a cancelCtx when
// an AfterFunc has been registered with the parent.
// It holds the stop function used to unregister the AfterFunc.

type stopCtx struct {
	Context
	stop func() bool
}

添加的代码如下：

if a, ok := parent.(afterFuncer); ok {
		//parent 实现了 afterFunc 函数
		c.mu.Lock()
                
		stop := a.AfterFunc(func() {
			child.cancel(false, parent.Err(), Cause(parent))
		})
           //stopCtx 作为 cancelCtx 的 parent
		c.Context = stopCtx{
			Context: parent,
			stop:    stop,
		}
		c.mu.Unlock()
		return
	}

假设存在这样一个 customCtx，实现了 Context 接口，并且也实现了 afterFuncer 接口。当使用 WithCancel 函数，并且 customCtx 作为参数。那么，就会执行上述代码。

也就是说，当 parent 不为 cancelCtx（广义上的含义）时，构建层级关系不是通过 children 字段

那么如何取消呢？下面会有介绍

在 removeChild 函数中添加了下面的内容。当 parent 为 stopCtx，则直接调用 s.stop 函数。

// removeChild removes a context from its parent.
func removeChild(parent Context, child canceler) {
   //新添加的
	if s, ok := parent.(stopCtx); ok {
		s.stop()
		return
	}
  //到此为止
  
	p, ok := parentCancelCtx(parent)
	if !ok {
		return
	}
	p.mu.Lock()
	if p.children != nil {
		delete(p.children, child)
	}
	p.mu.Unlock()
}

当一个 cancelCtx 的 Context 为 stopCtx，则其取消时，通过调用 stopCtx 函数来取消。来消除这种层级关系。

stopCtx 的作用可由下图表示：

stopCtx.drawio.png

5. 总结

在你看完上述源码解析之后，我的建议是你可以去写一点简单的代码感受一下 context 的用法，这样或许能够加深体会。在 context package document 中有许多例子，可以试着阅读一下。或者阅读这篇如何使用 context。

context 是并发安全的。
当使用 context 作为函数参数时，直接把它放在第一个参数的位置，并且命名为 ctx。另外，不要把 context 嵌套在自定义的类型里。
go1.21 引入了 AfterFunc API 可以实现 merge context。
context 之间的信号传递是通过构建层级关系实现的。
context 可以作为并发控制和超时控制的标准做法。

Go Package: Context 源码解析