Go Context底层原理

1
2
3
4
5
6

type Context interface {
	Deadline() (deadline time.Time, ok bool) // 用于子协程获取deadline
	Done() <-chan struct{} // 用于子协程获取管道的取消信号
	Err() error // 返回上下文被取消或超时时的错误信息
	Value(key any) any // 用于获取键值对
}

1
2
3
4
5
6
7
8
9

type emptyCtx int

var (  
    background = new(emptyCtx)  
)  
 
func Background() Context {  
    return background  
}

1
2
3
4

type valueCtx struct {  
    Context  // 指向父context的指针
    key, val any  // 键值对
}

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12

func WithValue(parent Context, key, val any) Context {
	if parent == nil {
		panic("cannot create context from nil parent")
	}
	if key == nil {
		panic("nil key")
	}
	if !reflectlite.TypeOf(key).Comparable() {
		panic("key is not comparable")
	}
	return &valueCtx{parent, key, val}
}

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32

func value(c Context, key any) any {  
    for {  
       switch ctx := c.(type) {  
       case *valueCtx:  
          if key == ctx.key {  
             return ctx.val  
          }  
          c = ctx.Context  
       case *cancelCtx:  
          if key == &cancelCtxKey {  
             return c  
          }  
          c = ctx.Context  
       case *timerCtx:  
          if key == &cancelCtxKey {  
             return &ctx.cancelCtx  
          }  
          c = ctx.Context  
       case *emptyCtx:  
          return nil  
       default:  
          return c.Value(key)  
       }    
    }
}

func (c *valueCtx) Value(key any) any {  
    if c.key == key {  
       return c.val  
    }  
    return value(c.Context, key)  
}

1
2
3
4
5
6
7
8

type cancelCtx struct {
	Context

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

1
2
3
4
5
6
7
8

func WithCancel(parent Context) (ctx Context, cancel CancelFunc) {
	if parent == nil {
		panic("cannot create context from nil parent")
	}
	c := newCancelCtx(parent)
	propagateCancel(parent, &c)
	return &c, func() { c.cancel(true, Canceled) }
}

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38

// propagateCancel arranges for child to be canceled when parent is.
func propagateCancel(parent Context, child canceler) {
	done := parent.Done()
	if done == nil {
		return // parent is never canceled
	}

	select {
	case <-done:
		// parent is already canceled
		child.cancel(false, parent.Err())
		return
	default:
	}

	if p, ok := parentCancelCtx(parent); ok {
		p.mu.Lock()
		if p.err != nil {
			// parent has already been canceled
			child.cancel(false, p.err)
		} else {
			if p.children == nil {
				p.children = make(map[canceler]struct{})
			}
			p.children[child] = struct{}{}
		}
		p.mu.Unlock()
	} else {
		atomic.AddInt32(&goroutines, +1)
		go func() {
			select {
			case <-parent.Done():
				child.cancel(false, parent.Err())
			case <-child.Done():
			}
		}()
	}
}

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15

func parentCancelCtx(parent Context) (*cancelCtx, bool) {
	done := parent.Done()
	if done == closedchan || done == nil {
		return nil, false
	}
	p, ok := parent.Value(&cancelCtxKey).(*cancelCtx)
	if !ok {
		return nil, false
	}
	pdone, _ := p.done.Load().(chan struct{})
	if pdone != done {
		return nil, false
	}
	return p, true
}

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27

func (c *cancelCtx) cancel(removeFromParent bool, err error) {
	if err == nil {
		panic("context: internal error: missing cancel error")
	}
	c.mu.Lock()
	if c.err != nil {
		c.mu.Unlock()
		return // already canceled
	}
	c.err = err
	d, _ := c.done.Load().(chan struct{})
	if d == nil {
		c.done.Store(closedchan)
	} else {
		close(d)
	}
	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)
	}
}