Context

context 数据结构

type Context interface {
    Deadline() (deadline time.Time, ok bool) // 返回 context 的过期时间
    Done() <-chan struct{}                   // 返回 context 中的 channel
    Err() error                              // 返回错误
    Value(key any) any                       // 返回 context 中的对应 key 的值
}

context.Background | context.TODO

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

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

func Background() Context {
    return background
}

func TODO() Context {
    return todo
}

emptyCtx

emptyCtx 对 Context 接口实现如下，emptyCtx 本质上类型为一个整型。我们观察其实现，发现 emptyCtx 没有实现 Context 的高级功能，像 cancel | timeout

Done 方法返回一个 nil channel，写入或者读取数据，均会陷入阻塞
Err 方法返回的错误永远为 nil
Value 方法返回的 value 同样永远为 nil
Deadline 标识当前 context 不存在过期时间

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
}

context.WithCancel

下面是 WithCancel 源码：

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 := &cancelCtx{}
    c.propagateCancel(parent, c)
    return c
}

校验父 context 非空
在 propagateCancel 方法中构建父子 ctx, 保证父 context 终止时，该 cancelCtx 也会被终止
将 cancelCtx 返回，返回用以终止该 cancelCtx 的闭包函数

下面是 cancelCtx 结构体：

type cancelCtx struct {
    Context   
  mu       sync.Mutex
  done     atomic.Value             
  children map[canceler]struct{}
  err      error
  cause    error
}

// A canceler is a context type that can be canceled directly. The
// implementations are *cancelCtx and *timerCtx.
type canceler interface {
    cancel(removeFromParent bool, err, cause error)
    Done() <-chan struct{}
}

下面是 propagateCancel 方法：

// propagateCancel arranges for child to be canceled when parent is.
// It sets the parent context of cancelCtx.
func (c *cancelCtx) propagateCancel(parent Context, child canceler) {
    c.Context = parent

    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 {
        // parent is a *cancelCtx, or derives from one.
        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()
        return
    }

    if a, ok := parent.(afterFuncer); ok {
        // parent implements an AfterFunc method.
        c.mu.Lock()
        stop := a.AfterFunc(func() {
            child.cancel(false, parent.Err(), Cause(parent))
        })
        c.Context = stopCtx{
            Context: parent,
            stop:    stop,
        }
        c.mu.Unlock()
        return
    }

    goroutines.Add(1)
    go func() {
        select {
        case <-parent.Done():
            child.cancel(false, parent.Err(), Cause(parent))
        case <-child.Done():
        }
    }()
}

propagateCancel 方法中：

parent 是不会被 cancel 的类型（如 emptyCtx），则直接返回
parent 已经被 cancel，则直接终止子 context，并以 parent 的 err 作为子 context 的 err， parent 的 Cause 作为子 context 的 cause
parent 是 cancelCtx 的类型，加锁，并将子 context 添加到 parent 的 children map 当中，函数终止
parent 不是 cancelCtx 类型但实现了 afterFuncer 接口，加锁，以 child cancel 为参数传入父 context 的 AfterFunc 函数中，构建父子 context, 此时父 context 为 stopCtx
假如 parent 不是 cancelCtx 类型，但存在 cancel 的能力（用户自定义实现的 context），则启动一个协程，通过多路复用的方式监控 parent 状态，倘若其终止，则同时终止子 context，并透传 parent 的 err

下面是 propagateCancel 函数中的 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 特有的协议）

cancel

下面是 cancel 源码，当我们执行 withCancel 返回的 cancelFunc 时，会调用 func() { c.cancel(true, Canceled, nil) }

// cancel closes c.done, cancels each of c's children, and, if
// removeFromParent is true, removes c from its parent's children.
// cancel sets c.cause to cause if this is the first time c is canceled.
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)
    }
}

进入方法主体，首先校验传入的 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 中移除

context.WithTimout | context.WithDeadline

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

WithTimeout 本质上会调用 context.WithDeadline 方法

func WithDeadlineCause(parent Context, d time.Time, cause error) (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{
        deadline: d,
    }
    c.cancelCtx.propagateCancel(parent, c)
    dur := time.Until(d)
    if dur <= 0 {
        c.cancel(true, DeadlineExceeded, cause) // 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, cause)
        })
    }
    return c, func() { c.cancel(true, Canceled, nil) }
}

校验 parent context 非空
校验 parent 的过期时间是否早于自己，若是，则构造一个 cancelCtx 返回即可（此时在给子 ctx 包装 timeCtx, 没意义，因为一旦父 context 退出，子 context 也会被cancel）
构造出一个新的 timerCtx
propagateCancel 构建父子 ctx
判断过期时间是否已到，若是，直接 cancel timerCtx，并返回 DeadlineExceeded ("context deadline exceeded") 错误
加锁
启动 time.Timer，设定一个延时时间，即达到过期时间后会终止该 timerCtx，并返回 DeadlineExceeded 的错误
解锁
返回 timerCtx，和一个封装了 cancel 逻辑的函数

下面是 timerCtx 源码：

// A timerCtx carries a timer and a deadline. It embeds a cancelCtx to
// implement Done and Err. It implements cancel by stopping its timer then
// delegating to cancelCtx.cancel.
type timerCtx struct {
    cancelCtx
    timer *time.Timer // Under cancelCtx.mu.

    deadline time.Time
}

复用了 cancelCtx 的能力, 并且加入了定时器，可以等时间到了再去 cancel

context.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

valueCtx.Value

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

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 withoutCancelCtx:
            if key == &cancelCtxKey {
                // This implements Cause(ctx) == nil
                // when ctx is created using WithoutCancel.
                return nil
            }
            c = ctx.c
        case *timerCtx:
            if key == &cancelCtxKey {
                return &ctx.cancelCtx
            }
            c = ctx.Context
        case backgroundCtx, todoCtx:
            return nil
        default:
            return c.Value(key)
        }
    }
}

Value 方法中，判断当前 valueCtx 的 key 是否等于用户传入的 key，是，则直接返回其 value。不是，调用 value 方法，从 parent context 中依次向上寻找

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

一个 valueCtx 实例只能存一个 kv 对，因此 n 个 kv 对会嵌套 n 个 valueCtx，造成空间浪费；
基于 k 寻找 v 的过程是线性的，时间复杂度 O(N)；

[Go] Golang(1.23.2) 源码走读 - Context