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