Golang 笔记（二）：Context 源码剖析（2）

// A canceler is a context type that can be canceled directly. The
// implementations are *cancelCtx and *timerCtx.
type canceler interface {
    cancel(removeFromParent bool, err error)
   Done() <-chan struct{}
}

func propagateCancel(parent Context, child canceler) {
 2    done := parent.Done()
 3    if done == nil {
 4        return // 父节点不可取消
 5    }
 6
 7    select {
 8    case <-done:
 9        // 父节点已经取消
10        child.cancel(false, parent.Err())
11        return
12    default:
13    }
14
15    if p, ok := parentCancelCtx(parent); ok { // 找到一个 cancelCtx 实例
16        p.mu.Lock()
17        if p.err != nil {
18            // 父节点已经被取消
19            child.cancel(false, p.err)
20        } else {
21            if p.children == nil {
22                p.children = make(map[canceler]struct{}) // 惰式创建
23            }
24            p.children[child] = struct{}{}
25        }
26        p.mu.Unlock()
27    } else {                                // 找到一个非 cancelCtx 实例
28        atomic.AddInt32(&goroutines, +1)
29        go func() {
30            select {
31            case <-parent.Done():
32                child.cancel(false, parent.Err())
33            case <-child.Done(): 
34            }
35        }()
36    }
37}

func (c *cancelCtx) cancel(removeFromParent bool, err error) {
 2    if err == nil { // 需要给定取消的理由，Canceled or DeadlineExceeded
 3        panic("context: internal error: missing cancel error")
 4    }
 5
 6    c.mu.Lock()
 7    if c.err != nil {
 8        c.mu.Unlock()
 9        return // 已经被其他 goroutine 取消
10    }
11
12    // 记下错误，并关闭 done
13    c.err = err
14    if c.done == nil {
15        c.done = closedchan
16    } else {
17        close(c.done)
18    }
19
20    // 级联取消
21    for child := range c.children {
22        // NOTE: 持有父 Context 的同时获取了子 Context 的锁
23        child.cancel(false, err)
24    }
25    c.children = nil
26    c.mu.Unlock()
27
28    // 子树根需要摘除，子树中其他节点则不再需要
29    if removeFromParent {
30        removeChild(c.Context, c)
31    }
32}

type timerCtx struct {
 2    cancelCtx
 3    timer *time.Timer // Under cancelCtx.mu
 4
 5    deadline time.Time
 6}
 7
 8func (c *timerCtx) Deadline() (deadline time.Time, ok bool) {
 9    return c.deadline, true
10}
11
12func (c *timerCtx) cancel(removeFromParent bool, err error) {
13  // 级联取消子树中所有 Context
14    c.cancelCtx.cancel(false, err)
15
16    if removeFromParent {
17    // 单独调用以摘除此节点，因为是摘除 c，而非 c.cancelCtx
18        removeChild(c.cancelCtx.Context, c)
19    }
20
21  // 关闭计时器
22    c.mu.Lock()
23    if c.timer != nil {
24        c.timer.Stop()
25        c.timer = nil
26    }
27    c.mu.Unlock()
28}

func WithDeadline(parent Context, d time.Time) (Context, CancelFunc) {
 2    if cur, ok := parent.Deadline(); ok && cur.Before(d) {
 3        // 祖先节点的时限更早
 4        return WithCancel(parent)
 5    }
 6
 7    c := &timerCtx{             
 8        cancelCtx: newCancelCtx(parent), // 使用一个新的 cancelCtx 实现部分 cancel 功能
 9        deadline:  d,
10    }
11    propagateCancel(parent, c) // 构建 Context 取消树，注意传入的是 c 而非 c.cancelCtx
12    dur := time.Until(d)       // 测试时限是否设的太近以至于已经结束了
13    if dur <= 0 {
14        c.cancel(true, DeadlineExceeded)
15        return c, func() { c.cancel(false, Canceled) }
16    }
17
18  // 设置超时取消
19    c.mu.Lock()
20    defer c.mu.Unlock()
21    if c.err == nil {
22        c.timer = time.AfterFunc(dur, func() {
23            c.cancel(true, DeadlineExceeded)
24        })
25    }
26    return c, func() { c.cancel(true, Canceled) }
27}

for ; ; time.Sleep(time.Second) {
2        select {
3        case <-context.Done():
4            return
5        default:
6        }
7
8      // 一些耗时操作
9}