Go-Context底层原理剖析

背景


context包主要的作用 是在并发的情况场景下去同步取消信号携带上文值

下面我们来看看各种功能对应的实现。

内容

context 包的代码并不长,context.go 文件总共不到 800 行,其中还有很多大段的注释,代码可能也就 200 行左右的样子,是一个非常值得研究的代码库。

默认的上下文

context 包中最常用的方法是 context.Backgroundcontext.TODO,这两个方法都会返回预先初始化好的私有变量 backgroundtodo

go 复制代码
func Background() Context {
  return backgroundCtx{}
}

type backgroundCtx struct{ emptyCtx }

func TODO() Context {
   return todoCtx{}
}

type todoCtx struct{ emptyCtx }

它们是指向私有结构体 context.emptyCtx,这是最简单、最原始的上下文类型:

go 复制代码
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.Backgroundcontext.TODO 也只是互为别名,没有太大的差别,只是在使用和语义上稍有不同:

context.Background 是上下文的默认值,所有其他的上下文都应该从它衍生出来;

context.TODO 应该仅在不确定应该使用哪种上下文时使用;在多数情况下,如果当前函数没有上下文作为入参,我们都会使用 context.Background 作为起始的上下文向下传递。


同步取消信号

context.WithCancel函数会返回一个取消的上下文,和一个取消函数(cancel()),当我们执行返回的取消函数时,当前上下文以及它的子上下文都会被取消,所有的 Goroutine 都会同步收到这个取消信号,如图所示:

看一下实现怎么做的?

直接看context.WithCancel函数是怎样写的

go 复制代码
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
}

这个函数的主要作用是初始化父子上下文的关系,主要的逻辑在propagateCancel函数里面,函数的内容为下:

go 复制代码
// propagateCancel 函数用于将取消信号从父上下文传播到子上下文
func (c *cancelCtx) propagateCancel(parent Context, child canceler) {
	// 将父上下文设置为当前上下文,以便于子上下文可以访问父上下文的属性。
	c.Context = parent
	// 获取父上下文的Done通道,用于监听父上下文的取消事件。
	done := parent.Done()
	// 如果父上下文没有设置Done通道,说明它永远不会被取消,直接返回。
	if done == nil {
		return // parent is never canceled
	}

	// 选择性地等待父上下文的Done通道被关闭,如果通道关闭,说明父上下文已经被取消,
	// 那么就调用子上下文的cancel方法,并将父上下文的错误和原因传递给子上下文。
	select {
	case <-done:
		// parent is already canceled
		child.cancel(false, parent.Err(), Cause(parent))
		return
	default:
	}

	// 如果父上下文是一个*cancelCtx类型或者从它派生出来的类型,
	// 那么就使用它自己的机制来传播取消信号。
	if p, ok := parentCancelCtx(parent); ok {
		p.mu.Lock()
		// 如果父上下文已经被取消,那么就传递取消信号给子上下文。
		if p.err != nil {
			child.cancel(false, p.err, p.cause)
		} else {
			// 如果父上下文的children字段还没有被创建,那么初始化它,
			// 并将子上下文添加到字段中。
			if p.children == nil {
				p.children = make(map[canceler]struct{})
			}
			p.children[child] = struct{}{}
		}
		p.mu.Unlock()
		return
	}

	// 如果父上下文实现了AfterFunc方法,那么就使用它来在父上下文结束时
	// 取消子上下文。
	if a, ok := parent.(afterFuncer); ok {
		c.mu.Lock()
		// 使用AfterFunc方法注册一个函数,该函数在父上下文结束时执行。
		stop := a.AfterFunc(func() {
			child.cancel(false, parent.Err(), Cause(parent))
		})
		// 更新当前上下文,使其包含停止函数的上下文。
		c.Context = stopCtx{
			Context: parent,
			stop:    stop,
		}
		c.mu.Unlock()
		return
	}

	// 如果以上条件都不满足,那么创建一个新的goroutine来监控父上下文和子上下文的Done通道。
	// 当任意一个通道关闭时,就调用子上下文的cancel方法。
	goroutines.Add(1)
	go func() {
		select {
		case <-parent.Done():
			child.cancel(false, parent.Err(), Cause(parent))
		case <-child.Done():
		}
	}()
}

总结一下流程:

  1. 当 parent.Done() == nil,parent 不会触发取消事件时,当前函数会直接返回;

  2. 当 parent上下文以及是一个CancelCtx,会判断 parent 上下文是否已经触发了取消信号;

    如果已经被取消,child 会立刻被取消;

    如果没有被取消,child 会被加入 parent 的 children 列表中,等待 parent 释放取消信号;

  3. 当父上下文第一个CancelCtx或者自定义的类型;

    运行一个新的 Goroutine 同时监听 parent.Done() 和 child.Done() 两个 Channel;

  4. 在 parent.Done() 关闭时调用 child.cancel 取消子上下文

然后我们看看cancel()函数的实现:

go 复制代码
// cancel 方法设置取消错误,并将其传播给所有子上下文。
func (c *cancelCtx) cancel(removeFromParent bool, err, cause error) {
	// 如果err参数为空,则抛出恐慌,因为取消错误是必须的。
	if err == nil {
		panic("context: internal error: missing cancel error")
	}
	// 如果cause参数为空,则将其设置为err,因为cause通常包含导致取消的详细信息。
	if cause == nil {
		cause = err
	}
	// 锁定cancelCtx的互斥锁,以安全地修改内部状态。
	c.mu.Lock()
	// 如果cancelCtx已经被取消,则直接解锁并返回,不做任何处理。
	if c.err != nil {
		c.mu.Unlock()
		return // already canceled
	}
	// 设置cancelCtx的错误和原因。
	c.err = err
	c.cause = cause
	// 获取并关闭cancelCtx的Done通道,以通知等待的goroutine。
	d, _ := c.done.Load().(chan struct{})
	if d == nil {
		c.done.Store(closedchan)
	} else {
		//当close当时候,我们可以从Done()方法接受到消息
		close(d)
	}
	// 遍历cancelCtx的子上下文,并调用它们的cancel方法,传播取消信号。
	for child := range c.children {
		// 注意:在持有父上下文的锁时获取子上下文的锁。
		child.cancel(false, err, cause)
	}
	// 清空cancelCtx的子上下文映射。
	c.children = nil
	// 解锁cancelCtx的互斥锁。
	c.mu.Unlock()

	// 如果参数removeFromParent为真,则从父上下文的children列表中移除当前上下文。
	if removeFromParent {
		removeChild(c.Context, c)
	}
}

它会取消当前的上下文,如果有孩子上下文,会把孩子上下文都取消了,context.WithDeadlinecontext.WithTimeout 也都能创建可以被取消的计时器上下文 context.timerCtx

它们是通过嵌入 context.cancelCtx 结构体继承了相关的变量和方法,通过持有的定时器 timer 和截止时间 deadline 实现了定时取消的功能,源码就不展示了。

然后我们来看context携带值怎么实现的,它是怎么确保的并发安全的呢?

携带上文值

我们直接查看context.WithValue方法的实现,它的作用是能从父上下文中创建一个子上下文并携带你传入的key,value值。

我们来查看一下函数源码

go 复制代码
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}
}

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)
}

可以看到valueCtx结构,是一个链表的结构,当前的上下文,会携带父上下文,和自己携带的key,value值。

然后我们来看看查询数据方法Value是怎么实现的?

go 复制代码
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
			}
	....
		case backgroundCtx, todoCtx:
			return nil
		default:
			return c.Value(key)
		}
	}
}

可以看到函数,可以看到每个goroutine携带值都会进行一次加上下文的操作,所以每个都有一个自己的上下文,所以写入数据是并发完全的,然后查询数据,是通过层级的形式去遍历上下文对比valueCtxkey是否跟查询的值相等,如果相等就返回,不相等就去父上下文对比key值,直到找到key,如图:

总结

  1. 同步取消信号,采用树结构去存储上下文的关系,实现了同步取消本身上下文和取消全部子上下文的功能;

  2. 通过类似链表结构来存储层级goroutine数据来实现了并发安全,减少锁的使用,但是查询的速度一般,O(n)的查询时间复杂度;

  3. 而且不能查询兄弟上下文的值,而且它不会限制key和value的值或者类型,并发场景下处理某些类型如map,slice是不安全的;

参考文献

【go语言设计与实现】 draveness.me/golang/docs...

【深度解密GO-context】 www.cnblogs.com/qcrao-2018/...

【Golang Context 是好的设计吗?】segmentfault.com/a/119000001...

相关推荐
王码码20359 小时前
Go语言的测试:从单元测试到集成测试
后端·golang·go·接口
王码码20359 小时前
Go语言中的测试:从单元测试到集成测试
后端·golang·go·接口
嵌入式×边缘AI:打怪升级日志9 小时前
使用JsonRPC实现前后台
前端·后端
小码哥_常10 小时前
从0到1:Spring Boot 中WebSocket实战揭秘,开启实时通信新时代
后端
lolo大魔王10 小时前
Go语言的异常处理
开发语言·后端·golang
IT_陈寒12 小时前
Python多进程共享变量那个坑,我差点没爬出来
前端·人工智能·后端
码事漫谈13 小时前
2026软考高级·系统架构设计师备考指南
后端
AI茶水间管理员14 小时前
如何让LLM稳定输出 JSON 格式结果?
前端·人工智能·后端
其实是白羊14 小时前
我用 Vibe Coding 搓了一个 IDEA 插件,复制URI 再也不用手动拼了
后端·intellij idea
用户83562907805114 小时前
Python 操作 Word 文档节与页面设置
后端·python