Golang Context 使用场景示例

context 用来在 goroutine 之间传递上下文信息,使用场景包括:取消 goroutine、传递共享的数据、防止 goroutine 泄漏等。

1. 取消 goroutine

我们先来设想一个场景:打开外卖的订单页,地图上显示外卖小哥的位置,而且是每秒更新 1 次。app 端向后台发起 websocket 连接(现实中可能是轮询)请求后,后台启动一个协程,每隔 1 秒计算 1 次小哥的位置,并发送给端。如果用户退出此页面,则后台需要“取消”此过程,退出 goroutine,系统回收资源。

后端可能的实现如下:

func Perform() {
    for {
        calculatePos()
        sendResult()
        time.Sleep(time.Second)
    }
}

如果需要实现“取消”功能,并且在不了解 context 功能的前提下,可能会这样做:给函数增加一个指针型的 bool 变量,在 for 语句的开始处判断 bool 变量是发由 true 变为 false,如果改变,则退出循环。

上面给出的简单做法,可以实现想要的效果,没有问题,但是并不优雅,并且一旦协程数量多了之后,并且各种嵌套,就会很麻烦。优雅的做法,自然就要用到 context。

    func Perform(ctx context.Context) {
        for {
            calculatePos()
            sendResult()
            select {
            case <-ctx.Done():
                // 被取消,直接返回
                return
            case <-time.After(time.Second):
                // block 1 秒钟 
            }
        }
    }

主流程可能是这样的:

ctx, cancel := context.WithTimeout(context.Background(), time.Hour)
go Perform(ctx)
// ……
// app 端返回页面,调用cancel 函数
cancel()

注意一个细节,WithTimeOut 函数返回的 context 和 cancelFun 是分开的。context 本身并没有取消函数,这样做的原因是取消函数只能由外层函数调用,防止子节点 context 调用取消函数,从而严格控制信息的流向:由父节点 context 流向子节点 context。

2. 传递共享的数据

对于 Web 服务端开发,往往希望将一个请求处理的整个过程串起来,这就非常依赖于 Thread Local(对于 Go 可理解为单个协程所独有) 的变量,而在 Go 语言中并没有这个概念,因此需要在函数调用的时候传递 context。

package main
import (
    "context"
    "fmt"
)
func main() {
    ctx := context.Background()
    process(ctx)
    ctx = context.WithValue(ctx, "traceId", "aizws-2021")
    process(ctx)
}

func process(ctx context.Context) {
    traceId, ok := ctx.Value("traceId").(string)
    if ok {
        fmt.Printf("process over. trace_id=%s\n", traceId)
    } else {
        fmt.Printf("process over. no trace_id\n")
    }
}

运行结果:

process over. no trace_id
process over. trace_id=aizws-2021

第一次调用 process 函数时,ctx 是一个空的 context,自然取不出来 traceId。第二次,通过 WithValue 函数创建了一个 context,并赋上了 traceId 这个 key,自然就能取出来传入的 value 值。

当然,现实场景中可能是从一个 HTTP 请求中获取到的 Request-ID。所以,下面这个样例可能更适合:

const requestIDKey int = 0
func WithRequestID(next http.Handler) http.Handler {
    return http.HandlerFunc(
        func(rw http.ResponseWriter, req *http.Request) {
            // 从 header 中提取 request-id
            reqID := req.Header.Get("X-Request-ID")
            // 创建 valueCtx。使用自定义的类型,不容易冲突
            ctx := context.WithValue(
                req.Context(), requestIDKey, reqID)
            // 创建新的请求
            req = req.WithContext(ctx)
            // 调用 HTTP 处理函数
            next.ServeHTTP(rw, req)
        }
    )
}
// 获取 request-id
func GetRequestID(ctx context.Context) string {
    ctx.Value(requestIDKey).(string)
}
func Handle(rw http.ResponseWriter, req *http.Request) {
    // 拿到 reqId,后面可以记录日志等等
    reqID := GetRequestID(req.Context())
    ...
}
func main() {
    handler := WithRequestID(http.HandlerFunc(Handle))
    http.ListenAndServe("/", handler)
}

3. 防止 goroutine 泄漏

前面那个例子里,goroutine 还是会自己执行完,最后返回,只不过会多浪费一些系统资源。这里改编一个“如果不用 context 取消,goroutine 就会泄漏的例子”。

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
        }
    }
    // ……
}

增加一个 context,在 break 前调用 cancel 函数,取消 goroutine。gen 函数在接收到取消信号后,直接退出,系统回收资源。

下一章:Golang Context 源码剖析

context 包的代码并不长,大约有 200 行,是一个非常精简的代码库。context 包中包含了一些结构体、接口以及函数的定义,下表中汇总了它们的名称、类型和用途。 1. contex ...