深入 GO Context「源码分析+项目案例」

写作背景

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

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

名词解释

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

学习 Context 看源码是最好的方法，Context 源码非常精简值得大家研究一番。我梳理下源码画几个类图方便大家理解（尝试在一张图中画全类图但太难看了容易绕晕，拆成 4 张吧）。

下面我解释下类图

3 个接口

Context 接口

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

stringer 接口

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

canceler 接口

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

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

4 个结构体

valueCtx 结构体

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

我给大家贴一段源码：

go 复制代码

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 接口的子级。

我给大家贴一段源码：

go 复制代码

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 函数创建一个可撤销的 parent 的子值。通过调用该函数，获得了一个 Context 值和一个"触发撤销信号的函数"。

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

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

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

再给大家贴一段源码：

go 复制代码

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 方法的结果值（一个接收通道）表达出来。

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

scss 复制代码

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() 撤销，协程收到信号量后打印日志。日志打印如下：

diff 复制代码

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

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

scss 复制代码

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防止主线程退出
}

上段代码打印日志如下：

diff 复制代码

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

WithCancelCause 函数也可以创建一个可撤销的 parent 的子值

该函数类似于 WithCancel 函数，只不过该函数第二个返回值是 CancelCauseFunc 而不是 CancelFunc。CancelCauseFunc 支持你自定义 Cause，其他和 WithCancel 一样的。

举一个简单案例：

scss 复制代码

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

上段代码日志打印如下：

diff 复制代码

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

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

timerCtx 结构体

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

我给大家贴一段源码：

go 复制代码

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 函数，代码很简单就不细讲了。

举一个简单案例：

scss 复制代码

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

上段代码日志打印如下：

diff 复制代码

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

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

emptyCtx 结构体

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

我给大家贴一段源码：

go 复制代码

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

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

我们在调用三方 API 时，不清楚对方接口的耗时情况并且你的业务处理要低延迟（尤其是在 qps 高的场景）耗时长会拉低你的业务处理速度导致的程序阻塞，我们业务就出现过这个问题。

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

go 复制代码

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
}

上游代码调用如下：

go 复制代码

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

日志打印如下：

sql 复制代码

=== 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），程序启动时需要启动消费者，消费消息处理业务逻辑，当程序重启或者停止时我需要平滑启停防止消息被丢弃。代码如下：

go 复制代码

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("程序退出完毕准备启动")
}

打印日志如下：

diff 复制代码

=== 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 获取业务需要的数据。

go 复制代码

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

对于调用者代码如下：

go 复制代码

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 携带固定参数场景，日志输出如下：

ini 复制代码

=== 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 还可以利用通道实现，代码如下：

go 复制代码

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 等待协程日志打印完毕
}

上段代码日志输出：

diff 复制代码

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

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

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

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

4、 Context 应该放第一个参数还是放在结构体里面？官方建议你放在第一个参数，我一般也是这么用的。但是你放到结构体里面也不是不行。

参考文献

context package - context - Go Packages