zl程序教程

您现在的位置是:首页 >  Javascript

当前栏目

Go 语言 context 优秀实践

2023-02-19 12:24:07 时间

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

输出结果:

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

5.总结

本文我们介绍 context​ 包的传值、超时和取消的使用方式,context​ 包的这三个功能,我们不仅可以用于跨 goroutine 的操作,而且还可以用于跨服务的操作。