深入底层读源码！一文详解context标准库

一、简介

        context上下文包是go中的标准库，主要用于在不同Goroutine之间传递取消信号、超时警告、截止时间信息以及其他请求范围的值。

        笔者这里的GO的版本是1.21.6，所以下面的代码都是GO1.21.6中的源码。

        context是一个接口，其中实现了四个方法

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

• Deadline方法返回请求的截止时间（deadline），以及一个布尔值表示是否存在截止时间。如果截止时间存在，则ok为true；否则，为false。
• Done方法返回一个类型为<-chan struct{}的通道。这个通道会在上下文被取消或超时时关闭，可以用于监听上下文的取消信号。
• Err方法返回上下文的错误信息。一般情况下，当上下文被取消或超时时，会返回相应的错误信息。
• Value方法获取与给定键关联的值。这个方法允许在上下文中存储和检索键值对，可以用于在请求处理过程中传递一些自定义的数据。

        这里具体解释done的作用机制。

func MyHandler(ctx context.Context, w http.ResponseWriter, r *http.Request) {
    done := ctx.Done()

    select {
    case <-done:
        // 上下文已取消
        // 进行相应的处理
    default:
        // 继续处理请求
        // ...
    }
}

        在一个协程中，首先通过Done()创建这样一个通道，这个通道就是用于检测该上下文是否被取消或超时，如果被取消或超时，就会执行相应的逻辑关闭这个通道。而这个通道在没有关闭的时候，select在没有其他操作的情况下会处于阻塞状态，因为这个通道内没有可读取的值。而一旦关闭（取消或超时，详见后文源码），该通道就可读而不会阻塞，不过读到的值会是零值，然后就可以执行相应的逻辑处理。
        由此实现了上下文对协程的并发控制。

二、Context

    1、emptyCtx

        数据结构

type emptyCtx int

        方法

        emtpyCtx在源码中由Background()和TODO()两个函数生成。

var (
    background = new(emptyCtx)
    todo = new(emptyCtx)
)
 
func Background() Context {
    return background
}
 
func TODO() Context {
    return todo
}

         下面是该数据结构的数据类型和其实现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 any) any {
    return 
}

• Deadline() (deadline time.Time, ok bool) 方法：返回零值时间和 false，意味着没有设置截止时间。
• Done() <-chan struct{} 方法：返回 nil，表示上下文不会被取消，这个通道不可读也不写。
• Err() error 方法：返回 nil，表示没有错误发生。
• Value(key any) any 方法：返回空值。

    2、cancelCtx

        数据结构

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
}

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

• Context ：嵌入 Context 接口，cancelCtx 结构体可以继承 Context 接口中定义的方法和行为，使得 cancelCtx 类型可以被视为 context.Context 接口类型的值。这使得WithCancel（）等操作的parent入参不一定要是一个emptyCtx。
• mu sync.Mutex：互斥锁（Mutex），用于保护该结构体的并发访问。在需要保护共享资源的临界区使用互斥锁，确保在同一时间只有一个 goroutine 能够访问或修改结构体的字段，避免竞态条件。
• done atomic.Value：原子值atomic.Value(对值的读取和写入操作是原子的，不会受到并发访问的干扰。)，用于存储一个通道（chan struct{}），是上下文是否已经完成、取消或超时的标志。通道是在第一次调用cancel()时才延迟创建，减少占用的内存，并在其中放入一个值，然后通过关闭通道来表示上下文已完成。
• children map[canceler]struct{}：存储该上下文的子上下文对象。在这里，canceler 是一个接口类型，用于表示可取消的对象（即实现了上面的两种cancel和Done两个方法）。通过维护该映射，父上下文可以管理和传播取消信号给其子上下文。同时在第一次调用cancel时，将会将 children 设置为 nil，以防止进一步的子上下文注册。
• err error：存储表示上下文被取消的错误信息。当上下文被取消时，可以将取消的原因记录在 err 字段中。
• cause error：存储导致上下文取消的错误。通常情况下，这个错误是由父上下文传播到子上下文，以记录导致取消的根本原因。

        方法

        cancelCtx有以下的创建方式，同时必须有一个父节点作为入参。

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

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

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

        WithCancel ()返回一个cancelCtx类型的上下文，同时和这个cancelCtx的取消函数。
        WithCancelCause（）则会额外返回一个引发取消的原因的err类型变量
        在之前的版本在创建c := &cancelCtx{}是会是

c := newCancelCtx(parent)

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

       现在直接创建一个对应的结构体，而注入父上下文的操作则放到了propagateCancel()函数中。

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

        propagateCancel用于在父上下文被取消时，将取消信号传播给子上下文，从而实现具体的协程控制。 逻辑如下：

        判断父上下文是否能被取消，如果不能被取消直接返回。
        判断父上下文是否已经取消，如果是直接取消子上下文。（这里子上下文即是调用该方法的c）
        找到并判断父上下文是否是一个cancelCtx，如果是则把当前上下文加入到他的子上下文map中。
        判断父上下文是否有AfterFunc方法，有的话就是一个timerCtx，然后注入context的stop为父上下文的stop（后文具体这个字段的含义）
        最后判断完了启动一个协程监听父上下文和子上下文的情况，如果父上下文取消了该上下文也取消。

        进一步探究判断是否是cancelCtx的方法：

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
}

        函数首先通过父上下文的 Done() 方法获取一个通道 done，然后检查这个通道是否已经关闭或者为 nil。如果是的话，表示父上下文已经被取消或者不存在，那么函数返回 nil 和 false。
        再从父上下文的 Value 中获取键为 &cancelCtxKey 的值，并将其转换为 *cancelCtx 类型的指针。如果获取成功，则将得到的 *cancelCtx 对象存储在变量 p 中（基于 cancelCtxKey 为 key 取值时返回 cancelCtx 自身，接下来的value函数会体现这一点），并将 ok 标记设置为 true。如果获取失败，则返回 nil 和 false，表示未找到有效的 *cancelCtx 对象。
        函数尝试从 p 中获取 done 通道，并与之前获取的 done 进行比较。如果两者不相等，也会返回 nil 和 false，表示未找到有效的 *cancelCtx 对象。最后，如果以上条件都满足，函数返回 p 和 true，表示成功从父上下文中获取到了有效的 *cancelCtx 对象。
        这里的判断方式是根据cancelCtx特有的key的取值是自身的方式来找的。

        接下来是接口中的方法：

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

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

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

        首先这里的value方法就是前面所说的，cancelCtx的value会返回自身，所以可以通过这个来判断是否是一个cancelCtx
        Done方法通过原子值来操作通道节约空间。
        Err不再赘述。
        需要注意的是cancelCtx并未实现Deadline 方法，只是带有实现该方法的接口，使用没有什么影响。

func main() {
	parent := context.Background()
	cancelCtx, cancel := context.WithCancel(parent)
	fmt.Println(cancelCtx)
	fmt.Println(parent.Deadline())
	fmt.Println(cancelCtx.Deadline())
	defer cancel()
}

        输出： 

context.Background.WithCancel
0001-01-01 00:00:00 +0000 UTC false
0001-01-01 00:00:00 +0000 UTC false

        cancelCtx调用Deadline()的输出就是emptyCtx调用Deadline()的输出。

        最后是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()
	if c.err != nil {
		c.mu.Unlock()
		return // already canceled
	}
	c.err = err
	c.cause = cause
	d, _ := c.done.Load().(chan struct{})
	if d == nil {
		c.done.Store(closedchan)
	} else {
		close(d)
	}
	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)
	}
}

        两个入参 removeFromParent 是一个 bool 值，表示当前 context 是否需要从父 context 的 children set 中删除；第二个 err 则是 cancel 后需要展示的错误，第三个是取消引发的原因。
        这里我们回头看WithCancel()函数

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

        给的入参就是true, Canceled, nil。
        前面的一些赋值操作就很朴素，后面的具体操作首先是对通道的操作：这里对原子值进行一个load()操作，如果这里没有被store()那显然就是一个nil，换言之没有被调用过，没有谁要判断在等待他通道的关闭，那就直接预定义一个closedchan：

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

func init() {
	close(closedchan)
}

        如果获取的通道 d 不为 nil，说明有其他地方在等待当前上下文的完成信号。那就关闭通道。

        然后是遍历操作，遍历所有的子上下文，然后关闭掉

        最后是看是否要把当前上下文从父上下文中删除：

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

        首先看父节点有没有stop方法，有就直接调用，stop直接进行停止操作。

type stopCtx struct {
	Context
	stop func() bool
}

        随后就是判断是不是cancelCtx，是就删除子上下文映射，不是就直接返回。

3、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
}

        timerCtx 包含了一个计时器和一个截止时间。嵌入了一个 cancelCtx 来实现 Done 和 Err 方法。它通过停止计时器然后委托给 cancelCtx.cancel 来实现取消操作。所以timerCtx就是在cancelCtx的基础上加了一个deadline和timer计时器。

        方法

        首先是创建方法：context.WithTimeout() 和 context.WithDeadline()

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

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")
	}
	if cur, ok := parent.Deadline(); ok && cur.Before(d) {
		// The current deadline is already sooner than the new one.
		return WithCancel(parent)
	}
	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()
	if c.err == nil {
		c.timer = time.AfterFunc(dur, func() {
			c.cancel(true, DeadlineExceeded, cause)
		})
	}
	return c, func() { c.cancel(true, Canceled, nil) }

        context.WithDeadline()操作类似cancelCtx的创建操作，最后创建一个计时器，如果时间超过就直接执行取消函数。

func AfterFunc(d Duration, f func()) *Timer {
	t := &Timer{
		r: runtimeTimer{
			when: when(d),
			f:    goFunc,
			arg:  f,
		},
	}
	startTimer(&t.r)
	return t
}

        在 runtimeTimer 结构体中设置定时器的触发时间 when，即当前时间加上时间间隔 d。同时，将需要执行的函数 f 封装成一个 goFunc，并将其作为参数 arg 传递给 runtimeTimer。（runtime是sleep库里面的结构体，这里就不详细介绍了，大致就理解成一个计时器，到了时间就去执行对应的函数）
        调用 startTimer 函数来启动定时器，以便在指定的时间到达时触发执行函数。
        context.WithTimeout() 和 context.WithDeadline()没啥区别，无非是前者要求输入一个持续时间，后者要求输出准确的过期时间。

        Deadline()展示一下过期内容

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

        timerCtx.cancel()

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

        复用cancelCtx的cancel。同时停止一下计时器。

func (t *Timer) Stop() bool {
	if t.r.f == nil {
		panic("time: Stop called on uninitialized Timer")
	}
	return stopTimer(&t.r)
}

func stopTimer(t *timer) bool {
	return deltimer(t)
}

4、valueCtx

        数据结构

type valueCtx struct {
	Context
	key, val any
}

        换言之较emptyCtx更新了value方法的实现

        方法

        创建方法：WithValue

func WithValue(parent Context, key, val any) Context {
	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}
}

        Value方法：

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

        如果在当前节点找到了对应的key就返回，否则向上（父上下文）继续查找。
        Value则是判断父上下文的具体类型，进行返回。是valueCtx则判断key；是cancelCtx就返回本身；是withoutCancelCtx返回nil；是timerCtx返回他包含的cancelCtx；backgroundCtx和TODO也返回nil。

        用法注意事项：

• 一个 valueCtx 实例只能存一个 kv 对，因此 n 个 kv 对会嵌套 n 个 valueCtx，造成空间浪费；
• 基于 k 寻找 v 的过程是线性的，时间复杂度 O(N)；
• 不支持基于 k 的去重，相同 k 可能重复存在，并基于起点的不同，返回不同的 v. 由此得知，valueContext 的定位类似于请求头，只适合存放少量作用域较大的全局 meta 数据.

5、withoutCtx

// WithoutCancel returns a copy of parent that is not canceled when parent is canceled.
// The returned context returns no Deadline or Err, and its Done channel is nil.
// Calling [Cause] on the returned context returns nil.
func WithoutCancel(parent Context) Context {
	if parent == nil {
		panic("cannot create context from nil parent")
	}
	return withoutCancelCtx{parent}
}

type withoutCancelCtx struct {
	c Context
}

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

WithoutCancel 方法返回一个 parent 的副本，当 parent 被取消时，该副本不会被取消。

返回的上下文不会返回任何截止时间 (Deadline) 或错误 (Err)，其 Done 通道为 nil。在返回的上下文上调用 Cause 方法将返回 nil。