Go的context包笔记

0前言

context包作为使用go进行server端开发的重要工具，其源码只有791行，不包含注释的话预计在500行左右，非常值得我们去深入探讨学习，于是在本篇笔记中我们一起来观察源码的实现，知其然更要知其所以然。（当前使用go版本为1.22.2）

1核心数据结构

整体的接口实现和结构体embed图

1.1Context接口

context接口定义了四个方法：

• Deadline方法返回context是否为timerctx，以及它的结束时间
• Done方法返回该ctx的done channel
• Err方法返回该ctx被取消的原因
• Value方法返回key对应的value

2emptyCtx

先来观察源码

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实现了context接口中的所有方法，对于每个方法返回的都是空值，它没有值、不能被取消以及没有截止时间，它只作为一个空context的载体，相当于所有ctx的祖先。

如何创建一个context？

context包提供了Background()方法和TODO()方法，都用于创建一个空的context。

func Background() Context {
	return backgroundCtx{}
}
func TODO() Context {
	return todoCtx{}
}
type backgroundCtx struct{ emptyCtx }

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

type todoCtx struct{ emptyCtx }

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

它们返回的都为空的context，虽然方法名不同，但是效果是一样的，那么什么时候使用background，什么时候使用TODO呢，这是官方给出的注释

“TODO 会返回一个非空的、为空的 [Context]。 代码应该在不清楚应该使用哪个 [Context] 或者 [Context] 尚未可用（因为周围的函数尚未被扩展以接受 [Context] 参数）时使用 context.TODO。”

“background返回一个非空的、空的上下文对象。它不会被取消，没有值，也没有截止时间。它通常被主函数、初始化和测试使用，作为进入请求的顶级上下文。”

3cancelctx

先来观察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
}

• cancelCtx 内嵌了 Context 作为其父 context，根据go语言的特性，cancelCtx结构体就隐式实现了context接口中的所有方法，可以被当作一个接口来被调用，但是其方法还需要被具体的实现赋值才能进行调用。并且可以得知，cancelCtx的父类一定也是一个Context。

• mu是cancelCtx的内置锁，用来协调并发场景下的资源获取

• done的实际类型为chan struct{}，通过atomic包来实现并发安全，可以用于反应该ctx的生命周期情况，done是懒汉式创建的，只有第一次调用Done()方法时才会被创建，在下文的Done方法中会提到

• children用于关联和子ctx的关系，当取消该ctx时，可以接连通知子ctx进行关闭，及时释放资源。

• err用于返回ctx关闭的原因，调用的是context包定义的内置error

  var Canceled = errors.New("context canceled")
  
  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 }
  

• cause返回的是该ctx失效更底层的原因，例如导致DeadlineExceeded err的具体原因是“database connection timeout”

  // Example use:
  //
  //	ctx, cancel := context.WithCancelCause(parent)
  //	cancel(myError)
  //	ctx.Err() // returns context.Canceled
  //	context.Cause(ctx) // returns myError
  

3.1Done方法的实现

func (c *cancelCtx) Done() <-chan struct{} {
	d := c.done.Load() //获取cancelCtx的done通道
	if d != nil { //假若已经创建过了done通道，则直接返回
		return d.(chan struct{}) 
	}
	c.mu.Lock() //上锁
	defer c.mu.Unlock()
	d = c.done.Load() //Double check是否创建done通道，因为在上锁前，可能其他goroutine调用了该ctx的Done方法。
	if d == nil { //如果仍然未创建
		d = make(chan struct{}) //创建该done channel
		c.done.Store(d) //存储
	}
	return d.(chan struct{})
}

通过代码可以看见，ctx的done只有当被调用过Done方法时才会被创建，那么为什么这样子设计呢？很容易想到主要目的就是为了节省了不必要的资源浪费，提高效率，在很多情况下创建context并不需要监听done通道，只有在需要时才被创建，符合go语言的设计理念，只有需要的时候才引入。

3.2value方法的实现

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

若参数key与cancelCtxKey相符，则返回当前ctx本身

否则，就向父层 层层寻找

源码表现的非常晦涩，为了具体知道这个Value方法是做什么的，我们先来看对于cancelCtxKey的定义

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

cancelCtxKey是一个私有的、唯一的标识符，它用于返回cancelCtx它本身。

对于该cancelCtxKey具体使用场景，下面还会讲到

3.3创建cancelCtx的WithCancel

//WithCancel返回一个带有新的Done通道的父context的副本。只有当父ctx被关闭，或者返回的cancel方法被调用时，该ctx的Done通道才会被关闭。
func WithCancel(parent Context) (ctx Context, cancel CancelFunc) {
	c := withCancel(parent)
	return c, func() { c.cancel(true, Canceled, nil) }
}

该方法返回了一个cancelCtx，以及关闭它的cancel方法。

在1.20版本中，新增了一个WithCancelCause方法，该方法返回了一个cancelctx和它的CancelCauseFunc，我们也来看一下

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

在代码方面基本和withCancel方法一致，但是返回的CancelCauseFunc可以用于给用户自定义ctx被取消的原因，例如

ctx, cancel := context.WithCancelCause(parent)
cancel(myError)
ctx.Err() // returns context.Canceled
context.Cause(ctx) // returns myError

接着来看withCancel

func withCancel(parent Context) *cancelCtx {
	if parent == nil { //若父ctx是nil，那么不能创建
		panic("cannot create context from nil parent")
	}
	c := &cancelCtx{}
	c.propagateCancel(parent, c) //传递cancel给新创建的ctx
	return c
}

主要来看propagateCancel做了什么：

//该方法主要用于建立父context和子context的联系，如果父context也是一个cancelCtx，它需要保证父context被取消时，子context也能跟着被取消。
func (c *cancelCtx) propagateCancel(parent Context, child canceler) {
	c.Context = parent //将父context内嵌入子context中

	done := parent.Done() //获取父ctx的Done通道
	if done == nil { //如果通道不存在，那么说明父context不是cancelCtx，不需要为它们两个之间建立联系，因为父context永远不会关闭。
		return // parent is never canceled
	}

	select {
	case <-done: //非堵塞地获取父done通道的状态，如果done通道以及被closed，那么这里会接受到一个零值，如果没有被closed，会执行default后面的语句。
		//接受到零值，说明done通道以及被关闭了，父context已经被取消，此时应该立即调用子context的cancel方法，取消子context。
		child.cancel(false, parent.Err(), Cause(parent))
		return
	default:
	}
	//否则，执行以下代码
    //此处判断父context是否是一个context包实现的cancelCtx，如何理解会在下文讲述该方法时继续说明
	if p, ok := parentCancelCtx(parent); ok {
		// 如果父context是一个cancelCtx，或者是从某个cancelCtx衍生出来的context
		p.mu.Lock() //加锁
		if p.err != nil {
			// 如果存在err，说明已经被取消，此时也应该调用child的cancel方法取消子ctx
			child.cancel(false, p.err, p.cause)
		} else {
			if p.children == nil { //如果这是p的第一个子cancelCtx，需要初始化map的内存
				p.children = make(map[canceler]struct{})
			}
			p.children[child] = struct{}{}//绑定子ctx到父ctx的map中，用于当父ctx取消时，能通知所有的child跟着取消。
		}
		p.mu.Unlock()
		return
	}
	//如果父ctx不是cancelctx并且实现了AfterFuncer接口，即实现了AfterFunc方法（该方法会在ctx被取消后唯一一次调用），那么就需要为父ctx再设置一个afterFunc方法，用于取消child并且传递err和cause
	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, //这里的stop方法可以用于当父ctx调用afterFunc的时候，取消父ctx对cancel函数的调用（看需求）
		}
		c.mu.Unlock()
		return
	}
	//下面的情况为增加一个goroutine，监听父ctx自己实现的Done channel（用户自定义的）
	goroutines.Add(1)
	go func() {
		select {
		case <-parent.Done()://如果监听到父ctx自定义实现的Done channel关闭时，就关闭child
			child.cancel(false, parent.Err(), Cause(parent))
		case <-child.Done(): //如果child先关闭，那么就立即释放协程，避免协程泄露
		}
	}()
}

propagateCancel方法主要的作用是，保证了对于父ctx被取消时，为了能及时取消子ctx，避免不必要的资源浪费，建立父ctx和子ctx之间的联系。使用流程图来表示该方法如下（省略了检查afterFunc）

在该方法中，比较难以理解的地方是第二个if,"if p, ok := parentCancelCtx(parent); ok"究竟做了什么，为此我们跟进源码查看：

// parentCancelCtx 返回父级对象的底层 cancelCtx。
// 它通过查找 parent.Value(&cancelCtxKey) 来找到最内层的 enclosing cancelCtx，然后检查 parent.Done() 是否与该 cancelCtx 匹配。（如果不匹配，则该 cancelCtx 可能已被封装在提供了不同 done 通道的自定义实现中，在这种情况下，我们不应该绕过它。）
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
}

三个if分别做了什么：

• 第一个if：获取parentCtx的done channel，并且查看情况，若已经关闭或者父Ctx不可取消，此时返回false，回到propagate方法中最终会cancel掉child

• 第二个if：可以理解为通过Value方法找到离parentCtx“最近的”cancelCtx p，一般情况下，如果parentCtx就是一个CancelCtx，这时候就是parentCtx它本身。如果p不存在，也就是没有可以取消的ctx，此时也会返回false。

• 第三个if：找到了p后，还读取了p的done channel，这时候一般情况，pdone 当然会 == done，因此最终会返回p和true，那么什么时候会不相等呢？为什么会不相等呢？为此，看下方的层次图来理解这个if

在这个情况假设下，ParentCtx2是从ParentCtx1衍生出来的，ParentCtx1是一个标准的CancelCtx，而ParentCtx2是一个用户自定义了ctx，它内层继承了ParentCtx1并且自己实现了Done()方法，这时候代码中的“p”找到的就是ParentCtx1，而parent是ParentCtx2，此时的done和pdone就是两个不同的channel了，这时候cihldctx应该监听哪一个done channel呢？答案是监听用户自定义实现的Ctx的chennel，因为我们不应该绕过用户实现的Done channel，这更加符合ctx到层次逻辑。假如这时候不去判断pdone == done，直接返回的指针就是ParentCtx1的指针了。

3.4Cancel方法实现

接下来我们来看cancel方法是如何实现的。

//cancel 关闭 c.done，cancel 每一个c的children。如果removeFromParent为true，将会把c从parentCtx的child中移除。
func (c *cancelCtx) cancel(removeFromParent bool, err, cause error) {
	if err == nil { //必须存在err
		panic("context: internal error: missing cancel error")
	}
	if cause == nil { //没有自定义设置cause，默认为err
		cause = err
	}
	c.mu.Lock() //上锁
	if c.err != nil { //double check
		c.mu.Unlock()
		return // already canceled
	}
	c.err = err
	c.cause = cause
	d, _ := c.done.Load().(chan struct{})
	if d == nil { //done没有被创建，直接存储context包内以及创建了的关闭的channel，不需要再次创建
		c.done.Store(closedchan)
	} else {
		close(d) //关闭done
	}
	for child := range c.children { //关闭每一个子context
		// 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) //从父child中移除自己
	}
}

代码比较简单，问题在于什么时候removFromParent为true呢？为什么为true？

回到withCancel方法中，我们可以看到返回的cancel方法中，此时removeFromParent就为true

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

当用户主动调用cancel()时，就会将子ctx从父ctx中的child删除。因为此时没有必要再在父ctx中接受父或祖先的cancel通知。而当调用cancel函数内部，对child执行的cancel就为false，这是因为后面设置了c.children = nil，这时候是从父ctx的方向关闭了子ctx对其的链接。

4afterFuncCtx

afterFuncCtx是1.20版本后引入的新ctx，它的作用是当ctx被取消后，能执行一次自定义的F函数，一般用于回收资源等。

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

可以看到afterFuncCtx embed了cancelCtx，在此基础上添加了once和f。once保证了f只会被执行一次。接着我们来看如何实现一个afterFuncCtx

4.1AfterFunc

func AfterFunc(ctx Context, f func()) (stop func() bool) {
	a := &afterFuncCtx{ //创建局部变量afterFuncCtx，记录cancel时需要被执行的func
		f: f,
	}
	a.cancelCtx.propagateCancel(ctx, a)
	return func() bool { //返回stop方法，如果需要不执行func，则调用stop
		stopped := false
		a.once.Do(func() { //尝试stop func函数，如果成功则stop会被设置为true
			stopped = true
		})
		if stopped { //stop后，a没有存在的意义，进行取消。
			a.cancel(true, Canceled, nil)
		}
		return stopped //true为成功取消，false表示f已经被执行或者正在被执行。
	}
}

这是一个闭包实现，闭包是指在函数内部定义的函数（如这里返回的 stop 函数），它会“捕获”并保存定义时可访问的所有外部变量。在 AfterFunc 方法中，虽然 a 是局部变量，但返回的 stop 方法引用了 a，形成了一个闭包，闭包会将 a 的内存保留在堆上，即使 AfterFunc 方法返回后，a 依然存在。所以当虽然没有返回a，但是返回的stop方法任然能调用a，a的生命周期超出了afterFunc方法。

以下是一个示例

package main

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

func main() {
	// 创建一个 2 秒后超时的 context
	ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second)
	defer cancel()

	// 注册 AfterFunc，context 完成时将调用清理操作
	stop := context.AfterFunc(ctx, func() {
		fmt.Println("清理操作正在执行...")
	})

	// 等待 3 秒
	time.Sleep(3 * time.Second)

	// 尝试停止清理操作
	stopped := stop()
	if stopped {
		fmt.Println("成功停止清理操作")
	} else {
		fmt.Println("清理操作已经开始或已停止")
	}

	// 等待 context 超时
	<-ctx.Done()

	fmt.Println("程序结束")
}

4timerCtx

接下来来看timerCtx

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

	deadline time.Time
}

timer直接内嵌了cancelCtx以实现Done和Err，并且新添了timer和deadline字段，deadline用于查看ctx的截止时间，timer用于完成过时取消ctx。

4.1WithDeadline

先来看WithTimeout方法

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

WithTimeout接受一个存活时间来创建一个timerCtx，可以看到最终都是调用了WithDeadline

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

WithDeadline又调用了WithDeadlineCause，返回了timerCtx和一个CancelFunc。

4.2WithDeadlineCause

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) { //如果parent的deadline更早，则直接返回parent的副本，不需要再创建timer
		// 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 {//检测是否在创建过程中已经过了ddl
		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() { //设置超时自动cancel
			c.cancel(true, DeadlineExceeded, cause)
		})
	}
	return c, func() { c.cancel(true, Canceled, nil) }//返回timerCtx和cancelfunc
}

比较好理解，所以接着往下看cancel方法

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

这里值得注意的地方只有第二行，它调用的是c.cancelCtx.cancel，取消的并非是父ctx，而是它本身，这里的作用是为了区分cancel方法的实现，c.cancelCtx是父ctx的一个副本，并不是父ctx，所以真正的parent是c.cancelCtx.Context。

5valueCtx

type valueCtx struct {
	Context
	key, val any
}

valueCtx内嵌了Context接口所以拥有该接口的所有方法，以及添加了k-v pair。

5.1WithValue

// 提供的键必须可比较，并且不应是字符串或任何其他内置类型，以避免在使用上下文时与其他包发生冲突。使用WithValue的用户应为其键定义自己的类型。为了避免在将值赋给接口{}时分配内存，上下文键通常具有具体的类型struct{}。或者，应将导出的上下文键变量的静态类型设置为指针或接口。
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() { //必须保证key是可比较的
		panic("key is not comparable")
	}
	return &valueCtx{parent, key, val}
}

5.2value

func value(c Context, key any) any {
	for {
		switch ctx := c.(type) { //对context进行类型断言
		case *valueCtx:
			if key == ctx.key { //如果key就是当前ctx的key，则直接返回val
				return ctx.val
			}
			c = ctx.Context //否则向父ctx查询
		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 //不存在key value
		default:
			return c.Value(key) //返回用户实现的Value
		}
	}
}

可以看见，valueCtx查询value的过程，类似于链表查询，它是自底向上的，并且时间复杂度为O（N），它并不适用于存放大量的kv数据，原因有以下：

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

感谢观看，参考博客：

Golang context 实现原理 (qq.com)

深入Go：Context-腾讯云开发者社区-腾讯云 (tencent.com)

Go context的使用和源码分析_&cancelctxkey-CSDN博客