必知必会系列-Context

定义

在Go服务器中，每个传入的请求都在自己的goroutine中处理。请求的处理程序经常启动额外的goroutine来访问后端服务，如数据库和RPC服务。处理一个请求的一组goroutine通常需要访问该请求相关的特定的值，比如最终用户的身份、授权令牌和请求的deadline等。当一个请求被取消或处理超时时，所有在该请求上工作的goroutines应该迅速退出，以便系统可以回收他们正在使用的任何资源。这时便诞生了Context，主要用于

1. goroutine 之间推出通知
2. 元数据传递 上下文 context.Context Go 语言中用来设置截止日期、同步信号，传递请求相关值的结构体。上下文与 Goroutine 有比较密切的关系，是 Go 语言中独特的设计，在其他编程语言中我们很少见到类似的概念。

context.Context 是一个接口，定义了四个需要实现的方法：

1. Deadline --- 返回 context.Context 被取消的时间，也就是完成工作的截止日期

2. Done --- 返回一个 Channel，这个 Channel 会在当前工作完成或者上下文被取消后关闭，多次调用 Done 方法会返回同一个 Channel

3. Err --- 返回 context.Context 结束的原因，它只会在 Done 方法对应的 Channel 关闭时返回非空的值

   1. 如果 context.Context 被取消，会返回 Canceled 错误
   2. 如果 context.Context 超时，会返回 DeadlineExceeded 错误

4. Value --- 从 context.Context 中获取键对应的值，对于同一个上下文来说，多次调用 Value 并传入相同的 Key 会返回相同的结果，该方法可以用来传递请求特定的数据；

主要方法

• context.Background
• context.TODO
• context.WithDeadline
• context.WithValue
• context.WithCancel

其中

• context.Background和 context.TODO 只是通过空方法实现了该接口，主要用于函数入参占位context.Background 是上下文默认值，所有其他上下文都应该从它衍生出来
• context.TODO 应该仅在不确定应该使用哪种上下文时使用。在多数情况下，如果当前函数没有上下文作为入参，我们都会使用 context.Background 作为起始的上下文向下传递

// Context's methods may be called by multiple goroutines simultaneously.
type Context interface {
    // Deadline returns the time when work done on behalf of this context
    // should be canceled. Deadline returns ok==false when no deadline is
    // set. Successive calls to Deadline return the same results.
    Deadline() (deadline time.Time, ok bool)

    // Done returns a channel that's closed when work done on behalf of this
    // context should be canceled. Done may return nil if this context can
    // never be canceled. Successive calls to Done return the same value.
    // The close of the Done channel may happen asynchronously,
    // after the cancel function returns.
    Done() <-chan struct{}

    // If Done is not yet closed, Err returns nil.
    // If Done is closed, Err returns a non-nil error explaining why:
    // Canceled if the context was canceled
    // or DeadlineExceeded if the context's deadline passed.
    // After Err returns a non-nil error, successive calls to Err return the same error.
    Err() error

    Value(key any) any
}

内部实现结构体

canceler 取消接口

• 有两个实现 cancelCtx timerCtx

// A canceler is a context type that can be canceled directly. The
// implementations are cancelCtx and timerCtx.
type canceler interface {
    // removeFromParent 表示是否与父节点断开联系
    cancel(removeFromParent bool, err error)
    Done() <-chan struct{}
}

可复用的通道 closedchan

// closedchan is a reusable closed channel.
var closedchan = make(chan struct{})

func init() {
    close(closedchan)
}

emptyCtx

• 空实现

// An emptyCtx is never canceled, has no values, and has no deadline. It is not
// struct{}, since vars of this type must have distinct addresses.
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 any) any {
    return nil
}

cancelCtx

• mu 用于 保证并发安全
• 第一次调用cancel方法 会将done中存储的 chan struct{} 关闭
• children 存放该节点下的子节点，用于取消信号的向下传递
• Err 用于记录取消原因 超时取消｜cancel方法调用

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
}

func (c *cancelCtx) Value(key any) any {
    // &cancelCtxKey is the key that a cancelCtx returns itself for.
    if key == &cancelCtxKey {
        return c
    }
    return value(c.Context, key)
}

func (c *cancelCtx) Done() <-chan struct{} {
    d := c.done.Load()
    if d != nil {
        return d.(chan struct{})
    }
    c.mu.Lock()
    defer c.mu.Unlock()
    d = c.done.Load() 
    // double checked locking 双重检查
    if d == nil {
        d = make(chan struct{})
        c.done.Store(d) 
    }
    return d.(chan struct{})
}

func (c *cancelCtx) Err() error {
    c.mu.Lock()
    err := c.err
    c.mu.Unlock()
    return err
}

// cancel closes c.done, cancels each of c's children, and, if
// removeFromParent is true, removes c from its parent's children.
func (c *cancelCtx) cancel(removeFromParent bool, err error) {
    if err == nil {
        panic("context: internal error: missing cancel error")
    }
    c.mu.Lock()
    if c.err != nil {
        c.mu.Unlock()
        return // already canceled
    }
    c.err = err
    d, _ := c.done.Load().(chan struct{})
    // 懒加载，如果为空，则设置为 复用的已经关闭的通道closedchan
    if d == nil {
        c.done.Store(closedchan)
    } else {
        // 如果不为 nil 则关闭该通道
        close(d)
    }
    
    // 取消子节点
    for child := range c.children {
        // NOTE: acquiring the child's lock while holding parent's lock.
        child.cancel(false, err)
    }
    c.children = nil
    c.mu.Unlock()
    
    // 将该节点从父节点移除
    if removeFromParent {
        removeChild(c.Context, c)
    }
}

重点讲一下 cancelCtx实现的 Done方法和cancel()方法

Done() 方法

加锁 从变量done中获取chan struct{}，如果为nil 则初始化一个，并存储在变量done中，保证该方法返回的幂等性。如果不为nil则直接返回

cancel()方法 cancel(removeFromParent bool, err error)

• 加锁 获取err字段，如果不为空，说明已经被取消了，提前返回
• 从变量done中获取chan struct{}，如果为nil 说明没有调用过Done方法【因为该变量是懒加载，在Done方法调用的时候才会初始化该变量】，直接赋值一个已经关闭的chan closedchan【局部变量，可复用】
• 如果获取的done变量不为空，说明done变量初始化过，则调用close方法关闭该chan
• 遍历children节点，依次调用子节点的取消方法
• 如果removeFromParent=true, 还需要在父节点的children删除该节点

timerCtx

• 内嵌cancelCtx结构体
• timer 定时器 用于实现超时后 调用cancel方法

// A timerCtx carries a timer and a deadline. It embeds a cancelCtx to
 // implement Done and Err.  It implements cancel by stopping its timer then
// delegating to cancelCtx.cancel.
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) cancel(removeFromParent bool, err error) {
    c.cancelCtx.cancel(false, err)
    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()
}

---

内部方法

newCancelCtx

初始化 带取消功能的 context

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
}

// newCancelCtx returns an initialized cancelCtx.
func newCancelCtx(parent Context) cancelCtx {
    return cancelCtx{Context: parent}
}

parentCancelCtx

• 从context中 寻找 cancelCtx
• 如果用户实现了自定义的Context接口，该函数返回false,nil

// parentCancelCtx returns the underlying *cancelCtx for parent.
// It does this by looking up parent.Value(&cancelCtxKey) to find
// the innermost enclosing *cancelCtx and then checking whether
// parent.Done() matches that *cancelCtx. (If not, the *cancelCtx
// has been wrapped in a custom implementation providing a
// different done channel, in which case we should not bypass it.)
func parentCancelCtx(parent Context) (*cancelCtx, bool) {
    done := parent.Done()
    // 如果父节点不是可取消节点 或者父节点已经调用了cancel方法 提前返回
    if done == closedchan || done == nil {
        return nil, false
    }
    // 如果用户实现了自定义的context接口，不做任何处理
    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
}

removeChild

• 从父节点中移除子节点

// removeChild removes a context from its parent.
// 如果父节点是可取消节点，才将子节点从父节点中移除
func removeChild(parent Context, child canceler) {
    p, ok := parentCancelCtx(parent)
    // 如果父节点不是一个可取消节点，提前返回
    if !ok {
        return
    }
    p.mu.Lock()
    if p.children != nil {
        delete(p.children, child)
    }
    p.mu.Unlock()
}

propagateCancel

• 建立子节点和父节点的关联关系，便于取消信号的传递

// propagateCancel arranges for child to be canceled when parent is.
func propagateCancel(parent Context, child canceler) {
    done := parent.Done()
    if done == nil {
        return // parent is never canceled
    }

    select {
    case <-done:
        // 父节点已经取消
        child.cancel(false, parent.Err())
        return
    default:
    }

    if p, ok := parentCancelCtx(parent); ok {
        p.mu.Lock()
        if p.err != nil {
            // 父节点已经取消
            child.cancel(false, p.err)
        } else {
            // 将子节点加入到父节点的children中
            if p.children == nil {
                p.children = make(map[canceler]struct{})
            }
            p.children[child] = struct{}{}
        }
        p.mu.Unlock()
    } else {
         // 如果父节点不是内部定义的cancelCtx 另外开协程 监听父节点的取消 & 处理子节点的取消
        atomic.AddInt32(&goroutines, +1)
        go func() {
            select {
            case <-parent.Done():
                child.cancel(false, parent.Err())
            case <-child.Done():
            }
        }()
    }
}

value方法

// 递归寻找key
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
            }
            c = ctx.Context
        case *timerCtx:
            if key == &cancelCtxKey {
                return &ctx.cancelCtx
            }
            c = ctx.Context
        case *emptyCtx:
            return nil
        default:
            return c.Value(key)
        }
    }
}

对外暴露的方法

WithValue

• WithValue 对应的key最好不要是string类型或者任何build-in类型，以避免key冲突。且最好是指针类型

// WithValue returns a copy of parent in which the value associated with key is
// val.
//
// Use context Values only for request-scoped data that transits processes and
// APIs, not for passing optional parameters to functions.
//
// The provided key must be comparable and should not be of type
 // string or any other built-in type to avoid collisions between
// packages using context. Users of WithValue should define their own
// types for keys. To avoid allocating when assigning to an
// interface{}, context keys often have concrete type
// struct{}. Alternatively, exported context key variables' static
// type should be a pointer or interface.
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}
}

// A valueCtx carries a key-value pair. It implements Value for that key and
// delegates all other calls to the embedded Context.
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)
}

WithCancel 可取消上下文

• 内部会创建 cancelCtx

• 通过propagateCancel 建立父节点和子节点的联系，便于取消信号的透传

• 返回取消方法 cancelCtx.cancel

// WithCancel returns a copy of parent with a new Done channel. The returned
// context's Done channel is closed when the returned cancel function is called
// or when the parent context's Done channel is closed, whichever happens first.
//
// Canceling this context releases resources associated with it, so code should
// call cancel as soon as the operations running in this Context complete.
func WithCancel(parent Context) (ctx Context, cancel CancelFunc) {
    if parent == nil {
        panic("cannot create context from nil parent")
    }
    c := newCancelCtx(parent)
    propagateCancel(parent, &c)
    return &c, func() { c.cancel(true, Canceled) }
}

WithTimeout 超时取消上下文【timeout】

本质调用WithDeadline 方法

func WithTimeout(parent Context, timeout time.Duration) (Context, CancelFunc) {
    return WithDeadline(parent, time.Now().Add(timeout)) 
}

WithDeadline 超时取消上下文

func WithDeadline(parent Context, d time.Time) (Context, CancelFunc)

WithDeadline 超时

• 如果父节点的超时时间小于当前设置的，以父节点为准
• 当前时间大于超时时间或者调用了CancelFunc，context.Done()通道关闭

// WithDeadline returns a copy of the parent context with the deadline adjusted
// to be no later than d. If the parent's deadline is already earlier than d,
// WithDeadline(parent, d) is semantically equivalent to parent. The returned
// context's Done channel is closed when the deadline expires, when the returned
// cancel function is called, or when the parent context's Done channel is
// closed, whichever happens first.
//
// Canceling this context releases resources associated with it, so code should
// call cancel as soon as the operations running in this Context complete.

func WithDeadline(parent Context, d time.Time) (Context, CancelFunc) {
    if parent == nil {
        panic("cannot create context from nil parent")
    }
    if cur, ok := parent.Deadline() ; ok && cur.Before(d) {
        // The current deadline is already sooner than the new one.
        return WithCancel(parent)
    }
    c := &timerCtx{
        cancelCtx: newCancelCtx(parent),
        deadline:  d,
    }
    // 在 parent 和 child 之间同步取消和结束的信号，
    // 保证在 parent 被取消时，child 也会收到对应的信号
    propagateCancel(parent, c)
    dur := time.Until(d)
    // 发现已经到超时时间，提前取消
    if dur <= 0 {
        c.cancel(true, DeadlineExceeded) // deadline has already passed
        return c, func() { c.cancel(false, Canceled) }
    }
    c.mu.Lock()
    defer c.mu.Unlock()
    if c.err == nil {
        // 注册定时器，指定时间后执行 cancel方法
        c.timer = time.AfterFunc(dur, func ()  {
            c.cancel(true, DeadlineExceeded)
        })
    }
    return c, func() { c.cancel(true, Canceled) }
}

---

相关知识点

atomic.Value

原子地设置和读取变量

• 不能用atomic.Value原子值存储nil
• 第一次向原子值存储值,决定了它今后能且只能存储该类型的值

// A Value provides an atomic load and store of a consistently typed value.
// The zero value for a Value returns nil from Load.
// Once Store has been called, a Value must not be copied.
//
// A Value must not be copied after first use.
type Value struct {
    v any
}

// ifaceWords is interface{} internal representation.
type ifaceWords struct {
    // 指向类型
    typ  unsafe.Pointer
    // 指向数据
    data unsafe.Pointer
}

Store 方法

• 通过unsafe.Pointer将现有的 和要写入的 值分别转成ifaceWords类型，得到这两个interface{}的原始类型（typ）和真正的值（data）
• 通过LoadPointer这个原子操作拿到当前Value中存储的类型。下面根据这个类型的不同，分3种情况处理：

1. （第一次写入） 使用CAS操作，先尝试将typ设置为uintptr(0)中间状态。如果失败，则说明其他协程抢先完成了赋值操作，继续for循环。如果设置成功，那证明当前协程抢到了这个"乐观锁"，则先写data字段，再设置typ字段
2. （第一次写入未完成） 继续循环，"忙等"直到第一次写入完成
3. （第一次写入已完成） 检查上一次写入的类型与这一次要写入的类型是否一致，如果不一致则抛出异常。反之，则直接把这一次要写入的值写入到data字段

var firstStoreInProgress byte

// Store sets the value of the Value to x.
// All calls to Store for a given Value must use values of the same concrete type.
// Store of an inconsistent type panics, as does Store(nil).
func (v *Value) Store(val any) {
    if val == nil {
        panic("sync/atomic: store of nil value into Value")
    }
    // 通过Unsafe转换为 interface的内部实现结构体
    vp := (*ifaceWords)(unsafe.Pointer(v))
    vlp := (*ifaceWords)(unsafe.Pointer(&val))
    for {
        typ := LoadPointer(&vp.typ)
        // tpe==nil 表示该字段第一次set
        if typ == nil {
            // Attempt to start first store.
            // Disable preemption so that other goroutines can use
            // active spin wait to wait for completion.
            runtime_procPin() // 关闭协程的抢占
            // 通过CAS 原子指令设置一个中间状态
            if !CompareAndSwapPointer(&vp.typ, nil, unsafe.Pointer(&firstStoreInProgress)) {
                runtime_procUnpin()
                continue
            }
            // Complete first store.
            StorePointer(&vp.data, vlp.data)
            StorePointer(&vp.typ, vlp.typ)
            runtime_procUnpin()
            return
        }
        // 如果第一次set还没有完成，继续等待
        if typ == unsafe.Pointer(&firstStoreInProgress) {
            // First store in progress. Wait.
            // Since we disable preemption around the first store,
            // we can wait with active spinning.
            continue
        }
        // 第一次设置完成，便只需要比较类型，然后原子设置data字段便可
        if typ != vlp.typ {
            panic("sync/atomic: store of inconsistently typed value into Value")
        }
        StorePointer(&vp.data, vlp.data)
        return
    }
}

Load 方法

1. 如果当前的typ是 nil 或者uintptr(0)，证明第一次写入还没有开始，或者还没完成，那就直接返回 nil
2. 反之，基于当前的typ和data构造出一个新的interface{}返回出去

// Load returns the value set by the most recent Store.
// It returns nil if there has been no call to Store for this Value.
func (v *Value) Load() (val any) {
    vp := (*ifaceWords)(unsafe.Pointer(v)) 
    typ := LoadPointer(&vp.typ)
    if typ == nil || typ == unsafe.Pointer(&firstStoreInProgress) {
        // First store not yet completed.
        return nil
    }
    data := LoadPointer(&vp.data)
    vlp := (*ifaceWords)(unsafe.Pointer(&val))
    vlp.typ = typ
    vlp.data = data
    return
}