Go context 源码理解

1 核心数据结构

1.1 context.Context

context数据结构,interface 定义了四个核心 api：

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

1.2 标准 error

var Canceled = errors.New("context canceled") context 被 cancel 时会报此错误；

var DeadlineExceeded error = deadlineExceededError{} context 超时时会报此错误.

type deadlineExceededError struct{}

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

2 emptyCtx

2.1 类的实现

type emptyCtx int 是一个空的 context，本质上类型为一个整型；

方法会返回一个公元元年时间以及 false 的 flag，标识当前 context 不存在过期时间；
func (*emptyCtx) Deadline() (deadline time.Time, ok bool) {
    return
}

func (*emptyCtx) Done() <-chan struct{} {
    return nil 方法返回一个 nil 值，用户无论往 nil 中写入或者读取数据，均会陷入阻塞；
}

func (*emptyCtx) Err() error {
    return nil 方法返回的错误永远为 nil；
}

func (*emptyCtx) Value(key any) any {
    return  返回的 value 同样永远为 nil.
}

2.2 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
}

3 cancelCtx

3.1 cancelCtx 数据结构

type cancelCtx struct {
    Context //embed 了一个 context 作为其父 context. 可见，cancelCtx 必然为某个 context 的子 context；
   
    // 内置了一把锁，用以协调并发场景下的资源获取；
    mu       sync.Mutex            // protects following fields
    
    done     atomic.Value          // of chan struct{}, created lazily, closed by first cancel call 。实际类型为 chan struct{}，即用以反映 cancelCtx 生命周期的通道；怎么懒加载的可以看下面贴的 Done 方法源码
    
    children map[canceler]struct{} // set to nil by the first cancel call  一个 set，指向 cancelCtx 的所有子 context；那什么是 canceler 呢，下面贴了
    
    err      error                 // set to non-nil by the first cancel call 记录了当前 cancelCtx 的错误. 必然为某个 context 的子 context；
}

// 可以看到 canceler 其实只暴露了两个方法。类似开闭原则。其实父亲只关心孩子的 cancel 和 Done ，其他的不需要关系，则直接不露
type canceler interface {
    cancel(removeFromParent bool, err error)
    Done() <-chan struct{}
}

3.2 Deadline 方法

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

3.3 Done 方法

func (c *cancelCtx) Done() <-chan struct{} {
    d := c.done.Load()
//    基于 atomic 包，读取 cancelCtx 中的 chan；倘若已存在，则直接返回；
    if d != nil {
        return d.(chan struct{})
    }
    c.mu.Lock()
    defer c.mu.Unlock()
    d = c.done.Load()
//    加锁后，【再次】检查 chan 是否存在，若存在则返回；（double check）
    if d == nil {
        d = make(chan struct{})
        // 初始化 chan 存储到 aotmic.Value 当中，并返回.（懒加载机制）
        c.done.Store(d)
    }
    return d.(chan struct{})
}

3.4 Err 方法

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

3.5 Value 方法

func (c *cancelCtx) Value(key any) any {
    if key == &cancelCtxKey {
        // 倘若 key 特定值 &cancelCtxKey，则返回 cancelCtx 自身的指针；cancelCtxKey 是一个全局变量。在下面的 parentCancelCtx 方法里会用到
        return c
    }
    
    return value(c.Context, key)
}

• 否则遵循 valueCtx 的思路取值返回，具体见下一章节！

3.6 context.WithCancel()

3.6.1 context.WithCancel()

func WithCancel(parent Context) (ctx Context, cancel CancelFunc) {
    if parent == nil {
        panic("cannot create context from nil parent")
    }
   
    c := newCancelCtx(parent) //注入父 context 构造好一个新的 cancelCtx；
    propagateCancel(parent, &c) //在 propagateCancel 方法内启动一个守护协程，以保证父 context 终止时，该 cancelCtx 也会被终止；propagateCancel具体代码，下面会贴出来
    
    return &c, func() { c.cancel(true, Canceled) } //在 propagateCancel 方法内启动一个守护协程，以保证父 context 终止时，该 cancelCtx 也会被终止；
}

3.6.2 newCancelCtx

1
2
3

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

3.6.3 propagateCancel

传递父子 context 之间的 cancel 事件
func propagateCancel(parent Context, child canceler) {
    done := parent.Done()
    if done == nil {
//倘若 parent 是不会被 cancel 的类型（如 emptyCtx），则直接返回；
        return // parent is never canceled
    }

    select {
    case <-done:
        // parent is already canceled
        // 倘若 parent 已经被 cancel，则直接终止子 context，并以 parent 的 err 作为子 context 的 err；
        child.cancel(false, parent.Err())
        return
    default:
    }

    注意，能走到这的ctx 其实都有 Cancel 能力了，但是还要判断一下是用户自己实现的，还是 cancelCtx 类型的
    
    if p, ok := parentCancelCtx(parent); ok {
		//假如 parent 是 cancelCtx 的类型，则加锁，具体怎么判断的下面我贴了 parentCancelCtx 源码
        p.mu.Lock()
        if p.err != nil {
            // parent has already been canceled
            child.cancel(false, p.err)
        } else {
            if p.children == nil {
                //第一次生孩子，先初始化下仓库
                p.children = make(map[canceler]struct{})
            }
            p.children[child] = struct{}{}// 将子 context 添加到 map 
        }
        p.mu.Unlock()
    } else {
        假如 parent 不是 cancelCtx 类型 但又存在 cancel 的能力（比如用户自定义实现的 context），
        atomic.AddInt32(&goroutines, +1)
        go func() {
            select {
            // 则启动一个协程，通过多路复用的方式监控 parent 状态，倘若其终止，则同时终止子 context，并透传 parent 的 err.
            case <-parent.Done():
                child.cancel(false, parent.Err())
            case <-child.Done():
            }
        }()
    }
}

// 校验 parent 是否为 cancelCtx 的类型
func parentCancelCtx(parent Context) (*cancelCtx, bool) {
    done := parent.Done()
    if done == closedchan || done == nil {
        //倘若 parent 的 channel 已关闭或者是不会被 cancel 的类型，则返回 false
        return nil, false
    }
    
//    倘若以特定的 cancelCtxKey 从 parent 中取值，取得的 value 是 parent 本身，则返回 true. （基于 cancelCtxKey 为 key 取值时返回 cancelCtx 自身的指针，这时候断言就能知道它是不是 cancelCtx）.
    p, ok := parent.Value(&cancelCtxKey).(*cancelCtx)
    if !ok {
        return nil, false
    }
    // OK，对上暗号了，你是 cancelCtxKey
    pdone, _ := p.done.Load().(chan struct{})
    if pdone != done {
        return nil, false
    }
    return p, true
}

3.6.4 cancelCtx.cancel

两个入参，第一个 removeFromParent 是一个 bool 值，表示当前 context 是否需要从父 context 的 children set 中删除；第二个 err 则是 cancel 后需要展示的错误；
func (c *cancelCtx) cancel(removeFromParent bool, err error) {
    if err == nil {
        panic("context: internal error: missing cancel error")
    }
    c.mu.Lock()
    if c.err != nil {
        //若非空说明已被 cancel，则解锁返回；
        c.mu.Unlock()
        return // already canceled
    }
    c.err = err
    d, _ := c.done.Load().(chan struct{})
    if d == nil {
        // 懒加载。还没声明过。则取一个全局的closedchan注入
        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)//所有孩子都要👋🏻
    }
    c.children = nil
    c.mu.Unlock()

    if removeFromParent {
        // 根据传入的 removeFromParent flag 判断是否需要手动把 cancelCtx 从 parent 的 children set 中移除.
        removeChild(c.Context, c) //c.Context就是 c 的父亲
    }
}

走进 removeChild 方法中，观察如何将 cancelCtx 从 parent 的 children set 中移除：

func removeChild(parent Context, child canceler) {
    p, ok := parentCancelCtx(parent)//小蝌蚪找爸爸
    if !ok {
        // 爸爸不是 cancelCtx，直接返回
        return
    }
    p.mu.Lock()
    if p.children != nil {
        // 爸爸弃孩子
        delete(p.children, child)
    }
    p.mu.Unlock()
}

关于 removeFromParent 为啥有时候传 false ，我问了 chatGPT：

在 go 的源码里，为啥 cancelCtx 的 cancel 方法第一个参数removeFromParent，有时候传 true 有时候传 false
===
在 Go 的源码中，cancelCtx 结构体的 cancel 方法中的 removeFromParent 参数用于指示在取消操作后是否从父上下文中移除该子上下文。

当 removeFromParent 参数为 true 时，表示在取消操作后应该从父上下文中移除该子上下文。这通常是因为子上下文已经不再需要与父上下文相关联，例如，该子上下文已经完成了其任务或不再需要超时控制。

而当 removeFromParent 参数为 false 时，表示取消操作后不应该从父上下文中移除该子上下文。这可能是因为子上下文仍然需要与父上下文相关联，或者有其他需要考虑的因素。
===
为啥会有 false 的情况存在？子 Cancel 了不应该直接从父移除吗?
===
当一个父上下文管理多个子上下文时，可能存在一种情况，即子上下文之间相互依赖或相互影响。在这种情况下，即使一个子上下文被取消了，父上下文仍然需要知道取消的原因，并可能采取一些适当的措施。

举个例子，假设一个 HTTP 服务器处理请求的过程中创建了一个子上下文来管理每个请求的超时。如果某个请求的子上下文超时了，虽然这个请求已经被取消，但服务器可能仍然需要将该请求的取消状态记录下来，以便在日志中进行记录或执行清理操作。

在这种情况下，即使子上下文被取消了，也不应该从父上下文中移除，因为父上下文仍然需要知道请求的状态并采取相应的行动。

4 timerCtx

4.1 数据结构

type timerCtx struct {
    cancelCtx // 在 cancelCtx 基础上又做了一层封装
    timer *time.Timer // Under cancelCtx.mu.

    deadline time.Time
}

4.2 timerCtx.Deadline()

1
2
3

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

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

4.3 timerCtx.cancel

func (c *timerCtx) cancel(removeFromParent bool, err error) {
    // 复用继承的 cancelCtx 的 cancel 能力，进行 cancel 处理；
    c.cancelCtx.cancel(false, err)
    if removeFromParent {
        removeChild(c.cancelCtx.Context, c)
    }
    c.mu.Lock()
    if c.timer != nil {
        //停止 time.Timer(我都要 GG 了，留着也没用)
        c.timer.Stop()
        c.timer = nil
    }
    c.mu.Unlock()
}

4.4 context.WithTimeout & context.WithDeadline

func WithTimeout(parent Context, timeout time.Duration) (Context, CancelFunc) {
// 本质上会调用 context.WithDeadline 方法：
    return WithDeadline(parent, time.Now().Add(timeout))
}

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. 校验 parent 的过期时间是否早于自己，若是，则构造一个 cancelCtx 返回即可；
        return WithCancel(parent)
    }
    c := &timerCtx{
        cancelCtx: newCancelCtx(parent),
        deadline:  d,
    }
    propagateCancel(parent, c) //启动守护方法，同步 parent 的 cancel 事件到子 context；
    dur := time.Until(d)
    if dur <= 0 {
//判断过期时间已到，直接 cancel timerCtx，并返回 DeadlineExceeded 的错误；
        c.cancel(true, DeadlineExceeded) // deadline has already passed
        return c, func() { c.cancel(false, Canceled) }
    }
    c.mu.Lock()
    defer c.mu.Unlock()
    if c.err == nil {
        c.timer = time.AfterFunc(dur, func() {
             //启动 time.Timer，设定一个延时时间，即达到过期时间后会终止该 timerCtx，并返回 DeadlineExceeded 的错误；
            c.cancel(true, DeadlineExceeded)
        })
    }
    return c, func() { c.cancel(true, Canceled) }
}

5 valueCtx

5.1 数据结构

type valueCtx struct {
    Context
    key, val any
}

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

5.2 valueCtx.Value()

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

假如我这没这个 key，则从 parent 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 *timerCtx:
            if key == &cancelCtxKey {
                return &ctx.cancelCtx
            }
            c = ctx.Context
        case *emptyCtx:
            // 最后也找不着你！
            return nil
        default:
            return c.Value(key)
        }
    }
}

5.3 valueCtx 用法小结

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

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

5.4 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() {
        // 从这可看出，只要是可比较的类型，都可以放入 ctx 做 key
        panic("key is not comparable")
    }
    return &valueCtx{parent, key, val}
}

Adopting CC BY-NC-SA 4.0 License Agreement