Go 语言 Context 的使用

昨天已经记录了如何用 context 控制并发。那么，context 还有哪些使用场景呢？首先我们先来看一下 context 结构体：

Context 结构体

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

定义了四个需要实现的方法，分别是：

• Deadline — 返回 context.Context 被取消的时间，也就是完成工作的截止日期；
• Done — 返回一个 Channel，这个 Channel 会在当前工作完成或者上下文被取消后关闭，多次调用 Done 方法会返回同一个 Channel；
• Err — 返回 context.Context 结束的原因，它只会在 Done 方法对应的 Channel 关闭时返回非空的值；
1. 如果 context.Context 被取消，会返回 Canceled 错误；
2. 如果 context.Context 超时，会返回 DeadlineExceeded 错误；
• Value — 从 context.Context 中获取键对应的值，对于同一个上下文来说，多次调用 Value 并传入相同的 Key 会返回相同的结果，该方法可以用来传递请求特定的数据；

默认 context

context 包中最常用的方法是 context.Background、context.TODO，这两个方法都会返回预先初始化好的私有变量 background 和 todo，它们会在同一个 Go 程序中被复用：

func Background() Context {
    return background
}

func TODO() Context {
    return todo
}

这两个私有变量都是通过 new(emptyCtx) 语句初始化的，它们是指向私有结构体 context.emptyCtx 的指针，这是最简单、最常用的 context,context.emptyCtx 实现了 context 中所有的方法，但没有任何的功能。

从源代码来看，context.Background 和 context.TODO 只是互为别名，没有太大的差别，只是在使用和语义上稍有不同：

context.Background 是 context 的默认值，所有其他的 context 都是从它衍生出来的；
context.TODO 仅在不确定应该使用哪种 context 时使用；
在多数情况下，如果当前函数没有 context 作为入参，我们都会使用 context.Background 作为起始的 context 向下传递。

取消信号

context.WithCancel 函数能够从 context.Context 中衍生出一个新的子 context 并返回用于取消该 context 的函数。一旦我们执行返回的取消函数，当前 context 以及它的子 context 都会被取消，所有的 Goroutine 都会同步收到这一取消信号。

context.WithCancel 函数的实现原理：

func WithCancel(parent Context) (ctx Context, cancel CancelFunc) {
    c := newCancelCtx(parent)
    propagateCancel(parent, &c)
    return &c, func() { c.cancel(true, Canceled) }
}

context.newCancelCtx 将传入的 context 包装成私有结构体 context.cancelCtx；
context.propagateCancel 会构建父子 context 之间的关联，当父 context 被取消时，子 context 也会被取消：context.propagateCancel 的作用是在 parent 和 child 之间同步取消和结束的信号，保证在 parent 被取消时，child 也会收到对应的信号，不会出现状态不一致的情况。
context 包中的另外两个函数 context.WithDeadline 和 context.WithTimeout 也都能创建可以被取消的计时器 context context.timerCtx：

func WithTimeout(parent Context, timeout time.Duration) (Context, CancelFunc) {
    return WithDeadline(parent, time.Now().Add(timeout))
}

func WithDeadline(parent Context, d time.Time) (Context, CancelFunc) {
    if cur, ok := parent.Deadline(); ok && cur.Before(d) { // **比较父 context 的截止日期与当前日期**
        return WithCancel(parent)
    }
    c := &timerCtx{
        cancelCtx: newCancelCtx(parent),
        deadline:  d,
    }
    propagateCancel(parent, c)
    dur := time.Until(d)
    if dur <= 0 {
        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 函数在创建 context.timerCtx 的过程中判断了父 context 的截止日期与当前日期，并通过 time.AfterFunc 创建定时器，当时间超过了截止日期后会调用 context.timerCtx.cancel 同步取消信号。

context.timerCtx 内部不仅通过嵌入 context.cancelCtx 结构体继承了相关的变量和方法，还通过持有的定时器 timer 和截止时间 deadline 实现了定时取消的功能：

type timerCtx struct {
    cancelCtx
    timer *time.Timer // Under cancelCtx.mu.

    deadline time.Time
}

func (c *timerCtx) Deadline() (deadline time.Time, ok bool) {
    return c.deadline, true
}

func (c *timerCtx) cancel(removeFromParent bool, err error) {
    c.cancelCtx.cancel(false, err)
    if removeFromParent {
        removeChild(c.cancelCtx.Context, c) // **停止持有的定时器减少不必要的资源浪费。**
    }
    c.mu.Lock()
    if c.timer != nil {
        c.timer.Stop()
        c.timer = nil
    }
    c.mu.Unlock()
}

传值

context 包中的 context.WithValue 能从父 context 中创建一个子 context，传值的子 context 使用 context.valueCtx 类型：

func WithValue(parent Context, key, val interface{}) Context {
    if key == nil {
        panic("nil key")
    }
    if !reflectlite.TypeOf(key).Comparable() {
        panic("key is not comparable")
    }
    return &valueCtx{parent, key, val}
}

context.valueCtx 结构体会将除了 Value 之外的 Err、Deadline 等方法代理到父 context 中，它只会响应 context.valueCtx.Value 方法，该方法的实现也很简单：

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

func (c *valueCtx) Value(key interface{}) interface{} {
    if c.key == key {
        return c.val
    }
    return c.Context.Value(key)
}

如果 context.valueCtx 中存储的键值对与 context.valueCtx.Value 方法中传入的参数不匹配，就会从父 context 中查找该键对应的值直到某个父 context 中返回 nil 或者查找到对应的值。

总结

Go 语言中的 context.Context 的主要作用是在多个 Goroutine 组成的树中同步取消信号以减少对资源的消耗和占用，同时它还能够设置截止日期（超时处理）、同步信号，传递请求相关值（设置中间件）的结构体。在真正使用传值的功能时我们应该非常谨慎，使用 context.Context 传递请求的所有参数一种非常差的设计，比较常见的使用场景是传递请求对应用户的认证令牌以及用于进行分布式追踪的请求 ID。

参考

Go 语言设计与实现