Go-Context底层原理剖析

棱棱2024-06-07 18:33

背景

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

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

内容

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

默认的上下文

context 包中最常用的方法是 context.Background、context.TODO，这两个方法都会返回预先初始化好的私有变量 background 和 todo：

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.Background 和 context.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.WithDeadline 和 context.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携带值都会进行一次加上下文的操作，所以每个都有一个自己的上下文，所以写入数据是并发完全的，然后查询数据，是通过层级的形式去遍历上下文对比valueCtx的key是否跟查询的值相等，如果相等就返回，不相等就去父上下文对比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...

相关推荐
customer081 小时前
【开源免费】基于SpringBoot+Vue.JS周边产品销售网站（JAVA毕业设计）
java·vue.js·spring boot·后端·spring cloud·java-ee·开源
Yaml42 小时前
智能化健身房管理：Spring Boot与Vue的创新解决方案
前端·spring boot·后端·mysql·vue·健身房管理
小码编匠3 小时前
一款 C# 编写的神经网络计算图框架
后端·神经网络·c#
AskHarries3 小时前
Java字节码增强库ByteBuddy
java·后端
佳佳_4 小时前
Spring Boot 应用启动时打印配置类信息
spring boot·后端
许野平5 小时前
Rust： 利用 chrono 库实现日期和字符串互相转换
开发语言·后端·rust·字符串·转换·日期·chrono
BiteCode_咬一口代码6 小时前
信息泄露！默认密码的危害，记一次网络安全研究
后端
齐 飞7 小时前
MongoDB笔记01-概念与安装
前端·数据库·笔记·后端·mongodb
热门推荐
01如何才能让手机厂商主动拥抱华为，接入鸿蒙系统？02【HarmonyOS】HUAWEI DevEco Studio 下载地址汇总03组基轨迹建模 GBTM的介绍与实现（Stata 或 R）04基于YOLOv10深度学习的CT扫描图像肾结石智能检测系统【python源码+Pyqt5界面+数据集+训练代码】深度学习实战、目标检测05(欧拉)openEuler系统添加网卡文件配置流程、(欧拉)openEuler系统手动配置ipv6地址流程、(欧拉)openEuler系统网络管理说明06【经验分享】Ubuntu22.04安装微信(linux官方版)07【TC3xx芯片】TC3xx芯片电源管理系统PMS详解08全面解析：构建基于深度学习的安全帽检测系统（UI界面+YOLO代码+数据集）09RAG 实践- Ollama+RagFlow 部署本地知识库10你还在傻傻的打开页面输用户名和密码？来我教你实现自动化