Go语言Context包源码学习
0前言
context包作为使用go进行server端开发的重要工具，其源码只有791行，不包含注释的话预计在500行左右，非常值得我们去深入探讨学习，于是在本篇笔记中我们一起来观察源码的实现，知其然更要知其所以然。（当前使用go版本为1.22.2）
1核心数据结构
整体的接口实现和结构体embed图
1.1Context接口
context接口定义了四个方法：
- Deadline方法返回context是否为timerctx，以及它的结束时间
- Done方法返回该ctx的done channel
- Err方法返回该ctx被取消的原因
- Value方法返回key对应的value
2emptyCtx
先来观察源码
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
}
emptyctx实现了context接口中的所有方法，对于每个方法返回的都是空值，它没有值、不能被取消以及没有截止时间，它只作为一个空context的载体，相当于所有ctx的祖先。
如何创建一个context？
context包提供了Background()方法和TODO()方法，都用于创建一个空的context。
func Background() Context {
return backgroundCtx{}
}
func TODO() Context {
return todoCtx{}
}
type backgroundCtx struct{ emptyCtx }
func (backgroundCtx) String() string {
return "context.Background"
}
type todoCtx struct{ emptyCtx }
func (todoCtx) String() string {
return "context.TODO"
}
它们返回的都为空的context，虽然方法名不同，但是效果是一样的，那么什么时候使用background，什么时候使用TODO呢，这是官方给出的注释
“TODO 会返回一个非空的、为空的 [Context]。 代码应该在不清楚应该使用哪个 [Context] 或者 [Context] 尚未可用（因为周围的函数尚未被扩展以接受 [Context] 参数）时使用 context.TODO。”
“background返回一个非空的、空的上下文对象。它不会被取消，没有值，也没有截止时间。它通常被主函数、初始化和测试使用，作为进入请求的顶级上下文。”
3cancelctx
先来观察cancelctx结构体的实现
type cancelCtx struct {
Context
mu sync.Mutex // protects following fields
done atomic.Value // of chan struct{}, created lazily, closed by first cancel call
children map[canceler]struct{} // set to nil by the first cancel call
err error // set to non-nil by the first cancel call
cause error // set to non-nil by the first cancel call
}
-
cancelCtx 内嵌了 Context 作为其父 context，根据go语言的特性，cancelCtx结构体就隐式实现了context接口中的所有方法，可以被当作一个接口来被调用，但是其方法还需要被具体的实现赋值才能进行调用。并且可以得知，cancelCtx的父类一定也是一个Context。
-
mu是cancelCtx的内置锁，用来协调并发场景下的资源获取
-
done的实际类型为chan struct{}，通过atomic包来实现并发安全，可以用于反应该ctx的生命周期情况，done是懒汉式创建的，只有第一次调用Done()方法时才会被创建，在下文的Done方法中会提到
-
children用于关联和子ctx的关系，当取消该ctx时，可以接连通知子ctx进行关闭，及时释放资源。
-
err用于返回ctx关闭的原因，调用的是context包定义的内置error
var Canceled = errors.New("context canceled") var DeadlineExceeded error = deadlineExceededError{} type deadlineExceededError struct{} func (deadlineExceededError) Error() string { return "context deadline exceeded" } func (deadlineExceededError) Timeout() bool { return true } func (deadlineExceededError) Temporary() bool { return true } -
cause返回的是该ctx失效更底层的原因，例如导致DeadlineExceeded err的具体原因是“database connection timeout”
// Example use: // // ctx, cancel := context.WithCancelCause(parent) // cancel(myError) // ctx.Err() // returns context.Canceled // context.Cause(ctx) // returns myError
3.1Done方法的实现
func (c *cancelCtx) Done() <-chan struct{} {
d := c.done.Load() //获取cancelCtx的done通道
if d != nil { //假若已经创建过了done通道，则直接返回
return d.(chan struct{})
}
c.mu.Lock() //上锁
defer c.mu.Unlock()
d = c.done.Load() //Double check是否创建done通道，因为在上锁前，可能其他goroutine调用了该ctx的Done方法。
if d == nil { //如果仍然未创建
d = make(chan struct{}) //创建该done channel
c.done.Store(d) //存储
}
return d.(chan struct{})
}
通过代码可以看见，ctx的done只有当被调用过Done方法时才会被创建，那么为什么这样子设计呢？很容易想到主要目的就是为了节省了不必要的资源浪费，提高效率，在很多情况下创建context并不需要监听done通道，只有在需要时才被创建，符合go语言的设计理念，只有需要的时候才引入。
3.2value方法的实现
func (c *cancelCtx) Value(key any) any {
if key == &cancelCtxKey {
return c
}
return value(c.Context, key)
}
若参数key与cancelCtxKey相符，则返回当前ctx本身
否则，就向父层 层层寻找
源码表现的非常晦涩，为了具体知道这个Value方法是做什么的，我们先来看对于cancelCtxKey的定义
// &cancelCtxKey is the key that a cancelCtx returns itself for.
var cancelCtxKey int
cancelCtxKey是一个私有的、唯一的标识符，它用于返回cancelCtx它本身。
对于该cancelCtxKey具体使用场景，下面还会讲到
3.3创建cancelCtx的WithCancel
//WithCancel返回一个带有新的Done通道的父context的副本。只有当父ctx被关闭，或者返回的cancel方法被调用时，该ctx的Done通道才会被关闭。
func WithCancel(parent Context) (ctx Context, cancel CancelFunc) {
c := withCancel(parent)
return c, func() { c.cancel(true, Canceled, nil) }
}
该方法返回了一个cancelCtx，以及关闭它的cancel方法。
在1.20版本中，新增了一个WithCancelCause方法，该方法返回了一个cancelctx和它的CancelCauseFunc，我们也来看一下
func WithCancelCause(parent Context) (ctx Context, cancel CancelCauseFunc) {
c := withCancel(parent)
return c, func(cause error) { c.cancel(true, Canceled, cause) }
}
在代码方面基本和withCancel方法一致，但是返回的CancelCauseFunc可以用于给用户自定义ctx被取消的原因，例如
ctx, cancel := context.WithCancelCause(parent)
cancel(myError)
ctx.Err() // returns context.Canceled
context.Cause(ctx) // returns myError
接着来看withCancel
func withCancel(parent Context) *cancelCtx {
if parent == nil { //若父ctx是nil，那么不能创建
panic("cannot create context from nil parent")
}
c := &cancelCtx{}
c.propagateCancel(parent, c) //传递cancel给新创建的ctx
return c
}
主要来看propagateCancel做了什么：
//该方法主要用于建立父context和子context的联系，如果父context也是一个cancelCtx，它需要保证父context被取消时，子context也能跟着被取消。
func (c *cancelCtx) propagateCancel(parent Context, child canceler) {
c.Context = parent //将父context内嵌入子context中
done := parent.Done() //获取父ctx的Done通道
if done == nil { //如果通道不存在，那么说明父context不是cancelCtx，不需要为它们两个之间建立联系，因为父context永远不会关闭。
return // parent is never canceled
}
select {
case <-done: //非堵塞地获取父done通道的状态，如果done通道以及被closed，那么这里会接受到一个零值，如果没有被closed，会执行default后面的语句。
//接受到零值，说明done通道以及被关闭了，父context已经被取消，此时应该立即调用子context的cancel方法，取消子context。
child.cancel(false, parent.Err(), Cause(parent))
return
default:
}
//否则，执行以下代码
//此处判断父context是否是一个context包实现的cancelCtx，如何理解会在下文讲述该方法时继续说明
if p, ok := parentCancelCtx(parent); ok {
// 如果父context是一个cancelCtx，或者是从某个cancelCtx衍生出来的context
p.mu.Lock() //加锁
if p.err != nil {
// 如果存在err，说明已经被取消，此时也应该调用child的cancel方法取消子ctx
child.cancel(false, p.err, p.cause)
} else {
if p.children == nil { //如果这是p的第一个子cancelCtx，需要初始化map的内存
p.children = make(map[canceler]struct{})
}
p.children[child] = struct{}{}//绑定子ctx到父ctx的map中，用于当父ctx取消时，能通知所有的child跟着取消。
}
p.mu.Unlock()
return
}
//如果父ctx不是cancelctx并且实现了AfterFuncer接口，即实现了AfterFunc方法（该方法会在ctx被取消后唯一一次调用），那么就需要为父ctx再设置一个afterFunc方法，用于取消child并且传递err和cause
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))
})
c.Context = stopCtx{
Context: parent,
stop: stop, //这里的stop方法可以用于当父ctx调用afterFunc的时候，取消父ctx对cancel函数的调用（看需求）
}
c.mu.Unlock()
return
}
//下面的情况为增加一个goroutine，监听父ctx自己实现的Done channel（用户自定义的）
goroutines.Add(1)
go func() {
select {
case <-parent.Done()://如果监听到父ctx自定义实现的Done channel关闭时，就关闭child
child.cancel(false, parent.Err(), Cause(parent))
case <-child.Done(): //如果child先关闭，那么就立即释放协程，避免协程泄露
}
}()
}
propagateCancel方法主要的作用是，保证了对于父ctx被取消时，为了能及时取消子ctx，避免不必要的资源浪费，建立父ctx和子ctx之间的联系。使用流程图来表示该方法如下（省略了检查afterFunc）
在该方法中，比较难以理解的地方是第二个if,"if p, ok := parentCancelCtx(parent); ok"究竟做了什么，为此我们跟进源码查看：
// parentCancelCtx 返回父级对象的底层 cancelCtx。
// 它通过查找 parent.Value(&cancelCtxKey) 来找到最内层的 enclosing cancelCtx，然后检查 parent.Done() 是否与该 cancelCtx 匹配。（如果不匹配，则该 cancelCtx 可能已被封装在提供了不同 done 通道的自定义实现中，在这种情况下，我们不应该绕过它。）
func parentCancelCtx(parent Context) (*cancelCtx, bool) {
done := parent.Done()
if done == closedchan || done == nil {
return nil, false
}
p, ok := parent.Value(&cancelCtxKey).(*cancelCtx)
if !ok {
return nil, false
}
pdone, _ := p.done.Load().(chan struct{})
if pdone != done {
return nil, false
}
return p, true
}
三个if分别做了什么：
-
第一个if：获取parentCtx的done channel，并且查看情况，若已经关闭或者父Ctx不可取消，此时返回false，回到propagate方法中最终会cancel掉child
-
第二个if：可以理解为通过Value方法找到离parentCtx“最近的”cancelCtx p，一般情况下，如果parentCtx就是一个CancelCtx，这时候就是parentCtx它本身。如果p不存在，也就是没有可以取消的ctx，此时也会返回false。
-
第三个if：找到了p后，还读取了p的done channel，这时候一般情况，pdone 当然会 == done，因此最终会返回p和true，那么什么时候会不相等呢？为什么会不相等呢？为此，看下方的层次图来理解这个if
在这个情况假设下，ParentCtx2是从ParentCtx1衍生出来的，ParentCtx1是一个标准的CancelCtx，而ParentCtx2是一个用户自定义了ctx，它内层继承了ParentCtx1并且自己实现了Done()方法，这时候代码中的“p”找到的就是ParentCtx1，而parent是ParentCtx2，此时的done和pdone就是两个不同的channel了，这时候cihldctx应该监听哪一个done channel呢？答案是监听用户自定义实现的Ctx的chennel，因为我们不应该绕过用户实现的Done channel，这更加符合ctx到层次逻辑。假如这时候不去判断pdone == done，直接返回的指针就是ParentCtx1的指针了。
3.4Cancel方法实现
接下来我们来看cancel方法是如何实现的。
//cancel 关闭 c.done，cancel 每一个c的children。如果removeFromParent为true，将会把c从parentCtx的child中移除。
func (c *cancelCtx) cancel(removeFromParent bool, err, cause error) {
if err == nil { //必须存在err
panic("context: internal error: missing cancel error")
}
if cause == nil { //没有自定义设置cause，默认为err
cause = err
}
c.mu.Lock() //上锁
if c.err != nil { //double check
c.mu.Unlock()
return // already canceled
}
c.err = err
c.cause = cause
d, _ := c.done.Load().(chan struct{})
if d == nil { //done没有被创建，直接存储context包内以及创建了的关闭的channel，不需要再次创建
c.done.Store(closedchan)
} else {
close(d) //关闭done
}
for child := range c.children { //关闭每一个子context
// NOTE: acquiring the child's lock while holding parent's lock.
child.cancel(false, err, cause)
}
c.children = nil
c.mu.Unlock()
if removeFromParent {
removeChild(c.Context, c) //从父child中移除自己
}
}
代码比较简单，问题在于什么时候removFromParent为true呢？为什么为true？
回到withCancel方法中，我们可以看到返回的cancel方法中，此时removeFromParent就为true
return c, func() { c.cancel(true, Canceled, nil) }
当用户主动调用cancel()时，就会将子ctx从父ctx中的child删除。因为此时没有必要再在父ctx中接受父或祖先的cancel通知。而当调用cancel函数内部，对child执行的cancel就为false，这是因为后面设置了c.children = nil，这时候是从父ctx的方向关闭了子ctx对其的链接。
4afterFuncCtx
afterFuncCtx是1.20版本后引入的新ctx，它的作用是当ctx被取消后，能执行一次自定义的F函数，一般用于回收资源等。
type afterFuncCtx struct {
cancelCtx
once sync.Once // either starts running f or stops f from running
f func()
}
可以看到afterFuncCtx embed了cancelCtx，在此基础上添加了once和f。once保证了f只会被执行一次。接着我们来看如何实现一个afterFuncCtx
4.1AfterFunc
func AfterFunc(ctx Context, f func()) (stop func() bool) {
a := &afterFuncCtx{ //创建局部变量afterFuncCtx，记录cancel时需要被执行的func
f: f,
}
a.cancelCtx.propagateCancel(ctx, a)
return func() bool { //返回stop方法，如果需要不执行func，则调用stop
stopped := false
a.once.Do(func() { //尝试stop func函数，如果成功则stop会被设置为true
stopped = true
})
if stopped { //stop后，a没有存在的意义，进行取消。
a.cancel(true, Canceled, nil)
}
return stopped //true为成功取消，false表示f已经被执行或者正在被执行。
}
}
这是一个闭包实现，闭包是指在函数内部定义的函数（如这里返回的 stop 函数），它会“捕获”并保存定义时可访问的所有外部变量。在 AfterFunc 方法中，虽然 a 是局部变量，但返回的 stop 方法引用了 a，形成了一个闭包，闭包会将 a 的内存保留在堆上，即使 AfterFunc 方法返回后，a 依然存在。所以当虽然没有返回a，但是返回的stop方法任然能调用a，a的生命周期超出了afterFunc方法。
以下是一个示例
package main
import (
"context"
"fmt"
"time"
)
func main() {
// 创建一个 2 秒后超时的 context
ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second)
defer cancel()
// 注册 AfterFunc，context 完成时将调用清理操作
stop := context.AfterFunc(ctx, func() {
fmt.Println("清理操作正在执行...")
})
// 等待 3 秒
time.Sleep(3 * time.Second)
// 尝试停止清理操作
stopped := stop()
if stopped {
fmt.Println("成功停止清理操作")
} else {
fmt.Println("清理操作已经开始或已停止")
}
// 等待 context 超时
<-ctx.Done()
fmt.Println("程序结束")
}
4timerCtx
接下来来看timerCtx
type timerCtx struct {
cancelCtx
timer *time.Timer // Under cancelCtx.mu.
deadline time.Time
}
timer直接内嵌了cancelCtx以实现Done和Err，并且新添了timer和deadline字段，deadline用于查看ctx的截止时间，timer用于完成过时取消ctx。
4.1WithDeadline
先来看WithTimeout方法
func WithTimeout(parent Context, timeout time.Duration) (Context, CancelFunc) {
return WithDeadline(parent, time.Now().Add(timeout))
}
WithTimeout接受一个存活时间来创建一个timerCtx，可以看到最终都是调用了WithDeadline
func WithDeadline(parent Context, d time.Time) (Context, CancelFunc) {
return WithDeadlineCause(parent, d, nil)
}
WithDeadline又调用了WithDeadlineCause，返回了timerCtx和一个CancelFunc。
4.2WithDeadlineCause
func WithDeadlineCause(parent Context, d time.Time, cause error) (Context, CancelFunc) {
if parent == nil {
panic("cannot create context from nil parent")
}
if cur, ok := parent.Deadline(); ok && cur.Before(d) { //如果parent的deadline更早，则直接返回parent的副本，不需要再创建timer
// The current deadline is already sooner than the new one.
return WithCancel(parent)
}
c := &timerCtx{
deadline: d,
}
c.cancelCtx.propagateCancel(parent, c)
dur := time.Until(d)
if dur <= 0 {//检测是否在创建过程中已经过了ddl
c.cancel(true, DeadlineExceeded, cause) // deadline has already passed
return c, func() { c.cancel(false, Canceled, nil) }
}
c.mu.Lock()
defer c.mu.Unlock()
if c.err == nil {
c.timer = time.AfterFunc(dur, func() { //设置超时自动cancel
c.cancel(true, DeadlineExceeded, cause)
})
}
return c, func() { c.cancel(true, Canceled, nil) }//返回timerCtx和cancelfunc
}
比较好理解，所以接着往下看cancel方法
4.3.cancel
func (c *timerCtx) cancel(removeFromParent bool, err, cause error) {
c.cancelCtx.cancel(false, err, cause)
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()
}
这里值得注意的地方只有第二行，它调用的是c.cancelCtx.cancel，取消的并非是父ctx，而是它本身，这里的作用是为了区分cancel方法的实现，c.cancelCtx是父ctx的一个副本，并不是父ctx，所以真正的parent是c.cancelCtx.Context。
5valueCtx
type valueCtx struct {
Context
key, val any
}
valueCtx内嵌了Context接口所以拥有该接口的所有方法，以及添加了k-v pair。
5.1WithValue
// 提供的键必须可比较，并且不应是字符串或任何其他内置类型，以避免在使用上下文时与其他包发生冲突。使用WithValue的用户应为其键定义自己的类型。为了避免在将值赋给接口{}时分配内存，上下文键通常具有具体的类型struct{}。或者，应将导出的上下文键变量的静态类型设置为指针或接口。
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() { //必须保证key是可比较的
panic("key is not comparable")
}
return &valueCtx{parent, key, val}
}
5.2value
func value(c Context, key any) any {
for {
switch ctx := c.(type) { //对context进行类型断言
case *valueCtx:
if key == ctx.key { //如果key就是当前ctx的key，则直接返回val
return ctx.val
}
c = ctx.Context //否则向父ctx查询
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 //不存在key value
default:
return c.Value(key) //返回用户实现的Value
}
}
}
可以看见，valueCtx查询value的过程，类似于链表查询，它是自底向上的，并且时间复杂度为O（N），它并不适用于存放大量的kv数据，原因有以下：
- 线性时间复杂度O（N），耗时太长
- 一个 valueCtx 实例只能存一个 kv 对，因此 n 个 kv 对会嵌套 n 个 valueCtx，造成空间浪费
- 不支持基于 k 的去重，相同 k 可能重复存在，并基于起点的不同，返回不同的 v. 由此得知，valueContext 的定位类似于请求头，只适合存放少量作用域较大的全局 meta 数据.
感谢观看，参考博客：
