Golang 之 context 深入解读

type Context interface {


    // 返回 context 是否设置了超时时间以及超时的时间点
    // 如果没有设置超时，那么 ok 的值返回 false
    // 每次调用都会返回相同的结果
    Deadline() (deadline time.Time, ok bool)


    // 如果 context 被取消，这里会返回一个被关闭的 channel
    // 如果是一个不会被取消的 context，那么这里会返回 nil
    // 每次调用都会返回相同的结果
    Done() <-chan struct{}


    // 返回 done() 的原因
    // 如果 Done() 对应的通道还没有关闭，这里返回 nil
    // 如果通道关闭了，这里会返回一个非 nil 的值：
    // - 若果是被取消掉的，那么这里返回 Canceled 错误
    // - 如果是超时了，那么这里返回 DeadlineExceeded 错误
    // 一旦被赋予了一个非 nil 的值之后，每次调用都会返回相同的结果
    Err() error


    // 获取 context 中保存的 key 对应的 value，如果不存在则返回 nil
    // 每次调用都会返回相同的结果
    Value(key interface{}) interface{}
}
复制

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 }
复制

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
}
复制

var (
    background = new(emptyCtx)
    todo       = new(emptyCtx)
)


func Background() Context {
    return background
}


func TODO() Context {
    return todo
}
复制

type valueCtx struct {
    // 父 ctx
    Context
    // kv
    key, val interface{}
}
复制

func WithValue(parent Context, key, val interface{}) 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}
}
复制

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

func (c *valueCtx) Value(key interface{}) interface{} {
    if c.key == key {
        return c.val
    }
    // 如果当前 ctx 没有找到，那么会去父 ctx 上面去找
    return c.Context.Value(key)
}
复制

type canceler interface {
    cancel(removeFromParent bool, err error)
    Done() <-chan struct{}
}
复制

type cancelCtx struct {
    // 父 ctx
    Context
    // 通过互斥锁来保证对下面三个 field 操作的安全性
    mu       sync.Mutex
    // 保存一个 chan struct{}，第一个 cancel() 调用会关闭这个通道
    done     atomic.Value
    // 维护所有子 canceler
    // 当前 ctx 被取消之后，它的所有子 canceler 都会被取消，并且这个属性会被置空
    children map[canceler]struct
    // 第一个 cancel() 调用会赋值 
    err      error
}
复制

func WithCancel(parent Context) (ctx Context, cancel CancelFunc) {
    if parent == nil {
        panic("cannot create context from nil parent")
    }
    // 初始化一个 cancelCtx 实例
    c := newCancelCtx(parent)
    // 将当前 ctx 挂载到父节点上，从而实现 cancel 操作的向下传播
    propagateCancel(parent, &c)
    return &c, func() { c.cancel(true, Canceled) }
}
复制

var cancelCtxKey int


func (c *cancelCtx) Value(key interface{}) interface{} {
    // 特殊路径，如果传入的 key 是 &cancelCtxKey，那么直接返回当前的 ctx
    if key == &cancelCtxKey {
        return c
    }
    // 否则去它的父 ctx 上面查找对应的 value
    return c.Context.Value(key)
}
复制

func (c *cancelCtx) cancel(removeFromParent bool, err error) {
    if err == nil {
        panic("context: internal error: missing cancel error")
    }
    c.mu.Lock()
    // 如果 c.err 不是 nil，代表已经被取消了，直接返回
    if c.err != nil {
        c.mu.Unlock()
        return
    }
    // 标记已经被取消
    c.err = err
    d, _ := c.done.Load().(chan struct{})
    // 关闭 channel，从而可以通知到其他协程
    // 如果 d == nil，代表之前没有调用过 Done() 方法，这里直接传入一个 closedchan（关闭的通道）
    if d == nil {
        c.done.Store(closedchan)
    } else {
        // 否则的话，需要关闭当前的 channel
        close(d)
    }
    // cancel 子节点
    for child := range c.children {
        child.cancel(false, err)
    }
    c.children = nil
    c.mu.Unlock()
    // 如果需要从父节点中移除当前节点，那么执行 remove 操作
    if removeFromParent {
        removeChild(c.Context, c)
    }
}


// removeChild 移除一个子节点
func removeChild(parent Context, child canceler) {
    p, ok := parentCancelCtx(parent)
    if !ok {
        return
    }
    p.mu.Lock()
    if p.children != nil {
        // 从 map 中移除当前元素
        delete(p.children, child)
    }
    p.mu.Unlock()
}
复制

func propagateCancel(parent Context, child canceler) {
    done := parent.Done()
    // done == nil 代表父 ctx 不是一个可以被 cancel 的 ctx，直接返回
    if done == nil {
        return
    }


    select {
    case <-done:
        // parent ctx 已经被 cancel 了，那么直接 cancel 当前 ctx
        child.cancel(false, parent.Err())
        return
    default:
    }
    // 判断父节点是不是一个（合法的） cancelCtx
    // yes
    if p, ok := parentCancelCtx(parent); ok {
        p.mu.Lock()
        // 双重检查
        if p.err != nil {
            // 父节点已经被 cancel 了，那么直接 cancel 当前 ctx
            // 第一个参数传 false 是因为当前节点还没有添加到父节点下面，自然没必要执行 remove 操作
            child.cancel(false, p.err)
        } else {
            if p.children == nil {
                p.children = make(map[canceler]struct{})
            }
            // 将当前节点挂载到父节点上面
            p.children[child] = struct{}{}
        }
        p.mu.Unlock()
    } else {
        atomic.AddInt32(&goroutines, +1)
        go func() {
            select {
            // 如果当前父节点取消，那么取消对应的子节点
            // 对应 parentCancelCtx 函数中 done == closedchan 的场景
            case <-parent.Done():
                child.cancel(false, parent.Err())
            // 子节点自己取消，退出 select
            case <-child.Done():
            }
        }()
    }
}


func parentCancelCtx(parent Context) (*cancelCtx, bool) {
    // 由于只有 cancelCtx 实现了 Done() 方法（emptyCtx 除外）
    // 通过继承机制，这里会返回向上查找的第一个 cancelCtx 的 channel
    done := parent.Done()
    // done == closedchan 代表已经父节点取消了
    // done == nil 代表不是一个 cancelCtx
    if done == closedchan || done == nil {
        return nil, false
    }
    // 可以联想一下 cancelCtx.Value() 以及 valueCtx.Value() 的实现
    // 它们都会递归往上层取值
    // 并且由于 key 是 &cancelCtxKey，它会向上取第一个 cancelCtx 实例
    p, ok := parent.Value(&cancelCtxKey).(*cancelCtx)
    // 如果最终断言失败，代表从当前节点一直往上层都找不到 cancelCtx
    // 直接返回失败
    if !ok {
        return nil, false
    }
    pdone, _ := p.done.Load().(chan struct{})
    // 这里检查两次拿到的 channel 是不是同一个
    // 如果不是，代表这个 *cancelCtx 被用户自定义的包装实现中提供了一个不同的 done channel
    // 如果是这种情况，那么会直接拦截，返回 false
    // 之所以要拦截，是因为在这里不清楚调用者会怎么使用自定义的 done channel
    // 由于 cancel 操作依赖这个 channel，如果依然把子节点放进当前节点中，可能会产生预期之外的效果
    if pdone != done {
        return nil, false
    }
    return p, true
}
复制

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()
    // 懒汉式
    if d == nil {
        d = make(chan struct{})
        c.done.Store(d)
    }
    return d.(chan struct{})
}
复制

func cancelled(ctx context.Context) bool {
    select {
    case <-ctx.Done():
        return true
    default:
        return false
    }
}
复制

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

type timerCtx struct {
    // 维护一个 cancelCtx
    cancelCtx
    // 通过 cancelCtx.mu 加锁保护
    timer *time.Timer
    // 超时时间
    deadline time.Time
}
复制

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

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) {
        // 如果父节点的截止日期比新的早
        // 那么新的截止日期实际上不会产生作用，因为它会被父节点提前 cancel
        // 所以这里直接返回一个 cancelCtx 即可
        return WithCancel(parent)
    }
    c := &timerCtx{
        cancelCtx: newCancelCtx(parent),
        deadline:  d,
    }
    // 挂载到父节点上，这样父节点取消的话，当前节点也能被取消
    propagateCancel(parent, c)
    dur := time.Until(d)
    // 如果已经超过截止时间了，那么直接取消
    if dur <= 0 {
        c.cancel(true, DeadlineExceeded)
        return c, func() { c.cancel(false, Canceled) }
    }
    c.mu.Lock()
    defer c.mu.Unlock()
    if c.err == nil {
        // 开启一个定时器，用于到时间自动取消
        c.timer = time.AfterFunc(dur, func() {
            c.cancel(true, DeadlineExceeded)
        })
    }
    return c, func() { c.cancel(true, Canceled) }
}
复制

const KEY_LOG = "LOG_ID"
const KEY_USER_ID = "USER_ID"


func TestWithValue(t *testing.T) {
    // 通过 WithValue() 生成一个保存 key-value 键值对的 ctx
    ctx := context.WithValue(context.Background(), KEY_LOG, "2021082900001")
    // 链式存入第二个 key
    ctx = context.WithValue(ctx, KEY_USER_ID, "112233")
    logId := GetLogID(ctx)
    t.Log(logId)
}


func GetLogID(ctx context.Context) string {
    // 通过 Value() 方法查找
    if logId := ctx.Value(KEY_LOG); logId != nil {
        return logId.(string)
    }
    return ""
}
复制

func TestWithCancel(t *testing.T) {
    ctx, cancel := context.WithCancel(ctx)
    // WithTimeout 可以实现超时自动调用 Cancel()
    // ctx, cancel := context.WithTimeout(context.Background(), 1*time.Second)


    go TakeTasks(ctx, "c1")
    // ctx2 是 ctx 的子 context，当 ctx 被取消之后，ctx2 也会被取消
    ctx2, _ := context.WithCancel(ctx)
    go TakeTasks(ctx2, "c2")


    time.Sleep(500 * time.Millisecond)
    // 也可以手动提前取消
    cancel()
    time.Sleep(100 * time.Millisecond)
}


func cancelled(ctx context.Context) bool {
    select {
    case <-ctx.Done():
        fmt.Println("finish taking tasks!")
        return true
    default:
        fmt.Println("continue!")
        return false
    }
}


func TakeTasks(ctx context.Context, flag string) {
    for {
        if cancelled(ctx) {
            break
        }
        fmt.Printf("%s taking tasks!\n", flag)
        time.Sleep(100 * time.Millisecond)
    }
}
复制