本文我们介绍 context 包的传值、超时和取消的使用方式,context 包的这三个功能,我们不仅可以用于跨 goroutine 的操作,而且还可以用于跨服务的操作。
1.介绍
Go 语言在 v1.7 引入 context 包,关于它的使用方式,我们在之前的文章中已经介绍过,感兴趣的读者朋友们可以翻阅。
本文我们介绍 context 包的最佳实践,包括传值、超时和取消。
2.传值
我们可以使用 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 传递的数据。
3.超时
我们可以使用 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。
4.取消
我们可以使用 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)
}
输出结果:
阅读上面这段代码,我们创建一个 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)
}
输出结果:
5.总结
本文我们介绍 context 包的传值、超时和取消的使用方式,context 包的这三个功能,我们不仅可以用于跨 goroutine 的操作,而且还可以用于跨服务的操作。