Go语言 context包源码学习

 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
39
40
41
42
43

// 睡眠5s，模拟长时间操作
func FuncB() (interface{}, error) {
	time.Sleep(5 * time.Second)
	return struct{}{}, nil
}

func FuncA() (interface{}, error) {

	var res interface{}
	var err error
	ch := make(chan interface{})

  // 调用FuncB()，将结果保存至 channel 中
	go func() {
		res, err = FuncB()
		ch <- res
	}()

  // 设置一个2s的定时器
	timer := time.NewTimer(2 * time.Second)
  
  // 监测是定时器先结束，还是 FuncB 先返回结果
	select {
    
    // 超时，返回默认值
	case <-timer.C:
		return "default", err
    
    // FuncB 先返回结果，关闭定时器，返回 FuncB 的结果
	case r := <-ch:
		if !timer.Stop() {
			<-timer.C
		}
		return r, err
	}

}

func main() {
	res, err := FuncA()
	fmt.Println(res, err)
}

 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

const LogKey = "LogKey"

// 模拟一个日志打印，每次从 Context 中取出 LogKey 对应的 Value 作为LogID
type Logger struct{}
func (logger *Logger) info(ctx context.Context, msg string) {
	logId, ok := ctx.Value(LogKey).(string)
	if !ok {
		logId = uuid.New().String()
	}
	fmt.Println(logId + " " + msg)
}
var logger Logger

// 日志打印 并 调用 FuncB
func FuncA(ctx context.Context) {
	logger.info(ctx, "FuncA")
	FuncB(ctx)
}

func FuncB(ctx context.Context) {
	logger.info(ctx, "FuncB")
}

// 获取初始化的，带有 LogID 的 Context，一般在程序入口做
func getLogCtx(ctx context.Context) context.Context {
	logId, ok := ctx.Value(LogKey).(string)
	if ok {
		return ctx
	}
	logId = uuid.NewString()
	return context.WithValue(ctx, LogKey, logId)
}

func main() {
	ctx = getLogCtx(context.Background())
	FuncA(ctx)
}

 1
 2
 3
 4
 5
 6
 7
 8
 9
10

type Context interface {

	Deadline() (deadline time.Time, ok bool)

	Done() <-chan struct{}

	Err() error

	Value(key interface{}) interface{}
}

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

 func Stream(ctx context.Context, out chan<- Value) error {
 	for {
    
    // 处理数据
 		v, err := DoSomething(ctx)
 		if err != nil {
 			return err
 		}
    
    // ctx.Done() 读取到数据，说明获取到了任务取消的信号
 		select {
 		case <-ctx.Done():
 			return ctx.Err()
    // 否则将结果输出，继续计算
 		case out <- v:
 		}
 	}
 }

1
2
3
4
5
6

func main() {
	ctx := context.Background()
	c := context.WithValue(ctx, "key", "value")
	v, ok := c.Value("key").(string)
	fmt.Println(v, ok)
}

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

var (
	background = new(emptyCtx)
	todo       = new(emptyCtx)
)

func Background() Context {
	return background
}

func TODO() Context {
	return todo
}

 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

type emptyCtx int

func (*emptyCtx) Deadline() (deadline time.Time, ok bool) {
	return
}

func (*emptyCtx) Done() <-chan struct{} {
	return nil
}

func (*emptyCtx) Err() error {
	return nil
}

func (*emptyCtx) Value(key interface{}) interface{} {
	return nil
}

func (e *emptyCtx) String() string {
	switch e {
	case background:
		return "context.Background"
	case todo:
		return "context.TODO"
	}
	return "unknown empty Context"
}

 1
 2
 3
 4
 5
 6
 7
 8
 9
10

func main() {
	ctx := context.Background()
	c := context.WithValue(ctx, "myKey", "myValue")

	v1 := c.Value("myKey")
	fmt.Println(v1.(string))

	v2 := c.Value("hello")
	fmt.Println(v2) //  nil
}

1
2
3
4

type valueCtx struct {
	Context
	key, val interface{}
}

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

func WithValue(parent Context, key, val interface{}) 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

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

 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

func doSomething(ctx context.Context) (interface{}, error) {
	res := make(chan interface{})
	go func() {
		fmt.Println("do something")
		time.Sleep(time.Second * 5)
		res <- "done"
	}()

	select {
	case <-ctx.Done():
		return nil, ctx.Err()
	case value := <-res:
		return value, nil
	}
}

func main() {
	ctx, cancel := context.WithCancel(context.Background())
	go func() {
		time.Sleep(time.Second * 2)
		cancel()
	}()
	res, err := doSomething(ctx)
	fmt.Println(res, err) // nil , context canceled
}

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

type canceler interface {
	cancel(removeFromParent bool, err error)
	Done() <-chan struct{}
}

// closedchan is a reusable closed channel.
var closedchan = make(chan struct{})

func init() {
	close(closedchan)
}

var Canceled = errors.New("context canceled")

1

type CancelFunc func()

1

var cancelCtxKey int

1
2
3
4
5
6
7
8

type cancelCtx struct {
	Context

	mu       sync.Mutex            // 保护如下字段，保证线程安全
	done     atomic.Value          // 保存 channel，懒加载，调用 cancel 方法时会关闭这个 channel
	children map[canceler]struct{} // 保存子节点，第一次调用 cancel 方法时会置为 nil
	err      error                 // 保存为什么被取消，默认为nil，第一次调用 cancel 会赋值
}

1
2
3
4
5
6

func (c *cancelCtx) Value(key interface{}) interface{} {
	if key == &cancelCtxKey {
		return c
	}
	return c.Context.Value(key)
}

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

func (c *cancelCtx) Done() <-chan struct{} {
	d := c.done.Load()
  
  // 如果已经有值了，直接返回
	if d != nil {
		return d.(chan struct{})
	}
  
  // 没有值，加锁赋值
	c.mu.Lock()
	defer c.mu.Unlock()
	d = c.done.Load()
	if d == nil {
		d = make(chan struct{})
		c.done.Store(d)
	}
	return d.(chan struct{})
}

1
2
3
4
5
6

func (c *cancelCtx) Err() error {
   c.mu.Lock()
   err := c.err
   c.mu.Unlock()
   return err
}

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

func newCancelCtx(parent Context) cancelCtx {
	return cancelCtx{Context: parent}
}

 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
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53

func propagateCancel(parent Context, child canceler) {
  
  // 如果 parent 不支持取消，那么就不支持取消传播，直接返回
	done := parent.Done()
	if done == nil {
		return 
	}

  // 到这里说明 done 不为 nil，parent 支持取消
  
	select {
	case <-done:
		// 如果 parent 此时已经取消了，那么直接告诉子节点也取消
		child.cancel(false, parent.Err())
		return
	default:
	}

  // 到这里说明此时 parent 还未取消
  
  // 如果 parent 是未取消的 cancelCtx 
	if p, ok := parentCancelCtx(parent); ok {
    
    // 加锁，防止并发更新
		p.mu.Lock()
    
    // 再次判断，因为有可能上一个获得锁的进行了取消操作。
    // 如果 parent 已经取消了，那么子节点也直接取消
		if p.err != nil {
			child.cancel(false, p.err)
		} else {
      // 把子Context 挂到父节点 parent cancelCtx 的 children字段下
      // 之后 parent cancelCtx 取消时，能通知到所有的 子Context 
			if p.children == nil {
				p.children = make(map[canceler]struct{})
			}
			p.children[child] = struct{}{}
		}
		p.mu.Unlock()
	} else {
    
   // parent 不是 cancelCtx 类型，可能是用户自己实现的Context
		atomic.AddInt32(&goroutines, +1)
    // 启动一个协程监听，如果 parent 取消了，子 Context 也取消
		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
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45

func (c *cancelCtx) cancel(removeFromParent bool, err error) {
  
  // err 是指取消的原因，必传，cancelCtx 中是 errors.New("context canceled")
	if err == nil {
		panic("context: internal error: missing cancel error")
	}
  
  // 涉及到保护字段值的修改，都需要加锁
	c.mu.Lock()
  
  // 如果该Context已经取消过了，直接返回。多次调用cancel，不会产生额外效果
	if c.err != nil {
		c.mu.Unlock()
		return 
	}
  
  // 给 err 赋值，这里 err 一定不为 nil
	c.err = err
  
  // close channel
	d, _ := c.done.Load().(chan struct{})
  // 因为c.done 是懒加载，有可能存在 nil 的情况
  // 如果 c.done 中没有值，直接赋值 closedchan；否则直接 close
	if d == nil {
		c.done.Store(closedchan)
	} else {
		close(d)
	}
  
  // 遍历当前 cancelCtx 所有的子Context，让子节点也 cancel
  // 因为当前的Context 会主动把子Context移除，子Context 不用主动从parent中脱离
  // 因此 child.cancel 传入的 removeFromParent 为false
	for child := range c.children {
		child.cancel(false, err)
	}
  // 将 children 置空，相当于移除自己的所有子Context
	c.children = nil
	c.mu.Unlock()
	
  // 如果当前 cancelCtx 需要从上层的 cancelCtx移除，调用removeChild方法
  // c.Context 就是自己的父Context
	if removeFromParent {
		removeChild(c.Context, c)
	}
}

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

func removeChild(parent Context, child canceler) {
  // parent 是否为未取消的 cancelCtx
	p, ok := parentCancelCtx(parent)
	if !ok {
		return
	}
  // 获取 parent cancelCtx 的锁，修改保护字段 children
	p.mu.Lock()
  // 将自己从 parent cancelCtx 的 children 中删除
	if p.children != nil {
		delete(p.children, child)
	}
	p.mu.Unlock()
}

 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

func parentCancelCtx(parent Context) (*cancelCtx, bool) {
  
  // 如果 parent context 的 done 为 nil， 说明不支持 cancel，那么就不可能是 cancelCtx
	// 如果 parent context 的 done 为 closedchan， 说明 parent context 已经 cancel 了
	done := parent.Done()
	if done == closedchan || done == nil {
		return nil, false
	}
  
  // 到这里说明支持取消，且没有被取消
  
	// 如果 parent context 属于原生的 *cancelCtx 或衍生类型，需要继续进行后续判断
	// 如果 parent context 无法转换到 *cancelCtx，则认为非 cancelCtx，返回 nil,fasle
	p, ok := parent.Value(&cancelCtxKey).(*cancelCtx)
	if !ok {
		return nil, false
	}
 
  // 经过上面的判断后，说明 parent context 可以被转换为 *cancelCtx，这时存在多种情况:
	//   - parent context 就是 *cancelCtx
	//   - parent context 是标准库中的 timerCtx
	//   - parent context 是个自己自定义包装的 cancelCtx
	//
	// 针对这 3 种情况需要进行判断，判断方法就是: 
	//   判断 parent context 通过 Done() 方法获取的 done channel 与 Value 查找到的 context 的 done channel 是否一致
	// 
	// 一致情况说明 parent context 为 cancelCtx 或 timerCtx 或 自定义的 cancelCtx 且未重写 Done()，
	// 这种情况下可以认为拿到了底层的 *cancelCtx
	// 
	// 不一致情况说明 parent context 是一个自定义的 cancelCtx 且重写了 Done() 方法，并且并未返回标准 *cancelCtx 的
	// 的 done channel，这种情况需要单独处理，故返回 nil, false
	pdone, _ := p.done.Load().(chan struct{})
	if pdone != done {
		return nil, false
	}
	return p, true
}

1
2
3
4
5

func slowOperationWithTimeout(ctx context.Context) (Result, error) {
		ctx, cancel := context.WithTimeout(ctx, 100*time.Millisecond)
    defer cancel()  // 如果未到截止时间，slowOperation就完成了，尽早调用 cancel() 释放资源
		return slowOperation(ctx)
}

1
2
3
4
5
6

type timerCtx struct {
   cancelCtx // 内嵌 cancelCtx
   timer *time.Timer // 受 cancelCtx.mu 互斥锁的保护

   deadline time.Time // 截止时间
}

1
2
3

func (c *timerCtx) Deadline() (deadline time.Time, ok bool) {
	return c.deadline, 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
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48

func WithDeadline(parent Context, d time.Time) (Context, CancelFunc) {
  
  // 传入的 parent 不能为 nil
	if parent == nil {
		panic("cannot create context from nil parent")
	}
  
  // parent 也有 deadline，并且比 d 还要早，直接依赖 parent 的取消传播即可
	if cur, ok := parent.Deadline(); ok && cur.Before(d) {
		// The current deadline is already sooner than the new one.
		return WithCancel(parent)
	}
  
  // 定义 timerCtx 接口
	c := &timerCtx{
		cancelCtx: newCancelCtx(parent),
		deadline:  d,
	}
  
  // 设置传播，如果parent 属于 cancelCtx，会挂载到 children 字段上
	propagateCancel(parent, c)
  
  // 距离截止时间 d 还有多久
	dur := time.Until(d)
	if dur <= 0 {
    // 已经到了截止时间，直接取消，同时从 parent 中取消挂载
    // 由于是超时，取消时的 err 是 DeadlineExceeded
		c.cancel(true, DeadlineExceeded) 
    
    // 再返回 c 和 CancelFunc，已经取消挂载了，此时的 CancelFunc 不会从 parent 中取消挂载
    // 后面再次调用 CancelFunc 不会产生任何效果了
    // 主动取消的话，err 是 Canceled
		return c, func() { c.cancel(false, Canceled) }
	}
  
  // 还没有到截止时间，定义一个定时器，过了 dur 会自动取消
	c.mu.Lock()
	defer c.mu.Unlock()
	if c.err == nil {
		c.timer = time.AfterFunc(dur, func() {
      // 由于是到了截止时间才取消，err 是 DeadlineExceeded
			c.cancel(true, DeadlineExceeded)
		})
	}
  
  // 返回 c 和 cancelFunc，主动取消的 err 是 Canceled
	return c, func() { c.cancel(true, Canceled) }
}

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

func (c *timerCtx) cancel(removeFromParent bool, err error) {
	c.cancelCtx.cancel(false, err)
  
  // 从父节点中移除
	if removeFromParent {
		removeChild(c.cancelCtx.Context, c)
	}
  
  // 把定时器停了，释放资源
  // 有可能还没到deadline，手动触发了 CancelFunc，此时把 timer 停了
	c.mu.Lock()
	if c.timer != nil {
		c.timer.Stop()
		c.timer = nil
	}
	c.mu.Unlock()
}

1
2
3

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