Go 语言 context 最佳实践
Go 语言 context 最佳实践
01
介绍
Go 语言在 v1.7 引入 context 包，关于它的使用方式，我们在之前的文章中已经介绍过，感兴趣的读者朋友们可以翻阅。
本文我们介绍 context 包的最佳实践，包括传值、超时和取消。
02
传值
我们可以使用 context 包的 func WithValue() 函数传递数据。
func main() {
ctx := context.WithValue(context.Background(), "ctxKey1", "ctxVal")
go func(ctx context.Context) {
// 读取 ctx 的 value
data, ok := ctx.Value("ctxKey1").(string)
if ok {
fmt.Printf("sub goroutine get value from parent goroutine, val=%s\n", data)
}
}(ctx)
time.Sleep(1 * time.Second)
}
输出结果：
sub goroutine get value from parent goroutine, val=ctxVal
阅读上面这段代码，我们使用 func WithValue() 函数创建一个 context，并且传递 key 为 ctxKey1 的数据。
我们知道 context 是并发安全的，所以我们每次使用 context 传递一个新数据，都需要使用 func WithValue() 函数创建一个新的 context，包装一下 parent context。
传递多个数据
...
ctx := context.WithValue(context.Background(), "ctxKey1", "ctxVal")
ctx = context.WithValue(ctx, "ctxKey2", "ctxVal2")
ctx = context.WithValue(ctx, "ctxKey3", "ctxVal3")
...
阅读上面这段代码，我们可以发现，如果使用 context 传递多个数据，就需要使用 func WithValue() 创建多个 context。
虽然通过使用 func WithValue() 创建多个 context 的方式，可以实现我们的需求，但是，它使代码不再优雅，并且性能也会降低。
怎么解决？
针对该场景，我们可以参考 gRPC 框架的 metadata 包的代码。定义一个 map，通过传递 map 类型的值，实现需要使用 context 传递多个数据的需求。
func main() {
ctxVal := make(map[string]string)
ctxVal["k1"] = "v1"
ctxVal["k2"] = "v2"
ctx := context.WithValue(context.Background(), "ctxKey1", ctxVal)
go func(ctx context.Context) {
// 读取 ctx 的 value
data, ok := ctx.Value("ctxKey1").(map[string]string)
if ok {
fmt.Printf("sub goroutine get value from parent goroutine, val=%+v\n", data)
}
}(ctx)
time.Sleep(1 * time.Second)
}
输出结果：
sub goroutine get value from parent goroutine, val=map[k1:v1 k2:v2]
修改传递数据
使用 context 包的 func WithValue() 函数传递的数据，不建议在传输过程中进行修改，如果遇到在传输过程中需要修改数据的场景，我们可以使用 COW 的方式处理，从而避免 data race。
func main() {
ctxVal := make(map[string]string)
ctxVal["k1"] = "v1"
ctxVal["k2"] = "v2"
ctx := context.WithValue(context.Background(), "ctxKey1", ctxVal)
go func(ctx context.Context) {
// 读取 ctx 的 value
data, ok := ctx.Value("ctxKey1").(map[string]string)
if ok {
ctxVal := make(map[string]string)
for k, v := range data {
ctxVal[k] = v
}
ctxVal["k3"] = "v3"
ctx = context.WithValue(ctx, "ctxKey1", ctxVal)
data, ok := ctx.Value("ctxKey1").(map[string]string)
if !ok {
fmt.Printf("sub goroutine get value from parent goroutine, val=%+v\n", nil)
}
fmt.Printf("sub goroutine get value from parent goroutine, val=%+v\n", data)
}
}(ctx)
time.Sleep(1 * time.Second)
}
输出结果：
sub goroutine get value from parent goroutine, val=map[k1:v1 k2:v2 k3:v3]
阅读上面这段代码，我们通过 COW（copy on write） 方式修改 context 传递的数据。
03
超时
我们可以使用 context 包的 func WithTimeout() 函数设置超时时间，从而避免请求阻塞。
func main() {
ctx, cancel := context.WithTimeout(context.Background(), 1*time.Millisecond)
defer cancel()
select {
case <-time.After(1 * time.Second):
fmt.Println("overslept")
case <-ctx.Done():
fmt.Println(ctx.Err())
}
输出结果：
context deadline exceeded
阅读上面这段代码，我们使用 func WithTimeout() 函数创建一个 1ms 取消的 context，使用 select ... case ... 读取 ctx.Done()，从而取消监听该 context 的 goroutine。
04
取消
我们可以使用 context 包的 func WithCancel() 函数取消操作，从而避免 goroutine 泄露。
func main() {
gen := func() <-chan int {
dst := make(chan int)
go func() {
var n int
for {
dst <- n
n++
}
}()
return dst
}
for n := range gen() {
fmt.Println(n)
if n == 5 {
break
}
}
time.Sleep(1 * time.Second)
}
输出结果：
0
1
2
3
4
5
阅读上面这段代码，我们创建一个 gen() 函数，启动一个 goroutine 生成整数，循环调用 gen() 函数输出生成的整数，当整数值为 5 时，停止循环，从输出结果看，没有发现问题。
但是，实际上该段代码会导致 goroutine 泄露，因为 gen() 函数一直在无限循环。
怎么解决？
我们可以使用 func WithCancel() 函数创建一个 context，作为 gen() 函数的第一个参数，当停止循环时，同时调用 context 的 CancelFunc 取消 gen() 函数启动的 goroutine。
func main() {
gen := func(ctx context.Context) <-chan int {
dst := make(chan int)
go func() {
var n int
for {
dst <- n
n++
}
}()
return dst
}
ctx, cancel := context.WithCancel(context.Background())
defer cancel()
for n := range gen(ctx) {
fmt.Println(n)
if n == 5 {
cancel()
break
}
}
time.Sleep(1 * time.Second)
}
输出结果：
0
1
2
3
4
5
05
总结
本文我们介绍 context 包的传值、超时和取消的使用方式，context 包的这三个功能，我们不仅可以用于跨 goroutine 的操作，而且还可以用于跨服务的操作。
参考资料：
- https://go.dev/talks/2014/gotham-context.slide#29
- https://go.dev/blog/context