图解Go语言Context-技术圈

今天与你分享下 Go 语言里面 context 包的相关知识。

一般新技术的出现都是为了解决现有技术存在的问题或者可以提供更优雅方便的实现方式，那我们就要想想 contex 包能解决什么问题？

关于 context 能解决什么问题，推荐你看看这两篇文章：

飞雪大佬的《Go 语言实战笔记（二十）| Go Context》^[1]

和

饶大的《深度解密 Go 语言之 context》

这两篇已经写得足够清晰明了，这里我主要想假设一些场景，更方便大家理解问题：

假设你开启了一个函数，你需要将一些常用的值传递给下游函数，但是不能通过函数参数传递，怎么办？
假设你开启了一个协程 A，协程 A 衍生出很多子协程，这些子协程又衍生出子协程，如果协程 A 所完成的任务“成果”不再需要，那我们怎么通知衍生出的子协程及时退出并释放占用的系统资源呢？
假设一个任务需要在 2s 内完成，如果超时，如何优雅地退出返回呢？
假设一个任务需要在中午 12 点完成，如果到点没有完成，又该如何优雅地退出呢？

好了，带着这些问题，我们接着往下看。

context 接口

理解 context 包，核心是需要理解 Context 接口：

type Context interface {

    Done() <-chan struct{}

    Err() error

    Deadline() (deadline time.Time, ok bool)

    Value(key interface{}) interface{}
}

这个接口有四个方法：

Done() 返回只读的 channel，在 goroutine 中，如果该 channel 可读，则意味着父 context 发起了取消操作或者是时间到期，理解这一点非常重要。

Err() 返回错误，表示 channel 被关闭的原因，被取消还是超时。

Deadline() 获取设置的截止时间，第一个是截止时间，表示到了这个点，context 会自动发起取消操作；第二个表示是否设置了截止时间。

Value() 方法获取 context 上绑定的值，是一个键值对，这个值一般是线程安全的。

Done() 是最常用的方法，经常与 select-case 配合使用，因为 context 取消的时候，我们就可以得到一个可读的 channel，以此来判断是否收到 context 取消的信号，最经典的用法可以在 context 包的代码中找到：

func Stream(ctx context.Context, out chan<- Value) error {
   for {
    v, err := DoSomething(ctx)
    if err != nil {
     return err
    }
    select {
    case <-ctx.Done():
     return ctx.Err()
    case out <- v:
    }
   }
  }

衍生 context

我们不需要自己实现 Context 接口，源码包已经为我们提供了两个实现接口的方法，分别是 Background() 和 TODO()；

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

func Background() Context {
 return background
}

func TODO() Context {
 return todo
}

context.Background() 返回空的 context，通常用在 main 函数里，作为根 context 衍生出子 context。

context.TODO() 也是返回空 context。主要用在还不清楚使用什么类型的 context 的时候，便于后期重构，先用它占个位。

它们两本质上是 emptyCtx 类型，不能被取消，没有值，也没有超时时间。

有了上面两个根 context，就可以衍生出子 context，源码包为我们提供了一系列 withXXX 函数用于生成子 context，分别是：

func WithCancel(parent Context) (ctx Context, cancel CancelFunc)

func WithDeadline(parent Context, deadline time.Time) (Context, CancelFunc)

func WithTimeout(parent Context, timeout time.Duration) (Context, CancelFunc)

func WithValue(parent Context, key, val interface{}) Context

这四个函数的第一个参数都是父 context，可以理解为基于父 context 生成子 context，即衍生子 context。

WithCancel() 基于父 context，返回子 context 和取消函数；
WithDeadline() 基于父 context，返回带截止时间的子 context 和取消函数；
WithTimeout() 基于父 context，返回带超时时间的子 context 和取消函数；
WithValue() 基于父 context，返回绑定键值对的子 context，没有取消函数；

前三个函数都会返回取消函数，需要注意的是只有创建该 context 的协程才能调用取消函数，且不推荐将取消函数作为参数传递。

我们可以调用取消函数取消一个 context，以及这个 context 下面所有的子 context。

通过这些函数，就能生成一棵 context 树，树的每个节点都可以有任意多个子节点，节点层级可以有任意多个。

context 树

我们一起来看下基于两个根 context 可以创建的 context 树是什么样。

两层树

rootCtx := context.Background()
childCtx := context.WithValue(rootCtx, "request_Id", "seekload")

上面的代码基于根 context 创建了两层 context 树，rootCtx 衍生出 childCtx，并携带键值对 {"request_Id" : "seekload"}。

三层树

rootCtx := context.Background()
childCtx := context.WithValue(rootCtx, "request_Id", "seekload")
childOfChildCtx, cancelFunc := context.WithCancel(childCtx)

基于两层树，childCtx 衍生出 childOfChildCtx，含有键值对并且具有取消功能。

多层树

rootCtx := context.Background()
childCtx1 := context.WithValue(rootCtx, "request_Id", "seekload")
childCtx2, cancelFunc := context.WithCancel(childCtx1)
childCtx3 := context.WithValue(rootCtx, "user_Id", "user_100")

上面的代码：

rootCtx 是根 context；
rootCtx 衍生出 childCtx1，并携带键值对 {"request_Id" : "seekload"}；
childCtx1 衍生出 childCtx2，可以取消 context；
rootCtx 衍生出 childCtx3，携带键值对 {"user_Id" : "user_100"}；

层级关系就像下面这样：

我们可以在任一结点 context 上创建子 context，比如从 childCtx1 衍生 childCtx4：

childCtx4 := context.WithValue(childCtx1, "token", "token_some")

层级关系就变成这样了：

Talk is cheap. Show me the code.

看到这里，可能你还是不知道怎么去用 context 包，接下来我们结合着示例展示下 withXXX 函数的使用方法。

如何使用

context.WithCancel()

context.WithCancel() 用于取消信号，直接来看例子：

func main() {
 ctx := context.Background()
 cancelCtx, cancelFunc := context.WithCancel(ctx)
 go task(cancelCtx)
 time.Sleep(time.Second * 3)
 cancelFunc()                        // 取消 context
 time.Sleep(time.Second * 1)         // 延时等待协程退出
 fmt.Println("number of goroutine: ",runtime.NumGoroutine()) // 协程数量
}

func task(ctx context.Context) {
 i := 1
 for {
  select {
  case <-ctx.Done(): // 接收取消信号
   fmt.Println("Gracefully exit")
   fmt.Println(ctx.Err()) // 取消原因
   return
  default:
   fmt.Println(i)
   time.Sleep(time.Second * 1)
   i++
  }
 }
}

输出：

1
2
3
Gracefully exit
context canceled
number of goroutine:  1

当调用 cancelFunc()，Done() 返回的 channel 变成可读，Err() 返回取消原因 “context canceled”，task() 函数执行 return 优雅地退出。

context.WithValue()

通过 context.WithValue() 可以在 goroutine 之间传递一些数据。

func main() {
 helloWorldHandler := http.HandlerFunc(HelloWorld)
 http.Handle("/hello", inejctRequestId(helloWorldHandler))
 http.ListenAndServe(":8080", nil)
}

func HelloWorld(w http.ResponseWriter, r *http.Request) {
 requestId := ""
 if m := r.Context().Value("requestId"); m != nil {
  if value, ok := m.(string); ok {
   requestId = value
  }
 }
 w.Header().Add("requestId", requestId)
 w.Write([]byte("Hello, world"))
}

func inejctRequestId(next http.Handler) http.Handler {
 return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
  requestId := uuid.New().String()
  ctx := context.WithValue(r.Context(), "requestId", requestId)
  req := r.WithContext(ctx)
  next.ServeHTTP(w, req)
 })
}

上面的代码，inejctRequestId() 是请求中间函数，通过 context.WithValue() 注入了键值对；HelloWorld() 是请求处理函数，从 context 获取到刚才绑定的 k-v。

go run 上面示例，然后执行：

curl -v localhost:8080/hello

会输出：

*   Trying ::1...
* TCP_NODELAY set
* Connected to localhost (::1) port 8080 (#0)
> GET /hello HTTP/1.1
> Host: localhost:8080
> User-Agent: curl/7.54.0
> Accept: */*
>
< HTTP/1.1 200 OK
< Requestid: e0b0544d-7993-4ff5-a2de-b29eacd3645a
< Date: Mon, 08 Feb 2021 13:13:06 GMT
< Content-Length: 12
< Content-Type: text/plain; charset=utf-8
<
* Connection #0 to host localhost left intact
Hello, world

从输出可以看到，返回里有 Requestid 信息。

context.WithTimeout()

context.WithTimeout() 可以设置一个超时时间，过期之后 channel done 会自动关闭，context 会被取消；超时之前可以调用取消函数手动取消 context。


func main() {
 ctx := context.Background()
 cancelCtx, cancel := context.WithTimeout(ctx, time.Second*3)
 defer cancel()
 go task(cancelCtx)
 time.Sleep(time.Second * 4)
}

func task(ctx context.Context) {
 i := 1
 for {
  select {
  case <-ctx.Done():
   fmt.Println("Gracefully exit")
   fmt.Println(ctx.Err())
   return
  default:
   fmt.Println(i)
   time.Sleep(time.Second * 1)
   i++
  }
 }
}

输出：

1
2
3
Gracefully exit
context deadline exceeded

上面的代码，context.WithTimeout() 设置了 3s 的超时时间，时间到了之后，context 自动取消，done channel 变成可读，Err() 返回取消原因，执行 return，task() 优雅地退出。

context.WithDeadline()

context.WithDeadline() 设置一个将来的时间点作为截止时间，时间到了之后，channel done 会自动关闭，context 会被取消；还未到截止时间可以调用取消函数手动取消 context。

func main() {
 ctx := context.Background()
 cancelCtx, cancel := context.WithDeadline(ctx, time.Now().Add(time.Second*3))
 defer cancel()
 go task(cancelCtx)
 time.Sleep(time.Second * 4)   // 延时，等待 task() 正常退出
}

func task(ctx context.Context) {
 i := 1
 for {
  select {
  case <-ctx.Done():
   fmt.Println("Gracefully exit")
   fmt.Println(ctx.Err())
   return
  default:
   fmt.Println(i)
   time.Sleep(time.Second * 1)
   i++
  }
 }
}

上面的代码设置的截止时间是 3s 钟之后的时间点，时间到了之后，context 自动取消，done channel 变成可读，Err() 返回取消原因，执行 return，task() 优雅地退出。

输出：

1
2
3
Gracefully exit
context deadline exceeded

相信你也能猜想到，其实 context.WithTimeout() 底层是通过 context.WithDeadline() 实现的，源码如下：

func WithTimeout(parent Context, timeout time.Duration) (Context, CancelFunc) {
 return WithDeadline(parent, time.Now().Add(timeout))
}

ps: 看完示例之后，建议你回头看看上面几节，相信对上面的内容理解会更加深刻。

总结

全文已经很好地回答了文章开始的四个问题，作为 Go 语言的核心功能之一，context 包已经得到广泛的应用，比如上面例子里提到的 http 包。在使用时有几个需要注意的地方：

context 是线程安全的，可在多个 goroutine 中传递；
使用 context 作为函数参数时，需作为第一个参数，并且命名为 ctx；
不要把 context 放在结构体中，要以参数的方式传递；
当不知道传递什么类型 context 时，可以使用 context.TODO()；
context 只能被取消一次，应当避免从已取消的 context 衍生 context；
只有父 context 和创建了该 context 的函数才能调用取消函数，避免传递取消函数 cancelFunc；

参考资料

[1]

《Go 语言实战笔记（二十）| Go Context》: https://www.flysnow.org/2017/05/12/go-in-action-go-context.html

推荐阅读

官发博文：进一步阐明关于 context 的最佳实践

福利

我为大家整理了一份从入门到进阶的Go学习资料礼包，包含学习建议：入门看什么，进阶看什么。关注公众号「polarisxu」，回复 ebook 获取；还可以回复「进群」，和数万 Gopher 交流学习。