context 源码浅读

context通常用来并发协调以及对 goroutine 的生命周期控制.

笔者在使用context的时候想要了解几种context的实现与相关操作，因此基于go 1.24 进行查看和尝试解读。
第一次输出文章有哪些错误和解读未到位的地方欢迎指出，望大家海涵
下面开始解读，首先是数据结构：

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

Deadline():

返回当前 context 的 deadline，等到对应的 deadline 对应的 context 会 cancel，第二个返回值如果该 context 没设置 deadline 则返回 false，否则返回 true

Done()：

返回一个 channel，用来标识对应 context 是否被 cancel，如果对应的 context 不会被 cancel 则返回 nil，可以用来控制 for 或者 goroutine，go 官方给的 example：

func Stream(context context.Context, out chan<- Value) error {
	 	for {
	 		v, err := DoSomething(context)
	 		if err != nil {
	 			return err
	 		}
	 		select {
	 		case <-context.Done():
	 			return context.Err()
	 		case out <- v:
	 		}
	 	}
	 }

上面例子代码，可以用 context 来控制 Stream 函数的返回，第一处 DoSomething 这里，如果 context 超时会提前返回错误，第二处 select，如果前面函数执行成功获取到 v，然后进入到 select 此时 context 突然超时也会直接返回。

Err() ：

返回一个 error，如果此时 context 对应的 channel 还未被关闭则返回 nil，否则返回 error：

DeadlineExceeded：超时/截止时间到了
Canceled：主动取消

Value(key any)：

返回与相关联的 context 的存储的 key 对应的 value 值，key 必须是可比较的类型：基本类型（string/int/bool 等）、指针、chan、接口（前提是动态值可比较）、以及元素也都可比较的 struct/array。

type contextKey int

var userIDKey contextKey 

v := context.Value(userIDKey)

此时别的包既拿不到 contextKey 这个类型名，也拿不到 userIDKey 这个变量名就算别的包也写了 type contextKey struct{}，那也是 另一个包里的另一个类型，在 Go 里它们不是同一个类型 Value 查找靠 **==**，类型不同不相等。

创建 context：

Background() 与 TODO():

这两个方法返回的本质都是一个 emptyCtx，emptyCtx是一个实现了 context 方法的空白 context，对 emptyCtx 调用 Done()，Err()，Value(key any)都会返回 nil，Deadline()返回 (time.Time{}, false)代表没有 Deadline。

Background 通常用于程序的根 context 创建，TODO 通常用于在程序中间不清楚怎么使用时创建。

cancelCtx

type cancelCtx struct {
	Context    //父context
	mu       sync.Mutex   //互斥锁        
	done     atomic.Value   //判断当前context是否被取消       
	children map[canceler]struct{}  //记录当前context的children context，以便当前context取消时一同cancel
	err      error    //返回当前 cancelCtx 的错误             
	cause    error    //cancel原因             
}

cancelCtx.Value(key any)

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

该方法有一个入参 key，该方法用于返回 context 中对应 key 的 value

cancelCtxKey 是 context 定义的一个特殊类型，当传入的 key 为&cancelCtxKey 则直接返回 context 本身
否则调用 value 从父 context 中查找

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

该方法没有入参，

从 context 中原子操作读取 done 如果取到的值不为 nil 则转换为 chan 直接返回
否则加锁，然后再次原子性读取（是为了防止在获取到锁之前，有其他的 goroutine 给当前的 context 的 done 赋值）
如果还是 nil 则创建一个 chan，并且原子操作给 context 赋值，然后解锁返回

Err() error

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

该方法没有入参，对 context 加锁，然后读取 err,解锁返回 err

WithCancel(parent Context)

该方法入参是父 context，该方法会调用 withCancel(parent Context)创建 cancelCtx，创建完返回新建的 cancelCtx 和取消该 context 的 cancel 函数。

withCancel(parent Context)

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

该方法入参是父 context，首先判断父 context，如果是 nil 则 panic，否则先初始化一个空白 cancelCtx，然后调用 c.propagateCancel(parent, c)给创建的 cancelCtx 赋值。

cancelCtx.propagateCancel(parent Context, child canceler)

func (c *cancelCtx) propagateCancel(parent Context, child canceler) {
	c.Context = parent 
	done := parent.Done()
	if done == nil {
		return 
	}
	select {
	case <-done:
		child.cancel(false, parent.Err(), Cause(parent)) 
		return
	default:
	}
	//1
	if p, ok := parentCancelCtx(parent); ok {
		p.mu.Lock()
		if p.err != nil {
			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
	}
	//2
	if a, ok := parent.(afterFuncer); ok {
		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
	}
	//3
	goroutines.Add(1)
	go func() {
		select {
		case <-parent.Done():
			child.cancel(false, parent.Err(), Cause(parent))
		case <-child.Done():
		}
	}()
}
type afterFuncer interface {
	AfterFunc(func()) func() bool
}

该方法有两个入参，分别是父 context 和子 context，实现了父 context 和子 context 的同步 cancel，

首先将 parent 赋值给子 context 的 context，然后判断父 context 是否可以取消，如果不可以被取消则直接返回
接着判断父 context 是否已经准备 cancel，如果父 context 准备 cancel 则直接将子 context 取消并返回
在第1部分判断父 context 是不是 cancelCtx 类型的 context，如果是则对父 context 进行加锁，判断父 context 是否被取消，如果取消则直接调用 cancel 取消子 context 解锁返回，否则惰性初始化 children 并将子 context 加入，然后解锁返回。
如果没有进入第1部分则进入第2部分的判断，在这里判断父 context 是否实现了回调函数，该回调函数在当 parent 结束时自动执行，如果实现了则对子 context 加锁，然后定义回调函数内容并拿到回调函数的返回值 func() bool, 该 func 用来移除/取消刚才注册的回调，如果回调函数还没执行则移除成功返回 true，否则返回 false，然后对子 context 赋值，解锁返回。
如果前面两部分都没进入，此时进入第三部分，启动一个 goroutine 同时监听父 context 的 Done 和子 cancelCtx 的 Done，如果父 context 先取消则直接对子 context 调用 cancel 取消，然后结束此协程；如果子 context 先取消则直接返回。

cancelCtx.cancel(removeFromParent bool, err, cause error)

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 
	}
	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 {
		child.cancel(false, err, cause)
	}
	c.children = nil
	c.mu.Unlock()

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

该方法的入参分别有三个：removeFromParent 用来判断否需要从父 context 的 children 中删除，err cancel 后显示的错误，cause 更为具体的原因

首先判断 err 是否为 nil，如果为 nil 则直接 panic
如果 cause 为 nil 则将 err 赋值给 cause
对 context 进行加锁，判断当前的 context 的 err 是否为空，不为空说明当前 context 已经被 cancel 则直接解锁返回。否则将 err 和 cause 赋值给当前 context
然后原子读取 context 的 done 并类型断言为 chan，如果 d 为 nil 则赋值 context 定义的一个可重入的封闭关闭管道，否则 close(d)
接着处理 context 的 child，遍历取消所有子 context，然后将 context 的 children 置为 nil，解锁
根据传入的 removeFromParent 来判断是否要执行 removeChild(c.Context, c)，返回

removeChild(parent Context, child canceler)：

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

该方法有两个入参，分别是父 context 和实现了 canceler 的子 context

如果 parent 是 stopCtx 类型，则直接执行 stop 函数取消 parent 的回调函数，直接返回
执行 parentCancelCtx 函数，返回 false 则直接返回
对 parent 加锁，判断 p.children != nil，然后从 p.children 中删除 child，解锁返回

parentCancelCtx(parent Context) ( * cancelCtx, bool)

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
}

该方法有一个入参父 context，

首先判断父 context 是否是已关闭或者是永远不会 cancel 的，如果是直接返回 false
然后去 parent 这条 ctx 链上，提取 parent 对应的 cancelCtx，如果不存在或存在但是对应的 done 跟 parent 的 done 不同则直接返回 false
前面都不执行的话最终返回 true

timerCtx

type timerCtx struct {

	cancelCtx 
	timer *time.Timer 
	deadline time.Time
}

timerCtx 嵌入了 cancelCtx 结构体，继承 cancelCtx 的能力，并新增 time.Timer 用于定时终止 context 和新增 deadline 字段用于字段 timerCtx 的过期时间.

timerCtx.Deadline()

1
2
3

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

该方法没有入参，返回当前 context 的截至时间和一个 bool 值

timerCtx.cancel(removeFromParent bool, err, cause error)

func (c *timerCtx) cancel(removeFromParent bool, err, cause error) {
	c.cancelCtx.cancel(false, err, cause)
	if removeFromParent {
		removeChild(c.cancelCtx.Context, c)
	}
	c.mu.Lock()
	if c.timer != nil {
		c.timer.Stop()
		c.timer = nil
	}
	c.mu.Unlock()
}

该方法有三个入参，removeFromParent 用来判断否需要从父 context 的 children 中删除，err cancel 后显示的错误，cause 更为具体的原因，

首先调用 cancel 方法取消 context 此时不会从 parent 中移除，然后根据入参决定要不要执行 removeChild
对 context 加锁，停止计时器并置为 nil，解锁返回

WithTimeout(parent Context, timeout time.Duration)

1
2
3

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

该方法有两个入参分别是父 context 和子 context 的持续时间，首先根据 timeout 转成一个绝对时间 deadline，本质上做了一个 WithDeadline 方法的调用

WithDeadline(parent Context, d time.Time)

1
2
3

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

该方法有两个入参分别是父 context 和子 context 的 deadline，用来调用 WithDeadlineCause 创建 timerCtx，返回一个 timerCtx 和取消 timerCtx 的 cancel 函数

WithDeadlineCause(parent Context, d time.Time, cause error)

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) {
		return WithCancel(parent)
	}
	c := &timerCtx{
		deadline: d,
	}
	c.cancelCtx.propagateCancel(parent, c)
	dur := time.Until(d)
	if dur <= 0 {
		c.cancel(true, DeadlineExceeded, cause) 
		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，子 context 的 deadline，cause 错误的原因

如果父 context 为空则直接 panic
获取父 context 的 deadline，获取成功并且早于子 context 的 deadline 则直接创建一个 cancelCtx
否则新建一个 timerCtx，然后对 c.cancelCtx 进行初始化，决定怎么实现 parent 的 cancel 时调用子 context 的 cancel
创建一个 dur 计算到 deadline 还有多久，判断是否已经过了 deadline，如果过了则直接 cancel 子 context
加锁，判断子 context 是否已经 cancel，没有的话创建一个*time.Timer，当 dur 过去后启动一个 goroutine 执行 cancel
解锁，返回封装好的 timerCtx 和取消 timerCtx 的 cancel 函数

valueCtx

type valueCtx struct {
	Context
	key, val any
}

一个 valueCtx 携带一个 key-value

WithValue(parent 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}
}

该方法的入参有三个，分别是父 parent，key 和 value

判断 parent 是否为空，为空则直接 panic
判断 key 是否为空，为空则直接 panic
判断 key 是否是可比较类型，不是则直接 panic
创建一个 valueCtx 赋值并返回

valueCtx.Value(key any)

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

该方法的入参有一个 key，

判断当前 context 的 key 存储的 key 是否与传入的 key 相等，相等的话返回存储的 val
否则调用 value(c.Context, key)，返回获取到的值

value(c Context, key any)

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 和 key，开启一个 for 循环，由下而上，依次对 key 进行匹配

先获取 context 的类型，根据不同类型不同处理方式
如果在过程中获取到 key 对应的 val 或者走到根 context，则直接返回

context源码解读到此结束，再次谢谢大家！