【Go进阶—并发编程】Context

Context 是 Go 应用开发常用的并发控制技术，它与 WaitGroup 最大的不同点是 Context 对于派生 goroutine 有更强的控制力，它可以控制多级的 goroutine。

尽管有很多的争议，但是在很多场景下使用 Context 都很方便，所以现在它已经在 Go 生态圈中传播开来了，包括很多的 Web 应用框架，都切换成了标准库的 Context。标准库中的 database/sql、os/exec、net、net/http 等包中都使用到了 Context。而且，如果遇到了下面的一些场景，也可以考虑使用 Context：

• 上下文信息传递 ，比如处理 http 请求、在请求处理链路上传递信息；
• 控制子 goroutine 的运行；
• 超时控制的方法调用；
• 可以取消的方法调用。

实现原理

接口定义

包 context 定义了 Context 接口，Context 的具体实现包括 4 个方法，分别是 Deadline、Done、Err 和 Value，如下所示：

type Context interface {
    Deadline() (deadline time.Time, ok bool)
    Done() <-chan struct{}
    Err() error
    Value(key interface{}) interface{}
}

Deadline 方法会返回这个 Context 被取消的截止日期。如果没有设置截止日期，ok 的值是 false。后续每次调用这个对象的 Deadline 方法时，都会返回和第一次调用相同的结果。

Done 方法返回一个 Channel 对象，基本上都会在 select 语句中使用。在 Context 被取消时，此 Channel 会被 close，如果没被取消，可能会返回 nil。当 Done 被 close 的时候，可以通过 ctx.Err 获取错误信息。

关于 Err 方法，你必须要记住的知识点就是：如果 Done 没有被 close，Err 方法返回 nil；如果 Done 被 close，Err 方法会返回 Done 被 close 的原因。

Value 返回此 ctx 中和指定的 key 相关联的 value。

Context 中实现了 2 个常用的生成顶层 Context 的方法：

• context.Background()：返回一个非 nil 的、空的 Context，没有任何值，不会被 cancel，不会超时，没有截止日期。一般用在主函数、初始化、测试以及创建根 Context 的时候。
• context.TODO()：返回一个非 nil 的、空的 Context，没有任何值，不会被 cancel，不会超时，没有截止日期。当你不清楚是否该用 Context，或者目前还不知道要传递一些什么上下文信息的时候，就可以使用这个方法。

事实上，它们两个底层的实现是一模一样的。绝大多数情况下可以直接使用 context.Background。

在使用 Context 的时候，有一些约定俗成的规则：

1. 一般函数使用 Context 的时候，会把这个参数放在第一个参数的位置。
2. 从来不把 nil 当做 Context 类型的参数值，可以使用 context.Background() 创建一个空的上下文对象，也不要使用 nil。
3. Context 只用来临时做函数之间的上下文透传，不能持久化 Context 或者把 Context 长久保存。把 Context 持久化到数据库、本地文件或者全局变量、缓存中都是错误的用法。
4. key 的类型不推荐字符串类型或者其它内建类型，否则容易在包之间使用 Context 时候产生冲突。使用 WithValue 时，key 的类型应该是自己定义的类型。
5. 常常使用 struct{} 作为底层类型定义 key 的类型。对于 exported key 的静态类型，常常是接口或者指针。这样可以尽量减少内存分配。

context 包中实现 Context 接口的 struct，除了用于 context.Background() 的 emptyCtx 外，还有 cancelCtx、timerCtx 和 valueCtx 三种。

cancelCtx

type cancelCtx struct {
    Context

    mu       sync.Mutex            // 互斥锁
    done     atomic.Value          // 调用 cancel 时会关闭的 channel
    children map[canceler]struct{} // 记录了由此 Context 派生的所有 child，此 Context 被 cancel 时会同时 cancel 所有 child
    err      error                 // 错误信息
}

WithCancel

cancelCtx 是通过 WithCancel 方法生成的。我们常常在一些需要主动取消长时间的任务时，创建这种类型的 Context，然后把这个 Context 传给长时间执行任务的 goroutine。当需要中止任务时，我们就可以 cancel 这个 Context，这样长时间执行任务的 goroutine，就可以通过检查这个 Context，知道 Context 已经被取消了。

WithCancel 返回值中的第二个值是一个 cancel 函数。记住，不是只有你想中途放弃，才去调用 cancel，只要你的任务正常完成了，就需要调用 cancel，这样，这个 Context 才能释放它的资源（通知它的 children 处理 cancel，从它的 parent 中把自己移除，甚至释放相关的 goroutine）。

看下核心源码：

func WithCancel(parent Context) (ctx Context, cancel CancelFunc) {
    if parent == nil {
        panic("cannot create context from nil parent")
    }
    c := newCancelCtx(parent)
    propagateCancel(parent, &c)
    return &c, func() { c.cancel(true, Canceled) }
}

func newCancelCtx(parent Context) cancelCtx {
    return cancelCtx{Context: parent}
}

func propagateCancel(parent Context, child canceler) {
    done := parent.Done()

    if p, ok := parentCancelCtx(parent); ok {
        p.mu.Lock()
        if p.err != nil {
            // parent 已经取消了，直接取消子 Context
            child.cancel(false, p.err)
        } else {
            // 将 child 添加到 parent 的 children 切片
            if p.children == nil {
                p.children = make(map[canceler]struct{})
            }
            p.children[child] = struct{}{}
        }
        p.mu.Unlock()
    } else {
        atomic.AddInt32(&goroutines, +1)
        // 没有 parent 可以“挂载”，启动一个 goroutine 监听 parent 的 cancel，同时 cancel 自身
        go func() {
            select {
            case <-parent.Done():
                child.cancel(false, parent.Err())
            case <-child.Done():
            }
        }()
    }
}

代码中调用的 propagateCancel 方法会顺着 parent 路径往上找，直到找到一个 cancelCtx，或者为 nil。如果不为空，就把自己加入到这个 cancelCtx 的 child，以便这个 cancelCtx 被取消的时候通知自己。如果为空，会新起一个 goroutine，由它来监听 parent 的 Done 是否已关闭。

当这个 cancelCtx 的 cancel 函数被调用的时候，或者 parent 的 Done 被 close 的时候，这个 cancelCtx 的 Done 才会被 close。

cancel 是向下传递的，如果一个 WithCancel 生成的 Context 被 cancel 时，如果它的子 Context（也有可能是孙，或者更低，依赖子的类型）也是 cancelCtx 类型的，就会被 cancel。

cancel

cancel 方法的作用是 close 自己及其后代的 done 通道，达到通知取消的目的。WithCancel 方法的第二个返回值 cancel 就是本函数。来看一下主要代码实现：

func (c *cancelCtx) cancel(removeFromParent bool, err error) {
    c.mu.Lock()
    // 设置 cancel 的原因
    c.err = err 
    // 关闭自身的 done 通道
    d, _ := c.done.Load().(chan struct{})
    if d == nil {
        c.done.Store(closedchan)
    } else {
        close(d)
    }
    // 遍历所有 children，逐个调用 cancel 方法
    for child := range c.children {
        child.cancel(false, err)
    }
    c.children = nil
    c.mu.Unlock()

    // 正常情况下，需要将自身从 parent 的 children 切片中删除
    if removeFromParent {
        removeChild(c.Context, c)
    }
}

timerCtx

type timerCtx struct {
    cancelCtx
    timer *time.Timer 

    deadline time.Time
}

timerCtx 在 cancelCtx 基础上增加了 deadline 用于标示自动 cancel 的最终时间，而 timer 就是一个触发自动 cancel 的定时器。timerCtx 可以由 WithDeadline 和 WithTimeout 生成， WithTimeout 实际调用了 WithDeadline，二者实现原理一致，只不过使用语境不一样：WithDeadline 是指定最后期限，WithTimeout 是指定最长存活时间。

WithDeadline

来看一下 WithDeadline 方法的实现：

func WithDeadline(parent Context, d time.Time) (Context, CancelFunc) {
    // 如果parent的截止时间更早，直接返回一个cancelCtx即可
    if cur, ok := parent.Deadline(); ok && cur.Before(d) {
        return WithCancel(parent)
    }
    c := &timerCtx{
        cancelCtx: newCancelCtx(parent),
        deadline:  d,
    }
    propagateCancel(parent, c) // 同cancelCtx的处理逻辑
    dur := time.Until(d)
    if dur <= 0 { //当前时间已经超过了截止时间，直接cancel
        c.cancel(true, DeadlineExceeded)
        return c, func() { c.cancel(false, Canceled) }
    }
    c.mu.Lock()
    defer c.mu.Unlock()
    if c.err == nil {
        // 设置一个定时器，到截止时间后取消
        c.timer = time.AfterFunc(dur, func() {
            c.cancel(true, DeadlineExceeded)
        })
    }
    return c, func() { c.cancel(true, Canceled) }
}

WithDeadline 会返回一个 parent 的副本，并且设置了一个不晚于参数 d 的截止时间，类型为 timerCtx（或者是 cancelCtx）。

如果它的截止时间晚于 parent 的截止时间，那么就以 parent 的截止时间为准，并返回一个类型为 cancelCtx 的 Context，因为 parent 的截止时间到了，就会取消这个 cancelCtx。如果当前时间已经超过了截止时间，就直接返回一个已经被 cancel 的 timerCtx。否则就会启动一个定时器，到截止时间取消这个 timerCtx。

综合起来，timerCtx 的 Done 被 Close 掉，主要是由下面的某个事件触发的：

• 截止时间到了；
• cancel 函数被调用；
• parent 的 Done 被 close。

和 cancelCtx 一样，WithDeadline（WithTimeout）返回的 cancel 一定要调用，并且要尽可能早地被调用，这样才能尽早释放资源，不要单纯地依赖截止时间被动取消。

valueCtx

type valueCtx struct {
    Context
    key, val interface{}
}

valueCtx 只是在 Context 基础上增加了一个 key-value 对，用于在各级协程间传递一些数据。

WithValue 基于 parent Context 生成一个新的 valueCtx，保存了一个 key-value 键值对。valueCtx 覆盖了 Value 方法，优先从自己的存储中检查这个 key，不存在的话会从 parent 中继续检查。