09.Golang 上下文Context 源码分析

Golang 上下文Context 源码分析

注意当前go版本代码为1.23

源码位置context.Context

介绍

context.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
type Context interface {
    // 返回 context 是否会被取消以及自动取消时间（即 deadline）可以表示 context 被取消的信号：当这个 channel 被关闭时，说明 context 被取消了。注意，这是一个只读的channel。
	Deadline() (deadline time.Time, ok bool)
    
	// 当 context 被取消或者到了 deadline，返回一个被关闭的 channel
    // Done可用于select语句：
    // 
    //  // Stream生成值并将其发送到out channel，直到DoSomething返回错误或ctx.Done关闭。
    //  func Stream(ctx context.Context, out chan<- Value) error {
    //      for {
    //          v, err := DoSomething(ctx)
    //          if err != nil {
    //              return err
    //          }
    //          select {
    //          case <-ctx.Done():
    //              return ctx.Err()
    //          case out <- v:
    //          }
    //      }
    //  }
    // 
    // 参见https://blog.golang.org/pipelines以了解更多关于如何使用Done channel进行取消的示例。
	Done() <-chan struct{}

	// 在 channel Done 关闭后，返回 context 取消原因
    // 如果Done尚未关闭，则返回nil。
    // 如果Done已经关闭，则返回一个非nil错误，解释为什么关闭：
    // 如果上下文被取消，则返回Canceled；
    // 如果上下文的截止日期已经过期，则返回DeadlineExceeded。
	Err() error

	// 获取 key 对应的 value
	Value(key interface{}) interface{}
}

默认上下文与emptyCtx

context 包中最常用的方法是 context.Background、context.TODO，这两个方法都会返回预先初始化好的私有变量 background 和 todo，它们会在同一个 Go 程序中被复用：

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
func Background() Context {
    return backgroundCtx{}
}

type backgroundCtx struct{ emptyCtx }


func TODO() Context {
	return todoCtx{}
}
type todoCtx struct{ emptyCtx }

type emptyCtx struct{}

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 any) any {
	return nil
}

context.emptyCtx 通过空方法实现了 context.Context 接口中的所有方法，它没有任何功能。

从源代码来看，context.Background 和 context.TODO 也只是互为别名，没有太大的差别，只是在使用和语义上稍有不同：

  • context.Background 是上下文的默认值，所有其他的上下文都应该从它衍生出来；
  • context.TODO 应该仅在不确定应该使用哪种上下文时使用；

在多数情况下，如果当前函数没有上下文作为入参，我们都会使用 context.Background 作为起始的上下文向下传递。

cancelCtx

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
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
// cancelCtx 结构体定义了一个可取消的上下文, 
// canceler接口的实现
type cancelCtx struct {
    // 继承自 Context 接口
    Context

    mu       sync.Mutex            // 用于保护以下字段的互斥锁
    done     atomic.Value          // 原子操作的值，保存一个通道（chan struct{}），懒加载，在第一次取消时创建并关闭
    children map[canceler]struct{} // 存储所有的子上下文，第一次取消时被置为nil
    err      error                 // 在第一次取消时设定的错误
    cause    error                 // 在第一次取消时设定的原因
}

// Value 方法用于获取上下文中键对应的值
func (c *cancelCtx) Value(key any) any {
    // 如果请求的键是特定于取消上下文的，就返回当前上下文
    if key == &cancelCtxKey {
        return c
    }
    // 否则，继续在父上下文中查找
    return value(c.Context, key)
}

// Done 方法返回一个通道，当上下文被取消时，该通道会关闭
func (c *cancelCtx) Done() <-chan struct{} {
    // 尝试加载已存在的done通道
    d := c.done.Load()
    if d != nil {
        return d.(chan struct{})
    }
    // 如果done通道尚未创建，则加锁创建
    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{})
}

// Err 方法返回上下文被取消的原因
func (c *cancelCtx) Err() error {
    c.mu.Lock()
    // 复制错误以防止在返回后被修改
    err := c.err
    c.mu.Unlock()
    return err
}


// cancel closes c.done, cancels each of c's children, and, if
// removeFromParent is true, removes c from its parent's children.
// cancel sets c.cause to cause if this is the first time c is canceled.
func (c *cancelCtx) cancel(removeFromParent bool, err, cause error) {
	if err == nil {
		panic("context: internal error: missing cancel error")
	}
	if cause == nil {
		cause = err
	}
	c.mu.Lock()
	if c.err != nil {
		c.mu.Unlock()
		return // already canceled // 已经被其他协程取消
	}
	c.err = err
	c.cause = cause
	d, _ := c.done.Load().(chan struct{})
	if d == nil {
		c.done.Store(closedchan)
	} else {
		close(d)
	}
	// 遍历它的所有子节点
	for child := range c.children {
	    // 递归地取消所有子节点
		child.cancel(false, err, cause)
	}
	// 将子节点置空
	c.children = nil
	c.mu.Unlock()

	if removeFromParent {
	    // 从父节点中移除自己 
		removeChild(c.Context, c)
	}
}

// removeChild removes a context from its parent.
func removeChild(parent Context, child canceler) {
	if s, ok := parent.(stopCtx); ok {
		s.stop()
		return
	}
	p, ok := parentCancelCtx(parent)
	if !ok {
		return
	}
	p.mu.Lock()
	if p.children != nil {
		delete(p.children, child)
	}
	p.mu.Unlock()
}

// 总体来看，cancel() 方法的功能就是关闭 channel：c.done；递归地取消它的所有子节点；从父节点从删除自己。达到的效果是通过关闭 channel，将取消信号传递给了它的所有子节点。goroutine 接收到取消信号的方式就是 select 语句中的读 c.done 被选中。



创建cancelCtx：WithCancel

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
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
func WithCancel(parent Context) (ctx Context, cancel CancelFunc) {
	c := withCancel(parent)
	return c, func() { c.cancel(true, Canceled, nil) }
}


func withCancel(parent Context) *cancelCtx {
	if parent == nil {
		panic("cannot create context from nil parent")
	}
	c := &cancelCtx{}
	c.propagateCancel(parent, c)
	return c
}

func (c *cancelCtx) propagateCancel(parent Context, child canceler) {
	// 将传入的 parent context 设置为当前 cancelCtx 的 context。
	// 这样，当前 cancelCtx 就继承了 parent context 的属性。
	c.Context = parent

	// 获取父 context 的 Done() channel。
	done := parent.Done()
	// 如果父 context 的 Done() channel 为 nil，表示父 context 永远不会被取消。
	// 因此，子 context 也无需监听父 context 的取消事件，直接返回。
	if done == nil {
		return // parent is never canceled
	}

	// 使用 select 语句非阻塞地检查父 context 是否已经被取消。
	select {
	case <-done:
		// 如果父 context 已经被取消，则立即取消子 context，并传递父 context 的错误信息和原因。
		child.cancel(false, parent.Err(), Cause(parent))
		return
	default:
		// 父 context 尚未取消，继续执行。
	}

	// 尝试将父 context 转换为 *cancelCtx 类型，或者看它是否派生自 *cancelCtx 类型。
	if p, ok := parentCancelCtx(parent); ok {
		// 父 context 是一个 *cancelCtx，或者派生自一个。
		// 对父 cancelCtx 的互斥锁进行加锁，以保护其内部状态。
		p.mu.Lock()
		// 如果父 cancelCtx 已经取消，则立即取消子 context，并传递父 cancelCtx 的错误信息和原因。
		if p.err != nil {
			// parent has already been canceled
			child.cancel(false, p.err, p.cause)
		} else {
			// 父 cancelCtx 尚未取消，将子 context 添加到父 cancelCtx 的子节点列表中。
			if p.children == nil {
				p.children = make(map[canceler]struct{})
			}
			p.children[child] = struct{}{}
		}
		// 解锁父 cancelCtx 的互斥锁。
		p.mu.Unlock()
		return
	}

	// 检查父 context 是否实现了 afterFuncer 接口。
	if a, ok := parent.(afterFuncer); ok {
		// 父 context 实现了 AfterFunc 方法。
		// 对当前 cancelCtx 的互斥锁进行加锁，以保护其内部状态。
		c.mu.Lock()
		// 使用父 context 的 AfterFunc 方法注册一个回调函数，该函数在父 context 被取消时执行。
		stop := a.AfterFunc(func() {
			// 在回调函数中，取消子 context，并传递父 context 的错误信息和原因。
			child.cancel(false, parent.Err(), Cause(parent))
		})
		// 将当前 cancelCtx 的 Context 设置为 stopCtx，
		// stopCtx 包装了父 context 并保存了 AfterFunc 返回的停止函数 stop。
		c.Context = stopCtx{
			Context: parent,
			stop:    stop,
		}
		// 解锁当前 cancelCtx 的互斥锁。
		c.mu.Unlock()
		return
	}

	// 如果父 context 不是 *cancelCtx 或实现了 afterFuncer 接口，则启动一个新的 goroutine 来监听父 context 的取消事件。
	goroutines.Add(1)
	go func() {
		// 使用 select 语句监听父 context 的 Done() channel 和子 context 的 Done() channel。
		select {
		case <-parent.Done():
			// 如果父 context 被取消，则取消子 context，并传递父 context 的错误信息和原因。
			child.cancel(false, parent.Err(), Cause(parent))
		case <-child.Done():
			// 如果子 context 被取消，则直接退出 goroutine。
		}
		// goroutine 结束后会隐式减少 goroutines 计数器，并不需要显式调用 goroutines.Done()
	}()
}

propagateCancel方法核心功能是实现 context 的取消传播机制。当一个父 context 被取消时，其所有子 context 也应该被取消。propagateCancel 函数负责建立父子 context 之间的取消关系。这样，调用上层 cancel 方法的时候，就可以层层传递，将那些挂靠的子 context 同时“取消”。

主要逻辑:

  1. 设置父 Context: 将传入的 parent Context 赋值给当前 cancelCtx 的 Context 字段，形成 context 的继承关系。

  2. 检查父 Context 是否可取消: 通过 parent.Done() 获取父 context 的 Done channel。如果为 nil，表示父 context 不可取消，直接返回。

  3. 快速检查父 Context 是否已取消: 使用 select 非阻塞地检查父 context 的 Done channel 是否已关闭，如果已关闭，则立即取消子 context。

  4. 尝试获取父 cancelCtx:尝试将父 context 转换为*cancelCtx类型，如果

    转换成功，则说明父 context 本身也是一个cancelCtx，或者派生自一个cancelCtx

    • 如果父 cancelCtx 已经取消: 直接取消子 context。
    • 如果父 cancelCtx 尚未取消: 将子 context 添加到父 cancelCtx 的子节点列表中，以便在父 cancelCtx 被取消时，能够通知到所有子节点。

  5. 检查父 Context 是否实现 AfterFunc接口。

    • 如果实现: 调用父 context 的 AfterFunc 方法，注册一个在父 context 被取消时取消子 context 的回调函数。

  6. 启动 goroutine 监听:如果父 context 不是cancelCtx类型，也没有实现*afterFuncer**接口，则启动一个新的 goroutine 来监听父 context 的取消事件。

    • goroutine 逻辑: 使用 select 监听父 context 的 Done channel 和子 context 的 Done channel，如果父 context 被取消，则取消子 context，如果子 context 被取消，则退出 goroutine。

timerCtx

timerCtx 基于 cancelCtx，只是多了一个 time.Timer 和一个 deadline。Timer 会在 deadline 到来时，自动取消 context。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
type timerCtx struct {
    cancelCtx
    timer *time.Timer // Under cancelCtx.mu.

    deadline time.Time
}

func (c *timerCtx) cancel(removeFromParent bool, err error) {
	// 直接调用 cancelCtx 的取消方法
	c.cancelCtx.cancel(false, err)
	if removeFromParent {
		// 从父节点中删除子节点
		removeChild(c.cancelCtx.Context, c)
	}
	c.mu.Lock()
	if c.timer != nil {
		// 关掉定时器，这样，在deadline 到来时，不会再次取消
		c.timer.Stop()
		c.timer = nil
	}
	c.mu.Unlock()
}

创建timerCtx ：WithTimeout与WithDeadline

context.WithDeadline 和 context.WithTimeout 都能创建可以被取消的计时器上下文 context.timerCtx：

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
54
55
56
57
58
59
60
61
62
func WithTimeout(parent Context, timeout time.Duration) (Context, CancelFunc) {
    return WithDeadline(parent, time.Now().Add(timeout))
}


// WithDeadline 实际上调用了 WithDeadlineCause，并将 cause 设置为 nil。
func WithDeadline(parent Context, d time.Time) (Context, CancelFunc) {
        return WithDeadlineCause(parent, d, nil)
}

// 参数：
//   parent: 要派生的父上下文。
//   d: 截止时间。
//   cause: 当截止时间到达时，设置到返回的 Context 中的取消原因。如果为 nil，则使用默认的 DeadlineExceeded 错误。
//
// 返回值：
//   Context: 新的上下文。
//   CancelFunc: 取消函数，调用它可以提前取消上下文。调用此函数不会设置取消原因，而是使用默认的 Canceled 错误。
func WithDeadlineCause(parent Context, d time.Time, cause error) (Context, CancelFunc) {
        // 如果父上下文为 nil，则无法创建新的上下文，因此抛出 panic。
        if parent == nil {
                panic("cannot create context from nil parent")
        }
        // 检查父上下文是否已经有截止时间，并且父上下文的截止时间是否早于当前指定的截止时间。
        // 如果是这样，说明父上下文会在当前指定的截止时间之前被取消，因此直接使用一个带有取消功能的父上下文即可。
        // 这样可以避免创建不必要的定时器和资源。
        if cur, ok := parent.Deadline(); ok && cur.Before(d) {
                // 父上下文的截止时间更早，直接返回一个基于父上下文的、可取消的上下文。
                return WithCancel(parent)
        }
        // 创建一个新的 timerCtx 结构体，用于管理截止时间和定时器。
        c := &timerCtx{
                deadline: d, // 设置截止时间。
        }
        // 将父上下文的取消信号挂到新的 timerCtx。
        // 这意味着如果父上下文被取消，新的 timerCtx 也会被取消。
        c.cancelCtx.propagateCancel(parent, c)
        // 计算距离截止时间还有多久。
        dur := time.Until(d)
        // 如果距离截止时间的时间小于等于 0，说明截止时间已经过去。
        if dur <= 0 {
                // 立即取消上下文，并设置取消原因为 DeadlineExceeded 或指定的 cause。
                c.cancel(true, DeadlineExceeded, cause)
                // 返回已取消的上下文和一个取消函数。
                // 该取消函数用于在需要时显式取消上下文，但不会设置取消原因。
                return c, func() { c.cancel(false, Canceled, nil) }
        }
        // 获取互斥锁，确保在设置定时器时没有并发修改。
        c.mu.Lock()
        defer c.mu.Unlock()
        // 如果上下文还没有被取消（c.err 为 nil）。
        if c.err == nil {
                // 创建一个新的定时器，当经过 dur 时间后，会执行指定的匿名函数。
                c.timer = time.AfterFunc(dur, func() {
                        // 定时器到期时，取消上下文，并设置取消原因为 DeadlineExceeded 或指定的 cause。
                        c.cancel(true, DeadlineExceeded, cause)
                })
        }
        // 返回新的上下文和一个取消函数。
        // 该取消函数用于在需要时显式取消上下文，并会停止定时器（如果存在）。
        return c, func() { c.cancel(true, Canceled, nil) }
}

传值：valueCtx

context 包中的 context.WithValue 能从父上下文中创建一个子上下文，传值的子上下文使用 context.valueCtx 类型，context.valueCtx 结构体会将除了 Value 之外的 Err、Deadline 等方法代理到父上下文中，它只会响应 context.valueCtx.Value 方法，该方法的实现也很简单：

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
// A valueCtx carries a key-value pair. It implements Value for that key and
// delegates all other calls to the embedded Context.
type valueCtx struct {
    Context
    key, val any
}

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

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 withoutCancelCtx:
          if key == &cancelCtxKey {
             // This implements Cause(ctx) == nil
             // when ctx is created using WithoutCancel.
             return nil
          }
          c = ctx.c
       case *timerCtx:
          if key == &cancelCtxKey {
             return &ctx.cancelCtx
          }
          c = ctx.Context
       case backgroundCtx, todoCtx:
          return nil
       default:
          return c.Value(key)
       }
    }
}

大致效果如下图：

valueCtx

和链表类似，只是它的方向相反：Context 指向它的父节点，链表则指向下一个节点。通过 WithValue 函数，可以创建层层的 valueCtx，存储 goroutine 间可以共享的变量。取值时，递归向上查找；

问题

1.context的主要作用：

  • 传递共享的数据 ，常见场景传递traceId或者requestId；

  • 取消 goroutine，以及避免 goroutine 泄漏：通过WithCancel等方法控制子context的协程运行；

    context.Context 的主要作用还是在多个 Goroutine 组成的树中同步取消信号以减少对资源的消耗和占用，虽然它也有传值的功能，但是这个功能我们还是很少用到，效率不高。

参考链接

1.Go 程序员面试笔试宝典

2.《Go学习笔记》

3.Go 语言设计与实现

Golang
#原创 #编程语言 #Golang #源码分析
09.Golang 上下文Context 源码分析
https://blog.longpi1.com/2024/12/25/09-Golang-上下文Context-源码分析/