基于 Go 1.18 源码分析

一、引言

1.1 什么是 Context?

ContextGo 1.7引入的一个标准库,官方 blog 里面介绍,最早是Google内部使用的一个库,主要用于在一个Request对应的多个Goroutine中传递数据,数据主要分为两种:

  1. 请求的基本信息,比如用户鉴权信息、请求的Request-ID等等。
  2. 请求的Deadline,如果请求被Cancel或者Timeout,能够控制多个Goroutine会返回。

整个 context.go 加上注释也就600行左右。核心就是Context type

type Context interface {
    // 获取 DeadLine 时间,使用 WithDeadline 和 WithTimeout 才有
    Deadline() (deadline time.Time, ok bool)

    // 返回一个代表context完成的管道,若是context无法关闭,done返回nil
    // WithCancel 安排 Done 在调用 cancel 时关闭;
    // WithDeadline 安排 Done 在截止日期到期时关闭; WithTimeout 安排 Done 在超时过后关闭。
    Done() <-chan struct{}

    // 若 done 没有关闭,err 返回 nil
    // 若 done is closed。如果是 cancel 就返回 Canceled => "context canceled"
    // 如果是超过 deadline 就返回 DeadlineExceeded => "context deadline exceeded"
    Err() error

    // 读取数据
    Value(key any) any
}

1.2 如何创建 Context

Go内置两个函数Background()TODO()用于创建Context

  • Background()是上下文的默认值,所有其他的上下文都应该从它衍生出来;
  • TODO() 应该仅在不确定应该使用哪种上下文时使用;

底层都是emptyCtx,本质没什么区别,不过一些代码检查的工具会检查是否有TODO函数。

type emptyCtx int

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

func Background() Context {
    return background
}

func TODO() Context {
    return todo
}

1.3 Conetext 的派生

Conetext 可以通过WithXXX来生成新的Context ,主要有4个函数来设置。

func WithCancel(parent Context) (ctx Context, cancel CancelFunc)
func WithDeadline(parent Context, deadline time.Time) (Context, CancelFunc)
func WithTimeout(parent Context, timeout time.Duration) (Context, CancelFunc)
func WithValue(parent Context, key, val interface{}) Context

WithCancel

WithCancel返回带有新done通道的Context

ctx0, cancel := context.WithCancel(context.Background())
deadline, ok := ctx0.Deadline()
fmt.Println(deadline, ok) // 0001-01-01 00:00:00 +0000 UTC ,  false
fmt.Println(ctx0.Err())   // nil

go func() {
    <-ctx0.Done()
    fmt.Println(ctx0.Err()) // context canceled
}()

cancel() // 结束以后 ctx0.Done() 变为可读状态

WithDeadline

返回一个带deadlineContext,如果父节点也有deadline,当前Contextdeadline以最先发生的情况为准,因为父节点Cancel的时候也会调用子节点Cancel

d := time.Now().Add(100000 * time.Second)
d1 := time.Now().Add(1 * time.Second)


ctx0, cancel := context.WithDeadline(context.Background(), d)
ctx1, cancel := context.WithDeadline(ctx0, d1)

// ctx1 以最先发生的情况为准 min(d, d1) 时间为 dealline

select {
case <-ctx0.Done():
    fmt.Println("ctx0 done : ", time.Now())
    fmt.Println(ctx0.Err()) // context deadline exceeded
case <-ctx1.Done(): 
    fmt.Println("ctx1 done : ", time.Now())
    fmt.Println(ctx0.Err()) // context deadline exceeded
}

cancel()                // 上面已经 done 了,这里执行 cancel 已经没意义了。
fmt.Println(ctx0.Err()) // context deadline exceeded

WithTimeout

WithTimeout底层就是调用WithDeadline,不在过多解释,也是以最先发生的情况为准

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

WithValue

WithValue用法比较简单

ctx := context.Background()
ctx1 := context.WithValue(ctx, "k", "v")
fmt.Println(ctx1.Value("k")) // v

二、Context 底层实现

2.1 底层依赖关系

image.png

2.2 emptyCtx

emptyCtx 是一个int新类型,空实现了Context的所有接口,主要给Background()TODO()使用,没有啥好说的。valueCtx也复用了emptyCtx除了Value的所有方法。

// An emptyCtx is never canceled, has no values, and has no deadline. It is not
// struct{}, since vars of this type must have distinct addresses.
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
}

func (e *emptyCtx) String() string {
    switch e {
    case background:
        return "context.Background"
    case todo:
        return "context.TODO"
    }
    return "unknown empty Context"
}

2.2 valueCtx

valueCtx有一个key和一个value来存储数据。

// A valueCtx carries a key-value pair. It implements Value for that key and
// delegates all other calls to the embedded Context.
type valueCtx struct {
    Context
    key, val any
}

调用WithValue会返回一个新的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}
}

image.png

关于Comparable这个里多说一下,valueCtxkey要是可比较(Comparable)的,就是支持key1 == key2这种写法。GoSliceMapFuunction都不支持==比较,

但是可以用reflect.DeepEqual做部分比较

这里的keyany,所以c.key == key其实是调用runtime.efaceeq来比较是否相等,这种需要对象的type都相等才可以。具体汇编点我

如果用String作为Key,可能导致被覆盖,但是用自定义的Struct就没有这个问题。

type myPrivateKey struct {}
ctx = context.WithValue(ctx, myPrivateKey{}, "abc")
ctx.Value(myPrivateKey{})

再来看下读取数据代码。其实就是一个循环,从下往上找到key相等的数据,然后返回。这个时间复杂度是O(n),所以往ctx里面塞了很多数据的话,读取速度会慢。这样设计WithValue的好处是,这样是并发安全的。

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 *timerCtx:
            if key == &cancelCtxKey {
                return &ctx.cancelCtx
            }
            c = ctx.Context
        case *emptyCtx:
            return nil
        default:
            return c.Value(key)
        }
    }
}

2.3 cancelCtx

cancelCtx 结构体如下,mu是保护并发设置childrenerr两个字段的,done是一个channel,在调用cancel的时候,业务方可以通过done感知到是否调用了Cancel()

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
}

创建 cancelCtx

WithCancel 会返回一个&cancelCtx和一个CancelFunc

type CancelFunc func()

func WithCancel(parent Context) (ctx Context, cancel CancelFunc) {
    if parent == nil {
        panic("cannot create context from nil parent")
    }
    c := newCancelCtx(parent) // new一个cancelCtx{}
    propagateCancel(parent, &c)
    return &c, func() { c.cancel(true, Canceled) }
}

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

在来看下propagateCancel这个函数

// propagateCancel arranges for child to be canceled when parent is.
// 这个函数是从 parent 往上找,看有没有 cancelCtx 或者 timeCtx
// 有的话把当前的 cancelCtx 加到父节点的 parent 里面去
func propagateCancel(parent Context, child canceler) {
    done := parent.Done()
    if done == nil { 
        // 如果 parent 是 emptyCtx 或者 valueCtx 就直接返回
        return // parent is never canceled
    }

    // 这里面检查一下父节点是不是已经 cancel 了
    // 如果 cancel 了,当前节点也需要 cancel 掉
    // 否则走 default 继续向下
    select {
    case <-done:
        // parent is already canceled
        child.cancel(false, parent.Err())
        return
    default:
    }

    // parentCancelCtx 就是向上递归找到 一个 cancelCtx 或者是 timeCtx.cancelCtx
    // ok 等于 true 表示递归找到了一个cancelCtx 的父节点,且这个节点没有 cancel 
    if p, ok := parentCancelCtx(parent); ok {
        p.mu.Lock()
        if p.err != nil {
            // 这里表示 父节点在加锁以后,被 Cancel 了,当前的 ctx 的 done 也要置为 cancel
            child.cancel(false, p.err)
        } else {
           // 父节点没有 cancel, 要把当前结点加到父节点的 children 中
           // 这样父节点 cancel 的时候就可以通知下面所有子节点去 cancel 掉
            if p.children == nil {
                p.children = make(map[canceler]struct{})
            }
            p.children[child] = struct{}{}
        }
        p.mu.Unlock()
    } else {
        // 走到这有两种情况。
        // 1. 自己实现了一个 Context,Done() != nil
        // 2. 在调用 parentCancelCtx 的瞬间, done 被 close了,这个时候 ok 也是false
        atomic.AddInt32(&goroutines, +1)
        go func() {
            select {
            case <-parent.Done():
                child.cancel(false, parent.Err())
            case <-child.Done():
            }
        }()
    }
}


func parentCancelCtx(parent Context) (*cancelCtx, bool) {
    done := parent.Done()
    if done == closedchan || done == nil {
        return nil, false
    }
    // parent.Value(&cancelCtxKey) 就是递归向上查到 cancelCtx 和 timerCtx. cancelCtx
    // 如果一直没有 cancelCtx 或者 timerCtx,最终返回为 nil
    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
}

读取数据

anyCtx.Value(&cancelCtxKey) 就是递归向上查到 cancelCtx 或者 timerCtx.cancelCtx,返回值类型是cancelCtx或者没查到就是nil

func (c *cancelCtx) Value(key any) any {
    if key == &cancelCtxKey {
        return c
    }
    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 *timerCtx:
            if key == &cancelCtxKey {
                return &ctx.cancelCtx
            }
            c = ctx.Context
        case *emptyCtx:
            return nil
        default:
            return c.Value(key)
        }
    }
}

Done()

Done() 比较简单,就是判断done是否为空,为空的话,就创建chan然后返回。

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

Cancel()

WithCancel返回的是一个func() { c.cancel(true, Canceled) },我们再看看cancel具体执行代码

func (c *cancelCtx) cancel(removeFromParent bool, err error) {
    if err == nil {
        panic("context: internal error: missing cancel error")
    }
    c.mu.Lock() // 加锁,设置 err ,还要读 children
    if c.err != nil {
        c.mu.Unlock()
        return // already canceled
    }
    c.err = err // 如果是 Cancel调用的话,这个就是 Canceled error
    d, _ := c.done.Load().(chan struct{})
    if d == nil { // 没有调用过 `Done()`,就执行了`Cancel`
        c.done.Store(closedchan)
    } else {
        close(d) // 关闭以后, <- ctx.Done() 就会返回
    }
    
    // 把子节点的所有 ctx 都取消掉
    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 { 
        // 把自己从父节点的 children 移除掉。
        // 因为父节点 Cancel 的时候,已经不需要再cancel这个节点了
        removeChild(c.Context, c)
    }
}

image.png

2.4 timerCtx

timerCtx复用了cancelCtx的大部分能力,然后多了一个deadline和一个TimertimerCtx 结构体如下:

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

创建 timerCtx

通过WithTimeoutWithDeadline都可以创建timerCtx

func WithTimeout(parent Context, timeout time.Duration) (Context, CancelFunc) {
    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")
    }
    
    // 如果 父节点的 Deadline 在 d 之前,那就不用再设置了
    // 直接调用 WithCancel 把当前节点挂到父节点之上就行了。
    if cur, ok := parent.Deadline(); ok && cur.Before(d) {
        // The current deadline is already sooner than the new one.
        return WithCancel(parent)
    }
    
    // new 一个 timerCtx , 然后也 new 一个 cancelCtx
    c := &timerCtx{
        cancelCtx: newCancelCtx(parent),
        deadline:  d,
    }
    
    // 向上递归找到类型为 cancelCtx 的父节点,把自己设置到 children 里面去
    propagateCancel(parent, c)
    dur := time.Until(d) // 算出还有多久到 deadline
    if dur <= 0 { // 已经到了,直接 cancel
        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 {
       // new 一个 Timer,过来 dur 时间去执行 cancel
        c.timer = time.AfterFunc(dur, func() {
            c.cancel(true, DeadlineExceeded)
        })
    }
    
    // 返回 ctx 和 cancelFun,这样可以在时间没到的时候,自己主动 cancel
    return c, func() { c.cancel(true, Canceled) }
}

DeadLine()

直接返回当前设置的deadline没啥好说的

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

cancel()

cancel就是调用cancelCtx.cancel,然后再关闭 timer,也没啥好说的。

func (c *timerCtx) cancel(removeFromParent bool, err error) {
    c.cancelCtx.cancel(false, err)
    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()
}

其他的func都是复用的cancelCtx的。

三、总结

平常其实用的最多还是valueCtxcancelCtxtimeCtx用的场景不是那么多。父节点cancel以后,所有子节点的 ctx 也都被cancel这个特性,新手刚刚开始用的时候很容易以为是Gobug