golang源码学习之context

本文不对context的具体用法讲解，仅对源码走读。

context一共分为6类：Background、TODO、WithCancel、WithDeadline、WithTimeout、WithValue，下面依次走读。

数据结构

// context/context.go
type Context interface {
    // 返回deadline和true, ok写死为true
    Deadline() (deadline time.Time, ok bool)
    // 用于cancel监听
    Done() <-chan struct{}
    // 返回错误消息
    Err() error
    // 获取WithValue存储的value
    Value(key interface{}) interface{}
}

Background、TODO

// context/context.go
var (
    background = new(emptyCtx)
    todo       = new(emptyCtx)
)

func Background() Context {
    return background
}

func TODO() Context {
    return todo
}

可以看出Background和TODO其实是一样的

// context/context.go
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 interface{}) interface{} {
    return nil
}

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

emptyCtx 的各个方法都返回的是零值。注意Done()返回nil,下面会利用它判断是否是Background和TODO。

WithCancel

// context/context.go
type cancelCtx struct {
    // 这里指的是parent
    Context

    mu       sync.Mutex            // protects following fields
    // 其实就是Done()返回值
    done     chan struct{}         // created lazily, closed by first cancel call
    // child集合
    children map[canceler]struct{} // set to nil by the first cancel call
    //错误消息
    err      error                 // set to non-nil by the first cancel call
}

func WithCancel(parent Context) (ctx Context, cancel CancelFunc) {
    // 创建cancelCtx
    c := newCancelCtx(parent)
    propagateCancel(parent, &c)
    return &c, func() { c.cancel(true, Canceled) }
}

func propagateCancel(parent Context, child canceler) {
    // 不具备cancel功能, Done()其实是返回一个叫done的成员变量， 而done是在cancel()里面赋值的
    // 这里指的是context.Background()、context.TODO()
    if parent.Done() == nil {
        return // parent is never canceled
    }
    // parent是否属于cancelCtx
    if p, ok := parentCancelCtx(parent); ok {
        p.mu.Lock()
        // parent已取消
        if p.err != nil {
            // parent has already been canceled
            child.cancel(false, p.err) // parent取消了，child直接执行取消
        } else {
            if p.children == nil {
                p.children = make(map[canceler]struct{})
            }
            // 如果parent没有取消，那么就将child加入到parent的children集合中
            p.children[child] = struct{}{}
        }
        p.mu.Unlock()
    } else { // 如果parent不属于cancelCtx
        go func() {
            select {
            case <-parent.Done():
                child.cancel(false, parent.Err()) // parent被cancel时，child执行cancel
            case <-child.Done():
            }
        }()
    }
}

// parent是否属于cancelCtx。
func parentCancelCtx(parent Context) (*cancelCtx, bool) {
    for {
        switch c := parent.(type) {
        case *cancelCtx: // 这里指context.WithCancel
            return c, true
        case *timerCtx: // 这里指context.WithDeadline、context.WithTimeout
            return &c.cancelCtx, true
        case *valueCtx: // 这里指 context.WithValue
            parent = c.Context
        default:
            return nil, false
        }
    }
}


func (c *cancelCtx) cancel(removeFromParent bool, err error) {
    if err == nil {
        panic("context: internal error: missing cancel error")
    }
    c.mu.Lock()
    // 多次cancel得到的err是一样的
    if c.err != nil {
        c.mu.Unlock()
        return // already canceled
    }
    c.err = err
    if c.done == nil {
        c.done = closedchan
    } else {
        close(c.done) //关闭channel
    }

    // 遍历孩子执行取消操作。
    for child := range c.children {
        // NOTE: acquiring the child's lock while holding parent's lock.
        //孩子执行取消操作, 为什么这里传false呢？ 应为我们在第一次做cancel时removeFromParent=true已经脱离parant了
        child.cancel(false, err)
    }
    c.children = nil
    c.mu.Unlock()

    if removeFromParent {
        // 这里的c.Context指的是parent
        removeChild(c.Context, c)
    }
}


func removeChild(parent Context, child canceler) {
    p, ok := parentCancelCtx(parent)

    // 因为我们是在parent属于cancelCtx的时候才将child加入的，所以这里直接return
    if !ok {
        return
    }
    p.mu.Lock()
    if p.children != nil {
        delete(p.children, child) // 移除child
    }
    p.mu.Unlock()
}

注意：只有cancelCtx才会有child,这里的cancelCtx包括WithCancel、WithDeadline、WithTimeout。

多次cancel得到的err是一样的。

context创建流程图

context

WithDeadline

// context/context.go

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

    deadline time.Time
}

func WithDeadline(parent Context, d time.Time) (Context, CancelFunc) {

    // 当parent为timerCtx 而且 deadline也比parent晚。那么就以WithCancel类型加入parent。
    // 因为child没有机会等到deadline。（timerCtx就比cancelCtx多了个时间）
    if cur, ok := parent.Deadline(); ok && cur.Before(d) {
        // The current deadline is already sooner than the new one.
        return WithCancel(parent)
    }
    c := &timerCtx{
        cancelCtx: newCancelCtx(parent),
        deadline:  d,
    }
    propagateCancel(parent, c)
    dur := time.Until(d)
    // 早于当前时间， 说明已经过了deadline, 直接执行cancel
    if dur <= 0 {
        c.cancel(true, DeadlineExceeded) // deadline has already passed
        return c, func() { c.cancel(true, Canceled) }
    }
    c.mu.Lock()
    defer c.mu.Unlock()
    if c.err == nil {
        // 经过dur后自动执行cancel
        c.timer = time.AfterFunc(dur, func() {
            c.cancel(true, DeadlineExceeded)
        })
    }
    return c, func() { c.cancel(true, Canceled) }
}

// ok 为固定的true
func (c *timerCtx) Deadline() (deadline time.Time, ok bool) {
    return c.deadline, true
}

流程见上图

timerCtx其实就在cancelCtx上面添加了个时间，它还是属于cancelCtx。

WithTimeout

// context/context.go
func WithTimeout(parent Context, timeout time.Duration) (Context, CancelFunc) {
    // WithTimeout其实也就是WithDeadline
    return WithDeadline(parent, time.Now().Add(timeout))
}

WithTimeout其实调用的就是WithDeadline

WithValue

// context/context.go
type valueCtx struct {
    Context
    key, val interface{}
}

func WithValue(parent Context, key, val interface{}) Context {
    if key == nil {
        panic("nil key")
    }
    if !reflect.TypeOf(key).Comparable() {
        panic("key is not comparable")
    }
    // 只是简单存储 k v
    return &valueCtx{parent, key, val}
}

func (c *valueCtx) Value(key interface{}) interface{} {
    if c.key == key {
        return c.val
    }
    // 如果本地获取不到会到parent中获取
    return c.Context.Value(key)
}

只是简单存储了K、V

<-ctx.Done()

全文没看见对cancelCtx的done写消息， 那为什么在cancel的时候select{ case <-ctx.Done(): }没有堵塞呢？

func init() {
    close(closedchan)
}

func (c *cancelCtx) Done() <-chan struct{} {
    c.mu.Lock()
    if c.done == nil {
        c.done = make(chan struct{})
    }
    d := c.done
    c.mu.Unlock()
    return d
}

func (c *cancelCtx) cancel(removeFromParent bool, err error) {
    
    ...
    if c.done == nil {
        c.done = closedchan
    } else {
        close(c.done) //关闭channel
    }
    ...

}

closedchan在init时候就进行了close。cancel的时候c.done总是处于close状态。
如果先<-ctx.Done()， 那么堵塞
如果先cancel， 那么<-ctx.Done()可执行， 并且返回的是false，说明chan是关闭的

结尾： context的代码整体来说还是比较易读的，但我肯定是写不出来啦~~~~