go context 源码

参考资料

(28 条消息) Go 语言 context 包源码剖析_binary~的博客-CSDN 博客

总结

• context 是 golang 支持的上下文，并发安全
• context 的作用：控制 goroutine 的生命周期，同步传参
• context 是一棵单节点树（链表），每个节点关联了父节点（最新的节点在根节点，先进后出）


源码


context 接口

type Context interface {
	Deadline() (deadline time.Time, ok bool) // 返回ctx被取消的时间。 当没有设置截止日期时，截止日期返回 ok==false
	Done() <-chan struct{}                   // 关闭信号
	Err() error                              // 如果 Done 尚未关闭，Err 返回 nil
	Value(key any) any                       // 读取context的数据，使用 WithValue 来存储数据，返回新 context

}

emptyCtx（空实现）

实现了 context 接口，但是没有实现功能；

其中 ==context.Background()== 和 ==context.TODO()== 就是 emtyCtx 结构

// 实现了context接口，但是没有实现功能
// An emptyCtx is never canceled, has no values, and has no deadline. It is not
// struct{}, since vars of this type must have distinct addresses.
type emptyCtx int

valueCtx（携带数据的 context）

valueCtx 目的就是为 Context 携带数据，他会继承 父 Context

该方法的实现就是从树的最底层向上找，直到找到或者到达根 Context 为止，context 树形结构如下


// valueCtx类的创建
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}
}

// 实现了Value方法，获取数据
func (c *valueCtx) Value(key any) any {
	if c.key == key {
		return c.val
	}
	// 递归从 父 context中获取数据
	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 *timerCtx:
			if key == &cancelCtxKey {
				return &ctx.cancelCtx
			}
			c = ctx.Context
		case *emptyCtx:
			return nil
		default:
			return c.Value(key)
		}
	}
}

cnacelCtx（带手动调用取消函数）

cnacelCtx 也会继承父 context

cancelCtx 是 Go 语言标准库中 context 包中的一个类型，表示一个可以被取消的上下文。它是 context.Context 接口的一个具体实现，用于在某些操作需要被取消时通知相关的 goroutine；

type cancelCtx struct {
	Context

	mu       sync.Mutex            // 一个互斥锁，用于保护接下来的字段
	done     atomic.Value          // 保存chan 类型的信道，该信道是懒加载的，即第一次调用cancel()函数时创建并关闭。
	children map[canceler]struct{} // 用于保存当前Context的子Context。第一次调用cancel()函数时，将其设置为nil
	err      error                 // 用于保存第一次调用cancel()函数时设置的错误信息
}

// 创建一个cancelCtx实例，返回一个cacelCtx实例，和一个CancelFunc函数
func WithCancel(parent Context) (ctx Context, cancel CancelFunc) {
	if parent == nil {
		panic("cannot create context from nil parent")
	}
	c := newCancelCtx(parent)
	propagateCancel(parent, &c)

	// 返回一个cancelCtx实例
	// 返回一个CancelFunc函数，函数调用时执行cancel：removeFromParent=true，err=errors.New("context canceled")
	return &c, func() { c.cancel(true, Canceled) }
}

// 方法实现了 context 包中上下文的取消功能，确保所有子级也被正确地取消。
// 关闭 c.done，取消 c 的所有子级（即其派生的上下文），并且如果 removeFromParent 为 true，则将 c 从其父级的子级列表中删除.
func (c *cancelCtx) cancel(removeFromParent bool, err error) {
	if err == nil { // 必须有err
		panic("context: internal error: missing cancel error")
	}
	c.mu.Lock()
	if c.err != nil { // 如果当前的ctx已经有err，表已退出，无须执行退出
		c.mu.Unlock()
		return // already canceled
	}
	c.err = err // 将退出的err赋值给当前ctx的err字段
	d, _ := c.done.Load().(chan struct{})
	if d == nil {
		c.done.Store(closedchan) // 如果当前ctx的done不存在，则赋值一个默认关闭的chan
	} else {
		close(d) // 关闭chan，通知其他协程
	}
	//遍历每一个children,取消所有孩子节点
	for child := range c.children {
		// NOTE: acquiring the child's lock while holding parent's lock.
		child.cancel(false, err) // 为什么这里是false？根节点和子节点断开即可，子节点和子节点之间的联系断不断没意义
	}
	c.children = nil
	c.mu.Unlock() //解锁

	if removeFromParent { //如果为true，则将当前节点和父节点断开链接
		removeChild(c.Context, c)
	}
}

用法

func main() {
	// 创建一个ctx
	parentCtx := context.Background()
	// 创建一个带cancel的ctx
	ctx, cancel := context.WithCancel(parentCtx)

	go func(ctx context.Context) {
		select {
		case <-ctx.Done(): // 当ctx的cancel执行之后，协程完成退出
			fmt.Println("context canceled")
		}

	}(ctx)
	time.Sleep(2 * time.Second)
	cancel() // 退出所有的ctx
}

timerCtx（支持超时自动取消 context）

timerCtx 基于 cancelCtx，继承了 cancelCtx 只是多了一个 time.Timer 和一个 deadline。

timerCtx 可以手动执行 Cancel 函数也可以在 deadline 到来时，自动取消 context。

WithTimeout 和 WithDeadline 都会创建一个 timerCtx，

不同的是 WithDeadline 表示是截止日期（几号几点结束），WithTimeout 是超时时间（多少时间后结束）

// 继承cancelCtx，并增加了 timer 和 deadline
type timerCtx struct {
	cancelCtx
	timer *time.Timer // 定时器

	deadline time.Time // 截止日期
}

// 该方法返回一个带有超时时间的timerCtx。与WithDeadline方法不同的是，WithTimeout方法的参数是一个持续时间（Duration），而不是一个时间点
func WithTimeout(parent Context, timeout time.Duration) (Context, CancelFunc) {
	return WithDeadline(parent, time.Now().Add(timeout))
}

// 该方法返回一个带有截止时间的timerCtx。如果在截止时间之前没有取消Context，定时器会自动触发取消操作
func WithDeadline(parent Context, d time.Time) (Context, CancelFunc) {
	if parent == nil {
		panic("cannot create context from nil 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,
	}
	//与父节点构建关联
	propagateCancel(parent, c)
	dur := time.Until(d)
	if dur <= 0 {
		c.cancel(true, DeadlineExceeded) // deadline has already passed
		return c, func() { c.cancel(false, Canceled) }
	}
	c.mu.Lock()
	defer c.mu.Unlock()
	if c.err == nil {
		//增加定时器，定时去取消
		c.timer = time.AfterFunc(dur, func() {
			c.cancel(true, DeadlineExceeded)
		})
	}
	return c, func() { c.cancel(true, Canceled) }
}

// 实现了cancel
func (c *timerCtx) cancel(removeFromParent bool, err error) {
	c.cancelCtx.cancel(false, err)
	if removeFromParent {
		// Remove this timerCtx from its parent cancelCtx's children.
		removeChild(c.cancelCtx.Context, c)
	}
	c.mu.Lock()
	if c.timer != nil {
		//如果定时器任务还未取消，停止定时器任务
		c.timer.Stop()
		c.timer = nil
	}
	c.mu.Unlock()
}

用法

func main() {
	// 2秒后超时关闭
	ctx, cancel := context.WithDeadline(context.Background(), time.Now().Add(2*time.Second))
	defer cancel() // return时自动关闭

	select {
	case <-ctx.Done(): // 2秒后会执行
		fmt.Println("Context cancelled")
	case <-time.After(3 * time.Second): // 3秒的定时器，因为ctx设置的2秒，所以不会执行
		fmt.Println("Time out")
	}
}

常见问题

为什么 context 是并发安全的

• context 在存数据时是通过 new 一个新的 ctx，并且关联父 ctx 的方式，取数据时遍历整颗树匹配数据，不存在数据并发读写的问题
• 详细看 WithValue() 和 Value()源码

context 的参数数据是如何存储的

使用 context.WithValue(ctx,key,value)，创建一个新的 valueCtx 节点，关联父节点