Go Context 详解

May 12, 2022

协程
Go

Go Context 详解

context用于在goroutine之间传递上下文信息

取消信号
超时信号
截止信号
K-V 数据

相同的context可以传递给多个goroutine，context是线程安全的。

在 Go 的 Server 中，通常每来一个请求，都会启动若干个goroutine同时工作，有些去数据库拿数据，有些调用下游服务获取相关数据……

而这些goroutine需要共享这些数据，比如登录的 token，处理请求的最大时间等等，当请求被取消或者被处理时，goroutine需要快速退出，因为他们的工作成果不再被需要了，在相关联的goroutine都退出之后，系统就可以回收相关的资源了。

而在 Go 中，我们不能去杀死协程，协程的关闭一般会使用channel+select方式来控制。

但是在某一些场景下，处理一个请求衍生了许多个协程，这些协程是互相关联的，需要共享一些全局变量，这个时候去使用channel+select就会比较麻烦，这时就可以通过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 塞到结构体里。直接将 Context 类型作为函数的第一参数，而且一般都命名为ctx。
不要向函数传入一个 nil 的 context，如果你实在不知道传什么，标准库给你准备好了一个 context：todo。
不要把本应该作为函数参数的类型塞到 context 中，context 存储的应该是一些共同的数据。例如：登陆的 session、cookie 等。
同一个 context 可能会被传递到多个 goroutine，别担心，context 是并发安全的。

取消 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()

注意一个细节，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
        }
    }
    // ……
}

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

context.Value 的查找过程是怎样的

context value.png

和链表有点像，只是它的方向相反：Context 指向它的父节点，链表则指向下一个节点。通过 WithValue 函数，可以创建层层的 valueCtx，存储 goroutine 间可以共享的变量。

查找的时候，会向上查找到最后一个挂载的 context 节点，也就是离得比较近的一个父节点 context