Go源码--context包

简介

Context 是go语言比较重要的且也是比较复杂的一个结构体，Context主要有两种功能:

取消信号：包括直接取消（涉及的结构体：cancelCtx ; 涉及函数：WithCancel）和携带截止日期的取消（涉及结构体：timeCtx,cancelCtx；涉及函数：WithTimeout/WithDeadline/WithDeadlineCause）；一般用来控制主从协程，当主协程执行取消操作时，当然希望从携程也执行取消操作（或者其他设定的操作），这样主协程可以对从协程的生命周期有控制权。
传递消息：包括跨API边界和进程间传递（涉及结构体：valueCtx；涉及函数：WithValue）；主从协程可以共享数据，这样在Context形成的带回溯的树状结构中，通过回溯任何协程可以根据key获取其他协程注册在valueCtx（继承了context）结构体中的value，从而做到在整个主从协程中共享数据。

由于上述涉及的结构体入参都包含Context 返回值也包含Context，所以继承自Context的结构体都可以是入参，这样上述函数就可以组合使用。WithValue一般和/WithCancel/WithTimeout/WithDeadline/WithDeadlineCause搭配使用，这样既可以控制从协程生命周期，又可以传递数据。这几个结构体的组合可以产生各种应用场景。我后续会介绍这种组合。下面我们先梳理下几种重要的接口和结构体吧。

重要结构体

Context接口

Context接口是梦开始的地方，拥有最崇高的地位。我们来看下其结构体

go 复制代码

type Context interface {
// Deadline 返回代表此上下文完成工作时应取消的时间。
	Deadline() (deadline time.Time, ok bool)
	Done() <-chan struct{}  // 返回 chan类型 如果返回的不是关闭的 则阻塞 否则 返回关闭的chan 不阻塞
	Err() error  // 返回Done不阻塞的原因 如果是 Canceled的则 返回cancel的原因 所以 这个参数不为空 不意味着是运行出错了
	Value(key any) any   // 返回 key对应的value，可以递归调用。用于查找整个context树上跟key匹配的value
}

emptyCtx 结构体

emptyCtx 采用空值实现了 Context的四个函数

go 复制代码

type emptyCtx struct{}

backgroundCtx 结构体

go 复制代码

type backgroundCtx struct{ emptyCtx }

其采用组合的方式来包含emptyCtx，从而可以调用其函数值，也就是说实现了Context这个结构体，这就是结构体嵌入的优势。todoCtx结构体跟backgroundCtx 结构一样。其主要用来作为树形结构的根节点，所以其实现的context函数都是空实现。

cancelCtx 结构体

介绍这个结构体前我们先来看下 canceler 接口

go 复制代码

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

canceler 的实现是 *cancelCtx和 *timerCtx，主要为context类型结构体进行功能补充，可以进行取消操作。

cancelCtx 结构体如下：

go 复制代码

// 可以被取消，当被取消，也会取消任何实现了canceler结构体的孩子Context
type cancelCtx struct {
	Context

	mu       sync.Mutex            // 保护下面的属性
	done     atomic.Value          // done存了chan struct{}，被懒惰创建（使用时才创建）
	children map[canceler]struct{} // 实现了 canceler接口的孩子map，key有用，value没用。主要用来组成树形结构，因为key是cancler类型的
								    // 所以能作为key（除了根节点的树节点）的只有cancelCtx，timeCtx和嵌入这两个任何一个的结构体；
	err      error                 // 
	cause    error                 // 
}

cancelCtx主要用来主协程控制从协程的生命周期，主协程调用cance()函数后其chan关闭，其子协程也会执行 Done()（关闭的chan不再阻塞）函数取消阻塞来执行后续逻辑。这里分两种情况，一种是只有一个With函数，子协程不产生新的context，也就是大家经常用的，第二种情况是在子协程使用With函数来产生孩子节点,这时children就会发挥作用，会形成带回溯的树形结构。下文会有详细讲解。

跟cancelCtx直接相关的函数主要是 WithCancel。

timerCtx 结构体

timerCtx结构体有定时器和截止时间两个参数。它嵌入了一个cancelCtx结构体来实现Done和Err方法。通过停止计时器（到时间了）然后委托给cancelCtx.cancel来实现cancel，也就是说timerCtx的cancel是通过调用cancelCtx的cancel方法来实现的。我们来看下其结构体：

go 复制代码

type timerCtx struct {
	cancelCtx // 内嵌取消context
	timer *time.Timer  // 定时器 用来定时取消操作

	deadline time.Time // 返回取消的时间戳
}

可以看到由于内嵌取消context其本身可以代表cancelCtx类型，可以断言成功。(这个特性可以将两者组合使用)

跟timerCtx直接相关的函数有 WithTimeout/WithDeadline/WithDeadlineCause三个。要再次强调下只要是context的实现结构体都可以作为With系列函数的入参，所以不同context变体的组合可以实现各种复杂的主从控制和信息传递操作。以上两个结构体介绍的都是主从控制方面的，现在我们来看下传递参数相关的结构体。

valueCtx 结构体

其结构体如下：

go 复制代码

type valueCtx struct {
	Context
	key, val any
}

这个结构体带来一个key-value对，其为key实现了Vaule函数，主要用来在树形结构中查找key对应的value。其他的方法调用都委托给了嵌入的Context。这样通过 valueCtx和timerCtx或者cancelCtx的组合就能即传递参数又控制子协程。当然vauleCtx也可以嵌入其他Context来实现参数的传递，例如我们常用的 go http包。

我们来简要绘制一个图来看下如上结构体之间的关系。

简要介绍了结构体接下来我们来介绍几种重要的函数，函数都是名如其人。

WithCancel

其英文解释如下：

// WithCancel returns a copy of parent with a new Done channel. The returned context's Done channel is closed 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.

翻译成中文是：

WithCancel返回带有Done 通道的上下文的副本，其实就是新建一个cancelCtx将其放入children key中。当返回的取消函数调用或者当父上下文通道已经关闭时，返回的上下文通道也会关闭。看哪个先发生吧。

其代码如下：

go 复制代码

func WithCancel(parent Context) (ctx Context, cancel CancelFunc) {
	c := withCancel(parent)  // 创造新cancelCtx 如果入参 ctx是cancelCtx类型 就将 新的cancelCtx放入其children key中
	return c, func() { c.cancel(true, Canceled, nil) } // 返回新 cancelCtx和 取消函数 
}

其使用例子如下：

go 复制代码

func TestWithCancel(t *testing.T) {
	var wg sync.WaitGroup
	ctx, concel := context.WithCancel(context.Background())
	wg.Add(1)
	go doSomthing(ctx, &wg)
	time.Sleep(3 * time.Second)
	concel()
	wg.Wait()
}

func doSomthing(ctx context.Context, wg *sync.WaitGroup) {
	defer wg.Done()
 	for {
		select {
		case <-ctx.Done():
			time.Sleep(2 * time.Second)

			fmt.Println("playing")
			return
		default:
			time.Sleep(1 * time.Second)
			fmt.Println("I am working")

		}

	}

}

其输出如下：

go 复制代码

I am working
I am working
playing

我们可以看到当还没有执行cancel时 doSoomthing 中 ctx.Done() 返回的chan是阻塞未关闭状态，当调用cancel()函数时，chan关闭，因为doSomthing()中入参是 context是一个接口，接口是入参，就会携带本身类型和其实现的结构体的指针。所以 ctx中的 chan就是唯一的，当主协程调用cancel关闭chan时，子协程ctx.Done函数返回关闭的结构体也就不再阻塞。还记得我在介绍cancelCtx结构体时，最后一段的描述吗，这时候因为就一个With函数不会产生新的 cancelCtx。

接下来我们来深入剖析下WithCancel源码。

WithCancel 的withCancel函数源码如下：

go 复制代码

func withCancel(parent Context) *cancelCtx {
	if parent == nil {
		panic("cannot create context from nil parent")
	}
	// 创建一个新的 结构体 作为 符合条件的父节点 child; 树结构
	c := &cancelCtx{}

	// 如果符合某条件 将 子节点插入树中
	c.propagateCancel(parent, c)
	return c
}

其中 propagateCancel函数源码如下：

go 复制代码

// propagateCancel 负责在父上下文被取消时，使子上下文也被取消。它设置了 cancelCtx 的父上下文(可以理解成树上的父子节点，制定一个节点的父节点和子节点)。
func (c *cancelCtx) propagateCancel(parent Context, child canceler) {
	c.Context = parent // 回溯 指向父节点 删除本子节点用 方便GC 和 通过Value方法回溯

	// 父节点 主要是context.Background() 节点 没有必要取消 它是根节点（设想下 链表的根节点是不是都是空的，上文中TestWithCancel函数 执行到 
	// 这边就返回了 因为parent 是 context.Background() 是根节点 且 其Done是空实现）
	done := parent.Done()  
	if done == nil {
		return // parent is never canceled
	}
    // 
	// 父节点已经取消 后续就不用再执行了 因为 按照 context 的设计理念 父节点取消 肯定是想子节点也取消
	select {
	case <-done:
		// parent is already canceled
		child.cancel(false, parent.Err(), Cause(parent))
		return
	default:
	}

	// 将 parent Context 类型的 转换成 cancelCtx类型 然后构建 树形传递链（一夫多子 一子一父） map可能共享 所以需要枷锁
	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
	}

	// 略过
}

可以看到上面有对孩子节点赋值的操作这是什么情况下发生的呢，接下来我们来看一个例子，是cancelCtx的另一个有传播树形链的使用场景。代码如下：

go 复制代码

func doSomthing2(ctx context.Context, wg *sync.WaitGroup) {
	defer wg.Done()
	for {
		select {
		case <-ctx.Done():
			fmt.Println("playing2")
			return
		default:
			time.Sleep(1 * time.Second)
			fmt.Println("I am working2")

		}

	}

}

func doSomthing(ctx context.Context, wg *sync.WaitGroup) {
	defer wg.Done()
	ctx1, cancel1:= context.WithCancel(ctx)  // ctx 当入参 产生新的 ctx1 这时 ctx下文是 ctx1 上文是 Backgroud产生的空context
	defer cancel1() 
	go doSomthing2(ctx1, wg)

	for {
		select {
		case <-ctx.Done():
 			fmt.Println("playing")
			return
		default:
			time.Sleep(1 * time.Second)
			fmt.Println("I am working")

		}

	}

}

func TestWithCancel(t *testing.T) {
	var wg sync.WaitGroup
	ctx, concel := context.WithCancel(context.Background())
	wg.Add(2)
	go doSomthing(ctx, &wg)
	time.Sleep(3 * time.Second)
	concel()
	wg.Wait()
}

运行结果是：

go 复制代码

I am working2
I am working
I am working
I am working2
I am working
playing
I am working2
playing2

可以看到这时候就形成了树形结构我们来看下形成的树：

如果在子协程中和子子协程中运用With函数进行嵌套最终可形成一颗很大的树，如图：

这时候就形成比较复杂的主从控制方案，如果valueCtx再嵌套上cancelCtx，就可以根据key在树上寻找value（因为指针是双向的，可以根据value函数递归查找，只能往上找不能往下），这样可以做到参数的向下传递。但我们一般常用的是主协程一个With函数，从协程就引用主协程的ctx 也不嵌套，就全局唯一一个ctx，如同第一个例子。

其结构图如下：

WithValue

WithValue 基于 valueCtx结构体，主要用来父context给子context传递数据，也就是说子context可以获得其父context的内容反之则不行。

一个例子：

go 复制代码

func doSomthing(ctx context.Context, wg *sync.WaitGroup) {
	defer wg.Done()
	for {
		select {
		case <-ctx.Done():
			fmt.Println("playing")
			return
		default:
			time.Sleep(1 * time.Second)
			fmt.Println("I am working")
			fmt.Println(ctx.Value("key"))
		}

	}

}

func TestWithCancel(t *testing.T) {
	var wg sync.WaitGroup
	ctx:=context.WithValue(context.Background(),"key","value")
	ctx, concel := context.WithCancel(ctx)
	wg.Add(1)
	go doSomthing(ctx, &wg)
	time.Sleep(3 * time.Second)
	concel()
	wg.Wait()
}

运行结果如下：

go 复制代码

I am working
value
I am working
value
I am working
value
playing

可以看到主协程传递的是 key value 可以在子协程通过 Value(key)函数来获取value。通过valuetx嵌入cancelCtx可以得到一颗树形结构，如上图cancelCtx形成的多节点树形结构，由于valueCtx 可以嵌入 cancelCtx所以可以加入其树形链条中。

其源码如下：

go 复制代码

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

源码比较简单，valueCtx本身一般不会用来传播，一般都是嵌入其他的结构体像cancelCtx，timerCtx 和其他的context，主要用来传递参数。子可以看所有的父参数，父不能看子，因为查找是递归查找。

valueCtx的主要功能是传递参数和查找参数，传递用的是 key,value。查找用的是Value函数我们来看下Value函数。

go 复制代码

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

// 用于在整个context链条（带回溯的树结构）上 寻找 匹配的值 是一个公共方法
func value(c Context, key any) any {
	for {
		switch ctx := c.(type) {
		case *valueCtx:
			if key == ctx.key {
				return ctx.val // 匹配上key 返回
			}
			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) // 回溯 递归查找
		}
	}
}

通过上述代码我们可以看到其之所以可以共享参数主要是由于有递归查找的原因所以子可以查找父传递的参数但是反之不行。

另一个用处是在http包内传递参数，感兴趣的可以查阅相关资料。

人为取消是不是有时候不爽啊，那就加入定时功能吧，所以下面的几个函数本质上就是在cancelCtx功能的扩展，我们看其结构体timerCtx就包含了cancelCtx，只不过在其基础上加了时间维度。下面来看看涉及的几个结构体。

WithDeadlineCause

WithDeadlineCause 是time相关功能的基础函数，WithDeadline和 WithTimeout都是基于这个函数。我们来看下其源码：

go 复制代码

// 定时取消操作
func WithDeadlineCause(parent Context, d time.Time, cause error) (Context, CancelFunc) {
	if parent == nil {
		panic("cannot create context from nil parent")
	}
	
	// 获取父context的取消时间，如果比传入时间早 就执行WithCancel取消逻辑
	if cur, ok := parent.Deadline(); ok && cur.Before(d) {
		// The current deadline is already sooner than the new one.
		return WithCancel(parent)
	}
	
	// 否则 就新建一个新timerCtx 并把取消时间赋给deadline
	c := &timerCtx{
		deadline: d,
	}
	// 构建 context树形结构 将c， parent父子化
	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()

	// dur时间过后 开始执行取消
	if c.err == nil {
		c.timer = time.AfterFunc(dur, func() {
			c.cancel(true, DeadlineExceeded, cause)
		})
	}
	// todo
	return c, func() { c.cancel(true, Canceled, nil) }
}

这个结构体不怎么使用我们一般使用它两个变种。如下

WithDeadline

变种自WithDeadlineCause

源码如下：

go 复制代码

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

只不过去掉了取消原因因为是定时取消原因不用说了

我们看一个例子

go 复制代码

func doWork(ctx context.Context,wg * sync.WaitGroup) {
	defer wg.Done()
	select {
	case <-time.After(3 * time.Second):
		// 模拟操作需要3秒完成
		fmt.Println("操作完成")
	case <-ctx.Done():
		// 当上下文取消时会进入这个分支
		fmt.Println("操作被取消:", ctx.Err())
	}
}
func TestWithDeadline(t *testing.T) {
	// 设置截止时间为当前时间后2秒
	deadline := time.Now().Add(2 * time.Second)
	ctx, cancel := context.WithDeadline(context.Background(), deadline)
	defer cancel() // 确保在主函数返回前取消上下文
	// 模拟操作

	var wg sync.WaitGroup
	wg.Add(1)
	go doWork(ctx,&wg)
	wg.Wait()
}

运行结果如下：

go 复制代码

操作被取消: context deadline exceeded

比较简单我们来看最后一个

WithTimeout

其源码如下：

go 复制代码

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

我们看到这个函数只是 TestWithDeadline一个特例，入参是一个时间戳，只不过后续需要转换成当前时间往后时间戳刻度作为定时时间。

所以上面的例子稍微修改下就可以使用

go 复制代码

func doWork(ctx context.Context,wg * sync.WaitGroup) {
	defer wg.Done()
	select {
	case <-time.After(3 * time.Second):
		// 模拟操作需要3秒完成
		fmt.Println("操作完成")
	case <-ctx.Done():
		// 当上下文取消时会进入这个分支
		fmt.Println("操作被取消:", ctx.Err())
	}
}
func TestWithTimeout(t *testing.T) {
	// 设置截止时间为当前时间后2秒
 	ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second)
	defer cancel() // 确保在主函数返回前取消上下文
	// 模拟操作

	var wg sync.WaitGroup
	wg.Add(1)
	go doWork(ctx,&wg)
	wg.Wait()
}

总结

如此 Context的所有变体和相关函数都介绍完毕了。我们来总结下。

Context接口是内核，cancelCtx内嵌了Context并实现了其部分函数，来执行一些取消操作，而timerCtx内嵌了 cancelCtx从而可以定时取消。
valueCtx内嵌了Context，其本身还有键值对参数，从而可以用来传递参数。其之所以不内嵌cancelCtx主要是为了提供一个基础组合结构体，用户就可以内嵌各种context变体来传递参数，使用范围更广，而不只是限制在cancel变体。
With函数入参和返回值都有Context所以各种变体可以进行组合。例如cancelCtx和valueCtx组合既可以控制从协程又可以向从协程传递参数。
With函数只有在子协程中也执行，才能构建Context树，否则With参数传递Contest就是全局唯一一个ctx，因为With函数Context入参是接口，所以实际上传递的是Context类型和其实现结构体的指针，所以单纯传递Context不会产生值复制。