使用context实现一对多的goroutine协作

参考实现

func main() {
	coordinateWithContext()
}

func coordinateWithContext() {
	total := 12
	var num int32
	// cxt实际上是一个指向cancelCtx结构体类型的指针，cancelCtx结构体内嵌了Context接口
	// cxt.done是一个chan struct{}的通道，当调用cancelFunc，就会将该通道关闭，此时<-cxt.Done()将不再阻塞
	// context.Background()返回了一个emptyCtx类型，该类型是int的别名类型，指针值实现了Context接口，代表根上下文。
	cxt, cancelFunc := context.WithCancel(context.Background())
	for i := 1; i <= total; i++ {
		// 当num值等于循环数量，表示任务完成，调用cancelFunc()，即WithCancel返回的匿名函数。这个匿名函数调用了cancelCtx.cancel()，将ctx.done关闭
		go addNum(&num, i, func() {
			if atomic.LoadInt32(&num) == int32(total) {
				cancelFunc()
			}
		})
	}
	<-cxt.Done()
}

// 每次调用尝试将numP加1，如果加失败了就自旋，否则就退出循环，返回时执行deferFunc
func addNum(numP *int32, id int, deferFunc func()) {
	defer func() {
		deferFunc()
	}()
	for i := 0; ; i++ {
		currNum := atomic.LoadInt32(numP)
		newNum := currNum + 1
		time.Sleep(time.Millisecond * 200)
		if atomic.CompareAndSwapInt32(numP, currNum, newNum) {
			break
		}
	}
}

源码解析

• Context接口如下：

type Context interface {
	Deadline() (deadline time.Time, ok bool)
	Done() <-chan struct{}
	Err() error
	Value(key interface{}) interface{}
}

• context.Background()返回的就是emptyCtx，代表根上下文，emptyCtx类型是int的别名类型，该类型的指针类型实现了Context接口，background具有唯一地址：

var (
	background = new(emptyCtx)
)

func Background() Context {
	return background
}

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
}

• context.WithCancel()做了如下几件事：
  • 使用一个Context接口类型的变量创建一个cancelCtx类型变量
  • 返回cancelCtx类型的指针值，还有一个匿名函数，这个匿名函数内部调用cancelCtx类型变量的cancel方法。
  • 匿名函数第一个参数值true表示取消时将cancelCtx从Context的孩子中去掉，第二个参数表示调用cancelCtx.Err()看到的错误就是Canceled

var Canceled = errors.New("context canceled")

func WithCancel(parent Context) (ctx Context, cancel CancelFunc) {
	c := newCancelCtx(parent)
	propagateCancel(parent, &c) // 父亲取消了把儿子也取消
	return &c, func() { c.cancel(true, Canceled) }
}

func newCancelCtx(parent Context) cancelCtx {
	return cancelCtx{Context: parent}
}

• concelCtx类型如下：

type cancelCtx struct {
	Context

	mu       sync.Mutex            // 互斥锁保护下面几个字段
	done     chan struct{}         // 没有初始化, 第一个cancel调用会关闭这个管道
	children map[canceler]struct{} // 第一个cancel调用会将其设置为nil
	err      error                 // 第一个cancel调用会将其设置为非nil
}

// 返回done
func (c *cancelCtx) Done() <-chan struct{} {
	c.mu.Lock()
	if c.done == nil {
		c.done = make(chan struct{})
	}
	d := c.done
	c.mu.Unlock()
	return d
}

// 返回err
func (c *cancelCtx) Err() error {
	c.mu.Lock()
	defer c.mu.Unlock()
	return c.err
}

var closedchan = make(chan struct{})

func init() {
	close(closedchan)
}

// 接收一个bool值和一个err值
// 第一次调用，将err值赋值给c.err，同时如果c.done未初始化，就直接赋一个已经关闭的管道给它;初始化了，就关闭
// 对于c.children，都进行取消，最终将c.children设置为nil
// 如果第一个参数removeFromParent是true，那么从父Context中将当前的cancelCtx去掉
// 第二次调用的话，检查到err已经被赋值，就直接返回。
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
	if c.done == nil {
		c.done = closedchan
	} else {
		close(c.done)
	}
	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)
	}
}

扩展

• 
除了WithCancel可以繁衍Context，还有WithDeadline、WithTimeout、WithValue。

  以上四个函数第一个参数都是Context类型。
• WithTimeout是对WithDeadline的封装，WithDeadline第二个参数接收的是绝对时间，WithTimeout第二个参数接收的是相对时间。定时撤销的错误类型是context.DeadlineExceeded。
• 首次撤销会赋err值，关闭done通道，并且深度优先的撤销子Context，最终将本Context断开与父值的联系。（因为本Context内嵌的父值）

WithValue

• WithValue得到的Context值是不可以撤销的，在撤销信号传播时，若遇到会直接跨过。并且会将信号传递给子值。
• 只有valueCtx可以存储数据，所以遇到其他的Ctx比如cancelCtx会直接跳过。
• 没有修改数据的接口，如果存储的是引用类型，比如切片可以从外部修改，需要自己保证并发安全

func WithValue(parent Context, key, val interface{}) Context {
	if key == nil {
		panic("nil key")
	}
	if !reflect.TypeOf(key).Comparable() {
		panic("key is not comparable")
	}
	return &valueCtx{parent, key, val}
}

type valueCtx struct {
	Context
	key, val interface{}
}

// 现在当前Context查询，查不到再去父Context查询
func (c *valueCtx) Value(key interface{}) interface{} {
	if c.key == key {
		return c.val
	}
	return c.Context.Value(key)
}