Go 上下文的理解与使用

发布时间 2023-08-29 10:02:15作者: 灯火消逝的码头

为什么需要 context

在 Go 程序中,特别是并发情况下,由于超时、取消等而引发的异常操作,往往需要及时的释放相应资源,正确的关闭 goroutine。防止协程不退出而导致内存泄露。如果没有 context,用来控制协程退出将会非常麻烦,我们来举一个例子。

假如说现在一个协程A开启了一个子协程B,这个子协程B又开启了另外两个子协程B1和B2来运行不同的任务,协程B2又开启了协程C来运行其他任务,现在协程A通知子协程B该退出了,这个时候我们需要完成这样的操作:A通知B退出,B退出时通知B1、B2退出,B2退出时通知C退出:

func TestChanCloseGoroutine(t *testing.T) {
	fmt.Printf("开始了,有%d个协程\n", runtime.NumGoroutine())

	var (
		chB  = make(chan struct{})
		chB1 = make(chan struct{})
		chB2 = make(chan struct{})
		chC  = make(chan struct{})
	)

	// 协程A
	go func() {
		// 协程B
		go func() {
			// 协程B1
			go func() {
				for {
					select {
					case <-chB1:
						return
					default:
					}
				}
			}()
			// 协程B2
			go func() {
				// 协程C
				go func() {
					for {
						select {
						case <-chC:
							return
						default:
						}
					}
				}()
				for {
					select {
					case <-chB2:
						// 通知协程C退出
						chC <- struct{}{}
						return
					default:
					}
				}
			}()

			for {
				select {
				case <-chB:
					chB1 <- struct{}{}
					chB2 <- struct{}{}
					return
				default:
				}
			}
		}()

		// 1秒后通知B退出
		time.Sleep(1 * time.Second)
		chB <- struct{}{}
		// A后续没有任务了,会自动退出
	}()

	time.Sleep(2 * time.Second)
	fmt.Printf("最终结束,有%d个协程\n", runtime.NumGoroutine())
}


// 结果
开始了,有2个协程
最终结束,有2个协程
// tips: Go Test 会启动两个额外的 goroutine 来运行代码,所以初始就会有2个 goroutine

通过 channel 来控制各个 goroutine 的关闭,程序看上去一点也不优雅。而且这才仅仅四个 goroutine ,就已经显得有些力不从心了,在真实的业务中,哪怕一个简单的 http 请求,都不可能启用四个 goroutine 就能够完成,且子协程的层级也绝非只有寥寥的三层!

context 是什么

context 在 Go 中是一个接口,它的定义如下:

type Context interface {
	Deadline() (deadline time.Time, ok bool)
	Done() <-chan struct{}
	Err() error
	Value(key any) any
}
  • Deadline 用来获取 ctx 的截止时间,如果没有截至时间,ok 将返回 false;
  • Done 里面是一个通道,当 ctx 被取消时,会返回一个关闭的 channel,如果该 ctx 永远都不会被关闭,则返回 nil;
  • Err 返回的 ctx 取消的原因,如果 ctx 没有被取消,会返回 nil。如果已经关闭了,会返回被关闭的原因,如果是被取消的会返回 canceled,超时的显示 deadline exceeded;
  • Value 会返回 ctx 中储存的值,会从当前 ctx 中一路向上追溯,如果整条 ctx 链中都没有找到值,则会返回nil。

context 的基本结构比较简单,里面也只有四个方法,如果到此没有理解四个方法也没有关系,下文会使用到这四个方法,届时将会很自然的掌握它们。

context 接口的实现

context 有四个不同的实现:emptyCtx、cancelCtx、timerCtx、valueCtx:

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
}

emptyCtx 是一个实现了 context 接口的整型,它不能储存信息,也不能被取消,它被当作根节点 ctx。cancelCtx、timerCtx、valueCtx 由于篇幅原因,这里不放出它们的源码,只解释它们的作用:cancelCtx 是一个可以主动取消的 ctx。timerCtx 也是一个可以主动取消的 ctx,不同于 cancelCtx,它还储存着额外的时间信息,当时间条件满足后,会自动取消该 ctx,利用这点,可以实现超时机制。valueCtx 比较简单,用来创建一个携带键值的 ctx。

context 的基本使用

创建一个根节点

创建根节点有两种方法:

ctx := context.Background()
ctx := context.TODO()

这两种方法其实本质上都是初始化了一个 emptyCtx:

var (
	background = new(emptyCtx)
	todo       = new(emptyCtx)
)

func Background() Context {
	return background
}

func TODO() Context {
	return todo
}

可以看到,在代码中,这两个函数其实是一模一样的,只是用于不同场景下:Background 推荐在主函数、初始化和测试中使用,TODO 用于不清楚使用哪个 context 时使用。根节点 ctx 不具备任何意义,也不能被取消。

创建一个子 ctx

可以通过WithCancel、WithDeadline、WithTimeout、WithValue 这四个主要的函数来创建子 ctx ,创建一个子 ctx 必须指定其归属的父 ctx,由此来形成一个上下文链,用来同步 goroutine 信号。来看一下它们的简单使用:

WithCancel 用来创建一个 cancelCtx,它可以被主动取消 :

func TestCtxWithCancel(t *testing.T) {
	ctx := context.Background()
	ctx, cancel := context.WithCancel(ctx)
	go func() {
		for {
			select {
			// 还记得前文提到的Done的方法吗
			// 当 ctx 取消时,ctx.Done()对应的通道就会关闭,case也就会被执行
			case <-ctx.Done():
				// ctx.Err() 会获取到关闭原因哦
				fmt.Println("协程关闭", ctx.Err())
				return
			default:
				fmt.Println("继续运行")
				time.Sleep(100 * time.Millisecond)
			}
		}
	}()

	// 等待一秒后关闭
	time.Sleep(1 * time.Second)
	cancel()
	// 等待一秒,让子协程有时间打印出协程关闭的原因
	time.Sleep(1 * time.Second)
}

// 结果
继续运行
继续运行
……
协程关闭 context canceled

WithDeadline 用来创建一个 timerCtx,当时间条件满足后,它会被自动取消 :

func TestCtxWithDeadline(t *testing.T) {
	ctx := context.Background()
	// 等待2秒后自动关闭
	ctx, cancel := context.WithDeadline(ctx, time.Now().Add(2*time.Second))
	defer cancel()
	// Deadline 前文也提到了,还记得吗?用来获取当前任务的截至时间
	if t, ok := ctx.Deadline(); ok {
		// time.DateTime 是 go1.20 版本的一个常量,其值是:"2006-01-02 15:04:05"
		fmt.Println(t.Format(time.DateTime))
	}
	go func() {
		select {
		case <-ctx.Done():
			// 手动关闭 context canceled
			// 自动关闭 context deadline exceeded
			fmt.Println("协程关闭", ctx.Err())
			return
		}
	}()

	time.Sleep(3 * time.Second)
}
// 结果
2023-05-10 18:00:36
协程关闭 context deadline exceeded


// 将最后的等待时间更改为一秒
func TestCtxWithDeadline(t *testing.T) {
	……
	time.Sleep(1 * time.Second)
}
// 结果
2023-05-10 18:01:45
协程关闭 context canceled

哪怕 WithDeadline 到达指定时间会自动关闭,但依然推荐使用 defer cancel() 。这是因为如果任务已经完成了,但是自动取消仍需要1天时间,那么系统就会白白浪费资源在这1天上。

WithTimeoutWithDeadline 同理,只不过是 WithTimeout 用来接受一个过期时间,而不是接受一个过期时间节点:

func TestCtxWithTimeout(t *testing.T) {
	ctx := context.Background()
	ctx, cancel := context.WithTimeout(ctx, 2*time.Second)
	defer cancel()
	go func() {
		select {
		case <-ctx.Done():
			fmt.Println("协程关闭", ctx.Err())
			return
		}
	}()

	time.Sleep(3 * time.Second)
}

// 结果
协程关闭 context deadline exceeded

WithValue 用来创建一个 valueCtx:

// 向上找到最近的上下文值
func TestCtxWithValue(t *testing.T) {
	ctx := context.Background()
	ctx1 := context.WithValue(ctx, "key", "ok")
	ctx2, _ := context.WithCancel(ctx1)
	// Value 会一直向上追溯到根节点,获取当前上下文携带的值,
	value := ctx2.Value("key")
	if value != nil {
		fmt.Println(value)
	}
}

// 结果
ok

这四个函数都是创建一个新的子节点,并不是直接修改当前 ctx,所以最后生成的 ctx 链有可能是这样的:

使用 ctx 退出 goroutine

回到开头提到的那个例子,我们使用 context 对其改造一下:

func TestCtxCloseGoroutine(t *testing.T) {
	fmt.Printf("开始了,有%d个协程\n", runtime.NumGoroutine())

	ctx := context.Background()

	// 协程A
	go func(ctx context.Context) {
		ctx, cancel := context.WithCancel(ctx)
		// 协程B
		go func(ctx context.Context) {
			// 协程B1
			go func(ctx context.Context) {
				for {
					select {
					case <-ctx.Done():
						return
					default:
					}
				}
			}(ctx)
			// 协程B2
			go func(ctx context.Context) {
				// 协程C
				go func(ctx context.Context) {
					for {
						select {
						case <-ctx.Done():
							return
						default:
						}
					}
				}(ctx)
				for {
					select {
					case <-ctx.Done():
						return
					default:
					}
				}
			}(ctx)

			for {
				select {
				case <-ctx.Done():
					return
				default:
				}
			}
		}(ctx)

		// 1秒后通知退出
		time.Sleep(1 * time.Second)
		cancel()
		// A后续没有任务了,会自动退出
	}(ctx)

	time.Sleep(2 * time.Second)
	fmt.Printf("最终结束,有%d个协程\n", runtime.NumGoroutine())
}

// 结果
开始了,有2个协程
最终结束,有2个协程

可以看到,和使用 channel 控制 goroutine 退出相比,context 大大降低了心智负担。context 优雅的实现了某一层任务退出,下层所有任务退出,上层任务和同层任务不受影响。

Go 语言最佳实践:每次 context 的传递都应该直接使用值传递,不应该使用指针传递。这样可以防止上下文的值被多个并发的 goroutine 修改而导致竞争问题。虽然使用值传递会导致一些微小的性能开销,因为每次传递上下文时都需要复制一份数据,但它提供了更好的并发安全性和程序可靠性。另外,由于上下文采用了值传递,也不应该向上下文中存入较大的数据,从而导致性能问题。