1
type Context interface {
2
    // Deadline 返回 context 被取消的时间
3
    // 返回 false 代表无 deadline，不会自动取消
4
    Deadline() (deadline time.Time, ok bool)
5

6
    // Done 返回一个 channel，当 context 被取消，即工作完成时，会关闭这个 channel
7
    // 若 context 不会被取消，Done 返回 nil
8
    // Done 关闭可能会与 cancel 函数的返回异步发生
9
    Done() <-chan struct{}
10

11
    // 在 Done 被关闭后，返回 context 的取消原因
12
    Err() error
13

14
    // Value 返回 context 中对应键中存储的值，若键不存在则返回 nil
15
    Value(key any) any
16
}

1
type cancelCtx struct {
2
     Context // 嵌入其父 ctx
3
     mu       sync.Mutex            // 互斥锁，保证线程安全
4
     done     atomic.Value          // 存储一个 chan struct{}, 懒加载，在调用 cancel() 时关闭
5
     children map[canceler]struct{} // 子 ctx 的集合，在 cancel() 被调用时设为 nil
6
     err      error                 // 在 cancel() 被调用时设置
7
     cause    error                 // 取消原因
8
}

1
// cancelCtxKey 是一个 int 指针类型的全局变量，用于判断当前 ctx 是否为 cancelCtx类型
2
// 返回最早的 cancelCtx
3
func (c *cancelCtx) Value(key any) any {
4
    // 返回自身
5
    if key == &cancelCtxKey {
6
        return c
7
    }
8
    return value(c.Context, key)
9
}
10

11
// 当找不到 &cancelCtxKey 时会继续调用 value(c.Context, key) 向上寻找
12
func value(c Context, key any) any {
13
    for {
14
        switch ctx := c.(type) {
15
        case *valueCtx:
16
            if key == ctx.key {
17
                return ctx.val
18
            }
19
            c = ctx.Context
20
        case *cancelCtx:
21
            if key == &cancelCtxKey {
22
                return c
23
            }
24
            c = ctx.Context
25
        case withoutCancelCtx:
26
            if key == &cancelCtxKey {
27
                return nil
28
            }
29
            c = ctx.c
30
        case *timerCtx:
31
            if key == &cancelCtxKey {
32
                return &ctx.cancelCtx
33
            }
34
            c = ctx.Context
35
        case backgroundCtx, todoCtx:
36
            return nil
37
        default:
38
            return c.Value(key)
39
        }
40
    }
41
}
42

43
// 懒创建并返回 done
44
func (c *cancelCtx) Done() <-chan struct{} {
45
    d := c.done.Load()
46
    if d != nil {
47
        return d.(chan struct{})
48
    }
49
    c.mu.Lock()
50
    defer c.mu.Unlock()
51
    d = c.done.Load()
52
    if d == nil {
53
        d = make(chan struct{})
54
        c.done.Store(d)
55
    }
56
    return d.(chan struct{})
57
}
58

59
func (c *cancelCtx) Err() error {
60
    c.mu.Lock()
61
    err := c.err
62
    c.mu.Unlock()
63
    return err
64
}

1
// cancelCtx 和 timerCtx 都实现了 canceler 接口
2
type canceler interface {
3
    cancel(removeFromParent bool, err, cause error)
4
    Done() <-chan struct{}
5
}

1
// WithCancel 返回一个带有新 Done channel 的派生 context
2
// 当父 context 被取消或当前 context 被取消时，Done 都会被关闭
3
func WithCancel(parent Context) (ctx Context, cancel CancelFunc) {
4
    c := withCancel(parent)
5
    return c, func() { c.cancel(true, Canceled, nil) }
6
}
7

8
func withCancel(parent Context) *cancelCtx {
9
    if parent == nil {
10
        panic("cannot create context from nil parent")
11
    }
12
    c := &cancelCtx{}
13
    c.propagateCancel(parent, c)        // 建立父子关系
14
    return c
15
}

1
// CancelCauseFunc 相比 CancelFunc 增加了一个取消原因，使用时需要传入一个 error 参数
2
// 可以调用 Cause 来获取具体的原因
3
type CancelCauseFunc func(cause error)
4

5
func WithCancelCause(parent Context) (ctx Context, cancel CancelCauseFunc) {
6
    c := withCancel(parent)
7
    return c, func(cause error) { c.cancel(true, Canceled, cause) }
8
}
9

10
// 如果 CancelCauseFunc(err) 设置过取消原因，则返回这个原因
11
// 否则返回 c.Err()
12
func Cause(c Context) error {
13
    if cc, ok := c.Value(&cancelCtxKey).(*cancelCtx); ok {
14
       cc.mu.Lock()
15
       defer cc.mu.Unlock()
16
       return cc.cause
17
    }
18
    return c.Err()
19
}

1
//
2
func WithoutCancel(parent Context) Context {
3
    if parent == nil {
4
        panic("cannot create context from nil parent")
5
    }
6
    return withoutCancelCtx{parent}
7
}
8

9
type withoutCancelCtx struct {
10
    c Context
11
}
12

13
func (withoutCancelCtx) Deadline() (deadline time.Time, ok bool) {
14
    return
15
}
16

17
func (withoutCancelCtx) Done() <-chan struct{} {
18
    return nil
19
}
20

21
func (withoutCancelCtx) Err() error {
22
    return nil
23
}
24

25
func (c withoutCancelCtx) Value(key any) any {
26
    return value(c, key)
27
}
28

29
func (c withoutCancelCtx) String() string {
30
    return contextName(c.c) + ".WithoutCancel"
31
}

1
func WithDeadline(parent Context, d time.Time) (Context, CancelFunc) {
2
    return WithDeadlineCause(parent, d, nil)
3
}
4

5
func WithDeadlineCause(parent Context, d time.Time, cause error) (Context, CancelFunc) {
6
    if parent == nil {
7
        panic("cannot create context from nil parent")
8
    }
9
    if cur, ok := parent.Deadline(); ok && cur.Before(d) {
10
        // 如果父节点的 deadline 更早
11
        return WithCancel(parent)
12
    }
13
    c := &timerCtx{
14
        deadline: d,
15
    }
16
    c.cancelCtx.propagateCancel(parent, c)
17
    dur := time.Until(d)
18
    if dur <= 0 {
19
        c.cancel(true, DeadlineExceeded, cause)
20
        return c, func() { c.cancel(false, Canceled, nil) }
21
    }
22
    c.mu.Lock()
23
    defer c.mu.Unlock()
24
    if c.err.Load() == nil {
25
        c.timer = time.AfterFunc(dur, func() {
26
            c.cancel(true, DeadlineExceeded, cause)
27
        })
28
    }
29
    return c, func() { c.cancel(true, Canceled, nil) }
30
}

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

5
func WithTimeoutCause(parent Context, timeout time.Duration, cause error) (Context, CancelFunc) {
6
    return WithDeadlineCause(parent, time.Now().Add(timeout), cause)
7
}

1
type valueCtx struct {
2
    Context
3
    key, val any
4
}
5

6
func (c *valueCtx) String() string {
7
    return contextName(c.Context) + ".WithValue(" +
8
        stringify(c.key) + ", " +
9
        stringify(c.val) + ")"
10
}
11

12
func stringify(v any) string {
13
    switch s := v.(type) {
14
    case stringer:
15
        return s.String()
16
    case string:
17
        return s
18
    case nil:
19
        return "<nil>"
20
    }
21
    return reflectlite.TypeOf(v).String()
22
}
23

24
func (c *valueCtx) Value(key any) any {
25
    if c.key == key {
26
        return c.val
27
    }
28
    return value(c.Context, key)
29
}
30

31
// 不断向上寻找直至找到 key
32
func value(c Context, key any) any {
33
    for {
34
        switch ctx := c.(type) {
35
        case *valueCtx:
36
            if key == ctx.key {
37
                return ctx.val
38
            }
39
            c = ctx.Context
40
        case *cancelCtx:
41
            if key == &cancelCtxKey {
42
                return c
43
            }
44
            c = ctx.Context
45
        case withoutCancelCtx:
46
            if key == &cancelCtxKey {
47
                return nil
48
            }
49
            c = ctx.c
50
        case *timerCtx:
51
            if key == &cancelCtxKey {
52
                return &ctx.cancelCtx
53
            }
54
            c = ctx.Context
55
        case backgroundCtx, todoCtx:
56
            return nil
57
        default:
58
            return c.Value(key)
59
        }
60
    }
61
}