你真的懂context吗？

前言

context是golang中的经典工具，主要在异步场景中用于控制并发以及控制goroutine，此外，context还有一定的数据存储能力。本人在学习过程中，在许多地方都看到了context，大致了解了使用方法，但是不了解底层使我非常难受，这篇文章就探究一下context的底层原理。

核心数据结构

context.Context

context.Context是一个接口，定义了context的常用方法：

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

1. Deadline():返回context的过期时间
2. Done():返回context中的channel
3. Err():返回错误
4. Value():返回context中对应key的值

标准error

var Canceled=errors.New("context canceled")
type deadlineExceededError struct{}
var DeadlineExceeded error = deadlineExceededError{}

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

emptyCtx

emptyCtx的实现

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 
}

1. emptyCtx是一个空的Context（实现了Context接口），本质类型是一个整形
2. Deadline方法会返回一个公元元年时间以及false的flag，标识当前Context不存在过期时间
3. Done方法返回一个nil值，用户无论往nil中写入或者读取数据，均会陷入阻塞
4. Err方法返回的错误永远为nil
5. Value方法返回的value永远为nil

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类型的一个实例

cancelCtx

数据结构

type cancelCtx struct {
    Context
   
    mu       sync.Mutex            
    done     atomic.Value        //由 `chan struct{}` 类型构成，延迟创建，由首次调用取消操作关闭  
    children map[canceler]struct{} //在首次调用cancel时设为空
    err      error     //在首次调用cancel时被设为非空值            
}

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

1. cancelCtx中embed（嵌入）了一个Context，这个Context是cancelCtx的父级，也就是说cancelCtx必然为某个Context的子Context
2. 内置了一把锁，用以协调并发场景下的资源获取
3. done：实际类型是chan struct{}，即用以反映cancelCtx生命周期的通道
4. children：一个set，指向cancelCtx的所有子Context
5. err：记录了当前cancelCtx的错误

同时定义了一个canceler接口，这个接口用于实现可取消的Context（如 *cancelCtx、*timerCtx 等内部类型）

同样，cancelCtx本身也是一个Context实例，也实现了Context接口

Deadline方法

cancelCtx未实现Deadline方法，仅是embed了一个带有Deadline方法的Context interface，因此倘如一个cancelCtx，会通过嵌入的Context实例的Deadline方法获取到其 Deadline（委托给父级）

Done方法

请结合注释看代码

func (c *cancelCtx)Done() <-chan struct{}{
    //这里不加锁快速读取当前是否有done channel，因为大多数情况下，done channel是存在的，如果存在，直接返回避免加锁开销
    d:=c.done.Load()  
    if d!=nil{
        return d.(chan struct{})
    }

    //如果done channel还没初始化，就需要安全的创建channel，用互斥锁mu保证：只有一个goroutine能进入创建逻辑，防止重复创建
    c.mu.Lock()
    defer c.mu.Unlock()

    //再次c.done.Load()，双重检查，因为可能有多个goroutine卡在拿锁这一步，当第一个goroutine创建完channel并释放锁后，其他的goroutine拿到锁后应该复用已经创建的channel，而不是再建一个
    d=c.done.Load()

    //如果done channel还不存在，则创建一个
    if d==nil{
        d=make(chan struct{})
        c.done.Store(d)
    }
    //返回channel
    return d.(chan struct{})
}

• Done方法也体现了一种优化策略：只有在真正需要某个资源或数据时，才去创建或加载它，如果一直不需要，就永远不加载。

• done channel不是创建cancelCtx时就跟着创建的，而是第一次调用Done()时才创建

• 如果你的代码从不监听ctx.Done(),这个channel就永远不会被创建-->节省内存和GC压力

Err方法

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

1. 加锁
2. 读取cancelCtx.err
3. 解锁
4. 返回结果

Value方法

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

1. 倘若key特定值&cancelCtxKey，则返回cancelCtx自身的指针
2. 否则遵循valueCtx的思路取值返回

context.WithCancel()

我们在日常编程中常用context.WithCancel()方法创建cancelCtx实例，那具体是怎么创建的呢？往下看

newCancelCtx()

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

• newCancelCtx()的作用是注入父Context，并返回一个新的cancelCtx

propagateCancel()

propagateCancel()的源码比较长，请结合注释阅读源码

// propagateCancel顾名思义，是用来传递父子Context之间的cancel事件的
func propagateCancel(parent Context, child canceler) {
    //获取父级Context的done channel
    done := parent.Done()

    //如果父级Context的done channel为nil，则说明父级是一个不会被cancel的类型(emptyCtx)，直接返回
    if done == nil {
        return 
    }

    //如果父级Context可以cancel，就监听父级的done channel
    select {
    //如果父级Context已经cancel，则直接终止子Context，并并以parent的err作为子Context的err
    case <-done:
        child.cancel(false, parent.Err())
        return
    default:
    }

    //这里判断父级Context是不是cancelCtx
    if p, ok := parentCancelCtx(parent); ok {
        //如果父级是cancelCtx，则加锁，并将子Context添加到父级的children map中
        p.mu.Lock()
        if p.err != nil {
            child.cancel(false, p.err)
        } else {
            if p.children == nil {
                p.children = make(map[canceler]struct{})
            }
            p.children[child] = struct{}{}
        }
        p.mu.Unlock()
    } else {
        //如果父级不是cancelCtx类型，但是又存在cancel的能力（比如用户自定义实现的Context），则启动一个协程，通过多路复用的方式监控父级状态，倘若其终止，则同时终止子Context，并传递父级的err
        atomic.AddInt32(&goroutines, +1)
        go func() {
            select {
            case <-parent.Done():
                child.cancel(false, parent.Err())
            case <-child.Done():
            }
        }()
    }
}

再来看看parentCancelCtx()

func parentCancelCtx(parent Context) (*cancelCtx, bool) {
    done := parent.Done()
    //倘若父级的channel是已经关闭的，或者父级是不会被cancel的类型，则返回false
    if done == closedchan || done == nil {
        return nil, false
    }

    //倘若以特定的cancelCtxKey从parent中取值，取得的value是父级本身，则返回true（基于cancelCtxKey为key取值时返回cancelCtx本身，是cancelCtx特有的协议）
    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
}

context.WithCancel()

func WithCancel(parent Context)(ctx Context,cancel CancelFunc){
    //检验父级Context是不是空，如果为空则引发panic
    if parent==nil{
        panic("cannot create context from nil parent")
    }

    //如果父级不为空，则开始创建cancelCtx，首先调用newCancelCtx，将该cancelCtx注入父级Context
    c := newCancelCtx(parent)

    //调用propagateCancel(),使得子Context开始监听父级Context的cancel事件
    propagateCancel(parent, &c)

    //将该cancelCtx返回，连带返回一个用于终止该cancelCtx的闭包函数
    return &c, func() { c.cancel(true, Canceled) }
}

cancelCtx.cancel()

我们经常调用cancelCtx.cancel()方法来取消一个cancelCtx，那么cancelCtx.cancel()具体做了哪些事情呢？

请结合注释食用源码

//cancelCtx.cancel方法有两个入参
// removeFromParent表示当前Context是否需要从父级Context的children set中删除
// err则是cancel后需要展示的错误
func (c *cancelCtx) cancel(removeFromParent bool, err error) {
    //如果传入的err为nil，则抛出一个panic
    if err == nil {
        panic("context: internal error: missing cancel error")
    }

    //加锁
    c.mu.Lock()

    //校验cancelCtx自带的err是否已经非空，如果非空，则说明已经cancel，则解锁返回
    if c.err != nil {
        c.mu.Unlock()
        return 
    }

    //将传入的err赋给cancelCtx.err
    c.err = err

    //获取cancelCtx.done channel
    d, _ := c.done.Load().(chan struct{})

    //如果cancelCtx.done channel为nil，则直接注入closedchan，closechan是一个预定义的已经关闭的channel，直接注入这个，则可以通知监听该Context的Context
    if d == nil {
        c.done.Store(closedchan)
    } else {
        //否则，关闭cancelCtx.done channel，也就是通知监听该Context的Context
        close(d)
    }

    //遍历cancelCtx的children set，一次将children context都进行cancel
    for child := range c.children {
        child.cancel(false, err)
    }

    //将cildren清空
    c.children = nil
    c.mu.Unlock()
    //根据removeFromParent判断是否手动把cancelCtx从parent的children set中移除
    if removeFromParent {
        removeChild(c.Context, c)
    }
}

timerCtx

数据结构

type timerCtx struct {
    cancelCtx
    timer *time.Timer
    deadline time.Time
}

timerCtx继承了cancelCtx的方法，而且新增了timer字段用于定时终止Context，和deadline字段用于设置timerCtx的过期时间。

timerCtx.Deadline()

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

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

context.WithTimeout() & context.WithDeadline()

我们可以利用context.WithTimeout()创建一个timerCtx。

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

context.WithTimeout本质会调用context.WithDeadline()

//WithDeadline有两个入参
// parent是timerCtx的父级Context
// d是timerCtx的过期时间
func WithDeadline(parent Context, d time.Time) (Context, CancelFunc) {
    //检验父级Context是不是空，如果为空则引发panic
    if parent == nil {
        panic("cannot create context from nil parent")
    }

    //检验父级的Deadline是否早于自己，如果是，则创建一个timerCtx返回即可，没有必要设置过期时间了
    if cur, ok := parent.Deadline(); ok && cur.Before(d) {
        
        return WithCancel(parent)
    }

    //创建timerCtx
    c := &timerCtx{
        cancelCtx: newCancelCtx(parent),
        deadline:  d,
    }

    //调用propagateCancel()方法，使得子Context开始监听父级Context的cancel事件
    propagateCancel(parent, c)
    
    //判断timerCtx的过期时间是否已经过了
    dur := time.Until(d)

    //如果过期时间已经过了，直接cancel timerCtx，并返回DeadlineExceeded的错误
    if dur <= 0 {
        c.cancel(true, DeadlineExceeded) 
        return c, func() { c.cancel(false, Canceled) }
    }

    //加锁
    c.mu.Lock()
    defer c.mu.Unlock()

    //启动time.Timer,设定一个延时时间，即达到过期时间后会终止该timerCtx，并返回DeadlineExceeded的错误
    if c.err == nil {
        c.timer = time.AfterFunc(dur, func() {
            c.cancel(true, DeadlineExceeded)
        })
    }

    //返回timerCtx和cancel闭包函数
    return c, func() { c.cancel(true, Canceled) }
}

timerCtx.cancel()

cancel方法用于关闭timerCtx

func (c *timerCtx) cancel(removeFromParent bool, err error) {
    //复用继承cancelCtx的cancel能力，进行cancel处理
    c.cancelCtx.cancel(false, err)

    //删除父级Context中的孩子中的timerCtx
    if removeFromParent {
        removeChild(c.cancelCtx.Context, c)
    }

    //加锁处理
    c.mu.Lock()
    //停止timerCtx的timer
    if c.timer != nil {
        c.timer.Stop()
        c.timer = nil
    }
    c.mu.Unlock()
}

valueCtx

数据结构

type valueCtx struct {
    Context
    key, val any
}

• valueCtx同样继承了一个父级Context
• 一个valueCtx中仅有一组kv对

valueCtx.Value()

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

• 假如当前valueCtx的key等于用户传入的key，则直接返回其value
• 假如不等，则调用value从父级Context中寻找

func value(c Context, key any) any {
    //开启一个for循环，由子及父的依次对key进行匹配
    for {
        //ctx := c.(type) 是类型 switch 的声明语法，它会根据接口变量 c 的实际动态类型，在每个 case 中将 ctx 自动转换为对应的具体类型。但是注意c必须是一个接口类型
        switch ctx := c.(type) {
        case *valueCtx:
            if key == ctx.key {
                return ctx.val
            }
            c = ctx.Context

        //对于某些特殊key进行特殊处理，比如如果key是某个可cancel的Context的地址，则返回该Context的指针，可以用于往上寻找第一个可cancel的Context
        case *cancelCtx:
            if key == &cancelCtxKey {
                return c
            }
            c = ctx.Context
        case *timerCtx:
            if key == &cancelCtxKey {
                return &ctx.cancelCtx
            }
            c = ctx.Context

        //如果是空Context，则返回nil
        case *emptyCtx:
            return nil
        default:
            return c.Value(key)
        }
    }
}

valueCtx用法

从源码中可以看出，valueCtx不适合视为存储介质，存放大量的kv数据，原因如下：

• 一个valueCtx实例仅能存一个kv对，因此如果你想存大量kv对，则需要创建多个valueCtx实例，浪费空间
• 基于k寻找v的过程是线性的，时间复杂度为O(n)
• 不支持k的去重，相同k可能存在不同v，并且基于你查找起点的不同，返回不同的v。由此得知，valueContext的定位类似于请求头，只适合存放少量作用域较大的全局meta数据

context.WithValue()

我们可以使用context.WithValue()创建一个valueCtx

func WithValue(parent Context, key, val any) Context {
    // 检验父级Context是不是空，如果为空则引发panic
    if parent == nil {
        panic("cannot create context from nil parent")
    }
    // 检验key是不是空，如果为空则引发panic
    if key == nil {
        panic("nil key")
    }
    // 检验key是不是可比较的,如果不能比较，引发panic
    if !reflectlite.TypeOf(key).Comparable() {
        panic("key is not comparable")
    }

    // 创建valueCtx并返回
    return &valueCtx{parent, key, val}
}

参考：小徐先生