go context 源码分析

WithCancel

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

// 包了一层，把 parent 赋值给子 context
func newCancelCtx(parent Context) cancelCtx {
    return cancelCtx{Context: parent}
}

type cancelCtx struct {
    Context

    mu       sync.Mutex            // protects following fields
    done     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
}

func propagateCancel(parent Context, child canceler) {
    if parent.Done() == nil {
        return // parent is never canceled
    }
  // 判断 parent 是否为 cancelCtx，timerCtx 也属于 cancelCtx，cancelCtx 有 child 
    if p, ok := parentCancelCtx(parent); ok {
        p.mu.Lock()
        if p.err != nil {
      // cancelCtx 被 cancel() 的时候，会给err属性赋值
            // parent has already been canceled
            child.cancel(false, p.err)
        } else { // parent 没有被 cancel()
            if p.children == nil {
                p.children = make(map[canceler]struct{})
            }
            p.children[child] = struct{}{} // 将 child 添加到 parent 中
        }
        p.mu.Unlock()
    } else { // parent 的状态未肯定 
        go func() {
            select {
            case <-parent.Done(): // 若是能够从 parent 中获取到值，说明 parent 被 cancel 了
                child.cancel(false, parent.Err())
            case <-child.Done():
            }
        }()
    }
}

// 类型判断
func parentCancelCtx(parent Context) (*cancelCtx, bool) {
    for {
        switch c := parent.(type) {
        case *cancelCtx:
            return c, true
        case *timerCtx:
            return &c.cancelCtx, true
        case *valueCtx:
            parent = c.Context
        default:
            return nil, false
        }
    }
}

分析下 WithCancel 返回的函数

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

func (c *cancelCtx) cancel(removeFromParent bool, err error) {
  // 被 cancel 的 cancelCtx 的 err 属性，必须被赋值
    if err == nil {
        panic("context: internal error: missing cancel error")
    }
    c.mu.Lock()
    if c.err != nil {
        c.mu.Unlock()
        return // 已经被 cancel
    }
    c.err = err
    if c.done == nil {
        c.done = closedchan // 注意此处！！！
    } else {
        close(c.done) // close 后，select 能够不断获取到默认值
    }
  // close 掉全部 child
    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 {
        removeChild(c.Context, c)
    }
}

func removeChild(parent Context, child canceler) {
  // 只有 cancelCtx 类型的 context，才有 child
    p, ok := parentCancelCtx(parent)
    if !ok {
        return
    }
    p.mu.Lock()
    if p.children != nil {
        delete(p.children, child)
    }
    p.mu.Unlock()
}


// closedchan is a reusable closed channel.
// closedchan init() 的时候就被 close 了，因此是能够经过 select 不断获取值的
var closedchan = make(chan struct{})
func init() {
    close(closedchan)
}