GO专家编程读书笔记 (5) 常见控制结构实现原理之Context

Golang 中 Context 与 WaitGroup 最大的不同在于，Context 可以控制树形结构的 goroutine，每一个 goroutine 具有相同的上下文。

由于 goroutine 派生出子 goroutine，而子 goroutine 又继续派生新的 goroutine。这种情况下使用 WaitGroup 就不太容易，因为子 goroutine 个数不容易确定，而使用 context 就可以很容易实现。

Context

Context 接口

Context 是一个接口，它定义了四个方法：

• Deadline() (deadline time.Time, ok bool)：返回一个 deadline，和一个是否已经设置 deadline 的布尔值。
• Done() <-chan struct{}：当 Context 被关闭后，返回的是一个被关闭的 channel；而没有关闭时，返回的是 nil。
• Err() error：当 Context 被关闭时，返回描述 Context 被关闭的原因；否则返回 nil。
• Value(key interface{}) interface{}：根据 key 查询 key 对应的 value，用于实现 ValueContext。

type Context interface {
    Deadline() (deadline time.Time, ok bool)

    Done() <-chan struct{}

    Err() error

    Value(key interface{}) interface{}
}

在 context 包中，定义了一个空的 Context 用于作为其他 Context 父节点或者全局的根节点。并且还实现了四种不同类型的 Context：

• Cancel Context：通过 WithCancel 创建。
• Deadline Context：通过 WithDeadline 创建。
• Timeout Context：通过 WithTimeout 创建。
• Value Context：通过 WithValue 创建。

context 包中，有四个结构体实现了 Context 接口：emptyCtx、cancelCtx、timerCtx、valueCtx。关系如下：

emptyCtx

context 包中定义了一个空的 context（emptyCtx），用于作为 context 的根节点，空的 context 只是简单的实现了 Context 接口，本身不包含任何值。

context 包中定义了一个公用的 emptCtx 全局变量，名为 background，可以使用 context.Background() 获取它：

var background = new(emptyCtx)
func Background() Context {
    return background
}

cancelCtx

数据结构

cancelCtx 的数据结构如下：

• children 中记录了由此 context 派生的所有 child，此 context 被 cancel 时会把其中的所有 child 都 cancel 掉。

type cancelCtx struct {
    Context

    mu       sync.Mutex            // protects following fields
    done     chan struct{}         // created lazily, closed by first cancel call
    children map[canceler]struct{} // set to nil by the first cancel call
    err      error                 // set to non-nil by the first cancel call
}

cancelCtx 与 deadline 和 value 无关，所以只需要实现 Done 和 Err 方法即可。

• 在实现 Done 方法时，直接返回成员变量 cancelCtx.done。
• 在实现 Err 方法时，直接返回成员变量 cancelCtx.err。在 cancelCtx 调用了 cancel 方法取消后，err 会被指定为一个全局变量 var Canceled = errors.New("context canceled")。

WithCancel

WithCancel() 方法做了三件事情：

• 初始化一个 cancelCtx 实例。
• 将 cancelCtx 实例添加到其父节点中：
  • 如果父节点也支持 cancel，也就是说其父节点肯定有 children 成员，那么把新 context 添加到 children 里即可。
  • 如果父节点不支持 cancel，就继续向上查询，直到找到一个支持 cancel 的节点，把新 context 添加到 children 里。
  • 如果所有的父节点均不支持 cancel，则启动一个协程等待父节点结束，然后再把当前 context 结束。
• 返回 cancelCtx 实例和 cancel 方法：

func WithCancel(parent Context) (ctx Context, cancel CancelFunc) {
    c := newCancelCtx(parent)	// 初始化一个 cancelCtx 实例
    propagateCancel(parent, &c)   //将自身添加到父节点
    return &c, func() { c.cancel(true, Canceled) }
}

cancel 实现

cancelCtx方法的核心在于 cancel 方法的实现，作用是关闭自己和后代，后代存储在 cancel.children 中。源码如下：

• 首先设置 err，关闭原因，这里这里被指定为了一个全局变量 var Canceled = errors.New("context canceled")。
• 接着遍历所有的 children，调用 cancel 方法。
• 将自己从 parent 中删除。

func (c *cancelCtx) cancel(removeFromParent bool, err error) {
    c.mu.Lock()

    c.err = err                          //设置一个error，说明关闭原因
    close(c.done)                     //将channel关闭，以此通知派生的context

    for child := range c.children {   //遍历所有children，逐个调用cancel方法
        child.cancel(false, err)
    }
    c.children = nil
    c.mu.Unlock()

    if removeFromParent {            //正常情况下，需要将自己从parent删除
        removeChild(c.Context, c)
    }
}

timerCtx

数据结构

timerCtx 继承于 cancelCtx，在 cancelCtx 基础上增加了 deadline 用于标示自动 cancel 的最终时间，而 timer 就是一个触发自动 cancel 的定时器。

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

    deadline time.Time
}

timerCtx 衍生出的 DeadlineContext 和 TimeoutContext 实现原理是一样的，只是 DeadlineContex 标识截止时间而 TimeoutContext 标识超时时间。

// DeadlineContext 和 TimeoutContext 的实现原理是一致的。
func WithTimeout(parent Context, timeout time.Duration) (Context, CancelFunc) {
    return WithDeadline(parent, time.Now().Add(timeout))
}

除了 cancelCtx 已经实现的 Done 和 Err 方法，timerCtx 还需要实现 Deadline 方法：

• 只需要返回 timerCtx.deadline 即可。

cancel 实现

timerCtx 的 cancel 的实现基本和 cancelCtx 的实现是一样的，需要进行的额外操作就是把定时器 timer 关闭。并且关闭原因也会不一样：

• 如果 deadline 到来之前手动关闭，则关闭原因与 cancelCtx 显示一致。
• 如果 deadline 到来时自动关闭，则原因为：context deadline exceeded。

valueCtx

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

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

由于 valueCtx 既不需要 cancel，也不需要 deadline，那么只需要实现 Value() 接口即可。

Value 接口

当前 context 查找不到 key 时，会向父节点查找，如果查询不到则最终返回 interface{}（孩子可以查询到父亲的 value 值）。

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