Context 在 Go 中的实践

背景

关于 Go 的 Context，我应该只写过一篇：Go 语言打印所有的 Context 值 ，而且是比较早期写的了，在不断地使用 Go 过程中，逐渐对于 Context 有了更多的认识，想扩展地写一下，但是觉得直接修改这篇文章不太合适，所以决定重新开一篇文章更加系统地讲一下。

Context 常见用法

首先我了解到 Context 并不是从 Go 诞生之初就存在的，而是在 Go 1.7 中开始加入的，加入的动机是想解决跨 goroutine 的协作控制，以及信息传递问题，所以从接口上看，可以看到 Context 的函数有：

func WithValue(parent Context, key, val any) Context
func WithCancel(parent Context) (ctx Context, cancel CancelFunc)
func WithDeadline(parent Context, d time.Time) (Context, CancelFunc)
func WithTimeout(parent Context, timeout time.Duration) (Context, CancelFunc)

这些接口就反应了 Context 的用处，这里就给一些示例介绍这些方法的使用。

传递信息

这是一段简单的通过 Context 传递信息的示例：

[root@liqiang.io]# cat example0.go
func main() {
    type favContextKey string
    f := func(ctx context.Context, k interface{}) {
        if v := ctx.Value(k); v != nil {
            fmt.Println("found value:", v)
            return
        }
        fmt.Println("key not found:", k)
    }
    k := favContextKey("language")
    ctx := context.WithValue(context.Background(), k, "Go")
    f(ctx, k)
    f(ctx, "language")
}

在这个简单的例子中，我往 Context 里面添加了一对 Key/Value，然后尝试从 Context 中将值取出来。这个例子的运行输出是：

[root@liqiang.io]# go run example0.go
found value: Go
key not found: language

从输出中可以发现，获取值传进去的都是 “language” 的字符串，但是一个是可以获取到值的，另外一个是获取不到值的，所以 Context 的 Key 是强类型对比的，关于这个话题我抽空再来聊一聊，如果感兴趣也可以看这个文档：Go Comparision Operators。

所以在平时使用的时候，一般想从一个 Context 中存取值都是通过 const 来做的，而这往往又是通过引入 SDK 的形式实现的，做法大概可以简化成这样：

[root@liqiang.io]# cat example1.go
func businessCode(ctx context.Context) {
    val := ctx.Value(sdk.CtxKeyHello)
    if val == nil{
        fmt.Println("value not found in context")
    } else {
        fmt.Println("value found in context ", val.(string))
    }
}

当然，也很多人使用另外一种方式来存取 Context 的值，那就是通过一个 Operator 的形式，例如暴露出来一个接口：

[root@liqiang.io]# cat example2.go
import "context"
type ContextOperator interface {
    WithHello(ctx context.Context, value interface{})
    GetHello(ctx context.Context)
}

将存取 Context 的值的操作都封装在一个接口中，这样就对用户的操作隐藏掉了很多细节，用户也无需关系你的 key 是什么了，但是这种方式还是一般以 SDK 的形式进行。

时序控制

在并发编程中，我们经常要面对的一个问题就是并发控制，例如我们开启了多个 goroutine 进行某项任务，但是，当其中任何一个 goroutine 执行的任务失败的时候，我们可能希望其他的 goroutine 也不需要执行了，直接返回，那么这个时候我们可能就可以考虑使用 Context 来控制，比如这段代码：

[root@liqiang.io]# cat main.go
func main() {
    ctx, cancel := context.WithCancel(context.Background())
    wg := sync.WaitGroup{}
    wg.Add(3)
    go func() {
        defer wg.Done()
        println("goroutine 1 start")
        <-ctx.Done()
        println("goroutine 1 done")
    }()
    go func() {
        defer wg.Done()
        println("goroutine 2 start")
        <-ctx.Done()
        println("goroutine 2 done")
    }()
    go func() {
        defer wg.Done()
        println("goroutine 3 start")
        time.Sleep(1 * time.Second)
        cancel()
        println("goroutine 3 fail, canceled")
    }()
    wg.Wait()
    println("main done")
}

这里我开启了三个 goroutine，其中第三个 goroutine 会调用 cancel 取消上下文，这样，当第一个和第二个 goroutine 收到消息之后，就知道应该退出了，然后它就会退出。

Context 的结构

那么，在使用 Context 的时候，不知道你有没有过一些疑问，例如如果我嵌套了多个 Cancel 的 Context，如果我 Cancel 了其中一个是否会导致所有的 Context 都被 Cancel？同样的，如果我收到一个 Context，我添加一个新的 Key/Value，如果这个 Context 中以前已经包含了对应的 key 了，那么以前的 key 是否会被替换？要想回答这些问题，我们还得回到源码中来。

我们先看几个我们常用的 Context 的实现：

cancelCtx

[root@liqiang.io]# cat context.go
func withCancel(parent Context) *cancelCtx {
    if parent == nil {
        panic("cannot create context from nil parent")
    }
    c := &cancelCtx{}
    c.propagateCancel(parent, c)
    return c
}
... ...
type cancelCtx struct {
    Context
    mu       sync.Mutex            // protects following fields
    done     atomic.Value          // of chan struct{}, created lazily, closed by first cancel call
    children map[canceler]struct{} // set to nil by the first cancel call
    err      error                 // set to non-nil by the first cancel call
    cause    error                 // set to non-nil by the first cancel call
}

可以看到 cancelCtx 是一个树形结构：

图 ：这个是图片说明
图片来源于 Understanding Context Package in Golang

timerCtx

[root@liqiang.io]# cat context.go
type timerCtx struct {
    cancelCtx
    timer *time.Timer // Under cancelCtx.mu.
    deadline time.Time
}

valueCtx

[root@liqiang.io]# cat context.go
type valueCtx struct {
    Context
    key, val any
}
func WithValue(parent Context, key, val any) Context {
    ... ...
    return &valueCtx{parent, key, val}
}
func (c *valueCtx) Value(key any) any {
    if c.key == key {
        return c.val
    }
    return value(c.Context, key)
}
func value(c Context, key any) any {
    for {
        switch ctx := c.(type) {
        case *valueCtx:
            ... ...
            c = ctx.Context
        case *cancelCtx:
            ... ...
            c = ctx.Context
        case withoutCancelCtx:
            ... ...
            c = ctx.c
        case *timerCtx:
            ... ...
            c = ctx.Context
        case backgroundCtx, todoCtx:
            return nil
        default:
            return c.Value(key)
... ...

从 valueCtx 的实现来看，Go 的 Context 的取值效率是非常低的，因为他是一个链表的操作，所以，如无必要，还是少往 Context 里面存取键值对。

跨进程传递

在前面介绍了 Context 在 Go 中的实现，以及在单个进程内的传递，但是，在实际使用中，我们可能需要的是跨进程的传递，例如在一个 RPC 请求中，我们经常会将 RequestID 之类的放在 Context，无论是同个进程还是跨进程，这个 RequestID 都是一致的。

跨进程传递的要素

那么在跨进程（RPC）场景下，Context 是如何传递的呢？其实本质上 Context 跨进程的传递就是三步：

• 确定要传递的 Context key 和 value；
• 将要传递的 key 和 value 附加到跨进程通信的消息中；
• 将跨进程通信消息中的 key 和 value 附加回 Context 中；

对于其中的每一步我们都有很多不一样的实践，这里做一个简单的探讨。

确定要传递的键值对

对于要传递的 Context key 和 value，常见的方式有两类，分别是全传递和选择传递：

• 全传递：对于 Context 的 key 和 value，不作区分，直接都传递了，这种比较少见；
• 选择传递：有选择地确定某些 key 和 value 需要传递，这个比较常见；

对于选择传递，实现的方式也五花八门，例如我曾经实现过的一个 RPC 框架，是通过 Getter 和 Setter 取特定的 key 进行传递，这样做的好处就是保证了传递的键值对的确定性，并对用户进行了限制，可以避免用户乱来，但是确定也比较明显，当需要支持一个新的键值对的时候，我们需要发一个新的版本。
在一个开源的 RPC 框架 kitex 中使用的是通过一个特定的 key 来保存要传递的键值对，然后直接将这些值全都传递下去，这样的做的好处也很明显，那就是无需关注用户需要传递什么值，只管无脑传递即可，但是缺点就是无法管控用户传递的键值对，对于服务治理平台来说有一定的麻烦（比如 key 冲突或者长度限制等）。

传输协议

当我们确定了要传递哪些键值对之后，如何传递这些键值对也是一个需要考虑的问题，这个实现也是因人而异，或者说是因协议而异，比如，如果你的协议是基于 HTTP 的，那么很自然，你会想到通过 HTTP Header 的形式来传递这些数据，虽然这很符合标准的协议，但是，有时它可能会带来一些问题，比如说支持的类型必须是 ASCII，长度限制等；如果你是自己实现的 RPC，那么也会有不同的选择，比如在字节的 RPC 演进中，也出现过将这些元数据放在请求和响应体里面的阶段，显然，这在服务治理上是非常麻烦的，比如 Mesh 场景下，你可能需要某些 Key 的值，那么你就需要去解请求体，这就很 tricky 了。

通过我个人的观察，比较通俗的做法就是将这些元数据放入一个 RPC 的元数据结构体中，例如 Kitex 中使用了一种叫做 TTHeader 的结构来承载。但是呢，即使是有这么个结构，不同的公司在实现的时候也会有不同的选择，例如，有的团队会选择一个 key 对应于元数据中的一个字段，也有的团队呢会选择奖所有的键值对封装到一个元数据的字段中；同时，还有的团队会选择某些 key 对应于元数据中的一个字段，剩余的键值对封装到一个元数据的字段中（我曾经实现过的一个 RPC 框架使用的就是这种）。

解析键值对

当我们收到一个跨进程的消息时，我们就会解析这个消息，并且将传递的键值对解析出来，这里的处理其实和第一步和第二步是紧密关联的，所以，基本上你第一步和第二步的选择就直接决定了第三步的处理方式，没有太大的变动空间，这里就不多做讨论了。

一种简单的实现

前面介绍了这么多，这里我就以 Context Operator 的方式为例，并且假设我是在两个 HTTP 服务之间传递 Context 为例，介绍一种简单的实现。首先我的 Context Operator 是这样的：

[root@liqiang.io]# cat example2.go
type ContextOperator interface {
    WithHello(ctx context.Context, value interface{})
    WithWorld(ctx context.Context, value interface{})
    GetHello(ctx context.Context)
    GetWorld(ctx context.Context)
    Marshal() ([]byte, error)
    Unmarshal([]byte) error
}

发送请求

当 Client 服务准备发送请求时，这时需要有一个 Middleware 能够进行 Context 的 Marshal，那么代码可能是这样的：

[root@liqiang.io]# cat example2.go
const contextHeader = "liqiang-io-context"
type Transport struct {
    rt http.RoundTripper
    co ContextOperator
}
func (t *Transport) RoundTrip(r *http.Request) (*http.Response, error) {
    ctx, err := t.co.Marshal(r.Context())
    if err == nil {
        r.Header.Add(contextHeader, string(ctx))
    }
    return http.DefaultTransport.RoundTrip(r)
}

这样就将 Context 的值放入到了 Http 的请求中，并且发送出去了，那么这个 Marshal实现可以简单地实现，就是一个字符串拼接：

[root@liqiang.io]# cat example2.go
func (o *ctxOp) Marshal(ctx context.Context) ([]byte, error) {
    var rtns []string
    if hello := o.GetHello(ctx); hello != nil {
        rtns = append(rtns, "hello="+*hello)
    }
    if world := o.GetWorld(ctx); world != nil {
        rtns = append(rtns, "world="+*world)
    }
    return []byte(strings.Join(rtns, string([]rune{0}))), nil
}

接受请求

然后在 HTTP 服务端这边接收到 HTTP 请求之后，就是相反的操作，对 Http Handler 进行了一层封装：

[root@liqiang.io]# cat example2.go
type wrapperHandler struct {
    h  http.Handler
    co ContextOperator
}
func (h *wrapperHandler) ServeHTTP(resp http.ResponseWriter, req *http.Request) {
    ctxVal := req.Header.Get(contextHeader)
    ctx, err := h.co.Unmarshal([]byte(ctxVal))
    if err == nil {
        req = req.WithContext(ctx)
    }
    h.h.ServeHTTP(resp, req)
}
func NewWrapperHandler(h http.Handler) http.Handler {
    return &wrapperHandler{
        h: h,
    }
}

同样的，Unmarshal 操作其实就是一个字符串的分解工作：

[root@liqiang.io]# cat example2.go
func (o *ctxOp) Unmarshal(content []byte) (context.Context, error) {
    vals := strings.Split(string(content), string([]rune{0}))
    ctx := context.Background()
    for _, val := range vals {
        switch val {
        case "hello":
            ctx = o.WithHello(ctx, val)
        case "world":
            ctx = o.WithWorld(ctx, val)
        }
    }
    return ctx, nil
}

cancel 和 timeout 传递

看上去上面的例子已经实现了 Context 的传递功能，但是实际上，这还不是 Go 中 Context 的全部，甚至于说 Go 中最常见的使用都无法满足，因为大多数情况下，大家都是使用的 Context 的 WithTimeout 和 WithCancel 功能。

cancel

WithCancel 功能可能看上去好理解，取消一个跨进程的请求就可以了，但是当你实际上操作的时候，会发现问题比你想象的更复杂一些，因为 Go 的默认网络库是不支持 Context 的，我们来看一下 net.Conn 的接口：

[root@liqiang.io]# cat net/net.go
type Conn interface {
    Read(b []byte) (n int, err error)
    Write(b []byte) (n int, err error)
    Close() error
}

可以发现，你没法通过 context 来控制 read 或者 write 的行为，因此，当你想传递一个 Cancel 给另一个进程时，从默认的网络接口来看是做不到的，所以，就有了一些骚操作，下面我举两个例子看看下一些常用的模式：

• goroutine 模式

因为 Go 中的 goroutine 很廉价，所以就有了这种操作：

[root@liqiang.io]# cat example3.go
import "context"
func passCancel1(ctx context.Context) {
    var networkDone chan struct{}
    go func() {
        // network operation 1
        networkDone <- struct{}{}
    }()
    select {
    case <-ctx.Done():
        // cancel
    case <-networkDone:
        // normal done
    }
    return
}

这种方式可能存在 goroutine 泄露，所以 goroutine 的处理需要很小心

• busy wait

另外一种就是给网络操作设置一个很短的超时时间，如果在 cancel 之前网络操作完成了，那就是正常，如果网络操作完成之前 cancel 掉，那么就是异常的，代码大概这么写：

[root@liqiang.io]# cat example4.go
func networkProcess(ctx context.Context) {
    for {
        conn.SetReadDeadline(time.Now().Add(time.Microsecond * 20))
        readBytes, err := conn.Read(buffer)
        if err != nil {
            return
        }
        if readBytes > 0 {
            return
        }
        select {
        case <-ctx.Done():
        default:
            continue
        }
    }
}

这种的问题其实也很明显，会出现大量 CPU 资源的浪费，其实从这个模式就可以看出这是因为网络操作不知道 Cancel Context 的缘故，所以在实践中，比较少使用跨进程的 Cancel Context，使用 timeout 的情况会比较常见。

timeout

Timeout Context 的传递本质上和 context 的值传递是类似的，但是 timeout 的传递又不太一样的地方在于从进程 A 传递到进程 B，进程 B 需要能够 timeout，这才是 timeout 传递的麻烦所在。在前面进程内 Context 的实现中，我们可以看到，timeoutCtx 本质上还是一个 Deadline Context，然后再看看 Deadline Context 又是如何实现的：

[root@liqiang.io]# cat context/context.go
func WithDeadline(parent Context, d time.Time) (Context, CancelFunc) {
    ... ...
    dur := time.Until(d)
    ... ...
    if c.err == nil {
        c.timer = time.AfterFunc(dur, func() {
            c.cancel(true, DeadlineExceeded)
        })
    }
    return c, func() { c.cancel(true, Canceled) }
}

可以看到，Deadline 是通过时间差来启动一个 timer，然后到时自动取消，那么看上去确实很简单，我们将这个 deadline 的值传给对端不就好了？

例如这样：

[root@liqiang.io]# cat exmaple2.go
func (o *ctxOp) Marshal(ctx context.Context) ([]byte, error) {
    var rtns []string
    if hello := o.GetHello(ctx); hello != nil {
        rtns = append(rtns, "hello="+*hello)
    }
    if world := o.GetWorld(ctx); world != nil {
        rtns = append(rtns, "world="+*world)
    }
    if dl, ok := ctx.Deadline(); ok {
        rtns = append(rtns, "deadline="+strconv.FormatInt(dl.UnixMilli(), 10))
    }
    return []byte(strings.Join(rtns, string([]rune{0}))), nil
}
func (o *ctxOp) Unmarshal(content []byte) (context.Context, context.CancelFunc, error) {
    vals := strings.Split(string(content), string([]rune{0}))
    ctx := context.Background()
    var cancel context.CancelFunc
    for _, val := range vals {
        switch val {
        case "hello":
            ctx = o.WithHello(ctx, val)
        case "world":
            ctx = o.WithWorld(ctx, val)
        case "deadline":
            i, err := strconv.ParseInt(val, 10, 64)
            if err == nil {
                t := time.Unix(i/1e9, i%1e9)
                ctx, cancel = context.WithDeadline(ctx, t)
            }
        }
    }
    return ctx, cancel, nil
}

这种实现看上去没啥问题，但是当放在分布式场景下时，就可能有问题了，因为这里用的是 UNIX Timestamp，这就要求两台机器之间的时间是完全同步的，如果不同步，就可能出现过早或者过晚结束的问题，有没有解决办法？办法肯定是有的，但是好坏是另说，在实践的时候，我们通常采用传递相对时间的方式，比如进程 A 设置的初始 Timeout 是 5000 ms，同时进程 A 自己处理已经使用了 500 ms，那么传递给进程 B 的超时时间就是 4500 ms，但是，因为考虑到网络延迟等因素，所以实际上进程 A 在进行响应等待的时候会设置长一点点的时间，例如 4550 ms（+ 50ms）来等待响应。

Ref

• context package
• Go Concurrency Patterns: Context