深入Golang的Context「源码分析+详细案例」

这是我写 go 源码系列第 4 篇，感兴趣可以阅读另外 3 篇

GO singleflight 你真的会用吗？「源码分析+详细案例」 - 掘金

深入 GO unsafe.Pointer & uintptr  - 掘金

深入GO之sync.Once，80%研发同学都不了解其中一行重要代码 - 掘金

GO原子操作「源码分析+案例」 - 掘金

写作背景

为什么会写 Context 呢？我在网上搜了一圈非常多文章在写，但没有讲具体的案例。我把这些年使用的场景和技巧总结下来，结合源码分析+项目实践学习会更高效。

本文源码和案例基于 GO 1.20.4 版本

名词解释

Context：中文翻译为上下文，在多个函数、方法、协程、跨 API、进程之间传递信息。它是 Go 语言标准库中的一个类型， Go 1.7 发布时才被加入到标准库的。

学习 Context 看源码是最好的方法，Context 源码非常精简值得大家研究一番。我梳理下源码画几个类图方便大家理解。

下面我解释下类图

3 个接口

Context 接口

Context 接口定义了跨 API 边界携带请求信息的标准方式。它可以携带截止时间（deadline）、撤销信号（cancellation signal）以及其他值。

stringer 接口

Stringer 接口定义了一个用于生成字符串的标准方式。String() 方法会返回一个字符串，通常用于调试和日志记录。

canceler 接口

canceler 接口是一个用于表示可撤销的上下文类型接口。实现了 canceler 接口的类型可以直接撤销，在需要撤销一系列相关操作时使用。

关于 3 个接口讲完了我提一个问题大家思考下？context 为啥需要定义 3 个接口而不是 1 个接口？答案在文章末尾给出。

4 个结构体

valueCtx 结构体

valueCtx 携带了一个键值对。它实现了 Value 方法，并将其他调用委托给嵌入的 Context

我给大家贴一段源码：

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

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 *timerCtx:
			if key == &cancelCtxKey {
				return ctx.cancelCtx
			}
			c = ctx.Context
		case *emptyCtx: // 遇到 emptyCtx 后退出
			return nil
		default:
			return c.Value(key)
		}
	}
}

往 Context 设置值的 WithValue 方法，这个方法通过传入的 parent 创建一个新的 Context，其中与 key 关联的值是 val。

从 Context 获取值的 Value 方法，该方法递归遍历 Context 获取 val，看源码 case *valueCtx 这部分如果你的 key 重复， val 在查找时会被最后这次覆盖哈。

cancelCtx 结构体

cancelCtx 表示可撤销的上下文。当 cancelCtx 被撤销时，它也会撤销任何实现了 canceler 接口的子级。

我给大家贴一段源码：

type cancelCtx struct {
	Context

	mu       sync.Mutex            // 锁，保护下面字段
	done     atomic.Value          // chan struct{} 的值，懒加载创建，由第一次撤销调用关闭
	children map[canceler]struct{} // 由第一次撤销调用设置为 nil
	err      error                 // 由第一次撤销调用设置为非 nil
	cause    error                 // 由第一次撤销调用设置为非 nil
}

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

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

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

通过调用 WithCancel 函数，获得了一个 Context 值和一个“触发撤销信号的函数”。

大家可能会比较疑惑 Context 是如何撤销的？

这里跟 Context 的 Done() 方法相关，该方法会返回一个 struct 的通道，但是这个通道并不是为了传递值的，而是让调用方感知到撤销的信号量。 当 Context 被撤销，该通道会被立刻关闭，当一个接受通道被关闭引用它的操作都会被立刻停止。

我们再来看 WithCancel 函数的第二个返回值“触发撤销信号的函数” cancel。

再给大家贴一段源码：

func (c *cancelCtx) cancel(removeFromParent bool, err, cause error) {
    // .... 这里省略多行代码，只保留关键代码
	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
    // ... 这里省略多行代码
}

可以调用 cancel 方法手动撤销，该方法内部会 for 循环遍历 c.children 调用 cancel 方法，一旦触发，撤销信号就会立即传达给 Context ，并由它的 Done 方法的返回值（通道）表达出来。

给大家展示一个我写的案例：

func contextCancel() {
	ctx, cancel := context.WithCancel(context.TODO())
	go func() {
		<-ctx.Done()
		fmt.Println("context 撤销")
	}()

	time.Sleep(time.Second * 3)
	cancel() // 手动撤销
	time.Sleep(time.Second * 1) // 休眠1s防止主线程退出程序结束日志未打印完成
}

func TestCtxCancel(t *testing.T) {
	contextCancel()
}

上段代码通过 WithCancel 方法创建一个子 Context，再创建一个协程，该协程基于调用表达式 cxt.Done() 的接收操作，感知撤销信号（必须要调用 cxt.Done() 方法才能感知撤销信号哈）。模拟业务逻辑操作休眠 3 秒，手动调用 cancel() 撤销，协程收到信号量后打印日志。日志打印如下：

=== RUN   TestCtxCancel
context 撤销
--- PASS: TestCtxCancel (4.00s)
PASS

除了感知”撤销信号量“以外，某一些场景感知撤销的原因也是有必要的，可以通过 Context 的  Err 方法获取撤销的具体原因。

func contextCancel() {
	ctx, cancel := context.WithCancel(context.TODO())
	go func() {
		<-ctx.Done()
		fmt.Println("context 撤销")
	}()

	time.Sleep(time.Second * 3)
	cancel()

	fmt.Printf("撤销原因=%v", ctx.Err())
	time.Sleep(time.Second * 1) // 休眠1s防止主线程退出
}

上段代码打印日志如下：

=== RUN   TestCtxCancel
context 撤销
撤销原因=context canceled--- PASS: TestCtxCancel (4.00s)
PASS

WithCancelCause 函数也可以创建一个可撤销的 Context 该函数类似于 WithCancel 函数，只不过该函数第二个返回值是 CancelCauseFunc 而不是 CancelFunc。CancelCauseFunc 支持你自定义 Cause，其他和 WithCancel 一样的。 举一个简单案例：

func contextWithCancelCause() {
	ctx, cancel := context.WithCancelCause(context.TODO())
	go func() {
		<-ctx.Done()
		fmt.Println("context 撤销")
	}()

	cancel(errors.New("数据库连接超时"))
	time.Sleep(time.Second * 1)
	fmt.Printf("撤销原因=%v,cause=%v\n", ctx.Err(), context.Cause(ctx))
	time.Sleep(time.Second * 1) // 休眠1s防止主线程退出
}

func TestCtxWithCancelCause(t *testing.T) {
	contextWithCancelCause()
}

上段代码日志打印如下：

=== RUN   TestCtxWithCancelCause
context 撤销
撤销原因=context canceled,cause=数据库连接超时
--- PASS: TestCtxWithCancelCause (2.00s)
PASS

WithCancelCause 函数你传入的 Cause 可以通过 context.Cause(ctx) 函数获取哈。

timerCtx 结构体

用于携带定时器和截止时间信息。它嵌入了一个 cancelCtx，以实现 Done 和 Err 方法。它通过停止定时器后委托给 cancelCtx.cancel 来实现撤销功能。

我给大家贴一段源码：

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

	deadline time.Time
}

func WithDeadline(parent Context, d time.Time) (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{
		cancelCtx: newCancelCtx(parent),
		deadline:  d,
	}
	propagateCancel(parent, c)
	dur := time.Until(d)
	if dur <= 0 {
		c.cancel(true, DeadlineExceeded, nil) // 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, nil)
		})
	}
	return c, func() { c.cancel(true, Canceled, nil) }
}

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

从源码可以发现它包含了一个定时器和截止时间信息，并利用了 cancelCtx 的功能来实现撤销操作。当撤销 timerCtx 时，它会首先停止定时器，然后委托给 cancelCtx.cancel 方法来执行撤销操作，以确保整个上下文树能够正确地被撤销。

官方提供 2 个函数可以创建 timerCtx，context.WithDeadline 函数和 context.WithTimeout 函数，代码很简单就不细讲了。

举一个简单案例：

func contextWithTimeout() {
	ctx, _ := context.WithTimeout(context.TODO(), time.Second*6)
	go func() {
		<-ctx.Done()
		fmt.Println("context 撤销")
	}()

	time.Sleep(time.Second * 9)
	fmt.Printf("插销原因=%v\n", ctx.Err())
}

上段代码日志打印如下：

=== RUN   TestCtxWithTimeout
context 撤销
插销原因=context deadline exceeded
--- PASS: TestCtxWithTimeout (9.00s)
PASS

“context deadline exceeded” 熟悉吗？是不是 http 、数据库、grpc 调用超时都会抛出这个错误

emptyCtx 结构体

它是一个用于表示空上下文的类型，通常用作默认的顶级上下文。「偷偷告诉你它是一个 int 类型，int 也能实现接口哈」。

我给大家贴一段源码：

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

func (e *emptyCtx) String() string {
	switch e {
	case background:
		return "context.Background"
	case todo:
		return "context.TODO"
	}
	return "unknown empty Context"
}

从源码看出 emptyCtx 虽然实现了接口但是里面是没有逻辑的。要创建 emptyCtx 对象不得不提 context.TODO 函数和 context.Background 函数。

context.TODO 它返回一个非 nil 但上下文为空的 Context（emptyCtx）。我总结了以下 2 点场景使用 TODO 函数。

1、  当代码中需要传递一个上下文，但目前尚不清楚应该使用哪个具体的上下文时，可以暂时使用 context.TODO。

2、  当一个函数尚未被修改以接受上下文参数，但可能在未来需要时，可以在函数调用中使用 context.TODO，以便将来很容易地添加上下文支持而无需修改调用方。

context.Background 它返回一个非 nil 但上下文为空的 Context（emptyCtx）。它永远不会被撤销，没有值，也没有截止时间。通常在主函数、初始化、测试以及作为传入请求的顶级上下文时使用。

虽然这两个函数我都举了使用场景大家可能还是比较疑惑，其实在开发中我也是混着用的。

Context 在项目中的应用场景

timerCtx 定时撤销的 Context

在下面这个场景我会创建定时撤销的 Context

我们在调用三方 API 时，不清楚对方接口的耗时情况（尤其是在 qps 高的场景）耗时长会拉低你的业务处理速度从而影响业务，让客户明显感知到。我们线上业务就出现过类似问题。

我举一个案例，该案例是某一个团队/三方提供的 API 耗时长，在我的业务中我请求该接口构造了一个定时 4s 的 Context，避免长耗时阻塞我的业务。

import (
	"context"
	"github.com/go-resty/resty/v2"
)

type GoogleClient struct {
	ctx context.Context
}

func NewGoogleClient(ctx context.Context) *GoogleClient {
	return &GoogleClient{ctx: ctx}
}

// GetGoogleData 模拟超时场景
func (c *GoogleClient) GetGoogleData() (string, error) {
	client := resty.New()
	_, err := client.R().SetContext(c.ctx).Get("https://www.google.com") // 模拟耗时长场景
	if err != nil {
		return "", err
	}

	return "你好 google", nil
}

上游代码调用如下：

func TestGetUser(t *testing.T) {
	ctx := context.Background()                            // 创建上下文
	ctx, cancel := context.WithTimeout(ctx, time.Second*4) // 创建一个4s撤销的上下文
	defer cancel()                                         // 撤销 context

	name, err := NewGoogleClient(ctx).GetGoogleData() // 获取google 数据
	if err != nil {
		panic(err)
	}
	fmt.Printf("name=%v", name)
}

日志打印如下：

=== RUN   TestGetUser
--- FAIL: TestGetUser (4.00s)
panic: Get "https://www.google.com": context deadline exceeded [recovered]
	panic: Get "https://www.google.com": context deadline exceeded

cancelCtx 可撤销的 Context

可撤销的 Context 在日常开发中非常常见，回想下在你们的开发中经常使用 go 关键字创建协程吗？当你创建一个协程，该协程在执行过程中你能全链路跟踪协程的状态吗？确定是正常退出还是阻塞了？所以在日常发开发中控制协程的退出或撤销很重要，这能使资源得到合理利用，避免内存泄露。

在某一个业务场景，需要消费队列消息（暂定为 kafka），程序启动时需要启动消费者，消费消息处理业务逻辑，当程序停止时需要关闭协程，通知消费者暂停消费。代码如下：

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

type KafkaConsumer struct {
	ctx    context.Context
	cancel context.CancelFunc
}

func NewKafkaConsumer(ctx context.Context) *KafkaConsumer {
	ctx, cancel := context.WithCancel(ctx)
	return &KafkaConsumer{ctx: ctx, cancel: cancel}
}

func (c *KafkaConsumer) Consume() error {
	for {
		select {
		case <-c.ctx.Done():
			fmt.Println("退出 kafka 退出消费...")
			return nil
		default:
			fmt.Println("消费了消息消息内容是 xxxx")
			time.Sleep(1 * time.Second)
		}
	}
}

// GracefulShutdown 平滑启停
func (c *KafkaConsumer) GracefulShutdown() {
	if c.cancel != nil {
		c.cancel()
	}
	// TODO 若有其逻辑需要关闭在后面写
}

func TestConsumer(t *testing.T) {
	consumer := NewKafkaConsumer(context.TODO())
	go func() {
		err := consumer.Consume()
		if err != nil {
			panic(err)
		}
	}()

	time.Sleep(3 * time.Second) // 模拟做了一些工作
	// 代码发布需要停止消费者
	// ......
	consumer.GracefulShutdown() // 平滑启停

	time.Sleep(1 * time.Second) // 给协程时间退出
	fmt.Println("程序退出完毕准备启动")
}

打印日志如下：

=== RUN   TestConsumer
消费了消息消息内容是 xxxx
消费了消息消息内容是 xxxx
消费了消息消息内容是 xxxx
退出 kafka 退出消费...
程序退出完毕准备启动
--- PASS: TestConsumer (4.00s)
PASS

请你思考一个问题，除了用 Context 通知协程退出还有其他方案吗？答案在文章末尾给出。

valueCtx 携带键值的 Context

可携带数据的 Context 在开发中使用非常广泛。

1、  全链路追踪场景用 Contex 携带 reqid、http 路由、ip 等字段实现全链路数据透传。

2、  定义一些标准接口为了减少参数透传，把一些固定的参数比如：租户 id、员工 id、登陆session 放到 Context 中。

3、  在日志输出场景，某些利于检索又是通用的字段（或数据）的也可以利用 Context 透传，在日志中间件中输出。

下面输出一些案例：

下面这段代码简单对 Context 进行了 2 次封装，对外露出了 ContextValue 类型是 map，key 和 value 都是 string 类型，使用方需要把透传的数据放入 ContextValue，调用 WithContext 构造一个 Context，它是携带了数据的 Context。当你需要获取数据时，调用 GetContextValue 、GetAccount、GetReqID  获取业务需要的数据。

import "context"

type ContextValueKey string

const (
	ContextValueKeyAccount ContextValueKey = "account"
	ContextValueKeyReqID   ContextValueKey = "reqId"
	contextKey             string          = "biz-ctx"
)

type ContextValue map[ContextValueKey]string

func WithContext(ctx context.Context, ctxVal ContextValue) context.Context {
	return context.WithValue(ctx, contextKey, ctxVal)
}

func GetContextValue(ctx context.Context) ContextValue {
	origin := getContext(ctx)
	cloned := make(ContextValue, len(origin))
	for k, v := range origin {
		cloned[k] = v
	}
	return cloned
}

func GetAccount(ctx context.Context) string {
	return getContext(ctx)[ContextValueKeyAccount]
}

func GetReqID(ctx context.Context) string {
	return getContext(ctx)[ContextValueKeyReqID]
}

func getContext(ctx context.Context) ContextValue {
	bc := ctx.Value(contextKey)
	if v, ok := bc.(ContextValue); ok {
		return v
	}
	return make(ContextValue)
}

对于调用者代码如下：

import (
	"context"
	"fmt"
	"testing"
)

func TestContext(t *testing.T) {
	bizCtx := make(ContextValue)
	bizCtx[ContextValueKeyAccount] = "123456"
	bizCtx[ContextValueKeyReqID] = "xxxjiuehdkjg"
	ctx := WithContext(context.TODO(), bizCtx)
	fmt.Printf("account=%v,reqid=%v\n", GetAccount(ctx), GetReqID(ctx))
	accountData := GetAccountData(ctx)
	fmt.Printf("accountData=%v\n", *accountData)
}

type AccountData struct {
	ID     string `json:"id"`     // 租户ID
	Name   string `json:"name"`   // 租户名称
	Status string `json:"status"` // 租户状态
	// ....
}

func GetAccountData(ctx context.Context) *AccountData {
	/*
		假设调用其他团队接口获取租户信息
	*/
	accountID := GetAccount(ctx)
	// 调用接口省略
   // http.Get()...
	// .....
	return &AccountData{ID: accountID, Name: "xx有限公司", Status: "inuse"}
}

模拟了日志输出场景、利用 Context 携带固定参数场景，日志输出如下：

=== RUN   TestContext
account=123456,reqid=xxxjiuehdkjg
accountData={123456 xx有限公司 inuse}
--- PASS: TestContext (0.00s)
PASS

总结

1、  关于第一个问题 context 为啥需要定义 3 个接口而不是 1 个接口？

我个人理解：相比于包含很多方法的大接口而言，小接口可以更加专注地表达某一种能力，同时也更容易被组合在一起，使用起来非常灵活，如果你经常研究源码你会发现源码有很多类似的小接口哈，比如： io 中的 ReadWriteCloser 接口和 ReadWriter ....等等，大家可以自己去研究。

2、  第二个问题我讲了使用 Context 通知协程退出，停止 Kafka 消费者消费，除了用 Context 还有其他方案吗？

除了使用 Context 还可以利用通道实现，代码如下：

func runner(exit chan bool) {
	for {
		select {
		case <-exit:
			fmt.Println("协程退出") // 收到退出协程信号，退出for循环
			return
		default:
			time.Sleep(time.Second * 2) // 模拟业务逻辑执行
			fmt.Println("休眠 2s 模拟业务逻辑执行")
		}
	}
}

func TestChanQuit(t *testing.T) {
	exit := make(chan bool, 1)
	go runner(exit)
	time.Sleep(time.Second * 3) // 休眠 3s
	exit <- true                // 模拟协程退出
	time.Sleep(time.Second * 4) // 休眠 4s 等待协程日志打印完毕
}

上段代码日志输出：

=== RUN   TestChanQuit
休眠 2s 模拟业务逻辑执行
休眠 2s 模拟业务逻辑执行
协程退出
--- PASS: TestChanQuit (7.00s)
PASS

3、  撤销信号如何在上下文树中传播？

我在前面讲过 WithCancel、WithDeadline、WithTimeout 都是被用来基于给定的 Context 创建可撤销的子值。这三个函数在被调用后都会返回两个值。第一个值就是可撤销的 Context，第二个值是用于“触发撤销信号的函数”，在“撤销函数”被调用后，它会执行取消操作，将取消信号传播给其所有子上下文。被取消的上下文会关闭其“通道”，并将其 Err 方法返回一个非空的错误。同时，它会调用所有子上下文的取消函数，向它们传播取消信号。子上下文收到取消信号后，会继续向它们的子上下文传播取消信号。这样就形成了一个递归的传播过程，直到所有相关的上下文都被取消为止。

WithDeadline、WithTimeout 创建的 Context 都是可以被手动撤销的哈。

4、  Context 应该放函数或者方法的第一个参数还是包装在结构体里面？官方建议你放在第一个参数，我一般也是这么用的。但如果你非要包装在结构体中，把结构体传入另一个函数或者方法也是可以的。

5、如果你需要全链路透传值，案例我用的 map ，我也建议你用 map，更灵活并且一次性构造 map 数据创建 Context，避免多次创建可以减少树的深度。

参考文献

context package - context - Go Packages

公众号地址

mp.weixin.qq.com/s/nU-RCDJ7L…