Go Context使用及源码解析

概述

context包提供了一种在Go程序中不同组件之间传递请求范围的值、取消信号和截止时间的方式。context包的设计目的是为了解决在Go的并发模型中，尤其是在goroutines之间，如何安全、高效地传递控制信息的问题。当 Context 被取消时，从其派生的所有 Context 也将被取消。

Context接口

Context接口定义了4个方法

type Context interface {

	//  返回 Context 被取消的时间点，以及一个布尔值 ok 表示是否设置了截止时间。如果没有设置deadline，ok==false
	Deadline() (deadline time.Time, ok bool)

	// 返回一个通道 chan struct{}，当 Context 被取消或超时时，该通道会被关闭。当context还未被关闭时，Done返回nil
	Done() <-chan struct{}

	// 返回 Context 关闭的原因
	Err() error

	// 从 Context 中获取与指定键关联的值。如果键不存在，则返回 nil。
	Value(key any) any
}

context类型

context包定义了几种context结构体

• emptyCtx 可以作为其他context的父节点，本身不包含任何值
• todoCtx 在不确定使用哪个Context时使用
• backgroundCtx 空context,不能被取消，其他context的父context
• cancelCtx 取消上下文，此context被canceled时，会把其中所有的child都cancel掉
• valueCtx 只是在Context基础上增加了一个key-value对，用于在各级协程间传递一些数据。
• timerCtx 在cancelCtx基础上增加了deadline用于表示自动cancel的最终时间，而timer就是一个触发自动cancel的定时器。由此，衍生出WithDeadline()和WithTimeout()。
• afterFuncCtx 在上下文取消时，触发回调函数,而且只会执行一次

func (a *afterFuncCtx) cancel(removeFromParent bool, err, cause error) {
    a.cancelCtx.cancel(false, err, cause)
    if removeFromParent {
       removeChild(a.Context, a)
    }
    a.once.Do(func() {
       go a.f()
    })
}

• stopCtx 用于在上下文取消时，停止某些异步操作或定时器的场景。
• WithoutCancelCtx 这个新上下文不会携带父上下文的取消信号

context 方法

context包提供了不同的方法用于创建不同类型的context

• WithCancel(parent Context) (ctx Context, cancel CancelFunc)

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

这里关键地方在于c.propagateCancel(parent,c)的处理。 判断父的context是否已经关闭，如果关闭，当前context立即关闭。如果没有关闭则将当前context加入到父context的children列表中

func (c *cancelCtx) propagateCancel(parent Context, child canceler) {
	c.Context = parent

	done := parent.Done()
	//done为nil，说明父context不支持取消，永远不会被取消，那么当前上下文不需要进行任何取消传播的逻辑，直接返回。
	if done == nil {
		return // parent is never canceled
	}

	//如果父context已经被取消，则当前context立即取消，并使用父context的错误信息和原因
	select {
	case <-done:
		// parent is already canceled
		child.cancel(false, parent.Err(), Cause(parent))
		return
	default:
	}

	//检查父上下文是否是 *cancelCtx
	if p, ok := parentCancelCtx(parent); ok {
		// parent is a *cancelCtx, or derives from one.
		p.mu.Lock()
		if p.err != nil {
			// parent has already been canceled, 如果父已经取消，子也立即取消
			child.cancel(false, p.err, p.cause)
		} else {
			//否则将当前的context添加到父的children列表中，这样当父context被取消时，所有子context被自动取消。
			if p.children == nil {
				//如果父context不存在children，make一个出来
				p.children = make(map[canceler]struct{})
			}
			p.children[child] = struct{}{}
		}
		p.mu.Unlock()
		return
	}

	//检查父上下文是否实现了 AfterFunc 方法
	if a, ok := parent.(afterFuncer); ok {
		// parent implements an AfterFunc method.
		c.mu.Lock()
		stop := a.AfterFunc(func() {
			child.cancel(false, parent.Err(), Cause(parent))
		})

		//将当前 cancelCtx 的上下文替换为 stopCtx
		c.Context = stopCtx{
			Context: parent,
			stop:    stop,
		}
		c.mu.Unlock()
		return
	}

	//如果父context既不是 *cancelCtx，也没有实现 afterFuncer 接口，那么启动一个 goroutine 中监控父上下文的 Done 通道
	//如果父上下文被取消，则取消子上下文，如果子上下文先完成，则退出goroutine
	goroutines.Add(1)
	go func() {
		select {
		case <-parent.Done():
			child.cancel(false, parent.Err(), Cause(parent))
		case <-child.Done():
		}
	}()
}

• Background() Context
• TODO() Context
• WithoutCancel(parent Context) Context
• WithDeadline(parent Context, d time.Time) (Context, CancelFunc)
• WithDeadlineCause(parent Context, d time.Time, cause error) (Context, CancelFunc)
• WithTimeout(parent Context, timeout time.Duration) (Context, CancelFunc)
• WithTimeoutCause(parent Context, timeout time.Duration, cause error) (Context, CancelFunc)
• WithValue(parent Context, key, val any) Context
• AfterFunc(ctx Context, f func()) (stop func() bool)

AfterFunc方法，用于在上下文取消时执行一个回调函数。

源码

func AfterFunc(ctx Context, f func()) (stop func() bool) {
        // 创建一个afterFuncCtx类型的context,在cancel时会调用f方法
	a := &afterFuncCtx{
		f: f,
	}
	a.cancelCtx.propagateCancel(ctx, a)
       
	return func() bool {
		stopped := false
                //这里通过once控制只会调用一次，如果stop先调用了，就不会执行afterFuncCtx的回调函数f。因此可以通过返回的stop接口来控制是否需要回调函数的执行
		a.once.Do(func() {
			stopped = true
		})
		if stopped {
			a.cancel(true, Canceled, nil)
		}
		return stopped
	}
}

使用例子

AfterFunc 使用例子

func main() {
	// 创建一个带有取消功能的上下文
	ctx, cancel := context.WithCancel(context.Background())

	// 使用 AfterFunc 注册一个回调函数，在上下文取消时执行
	cancelFunc := context.AfterFunc(ctx, func() {
		fmt.Println("Context canceled, cleaning up...")
	})

	time.Sleep(2 * time.Second)

	// 可选：如果需要，可以在上下文取消之前调用 cancelFunc 来取消注册的回调函数,如果打开，将不会执行回调函数
	// cancelFunc()

	// 取消上下文，这将触发 AfterFunc 注册的函数
	cancel()

	// 等待回调函数执行
	time.Sleep(1 * time.Second)
}

WithTimeout例子

func workerWithTimeout(ctx context.Context, duration time.Duration) {
	// 从上下文中获取值
	if v := ctx.Value("language"); v != nil {
		fmt.Printf("context language value : %v\n", v)
	} else {
		fmt.Printf("no language key found in context\n")
	}

	select {
	case <-time.After(duration):
		fmt.Printf("任务执行完成\n")
	case <-ctx.Done():
		fmt.Printf("任务执行超时被取消,%v\n", ctx.Err())
	}
	wg.Done()
}

func main() {
	ctx := context.Background()
	//设置超时时间5s
	ctx, cancel := context.WithTimeout(ctx, 5*time.Second)
	defer cancel()
	// 创建一个带有值的上下文
	ctx = context.WithValue(ctx, "language", "Go")

	wg.Add(1)
	go workerWithTimeout(ctx, 6*time.Second)

	// 等待子任务完成
	wg.Wait()
}

context uml图