context

context译名为上下文，其设计的初衷是用来跨多个协程尤其是父子协程间传递信号和一些简单的数据。它通常在处理 HTTP 请求、任务调度、数据库查询等场景中使用，尤其是在微服务架构中，gRPC 就通过context来跨进程跨网络来做元数据传递，链路控制等操作。

package main

import (
    "context"
    "fmt"
    "os/signal"
    "syscall"
    "time"
)

func main() {
    ctx, stop := signal.NotifyContext(context.Background(), syscall.SIGKILL, syscall.SIGINT, syscall.SIGTERM)
    defer stop()
    for {
       select {
       case <-ctx.Done():
          fmt.Println("terminate")
          return
       default:
       }
       fmt.Println("running")
       time.Sleep(100 * time.Millisecond)
    }
}

上面这个案例中，通过context来传递进程信号，当接收到信号后，程序就会自行退出，这也是context的一个应用场景。

结构

context.context并不是一个具体的实现，而是一个定义了一组方法的接口

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

• Deadline()，返回一个截至日期，和布尔值表示是否已经结束
• Done()，返回一个管道，用于通知结束消息
• Err()，返回上下文关闭的原因
• Value()，根据 key 获取指定的值

标准库中总共有下面几个可用的上下文

• Background，空白的上下文，通常用于上下文的根节点
• WithCancel，WithCancelCause，可取消的上下文
• WithDeadline，WithDeadlineCause，带截至时间的上下文
• WithTimeout，WithTimeoutCause，带超时时间的上下文
• WithValue，可以传值的上下文

具体到实现就主要是

• timerCtx
• cancelCtx
• emptyCtx

所以它的核心功能就四个点

• 取消
• 截止
• 传值
• 传播

弄懂了这几个，基本上就明白了context的运作原理。

取消

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
}


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

cancelCtx的核心就在于propagateCancel这个方法，它负责将可取消这一行为传播到父子上下文中。

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

1. 首先它会调用检查父 context 能否被取消，不能的话就直接返回

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

2. 返回检查父 context 是否已经被取消了，是的话则取消所有的子 context

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

3. 尝试将父 context 转换为cancelCtx类型，如果成功了的话就会将当前 context 加入到父 context 的children中。

   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
   }

4. 尝试将其转换为afterFuncer类型，如果成功了的话，就会将取消当前 context 的方法注册到父 context 的AfterFunc中

   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
   }

5. 如果还是不行，那么就单独开一个协程来监听Done管道，当收到信号时，就会取消子 context

   go func() {
       select {
       case <-parent.Done():
          child.cancel(false, parent.Err(), Cause(parent))
       case <-child.Done():
       }
   }()

然后由cancelCtx.cancel方法最终来负责取消子 context

func (c *cancelCtx) cancel(removeFromParent bool, err, cause error) {
	...
}

它的流程如下

1. 检查当前 context 是否已经被取消

   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
   }

2. 关闭done管道，发送关闭通知

   c.err = err
   c.cause = cause
   d, _ := c.done.Load().(chan struct{})
   if d == nil {
   	c.done.Store(closedchan)
   } else {
   	close(d)
   }

3. 遍历通知子 context

   for child := range c.children {
   	child.cancel(false, err, cause)
   }
   c.children = nil
   c.mu.Unlock()

4. 最根据传入的参数判断是否需要从父 context 中删除

   if removeFromParent {
   	removeChild(c.Context, c)
   }

截止

WithTimeout和WithDeadline都是具有截至时间的上下文，两者都是一个类型，仅仅只是使用的语义不一样，且都基于cancelCtx

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

	deadline time.Time
}

WithDeadlineCause负责创建具有截至时间的上下文

func WithDeadlineCause(parent Context, d time.Time, cause error) (Context, CancelFunc) {
	...
}

它的流程如下

1. 检查截至日期，如果父截止日期比当前截止日期早，那么父 context 肯定比当前 context 先取消，则直接创建cancelCtx类型的上下文

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

2. 构建timeCtx，并传播到子 context

   c := &timerCtx{
   	deadline: d,
   }
   c.cancelCtx.propagateCancel(parent, c)

3. 计算当前截止时间，如果已经截止了就直接取消

   dur := time.Until(d)
   if dur <= 0 {
   	c.cancel(true, DeadlineExceeded, cause) // deadline has already passed
   	return c, func() { c.cancel(false, Canceled, nil) }
   }

4. 如果没有截止的话，就通过time.AfterFunc设置在截至时间取消当前 context

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

对于timerCtx而言，它取消上下文方法就只是将timer停止，然后顺便停止子 context

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

传值

valueCtx可以在上下文中传值，获取值

type valueCtx struct {
    Context
    key, val any
}

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函数，其它的 context 具体实现也会进入到该函数来取值，比如

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

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

value函数是一个大的 for 循环，通过不断向上递归来寻找指定 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：如果 key 匹配当前 Context 的 key，则返回与 key 对应的 val。
• *cancelCtx：如果 key 匹配 cancelCtxKey，则返回当前的 cancelCtx，表示返回其自身。
• withoutCancelCtx：表示没有取消功能的 Context，如果 key 匹配 cancelCtxKey，则返回 nil。
• *timerCtx：如果 key 匹配 cancelCtxKey，返回与其关联的 cancelCtx。
• backgroundCtx 和 todoCtx：通常是没有携带任何额外值的特殊类型的 Context，遇到这两种类型时直接返回 nil。
• 如果是未知类型，则继续调用Value方法寻找。

小结

这几个 context 的核心就是cancelCtx，通过它来传播取消信号，其它的 context 类型以及一些第三方的 context 类型也都是一层套一层，在此基础上实现各式各样的功能。