context源码学习

特别提醒，本文所涉及的源码是go1.21.0 darwin/amd64

1. 什么是Context

在多个goroutine之间传递上下文的对象，传递的信息包括取消信号、截止时间以及其他一些跨api边界的值

2. 核心数据结构

2.1 Context

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

Context定义为Interface，定义了4个核心函数：

• Dealine：返回context的过期时间和一个bool值判断是否设置了deadline
• Done：返回context中的channel，该channel会在当前工作完成或者上下文被取消后关闭，多次调用返回的是同一个channel，如果该context是不能被取消的，则会返回nil
• Err：返回context结束的原因
• Value：返回contxet中存储的key对应的值

2.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超时时会返回该类错误

3 具体实现

3.1 emptyCtx

type emptyCtx struct{}

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，之前的一些版本中是直接将它定义为一个int类型

3.2 backgroundCtx

type backgroundCtx struct{ emptyCtx }

func (backgroundCtx) String() string {
    return "context.Background"
}

func Background() Context {
    return backgroundCtx{}
}

Background()返回一个非nil但是为空的context。它不能被取消，没有任何值，也没有截止时间。通常被主函数、初始化、测试以及作为传入请求的顶层context使用。

3.3 todoCtx

type todoCtx struct{ emptyCtx }

func (todoCtx) String() string {
    return "context.TODO"
}

func TODO() Context {
    return todoCtx{}
}

TODO()同样返回一个非nil但是为空的context。在代码不清楚应该使用哪个上下文或者上下文尚不可用时，应该使用context.TODO

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

cancelCtx是取消机制的关键

• 其结构体嵌入了一个`Context，作为其父类

• mu：内置一把锁，用于多goroutine并发场景下的资源获取

• done：原子类型的值，实际类型为chan struct{}，用来反应cancelCtx的生命周期
• children：一个集合，保存了所有该cancelCtx的子context

• err：记录cancelCtx发生的错误

• cause：记录calcelCtx被取消的原因

3.4.1 Value方法

// &cancelCtxKey is the key that a cancelCtx returns itself for.
var cancelCtxKey int

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

如果key的值为&cancelCtxKey，则返回cancelCtx本身的指针，否则返回对应key存储的value值

3.4.2 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{})
}

• 前面提到done属性是一个原子类型的值,因此通过atomic包中的Load()函数获取它的值,如果它已经存在则直接返回
• 如果不存在则先加锁，然后创建一个新的chan struct{}，然后通过Store()赋值给done，可以看到这是一个懒加载机制，在第一次调用c.Done()的时候，该属性才被创建，返回值，然后解锁

3.4.3 Err方法

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

加锁、读取值、解锁、返回结果

3.4.4 propagateCancel方法

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():
        }
    }()
}

该方法比较关键，会构建父子上下文之间的关联，当父上下文被取消时，子上下文也会被取消，不会出现状态不一致的情况

• 首先将parent赋值给Context属性
• 如果父上下文不会取消，则直接返回
• 如果父上下文已经取消，则直接终止子上下文
• 如果父上下文是可以取消的上下文类型，即cancelCtx，则先加锁、然后将其加入到父上下文的children中，最后解锁
• 如果父上下文实现了AfterFunc方法，则在上下文被取消后，将子上下文也取消，通过调用一个stop()函数来实现的
• 如果都没满足前面的条件，则启动一个协程监控parent状态，倘若父上下文终止，则终止子上下文

进一步观察parentCancelCtx是如何校验parent是否为cancelCtx类型的

var closedchan = make(chan struct{})
func init() {
    close(closedchan)
}

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
}

• closedchan表示一个已经关闭的通道
• 倘若parent的channel已经被关闭或者是不会被cancel的类型，则直接返回false
• 倘若用&cancelCtxKey能取到值并且得到的值是parent本身,返回true（cancelCtx的约定，key为&cancelCtxKey的时候，返回的是cancelCtx本身的指针）

3.4.5 String方法

func (c *cancelCtx) String() string {
    return contextName(c.Context) + ".WithCancel"
}

返回cancelCtx的名字

3.4.6 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)
    }
}

• 该方法接收三个参数，removeFromParent是一个bool值，表示需要从当前context的父context的children中删除
• 加锁，判断自己的err是否为空，如果不为空，说明已经被cancel，直接返回
• 处理 cancelCtx 的 channel，若channel 此前未初始化，则直接注入一个 closedChan，否则关闭该 channel
• 遍历该cancelCtx的子context，依次将其cancel
• 解锁，最后根据传入的参数removeFromParent，判断是否需要手动把 cancelCtx 从 parent 的children 集合中移除

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

func removeChild(parent Context, child canceler) {
    if s, ok := parent.(stopCtx); ok {
        s.stop()
        return
    }
    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集合）
• 加锁，将集合中key为child的删除，解锁

3.5 timerCtx

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

    deadline time.Time
}

timerCtx在cancelCtx的基础上，又做了一层封装，新增了一个time.Timer用于定时终止context，另外新增一个deadline用于记录过期时间

3.5.1 Deadline方法

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

返回过期时间

3.5.2 String方法

func (c *timerCtx) String() string {
    return contextName(c.cancelCtx.Context) + ".WithDeadline(" +
        c.deadline.String() + " [" +
        time.Until(c.deadline).String() + "])"
}

返回context的名字加上过期时间

3.5.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方法，进行处理
• 判断是否需要从 parent 的 children 集合 中移除，若是则进行处理
• 加锁、停止time.Timer、解锁返回

3.6 valueCtx

type valueCtx struct {
    Context
    key, val any
}

valueCtx同样继承了Context，同时一个valueCtx只能存储一对kv

3.6.1 String方法

func (c *valueCtx) String() string {
    return contextName(c.Context) + ".WithValue(type " +
        reflectlite.TypeOf(c.key).String() +
        ", val " + stringify(c.val) + ")"
}

返回valueCtx的名字以及存储的kv值

3.6.2 Value方法

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

• 如果参数的key等于他所保存的key则返回所保存的value
• 否则去parent的Context去找

进一步观察是如何在parent中寻找的

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

• 启动一个 for 循环，一直往上去父上下文中寻找
• 其中不同的上下文类型会有不同的处理方式；
• 找到匹配的 key，则将该value 进行返回.

4 创建context

4.1 cancelCtx

有两种方法，一个不带cause，一个带有cause

func WithCancel(parent Context) (ctx Context, cancel CancelFunc) {
    c := withCancel(parent)
    return c, func() { c.cancel(true, Canceled, nil) }
}

func WithCancelCause(parent Context) (ctx Context, cancel CancelCauseFunc) {
    c := withCancel(parent)
    return c, func(cause error) { c.cancel(true, Canceled, cause) }
}

4.2 withoutCancelCtx

func WithoutCancel(parent Context) Context {
    if parent == nil {
        panic("cannot create context from nil parent")
    }
    return withoutCancelCtx{parent}
}

4.3 Deadline

同样有两种方法，一个不带cause，一个带有cause

func WithDeadline(parent Context, d time.Time) (Context, CancelFunc) {
    return WithDeadlineCause(parent, d, nil)
}

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

4.4 Timeout

同样有两种方法，一个不带cause，一个带有cause，最后都是通过WithDeadlineCause来创建的

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

func WithTimeoutCause(parent Context, timeout time.Duration, cause error) (Context, CancelFunc) {
    return WithDeadlineCause(parent, time.Now().Add(timeout), cause)
}

4.5 ValueCtx

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