golang-context详解

前言

context是golang中的经典工具，主要在异步场景中用于实现并发协调以及对goroutine的生命周期控制，除此之外，context还兼具一定的数据存储能力

1. 核心数据结构

1.1 context.Context

Context为interface定义了4个核心api

• DealLine: 返回context的过期时间
• Done：返回context中的chan
• Err：返回错误
• Value：返回context中存储对应key的值

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

1.2 标准error

var Canceled = errors.New("context canceled") 

var DeadlineExceeded error = deadlineExceededError{}

type deadlineExceededError struct{}

func (deadlineExceededError) Error() string   { return "context deadline exceeded" }
func (deadlineExceededError) Timeout() bool   { return true }
func (deadlineExceededError) Temporary() bool { return true }

• Canceled：context 被 cancel 时会报此错误；
• DeadlineExceeded：context 超时时会报此错误.

2. emptyCtx

2.1 实现

type emptyCtx int

func (*emptyCtx) Deadline() (deadline time.Time, ok bool) {
    return
}

func (*emptyCtx) Done() <-chan struct{} {
    return nil
}

func (*emptyCtx) Err() error {
    return nil
}

func (*emptyCtx) Value(key any) any {
    return nil
}

• emptyCtx：一个空的context，本质是一个整形
• DealLine: 返回公元元年的时间以及false，标识当前context不存在过期时间
• Done：返回一个nil(空读写阻塞，写关闭异常，读关闭空零)
• Err：永远返回nil
• Value：返回的value永远为nil

2.2 context.Background()&context.TODO()

var (
    background = new(emptyCtx)
    todo       = new(emptyCtx)
)

func Background() Context {
    return background
}

func TODO() Context {
    return todo
}

我们所常用的 context.Background() 和 context.TODO() 方法返回的均是 emptyCtx 类型的一个实例.

3. cancleCtx

3.1 cancelCtx的数据结构

type cancelCtx struct {
    Context

    mu       sync.Mutex            // protects following fields
    done     atomic.Value          // of 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
    cause    error                 // set to non-nil by the first cancel call
}

type canceler interface {
    cancel(removeFromParent bool, err, cause error)
    Done() <-chan struct{}
}

• Context：内置了一个context为父context，可见。
• mu：内置了一把锁，用以协调并发场景下的资源获取
• 实际类型为chan struct{} 用以反映cancelCtx生命周期的通道
• children：一个set指向cancelCtx的所有子context
• err：记录当前cancelCtx的错误

3.2 Deadline 方法

cancelCtx 未实现该方法，仅是 embed 了一个带有 Deadline 方法的 Context interface，因此倘若直接调用会报错.

3.3 Done()

func (c *cancelCtx) Done() <-chan struct{} {
    d := c.done.Load()
    if d != nil {
        return d.(chan struct{})
    }
    c.mu.Lock()
    defer c.mu.Unlock()
    d = c.done.Load()
    if d == nil {
        d = make(chan struct{})
        c.done.Store(d)
    }
    return d.(chan struct{})
}

• 基于atomic包，读取cancelCtx中done chan倘若已存在则直接返回
• 加锁后，检查chan是否存在，存在返回 双重检查
• 初始化chan存储到done中返回 懒加载

3.4 Err()

func (c *cancelCtx) Err() error {
    c.mu.Lock()
    err := c.err
    c.mu.Unlock()
    return err
}

• 加锁
• 读取 cancelCtx.err
• 解锁
• 返回

3.5 Value()

func (c *cancelCtx) Value(key any) any {
    if key == &cancelCtxKey {
        return c
    }
    return value(c.Context, key)
}

• 指定特定的key时返回cancelCtx本身的指针
• 否则按照valueCtx 的思路取值返回

3.6 context.WithCancel()

3.6.1 context.WithCancel()

type CancelFunc func()
// 创建cancelCtx活得取消方法
func WithCancel(parent Context) (ctx Context, cancel CancelFunc) {
    c := withCancel(parent)
    return c, func() { c.cancel(true, Canceled, nil) }
}

func withCancel(parent Context) *cancelCtx {
    if parent == nil {
        panic("cannot create context from nil parent")
    }
    c := newCancelCtx(parent) // 创建cancelCtx
    propagateCancel(parent, c) // 
    return c
}

• 检验父context是否为空
• 注入父 context 构造好一个新的 cancelCtx
• 在 propagateCancel 方法内启动一个守护协程，以保证父 context 终止时，该 cancelCtx 也会被终止
• 将 cancelCtx 返回，连带返回一个用以终止该 cancelCtx 的闭包函数

3.6.2 newCancelCtx()

// 构造方法
func newCancelCtx(parent Context) *cancelCtx {
    return &cancelCtx{Context: parent}
}

• 注入父 context 后，返回一个新的 cancelCtx.

3.6.3 propagateCancel()

// propagateCancel arranges for child to be canceled when parent is.
func propagateCancel(parent Context, child canceler) {
    done := parent.Done()
    if done == nil {
        return // parent is never canceled
    }

    select {
    case <-done:
        // parent is already canceled
        child.cancel(false, parent.Err(), Cause(parent))
        return
    default:
    }

    if p, ok := parentCancelCtx(parent); ok {
        p.mu.Lock()
        if p.err != nil {
            // parent has already been canceled
            child.cancel(false, p.err, p.cause)
        } else {
            if p.children == nil {
                p.children = make(map[canceler]struct{})
            }
            p.children[child] = struct{}{}
        }
        p.mu.Unlock()
    } else {
        goroutines.Add(1)
        go func() {
            select {
            case <-parent.Done():
                child.cancel(false, parent.Err(), Cause(parent))
            case <-child.Done():
            }
        }()
    }
}

顾名思义，propagateCancel用以传递父子 context 之间的 cancel 事件：

• 倘若 parent 是不会被 cancel 的类型（如 emptyCtx），则直接返回
• 倘若 parent 已经被 cancel，则直接终止子 context，并以 parent 的 err 作为子 context 的 err
• 假如 parent 是 cancelCtx 的类型，则加锁，并将子 context 添加到 parent 的 children map 当中
• 假如 parent 不是 cancelCtx 类型，但又存在 cancel 的能力（比如用户自定义实现的 context），则启动一个协程，通过多路复用的方式监控 parent 状态，倘若其终止，则同时终止子 context，并透传 parent 的 err
  parentCancelCtx 是如何校验 parent 是否为 cancelCtx 的类型

func parentCancelCtx(parent Context) (*cancelCtx, bool) {
    done := parent.Done()
    if done == closedchan || done == nil {
        return nil, false
    }
    p, ok := parent.Value(&cancelCtxKey).(*cancelCtx)
    if !ok {
        return nil, false
    }
    pdone, _ := p.done.Load().(chan struct{})
    if pdone != done {
        return nil, false
    }
    return p, true
}

• 倘若 parent 的 channel 已关闭或者是不会被 cancel 的类型，则返回 false
• 倘若以特定的 cancelCtxKey 从 parent 中取值，取得的 value 是 parent 本身，则返回 true. （基于 cancelCtxKey 为 key 取值时返回 cancelCtx 自身，是 cancelCtx 特有的协议）

3.6.4 cancelCtx.cancel()

func (c *cancelCtx) cancel(removeFromParent bool, err, cause error) {
    if err == nil {
        panic("context: internal error: missing cancel error")
    }
    if cause == nil {
        cause = err
    }
    c.mu.Lock()
    if c.err != nil {
        c.mu.Unlock()
        return // already canceled
    }
    c.err = err
    c.cause = cause
    d, _ := c.done.Load().(chan struct{})
    if d == nil {
        c.done.Store(closedchan)
    } else {
        close(d)
    }
    for child := range c.children {
        // NOTE: acquiring the child's lock while holding parent's lock.
        child.cancel(false, err, cause)
    }
    c.children = nil
    c.mu.Unlock()

    if removeFromParent {
        removeChild(c.Context, c)
    }
}

• cancelCtx.cancel 方法有两个入参，removeFromParent 是一个 bool 值，表示当前 context 是否需要从父 context 的 children set 中删除
• 进入方法主体，首先校验传入的 err 是否为空，若为空则 panic
• 加锁
• 校验 cancelCtx 自带的 err 是否已经非空，若非空说明已被 cancel，则解锁返回
• 将传入的 err 赋给 cancelCtx.err
• 处理 cancelCtx 的 channel，若 channel 此前未初始化，则直接注入一个 closedChan，否则关闭该 channel
• 遍历当前 cancelCtx 的 children set，依次将 children context 都进行 cancel
• 解锁
• 根据传入的 removeFromParent flag 判断是否需要手动把 cancelCtx 从 parent 的 children set 中移除
  如何将 cancelCtx 从 parent 的 children set 中移除

func removeChild(parent Context, child canceler) {
    p, ok := parentCancelCtx(parent)
    if !ok {
        return
    }
    p.mu.Lock()
    if p.children != nil {
        delete(p.children, child)
    }
    p.mu.Unlock()
}

• 如果 parent 不是 cancelCtx，直接返回（因为只有 cancelCtx 才有 children set）
• 加锁
• 从parent的children set中删除对应的child
• 解锁返回

4. timerCtx

4.1 数据结构

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

    deadline time.Time
}

timerCtx 在 cancelCtx 基础上又做了一层封装，除了继承 cancelCtx 的能力之外，新增了一个 time.Timer 用于定时终止 context；另外新增了一个 deadline 字段用于字段 timerCtx 的过期时间.

4.2 Deadline()

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

context.Context interface 下的 Deadline api 仅在 timerCtx 中有效，由于展示其过期时间.

4.3 cancel()

func (c *timerCtx) cancel(removeFromParent bool, err, cause error) {
    c.cancelCtx.cancel(false, err, cause)
    if removeFromParent {
        // Remove this timerCtx from its parent cancelCtx's children.
        removeChild(c.cancelCtx.Context, c)
    }
    c.mu.Lock()
    if c.timer != nil {
        c.timer.Stop()
        c.timer = nil
    }
    c.mu.Unlock()
}

• 复用继承的 cancelCtx 的 cancel 能力，进行 cancel 处理
• 判断是否需要手动从 parent 的 children set 中移除，若是则进行处理
• 加锁
• 停止time.Timer
• 解锁返回

4.4 context.WithDeadline()&context.WithTimeout()

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

context.WithTimeout 方法用于构造一个 timerCtx，本质上会调用 context.WithDeadline 方法,传入的为多长时间后过期的一个时间间隔

func WithDeadline(parent Context, d time.Time) (Context, CancelFunc) {
    if parent == nil {
        panic("cannot create context from nil parent")
    }
    if cur, ok := parent.Deadline(); ok && cur.Before(d) {
        // The current deadline is already sooner than the new one.
        return WithCancel(parent)
    }
    c := &timerCtx{
        cancelCtx: newCancelCtx(parent),
        deadline:  d,
    }
    propagateCancel(parent, c)
    dur := time.Until(d)
    if dur <= 0 {
        c.cancel(true, DeadlineExceeded, nil) // deadline has already passed
        return c, func() { c.cancel(false, Canceled, nil) }
    }
    c.mu.Lock()
    defer c.mu.Unlock()
    if c.err == nil {
        c.timer = time.AfterFunc(dur, func() {
            c.cancel(true, DeadlineExceeded, nil)
        })
    }
    return c, func() { c.cancel(true, Canceled, nil) }
}

• 检验 parent context 非空
• 校验 parent 的过期时间是否早于自己，若是，则构造一个 cancelCtx 返回就行
• 构造出一个新的 timerCtx
• 启动守护方法，同步 parent 的 cancel 事件到子 context
• 判断过期时间是否已到，若是，直接 cancel timerCtx，并返回 DeadlineExceeded 的错误
• 加锁
• 校验err是否为空，如果为空证明未被 cancel，启动 time.Timer，设定一个延时时间，即达到过期时间后会终止该 timerCtx，并返回 DeadlineExceeded 的错误
• 解锁
• 返回 timerCtx，已经一个封装了 cancel 逻辑的闭包 cancel 函数

5. valueCtx()

5.1 数据结构

type valueCtx struct {
    Context
    key, val any
}

• valueCtx 同样继承了一个 parent context
• 一个 valueCtx 中仅有一组 kv 对

5.2 Value()

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

• 当前 valueCtx 的 key 等于用户传入的 key，则直接返回其 value
• 否则则从 parent context 中依次向上寻找

func value(c Context, key any) any {
    for {
        switch ctx := c.(type) {
        case *valueCtx:
            if key == ctx.key {
                return ctx.val
            }
            c = ctx.Context
        case *cancelCtx:
            if key == &cancelCtxKey {
                return c
            }
            c = ctx.Context
        case *timerCtx:
            if key == &cancelCtxKey {
                return ctx.cancelCtx
            }
            c = ctx.Context
        case *emptyCtx:
            return nil
        default:
            return c.Value(key)
        }
    }
}

• 启动一个 for 循环，由下而上，由子及父，依次对 key 进行匹配
• 其中 cancelCtx、timerCtx、emptyCtx 类型会有特殊的处理方式
• 找到匹配的 key，则将该组 value 进行返回

5.3 valueCtx 用法小结

阅读源码可以看出，valueCtx 不适合视为存储介质，存放大量的 kv 数据，原因有三：

• 一个 valueCtx 实例只能存一个 kv 对，因此 n 个 kv 对会嵌套 n 个 valueCtx，造成空间浪费；
• 基于 k 寻找 v 的过程是线性的，时间复杂度 O(N)；
• 不支持基于 k 的去重，相同 k 可能重复存在，并基于起点的不同，返回不同的 v. 由此得知，valueContext 的定位类似于请求头，只适合存放少量作用域较大的全局 meta 数据.

5.4 WithValue()

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

• parent context 为空，panic；
• key 为空 panic；
• key 的类型不可比较，panic；
• parent context 以及 kv对，返回一个新的 valueCtx.