2023-08-18 18:09

关注

context包介绍 | 青训营

context 包- 包核心方法

context 包的核心 API 有四个

context.WithValue: 设置键值对，并且返回一个新的 context 实例
context.WithCance
context.WithDeadline
cntext.WithTimeout: 三者都返回一个可取消的 context 实例，和取消函数

注意: context 实例是不可变的，每一次都是新创建的。

ctx := context.Background()一般是链路起点，或者调用起点 ctx := context.TODO()在你不确定 context 该用什么，用TODO()

func TestContext(t *testing.T) {
	// 一般是链路起点，或者调用起点
	ctx := context.Background()

	// 在你不确定 context 该用什么， 用TODO()
	//ctx := context.TODO()

	ctx = context.WithValue(ctx, mykey{}, "my-value")
}

func TestContext_WithCancel(t *testing.T) {
	ctx := context.Background()
	ctx, cancel := context.WithCancel(ctx)

	go func() {
		time.Sleep(time.Second)
		cancel()
	}()
	//defer cancel()

	<-ctx.Done()
	t.Log("hello cancel: ", ctx.Err())
}

func TestContext_WithDeadline(t *testing.T) {
	ctx := context.Background()
	ctx, cancel := context.WithDeadline(ctx, time.Now().Add(time.Second*3))
	deadline, _ := ctx.Deadline()
	t.Log("deadline: ", deadline)
	defer cancel()

	<-ctx.Done()
	t.Log("hello deadline: ", ctx.Err())
}

func TestContext_WithTimeout(t *testing.T) {
	ctx := context.Background()
	ctx, cancel := context.WithTimeout(ctx, time.Second*3)
	deadline, _ := ctx.Deadline()
	t.Log("deadline: ", deadline)
	defer cancel()

	<-ctx.Done()
	t.Log("hello timeout: ", ctx.Err())
}

context 包- Context 接口

Context 接口核心 API 有四个

Deadline: 返回过期时间，如果 ok 为 false，说明没有设置过期时间。略常用
Done: 返回一个 channel，一般用于监听 Context 实例的信号，比如说过期，或者正常关闭。常用
Err: 返回一个错误用于表达 Context 发生了什么。Canceled =>正常关闭，DeadlineExceeded => 过期超时。比较常用
Value: 取值。非常常用

context 包—安全传递数据

context 包我们就用来做两件事

安全传递数据
控制链路

安全传递数据，是指在请求执行上下文中线程安全地传递数据，依赖于 WithValue 方法。因为 Go 本身没有 thread-local 机制，所以大部分类似的功能都是借助于 context 来实现的。例子:

链路追踪的 trace id
AB测试的标记位
压力测试标记位
分库分表中间件中传递 sharding hint
ORM 中间件传递 SQL hint
Web 框架传递上下文

context 包-父子关系

context 的实例之间存在父子关系

当父亲取消或者超时，所有派生的子context 都被取消或者超时
当找 key 的时候，子 context 先看自己有没有，没有则去祖先果面找

控制是从上至下的，查找是从下至上的。

func TestContext_Parent(t *testing.T) {
	ctx := context.Background()
	parent := context.WithValue(ctx, "my-key", "parent my-value")
	child := context.WithValue(parent, "my-key", "child my-value")
	child2, cancel := context.WithTimeout(parent, time.Second)
	child3 := context.WithValue(parent, "new-key", "child my-value")
	defer cancel()

	t.Log("parent my-key", parent.Value("my-key"))
	t.Log("child my-key", child.Value("my-key"))
	t.Log("child2 my-key", child2.Value("my-key"))

	t.Log("child3 new-key", child3.Value("new-key"))
	t.Log("parent new-key", parent.Value("new-key"))
}

因为父 context 始终无法拿到子 context 设置的值，所以在逼不得已的时候我们可以在父 context 里面放一个 map，后续都是修改这个 map。

func TestContext_Parent(t *testing.T) {
	parent2 := context.WithValue(ctx, "map", map[string]interface{}{})
	child4, cancel := context.WithTimeout(parent2, time.Second)
	defer cancel()
	m := child4.Value("map").(map[string]interface{})
	m["key1"] = 123
	t.Log("parent2 map: ", parent2.Value("map"))
}

如何使用 context

context 使用起来非常方便。源码里对外提供了一个创建根节点 context 的函数：

func Background() Context

background 是一个空的 context，它不能被取消，没有值，也没有超时时间。

有了根节点 context，又提供了四个函数创建子节点 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 会在函数传递间传递。只需要在适当的时间调用 cancel 函数向 goroutines 发出取消信号或者调用 Value 函数取出 context 中的值。

在官方博客里，对于使用 context 提出了几点建议：

Do not store Contexts inside a struct type; instead, pass a Context explicitly to each function that needs it. The Context should be the first parameter, typically named ctx.
Do not pass a nil Context, even if a function permits it. Pass context.TODO if you are unsure about which Context to use.
Use context Values only for request-scoped data that transits processes and APIs, not for passing optional parameters to functions.
The same Context may be passed to functions running in different goroutines; Contexts are safe for simultaneous use by multiple goroutines.

我翻译一下：

不要将 Context 塞到结构体里。直接将 Context 类型作为函数的第一参数，而且一般都命名为 ctx。
不要向函数传入一个 nil 的 context，如果你实在不知道传什么，标准库给你准备好了一个 context：todo。
不要把本应该作为函数参数的类型塞到 context 中，context 存储的应该是一些共同的数据。例如：登陆的 session、cookie 等。
同一个 context 可能会被传递到多个 goroutine，别担心，context 是并发安全的。

传递共享的数据

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

package main

import (
    "context"
    "fmt"
)

func main() {
    ctx := context.Background()
    process(ctx)

    ctx = context.WithValue(ctx, "traceId", "qcrao-2019")
    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-2019

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

取消 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。

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