likes
comments
collection
share

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

作者站长头像
站长
· 阅读数 12

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

写作背景

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

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

名词解释

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

学习 Context 看源码是最好的方法,Context 源码非常精简值得大家研究一番。我梳理下源码画几个类图方便大家理解。 深入Golang的Context「源码分析+详细案例」

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

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

深入Golang的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() (stringerror) {
	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 bool1)
	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 应该放函数或者方法的第一个参数还是包装在结构体里面?官方建议你放在第一个参数,我一般也是这么用的。但如果你非要包装在结构体中,把结构体传入另一个函数或者方法也是可以的。 深入Golang的Context「源码分析+详细案例」

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

参考文献

context package - context - Go Packages

公众号地址

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