Golang Context源码解析

Posted 2022-08-28 wx61307a0120efd

        自从接触Go语言以来，Context就一直在用，不过一直弄不清楚这东西究竟有啥用，平常都是CV代码，最多就是稀里糊涂用个Background()。不过最近在开发Golang反向代理的时候遇到一个传值问题，因为调用的是底层的库不支持传递变量，所以甚是困惑，后来在stackoverflow搜索了一番，发现已有类似问题，给出的答案很简单，就是通过Context的WithValue传值，而且底层库自身的参数也包含Context，所以困惑的问题一下就解决了。回过头来想一下，如此简单明显的问题都不会，说明对Golang了解的过于片面，这可不是一个优秀工程师应有的表现。然后就看了下Context包的源码。

        Context包其实写的挺简单的，代码量也不多，很适合阅读。看完之后给我感觉就是Context就像一个共享内存，在不同的goroutine之间共享一些信息，比如值、或者信号等等，设计比较简洁，代码也比较有参考价值。

        下面直接看代码：

package context

import (
    "errors"
    "internal/reflectlite"
    "sync"
    "sync/atomic"
    "time"
)

/// Context 是一个接口，包含四个方法Deadline(), Done(), Err(),Value()，Context主要用于传递信号、值等。
type Context interface     
    /// 如果有设置当前工作完成的截止时间的话，则返回该时间，否则返回空，timerCtx 内部实现就有个变量保存了截止时间，另外需要保证该接口调用的幂等性。
    Deadline() (deadline time.Time, ok bool)

    /// Done()返回一个用于通知工作结束的chan,该通道一般结合select使用，如果该context没有取消函数，则返回nil
    /// 当调用cancel的时候，该通道会被关闭，则所有监听该通道的子协程都会收到closed信号
    Done() <-chan struct

    /// 返回Done()被关闭的原因，未被closed的时候返回nil
    Err() error

    /// 返回context中该key对应的值，当前context找不到则从派生路径上找，都找不到返回nil
    Value(key interface) interface



/// 默认的context被取消错误
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 


/// emptyCtx 顾名思义，就是一个空的context，它实现了Context，不过啥都没做，相当于Context根节点
type emptyCtx int

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 interface) interface 
    return nil


func (e *emptyCtx) String() string 
    switch e 
    case background:
        return "context.Background"
    case todo:
        return "context.TODO"
    
    return "unknown empty Context"


var (
    /// Context包定义了两个私有指针变量background和todo，这两者都是emptyCtx
    background = new(emptyCtx)
    todo       = new(emptyCtx)
)


/// 返回私有变量background，主要用于作为顶层的context，类似于c++的基类
func Background() Context 
    return background



/// 返回私有变量todo，这东西和Background()唯一的区别就是命名不同
func TODO() Context 
    return todo



/// 自定义取消函数类型 当调用该类型函数的时候需要停止当前work,并无需等待。另外该类型的实现可能同时被多个协程调用，需要考虑并发的问题，另外首次调用之后，后续的调用直接返回即可。
/// 这个不难理解，如果一个context已经取消了，再次取消应该直接返回
type CancelFunc func()

/// WithCancel 返回一个cancelCtx 和一个取消函数，cancelCtx应该是除了context包中Context之外最重要的结构了，其它的WithDeadline、WithTimeout实际上都是基于WithCancel扩展的
/// 实现上就是初始化一个cancelCtx，然后和parent context绑定，绑定的目的是当parent取消的时候，能够及时取消子context，具体实现看下面代码
func WithCancel(parent Context) (ctx Context, cancel CancelFunc) 
    /// 不允许从nil context派生
    if parent == nil  
        panic("cannot create context from nil parent")
    
    /// 这一步就是初始化cancelCtx
    c := newCancelCtx(parent) 
    
    /// 绑定到父context上
    propagateCancel(parent, &c) 

    /// 返回一个cancel函数，这个函数就是调用cancel取消自己，同时第一个参数表明需要将ctx从parent context的孩子队列中移除，移除就是解除和parent context的绑定关系
    return &c, func()  c.cancel(true, Canceled)  



/// 初始化一个cancelCtx
func newCancelCtx(parent Context) cancelCtx 
    return cancelCtxContext: parent



/// 保存go协程创建的数量，测试用
var goroutines int32


/// propagate 这个单词英文意思是繁衍，增殖; 使遗传; 扩散; 使蔓延，在这里指绑定父context和子context，绑定的目的是使得父context的取消能够影响到子context，否则父context取消不能及时通知子context
func propagateCancel(parent Context, child canceler) 
    done := parent.Done() 
    if done == nil 
        /// 如果父context 没有cancel channel，说明父context不会执行cancel，所以这里不用绑定了
        return // parent is never canceled
    

    select 
    /// 如果父context已经取消了，则将子context也取消掉，然后就可以返回了
    case <-done:
        // parent is already canceled 
        child.cancel(false, parent.Err())
        return
    default:
    

    /// 这一步是获取派生路径上的cancelCtx，如果有的话，则绑定到父context的children队列
    if p, ok := parentCancelCtx(parent); ok  
        p.mu.Lock() /// 加把锁保护下，防止并发
        if p.err != nil  /// 再次检查父context是否取消，如果取消了则相应的也应该将子context取消
            // parent has already been canceled
            child.cancel(false, p.err)
         else 
            // 如果父context没有取消则将子context加入到父context的children队列里面，这样当父context取消的时候，它能够取消所有子context
            if p.children == nil  /// 初始化队列
                p.children = make(map[canceler]struct)
            
            p.children[child] = struct /// 添加到队列
        
        p.mu.Unlock()
     else 
        /// 走到这里有两种情形，一种父context是cancelCtx，不过已经取消了，另一种是父context不是cancelCtx,这种情况有可能是用户自定义的cancelCtx，这种情况下就不用加到父contex队列里面去了，因为不知道父context有没有队列这个属性
        atomic.AddInt32(&goroutines, +1)
        go func() 
            select 
            case <-parent.Done(): // 如果父context先取消，则接着取消子context
                child.cancel(false, parent.Err())
            case <-child.Done():  // 如果子context先取消 则直接返回就OK了
            
        ()
    


// 顾名思义，这个专门给cancelCtx用的key
var cancelCtxKey int 

// 获取派生路径上父cancelCtx
func parentCancelCtx(parent Context) (*cancelCtx, bool) 
    done := parent.Done()

    // 如果父context已取消或者没有父context的Done channel为空，则直接返回
    if done == closedchan || done == nil 
        return nil, false
    
    /// 获取父辈派生路径上最近的cancelCtx, 这里使用了cancelCtxKey, 只有cancelCtx类型才实现了它
    /// 可以思考下这里为啥不用 p, ok := parent.(*cancelCtx)???，我觉得这样写也可以，只不过不符合原本的context设计规范，就像上面的propagateCancel，无非是加入到childen列表还是另起一个协程监听的区别，如果理解不正确，欢迎指正
    p, ok := parent.Value(&cancelCtxKey).(*cancelCtx)
    if !ok 
        return nil, false
    
    p.mu.Lock()
    // 判断和父context是否一样，不一样说明是父节点的父节点
    ok = p.done == done
    p.mu.Unlock()
    if !ok 
        return nil, false
    
    // 走到这里说明找到了派生路径上最近的一个cancelCtx类型
    return p, true



// 从父context中移除孩子节点自己，其实就是解除和parent的绑定
func removeChild(parent Context, child canceler) 
    /// 如果父节点不是cancelCtx类型，说明父context没有childen队列，因此无需删除直接返回
    p, ok := parentCancelCtx(parent)
    if !ok 
        return
    
    p.mu.Lock()
    /// 存在父节点，从父节点中删除自己
    if p.children != nil 
        delete(p.children, child)
    
    p.mu.Unlock()



/// canceler 接口，包含两个方法，cancelCtx实现了它
type canceler interface 
    cancel(removeFromParent bool, err error) // cancel函数，第一个参数表明是否需要从父context的孩子队列中移除，一般情况下如果child context主动取消，这时候就需要从父context中移除自身，使得父context取消的时候无需再取消子context

    // 返回Done chan
    Done() <-chan struct


/// 顾名思义，创建一个取消的channel,啥时候取消呢，包初始化的时候
var closedchan = make(chan struct)

/// 初始化操作直接将私有变量closedchan 关掉, 话说所有cancelCtx返回的cancel后，Done()返回同一个channel
func init() 
    close(closedchan)



/// cancelCtx 实现很好理解，首先需要一个变量保存父context，其次需要一个变量保存子context列表，另外还有个channel用于通知，对于这些变量，还需要加把锁保护下
type cancelCtx struct 
    Context

    /// 加把锁 并发保护下
    mu       sync.Mutex      
    /// 懒人模式，需要的时候才创建   
    done     chan struct          

    /// 子context列表 
    children map[canceler]struct // set to nil by the first cancel call 

    /// cancel的原因
    err      error                 


/// cancelCtx Value()实现很巧妙，可以在自己的代码里借鉴下
func (c *cancelCtx) Value(key interface) interface 
    if key == &cancelCtxKey  // 这里用到一个私有变量，只有cancelCtx类型才会返回自己
        return c
    
    return c.Context.Value(key) 


/// 返回cancelCtx Done通道
func (c *cancelCtx) Done() <-chan struct 
    c.mu.Lock()
    if c.done == nil  // 没有初始化则初始化，返回channel,调用cancel后才会关闭
        c.done = make(chan struct)
    
    d := c.done
    c.mu.Unlock()
    return d


/// 设置cancel err
func (c *cancelCtx) Err() error  
    c.mu.Lock()
    err := c.err
    c.mu.Unlock()
    return err


type stringer interface 
    String() string


func contextName(c Context) string 
    if s, ok := c.(stringer); ok 
        return s.String()
    
    return reflectlite.TypeOf(c).String()


func (c *cancelCtx) String() string 
    return contextName(c.Context) + ".WithCancel"



/// cancelCtx的cancel方法 主要逻辑就是检查自身是否关闭，没关闭则关闭一下，另外看是否需要将自身从父context移除
func (c *cancelCtx) cancel(removeFromParent bool, err error) 
    // 调用cancel方法必须传递一个err进来
    if err == nil 
        panic("context: internal error: missing cancel error")
    
    c.mu.Lock()       // 先抢占下锁，防止并发调用
    if c.err != nil  // 看看是否被其它协程取消了，如果是的话，直接返回
        c.mu.Unlock()
        return // already canceled
    
    // 取消的时候 err赋值
    c.err = err
    
    // 这里为nil说明还没有人监听channel,这里直接返回默认的close channel, 这样其它子协程即便监听也会立刻返回
    if c.done == nil  
        c.done = closedchan
     else  // 如果不为nil，说明有人监听了，这里直接关闭
        close(c.done)
    

    for child := range c.children      
        /// 将自身衍生的所有子context都取消掉，传递false，表示不需要从父context中移除自己，因为父context自己做了这步操作
        child.cancel(false, err)
    

    // 清空子context
    c.children = nil
    c.mu.Unlock()

    // 是否从父context中移除，通常cancel自己的话会传递true，标明将自己从父context中移除
    if removeFromParent 
        removeChild(c.Context, c)
    



/// WithDeadline返回一个timerCtx，本质上也是一个cancelCtx，相较于cancelCtx必须主动调用cancel之外，timerCtx还多了一种cancel调用方法，即设定一个定时器，到了指定时间自动取消
func WithDeadline(parent Context, d time.Time) (Context, CancelFunc) 
    /// parent context 不允许为nil
    if parent == nil 
        panic("cannot create context from nil parent")
    

    /// 如果父context的截止时间比当前设定的要早，则直接返回一个cancelCtx将自身和parent 绑定下，这是因为父context一定先取消，这时候只需要和父context绑定等ta调cancel取消自己就好了
    if cur, ok := parent.Deadline(); ok && cur.Before(d) 
        // The current deadline is already sooner than the new one.
        return WithCancel(parent)
    

    /// 如果父context取消的时间比自身还要晚，那就不能等待父context取消自己，需要自己主动起一个定时器取消，所以这里新建一个timerCtx
    c := &timerCtx
        cancelCtx: newCancelCtx(parent),
        deadline:  d,
    

    /// 和cancelCtx一样，绑定到父context上
    propagateCancel(parent, c)

    dur := time.Until(d)
    /// 如果已经到了截止时间，则直接取消就好了
    if dur <= 0 
        c.cancel(true, DeadlineExceeded) 
        /// 这里传递false是因为上一步传了true
        return c, func()  c.cancel(false, Canceled)  
    

    c.mu.Lock()
    defer c.mu.Unlock()
    /// 再次检查是否被取消，当父context取消的时候这里err不为nil
    if c.err == nil  
        c.timer = time.AfterFunc(dur, func()  /// 起一个定时器，超时则调用cancel 函数
            c.cancel(true, DeadlineExceeded)
        )
    
    return c, func()  c.cancel(true, Canceled) 



/// 内部封装了一个cancelCtx，另外需要一个定时器和一个时间变量保存相关配置
type timerCtx struct 
    cancelCtx
    timer *time.Timer // Under cancelCtx.mu.

    deadline time.Time


/// 返回截止时间
func (c *timerCtx) Deadline() (deadline time.Time, ok bool) 
    return c.deadline, true


func (c *timerCtx) String() string 
    return contextName(c.cancelCtx.Context) + ".WithDeadline(" +
        c.deadline.String() + " [" +
        time.Until(c.deadline).String() + "])"


/// timerCtx 和cancelCtx 取消唯一不同的地方在于timerCtx还需要关闭定时器
func (c *timerCtx) cancel(removeFromParent bool, err error) 

    /// 先通知子context关闭
    c.cancelCtx.cancel(false, err)
    
    /// 自己主动调用情况下 则需要将自身从父context中移除
    if removeFromParent  
        removeChild(c.cancelCtx.Context, c)
    

    c.mu.Lock()
    if c.timer != nil  /// 关闭定时器
        c.timer.Stop()
        c.timer = nil
    
    c.mu.Unlock()


/// WithTimeout 这个实现其实和WithDeadline一样，区别在于Deadline可以指定任意时间，WithTimeout则是在当前时间加上timeout
func WithTimeout(parent Context, timeout time.Duration) (Context, CancelFunc) 
    return WithDeadline(parent, time.Now().Add(timeout))



/// WithValue返回一个存储K,V的valueCtx, 这里需要key是可比较的，并且不建议使用内置的string等类型作为key，因为那样会导致冲突，另外
/// 不应该把WithValue用来传递普通参数。通常情况下，WithValue 用于传递一些跨API的数据，比如用于链路跟踪的traceID等等
/// valueCtx实现也很简单，主要就是一个K，V加上一个parent context
func WithValue(parent Context, key, val interface) Context     
    /// 同样不允许parent为nil
    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 &valueCtxparent, key, val


/// valueCtx实现更简单，就包含一个parent ctx和一对K,V，因为只有一对，所以也不需要锁保护
type valueCtx struct 
    Context
    key, val interface


// stringify tries a bit to stringify v, without using fmt, since we dont
// want context depending on the unicode tables. This is only used by
// *valueCtx.String().
func stringify(v interface) string 
    switch s := v.(type) 
    case stringer:
        return s.String()
    case string:
        return s
    
    return "<not Stringer>"


func (c *valueCtx) String() string 
    return contextName(c.Context) + ".WithValue(type " +
        reflectlite.TypeOf(c.key).String() +
        ", val " + stringify(c.val) + ")"


/// 查找key对应的value，找不到则从派生路径上去找
func (c *valueCtx) Value(key interface) interface 
    if c.key == key 
        return c.val
    
    return c.Context.Value(key)

参考文档：

1. ​​Go 语言并发编程与 Context | Go 语言设计与实现​​