Go-Context底层原理剖析

Go-Context底层原理剖析

引用

CSDN

1.

https://blog.csdn.net/CTZL123456/article/details/139534295

Go语言的Context包是处理并发场景中取消信号和上下文值传递的重要工具。本文将深入剖析Context包的底层实现原理，包括默认上下文、同步取消信号和携带上下文值等功能，帮助开发者更好地理解和使用Context包。

背景

Context包的主要作用是在并发场景下同步取消信号和携带上下文值。下面我们来看看各种功能对应的实现。

内容

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

默认的上下文

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

func Background() Context { return backgroundCtx{} }
type backgroundCtx struct{ emptyCtx }

func TODO() Context { return todoCtx{} }
type todoCtx struct{ emptyCtx }

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

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函数是怎样写的：

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函数里面，函数的内容为下：

// 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()函数的实现：

// 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值。

我们来查看一下函数源码：

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是怎么实现的？

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是不安全的；