细数 Context 使用场景

原创

Seekload 2022-11-10 12:49:43 ©著作权

文章标签 java python go 编程语言 golang 文章分类 运维

©著作权归作者所有：来自51CTO博客作者Seekload的原创作品，请联系作者获取转载授权，否则将追究法律责任

你好，我是四哥。

前一篇文章从源码的角度详细介绍了 Context 的实现原理，但是还没有提到 Context 的使用场景，今天我们一起来看下：

1.请求链路传值。

传值使用方式如下：

func func1(ctx context.Context) {
 ctx = context.WithValue(ctx, "k1", "v1")
 func2(ctx)
}

func func2(ctx context.Context) {
 fmt.Println("func2:",ctx.Value("k1").(string))
 ctx = context.WithValue(ctx, "k2", "v2")
 func3(ctx)
}

func func3(ctx context.Context) {
 fmt.Println("func3:",ctx.Value("k1").(string))
 fmt.Println("func3:",ctx.Value("k2").(string))
}

func main() {
 ctx := context.Background()
 func1(ctx)
}

我们在 func1() 通过函数 WithValue() 设置了一个键值对 k1-v1，在 func2() 可以获取到 func1() 设置的键值对，如果调用 func3() 时把这个 ctx 继续传入的话，在 func3() 中依然还是可以获取到 k1-v1。

但是在 func1() 中获取不到 func2() 设置的键值对 k2-v2，因为 context 只能自上而下携带值，这点需要注意。

2.取消耗时操作，及时释放资源。

使用 channel + select 的机制：

func func1() error {
 respC := make(chan int)        // 起消息通知作用
 // 处理逻辑
 go func() {
  time.Sleep(time.Second * 3)   // 模拟处理业务逻辑
  respC <- 1
  close(respC)
 }()

 // 判断是否超时
 select {
 case r := <-respC:
  fmt.Printf("Resp: %d\n", r)
  return nil
 case <-time.After(time.Second * 2):     // 超过设置的时间就报错
  fmt.Println("catch timeout")
  return errors.New("timeout")
 }
}

func main() {
 err := func1()
 fmt.Printf("func1 error: %v\n", err)
}

上面的方式平时也会用到，通过 context 怎么实现呢？

下面来看下如何使用 context 进行主动取消、超时取消。

主动取消：

func func1(ctx context.Context, wg *sync.WaitGroup) error {
 defer wg.Done()
 respC := make(chan int)
 go func() {
  time.Sleep(time.Second * 5)    // 模拟业务逻辑处理
  respC <- 10
 }()
 // 取消机制
 select {
 case <-ctx.Done():
  fmt.Println("cancel")
  return errors.New("cancel")
 case r := <-respC:
  fmt.Println(r)
  return nil
 }
}

func main() {
 wg := &sync.WaitGroup{}
 ctx, cancel := context.WithCancel(context.Background())
 wg.Add(1)
 go func1(ctx, wg)
 time.Sleep(time.Second * 2)
 cancel()   // 主动取消
 wg.Wait()  // 等待 goroutine 退出
}

超时取消：

func func1(ctx context.Context) {
 resp := make(chan int)
 go func() {
  time.Sleep(time.Second * 5)     // 模拟处理逻辑
  resp <- 1
 }()
 // 超时机制
 select {
 case <-ctx.Done():
  fmt.Println("ctx timeout")
  fmt.Println(ctx.Err())
 case <-resp:
  fmt.Println("done")
 }
 return
}
func main() {
 ctx, cancel := context.WithTimeout(context.Background(), time.Second*2)
 defer cancel()
 func1(ctx)
}

3.防止 goroutine 泄露。

引自【深度解密 Go 语言之 context^[1]】

func gen() <-chan int {
 ch := make(chan int)
 go func() {
  var n int
  for {
   ch <- n
   n++
   time.Sleep(time.Second)
  }
 }()
 return ch
}

这是一个可以生成无限整数的协程，但如果我只需要它产生的前 5 个数，那么就会发生 goroutine 泄漏：

func main() {
 for n := range gen() {
  fmt.Println(n)
  if n == 5 {
   break
  }
 }
 // ……
}

当 n == 5 的时候，直接 break 掉。那么 gen 函数的协程就会执行无限循环，永远不会停下来。发生了 goroutine 泄漏。

用 context 改进这个例子：

func gen(ctx context.Context) <-chan int {
 ch := make(chan int)
 go func() {
  var n int
  for {
   select {
   case <-ctx.Done():
    return
   case ch <- n:
    n++
    time.Sleep(time.Second)
   }
  }
 }()
 return ch
}

func main() {
 ctx, cancel := context.WithCancel(context.Background())
 defer cancel() // 避免其他地方忘记 cancel，且重复调用不影响

 for n := range gen(ctx) {
  fmt.Println(n)
  if n == 5 {
   cancel()
   break
  }
 }
 // ……
}