通过源码详解协程并发控制包go context

摘要

Context是go语言上下文管理的包，主要功能是通过给协程传递上下文信息以实现协程的并发控制，有句话说得好，“如果你没法控制协程在什么时候结束，那么这个协程就不应该被创建”，contex就是用于实现并发控制的。多个Context按树形或者链表的结构向前连接。Context是一个接口类型，实现了接口下的方法的类型，都可以认为是Context类型。包括emptyCtx、valueCtx、cancelCtx、timerCtx，它们是Context的具体实现。其中：

• emptyCtx是一个空的Context，不能存储数据、不能被撤销，常用作Context树的根节点。
• valueCtx是用来存储键值对数据的Context；
• cencelCtx是能够能够“手动”执行撤销的Context；
• timerCtx是能设置定时取消的也可“手动”撤销的Context。

1. context接口

type Context interface {

    Deadline() (deadline time.Time, ok bool) // 返回context过期时间的方法。不一定所有context都有过期时间，如果没有，ok将为false

    Done() <-chan struct{} // 返回一个只读的channel，可以通过监听该channel，去判断context是否被撤销

    Err() error // 如果context未取消，返回nil；如果已经取消，返回取消原因。

    Value(key interface{}) interface{} // 不同的context功能不能，valueCtx的返回值存储值的信息，cencelCtx、timerCtx则用于向上寻找最近的可撤销的context。
}

2. 实现context接口的类型

2.1 emptyCtx

emptyCtx是一个等价于int类型，是一个空的context，实现了Context接口。但它不能设置超时时间，不能存储消息。常用作Context的根节点（Context之间按树形或者链式结构向前连接）。

type emptyCtx int
// 不能设置过期时间，所有没过期时间，ok为默认值false（bool类型的零值）
func (*emptyCtx) Deadline() (deadline time.Time, ok bool) {
    return
}
// 不会被取消，所以为nil
func (*emptyCtx) Done() <-chan struct{} {
    return nil
}
// 不会被取消，所以为nil
func (*emptyCtx) Err() error {
    return nil
}
// 不能存储值，所以为nil
func (*emptyCtx) Value(key interface{}) interface{} {
    return nil
}
// 显示当前的context是background还是todo
func (e *emptyCtx) String() string {
    switch e {
    case background:
        return "context.Background"
    case todo:
        return "context.TODO"
    }
    return "unknown empty Context"
}
// background和todo本质上一样，只是语义上不一样。
// background常在主函数、初始化、测试中使用。
// todo 是在不确定使用什么context的时候才会使用
var (
    background = new(emptyCtx)
    todo       = new(emptyCtx)
)

func Background() Context {
    return background
}

func TODO() Context {
    return todo
}

2.2 valueCtx

能够存储键值对数据信息的Context,如下所示，由存储父context的接口、存储键值对数据的接口组成。

// 结构体类型，能够存储一对值
type valueCtx struct {
    Context
    key, val interface{}
}

如何创建valueCtx？使用WithValue，传入父context结点、要存储的k、v键值对数据。就会创建一个该父context的子结点。

// 在context（valueCtx）链表尾部添加值
func WithValue(parent Context, key, val interface{}) Context {
    if key == nil {
        panic("nil key")
    }
    // 判断Key是否可比较，不能比较则拒绝创建
    if !reflect.TypeOf(key).Comparable() {
        panic("key is not comparable")
    }
    // 在基于父context创建值
    return &valueCtx{parent, key, val}
}

添加的时候，要求传入的k是能够比较的值，并且也不是直接插入到Context（valueCtx）中的，而是重新创建一个新的context（valueCtx）存储（k,v）并插入到链表尾部。如下图所示。

如何查找值？按key值从当前结点开始，沿着父节点依次查找，直至找到要查找的key或者结遍历到了根结点。

func (c *valueCtx) Value(key any) any {  
    if c.key == key {  
        return c.val  
    }  
    return value(c.Context, key)  
}  
// 输入key查找value
func value(c Context, key any) any {  
// for循环依次遍历，没找到就向父结点去查找目标key
    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)  // 说明是自定义的context，因此直接返回value
        }  
    }  
}

实践一下，下面代码输出是什么？

func c1(ctx context.Context) context.Context {  
    c1Ctx := context.WithValue(ctx, "c1Key", "is c1")  
    return c1Ctx  
}  
func c2(ctx context.Context) context.Context {  
    c2Ctx := context.WithValue(ctx, "c2Key", "is c2")  
    return c2Ctx  
}  
func c3(ctx context.Context) context.Context {  
    c3Ctx := context.WithValue(ctx, "c3Key", "is c3")  
    return c3Ctx  
}  
  
func main() {  
    ctx := context.Background()  
    ctx1 := c1(ctx)  
    ctx2 := c2(ctx1)  
    ctx3 := c3(ctx2)  
    fmt.Println(ctx2.Value("c1Key"))  
    fmt.Println(ctx2.Value("c3Key"))  
    fmt.Println(ctx3.Value("c3Key"))  
}

结果：

2.3 cancelCtx

cancelCtx是可以撤销的context，它由存储父context的接口、并发控制的锁mu、接受终止信号的chan done、存储孩子结点的map children、存储撤销原因的err组成。协程通过监听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  
    cause error // set to non-nil by the first cancel call  
}

cancelCtx的Context的接口实现，只完成了Value、Done、Err，其中Value用于向上查找最近的可撤销的context，Done返回一个只读的channel，通过监听该channel，可以知道该context是否被撤销，Err用于返回撤销原因。Deadline未实现，直接继承父Context的Deadline。

// 该函数其实是用来判断
func (c *cancelCtx) Value(key any) any {  
    if key == &cancelCtxKey {  
        return c  
    }  
    return value(c.Context, key)  
}  
// 返回一个只读的channel，通过监听该channel，能够知道该context是否被撤销
func (c *cancelCtx) Done() <-chan struct{} {  
    // 判断该done内是否已经有信息存入，有则说明该context已经被其他协程撤销，直接返回结果
    d := c.done.Load()  
    if d != nil {  
        return d.(chan struct{})  
    }
    // 加锁操作
    c.mu.Lock() 
    defer c.mu.Unlock()
    // 因为申请锁过程中，该context状态可能被改变，所以需要再次检查
    d = c.done.Load()  
    // 为空说明该done还为开辟channel空间。
    if d == nil {  
        d = make(chan struct{})  
        c.done.Store(d)  // 创建一个channel，存入done中
    }  
    return d.(chan struct{})  
}  
// 因为该err状态可能被其它协程改变，所有需要加锁读取
func (c *cancelCtx) Err() error {  
    c.mu.Lock()  
    err := c.err  
    c.mu.Unlock()  
    return err  
}

如何创建一个cancelCtx？使用WithCancel，输入父context节点，得到可撤销的context、和撤销函数cancel，通过执行撤销函数，就可将协程以及子协程下的全部context撤销。

type CancelFunc func()
// 创建cancel
func WithCancel(parent Context) (ctx Context, cancel CancelFunc) {
    c := newCancelCtx(parent)
    propagateCancel(parent, &c)
    return &c, func() { c.cancel(true, Canceled) }
}

// 将parent作为父节点context生成一个新的子节点
func newCancelCtx(parent Context) cancelCtx {
    return cancelCtx{Context: parent}
}
// 用于建立父context与当前可撤销context的连接，以确保父context撤销，它的全部孩子也会撤销。
func propagateCancel(parent Context, child canceler) {
// parent.Done()返回nil表明父节点不能被撤销，所以就不需要继续下去了
     done := parent.Done()  
        if done == nil {  
            return // parent is never canceled  
     }
     // 父节点已经被取消，则马上把当前context也取消
    select {  
        case <-done:  
        // parent is already canceled  
            child.cancel(false, parent.Err(), Cause(parent))  
            return  
        default:  
    }
    // 判断父亲是不是cancelCtx类型，是则将当前context加入到父亲的孩子集合中
    if p, ok := parentCancelCtx(parent); ok {
        p.mu.Lock()
        if p.err != nil {
            // parent has already been canceled
            child.cancel(false, p.err)
        } else {
            if p.children == nil {
                p.children = make(map[canceler]struct{})
            }
            // 将当前子节点加入最近cancelCtx祖先节点的children中
            p.children[child] = struct{}{}
        }
        p.mu.Unlock()
    } else {
    // 若父亲不是cancelCtx、但父亲也是可取消的类型，则开启协程，监听父context取消信号，当取消信号触发，则马上取消孩子context
        go func() {
            select {
            case <-parent.Done():
                child.cancel(false, parent.Err())
            case <-child.Done():
            }
        }()
    }
}
// 用来判断父亲结点是不是可撤销的context
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  
}

撤销函数是怎样的？撤销函数需要传入撤销的错误及原因，如果原因为空，则默认为err，而err不能为空，是必传参数。运行过程如下：

1. 确认传入的err不为空；
2. 检查当前context是否已经被其它协程撤销，已经被撤销则不再进行操作；
3. context的err、cause赋值；
4. 取出context中done，并关闭该channel，读取已经关闭的channel将会得到一个0值，不再阻塞，用于通知监听的协程，context已经被关闭；
5. 撤销该context的所有孩子；
6. 从父结点移除该channel的信息。

func (c *cancelCtx) cancel(removeFromParent bool, err, cause error) {  
    if err == nil {  
        panic("context: internal error: missing cancel error")  
    }  
    if cause == nil {  
        cause = err  
    }  
    // 如果context的已经被其它协程撤销，则直接退出，不再继续
    c.mu.Lock()  
    if c.err != nil {  
        c.mu.Unlock()  
        return // already canceled  
    }  
    // 赋值
    c.err = err  
    c.cause = cause  
    // 拿出context中done，并关闭它
    d, _ := c.done.Load().(chan struct{})  
    if d == nil {  
        c.done.Store(closedchan)  // closedchan就是一个关闭的channel，和下面一样
    } else {  
        close(d)  
    }
    // 关闭所有孩子的context
    for child := range c.children {  
        // NOTE: acquiring the child's lock while holding parent's lock.  
        child.cancel(false, err, cause)  
    } 
    // 置空，以让gc垃圾回收
    c.children = nil  
    c.mu.Unlock()  
    // 从父节点删除该孩子信息
    if removeFromParent {  
        removeChild(c.Context, c)  
    }  
}

func removeChild(parent Context, child canceler) {  
    // 父结点是可撤销的context，才有存放孩子信息的map，才需要去做移除操作
    p, ok := parentCancelCtx(parent)  
    if !ok {  
        return  
    }  
    p.mu.Lock()  
    if p.children != nil {  
        delete(p.children, child)  
    }  
    p.mu.Unlock()  
}

2.4 timerCtx

能够根据时间或者事件取消的Context。其结构是对cencelCtx封装，并添加了个生命时间，到了deadline 会自动取消。

// 时间控制Context
type timerCtx struct {
    cancelCtx
    timer *time.Timer // Under cancelCtx.mu.

    deadline time.Time
}

它只实现了Deadline方法，其它是都是直接继承cancelCtx，并重写了cancelCtx的撤销操作。添加了时间计数器的处理。

// 返回Context的截至日期
func (c *timerCtx) Deadline() (deadline time.Time, ok bool) {
    return c.deadline, true
}
// 执行撤销操作
func (c *timerCtx) cancel(removeFromParent bool, err error) {
    // 调用父cancenCtx的撤销方法
    c.cancelCtx.cancel(false, err)
    if removeFromParent {
        // Remove this timerCtx from its parent cancelCtx's children.
        removeChild(c.cancelCtx.Context, c)
    }
    // 该context已经撤销，终止计时
    c.mu.Lock()
    if c.timer != nil {
       // 取消计时器
        c.timer.Stop()
        c.timer = nil
    }
    c.mu.Unlock()
}

如何创建timerCtx？两种方法：
WithDeadLine：输入父context、终止时间，创建一个该父context的孩子context，并返回创建好的context以及可手动撤销的函数CancelFunc。具体运行逻辑如下：

1. 父结点是不是为空；
2. 父结点是不是timerCtx，且撤销时间晚于当前要创建的timerCtx，是则说明没必要再为当前timerCtx创建计时器，直接创建cancelCtx就可；
3. 创建timerCtx，将其与父context关联起来，确保父亲撤销，当前结点也会被撤销；
4. 判断设置过期的时间是否合理，如果设置的过期时间在当前看，已经过期了，则直接撤销创建的timerCtx；
5. 创建计时器，在计时器达到deadline后，执行撤销函数；
6. 返回创建的timerCtx与撤销函数，支持用户手动撤销。

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,
	}
	propagateCancel(parent, c)
	dur := time.Until(d)
	if dur <= 0 {
		c.cancel(true, DeadlineExceeded, nil) // 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() {
			c.cancel(true, DeadlineExceeded, nil)
		})
	}
	return c, func() { c.cancel(true, Canceled, nil) }
}

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