传递共享的数据

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

package main

import (
    "context"
    "fmt"
)

func main() {
    ctx := context.Background()
    process(ctx)
    ctx = context.WithValue(ctx, "traceId", "qcrao-2020")
    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=qcrao-2020

第一次调用 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)
}

定时取消

某个应用需要获取网络上的数据，这需要花费一定的时间，若碰到网络延迟、机器负载过高等问题时，会导致获得的数据时间过长，导致请求阻塞，严重地会引起雪崩。这时候就可以用 context 的定时取消功能：

package main

import (
    "context"
    "fmt"
    "time"
)

func main() {
    ctx, cancel := context.WithTimeout(context.Background(), 1 * time.Second)
    defer cancel() // 避免其他地方忘记 cancel，且重复调用不影响
    ids := fetchWebData(ctx)
    fmt.Println(ids)
}
func fetchWebData(ctx context.Context) (res []int64) {
    select {
    case <- time.After(3 * time.Second):
        return []int64{100, 200, 300}
    case <- ctx.Done():
        return []int64{1, 2, 3}
    }
}

在 main 函数里，首先创建了一个定时 1s 的 context，到时间后会自动调用 cancel 函数，接着调用 fetchWebData 函数获取网络数据，最后打印返回的 ids。

在 fetchWebData 函数里，则通过设置 3s 的定时器，表示处理的时长，正常会返回 [100 200 300]，若 context 被取消，则返回默认值 [1 2 3]。

运行程序，结果如下：

[1 2 3]

若将 main 函数里的 context 超时时间改成 5s，则最终打印：

[100 200 300]

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

防止 goroutine 泄漏

上一节的例子里，如果不加 context，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 函数在接收到取消信号后，直接退出，系统回收资源。