Go-context源码解析

首先我们简单的来看一个例子，如下:（学好这个例子，我们就可以说完全掌握住context了，并且能重构一个context

func main() {
	ctx, cancel := context.WithCancel(context.Background())
	ctxV := context.WithValue(ctx, 1, "Hello World")

	go func(ctx context.Context) {
		val := ctx.Value(1)
		for {
			time.Sleep(time.Millisecond * 100)
			select {
			case <-ctx.Done():
				return
			default:
				fmt.Println("I am bot one, And i acclaim", val)
			}
		}
	}(ctxV)

	go func(ctx context.Context) {
		ctx2, _ := context.WithCancel(ctx)
		go func(ctx context.Context) {
			for {
				time.Sleep(time.Millisecond * 100)
				select {
				case <-ctx.Done():
					return
				default:
					fmt.Println("I am bot three")
				}
			}
		}(ctx2)
		for {
			time.Sleep(time.Millisecond * 100)
			select {
			case <-ctx.Done():
				return
			default:
				fmt.Println("I am bot two")
			}
		}
	}(ctx)

	time.Sleep(time.Second)
	fmt.Println(runtime.NumGoroutine()) // 猜猜这里输出的数字是多少?
	cancel()
	time.Sleep(time.Second)
	fmt.Println(runtime.NumGoroutine()) // 猜猜这里输出的数字是多少?
}

我们可以看到在结尾我们分别在cancel之前以及之后输出了当前的goruntine数量，那么前后有什么区别呢?答案是从4->1。为什么会这样呢?我们来解读一下context.WithCancel(context.Background())这部分函数

var (
	background = new(emptyCtx)
	todo       = new(emptyCtx)
)
func Background() Context {
	return background // 返回 最大节点（类似Object对象
}

func WithCancel(parent Context) (ctx Context, cancel CancelFunc) {
	if parent == nil { // 上级节点不为空，说明context都可以往上追溯，最终追溯到background节点或者toda节点
		panic("cannot create context from nil parent")
	}
	c := newCancelCtx(parent)                      // 实例化一个客户端节点
	propagateCancel(parent, &c)                    // 添加该子节点到父节点，及c->parent
	return &c, func() { c.cancel(true, Canceled) } // 返回子节点，以及对应的释放函数
}

同样的，context标准库支持多级封装，比如: context.WithCancel(ctx)这个函数之中的ctx可以是上一个context.WithCancel(ctx)返回的结果，从而形成了一颗树。当然，它们都拥有同一个最终的树根，当树根使用了cancel释放函数，同样的整棵树都将不复存在；那么将子节点链接到父节点这个过程就至关重要，它是由propagateCancel(parent, &c)来完成的

func propagateCancel(parent Context, child canceler) {
	done := parent.Done()
	if done == nil { // 确保父节点可以链接上子节点，如果父节点没有cancel函数，则说明子节点可以独立于父节点之外
		return // parent is never canceled
	}

	select { // 简单查看父节点是否已经销毁
	case <-done:
		// parent is already canceled
		child.cancel(false, parent.Err())
		return
	default:
	}

	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{})
			}
			p.children[child] = struct{}{} // 将子节点链接到父节点
		}
		p.mu.Unlock()
	} else {
		atomic.AddInt32(&goroutines, +1)
		go func() { // 启动一个监听器，如果父节点销毁，则子节点同样销毁
			select {
			case <-parent.Done():
				child.cancel(false, parent.Err())
			case <-child.Done():
			}
		}()
	}
}

好了，到这里。我相信大家的疑问依然没有解决，我们还差了一个关键函数Done()，我们可以发现任何销毁函数都与它有关，我们便跟一下它的设计:

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{})
	if d == nil {
		c.done.Store(closedchan) // 存入done之中closedchan（关闭隧道）
	} else {
		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 { // 同时将父节点上的记录从map结构中删除
		removeChild(c.Context, c) // 也就是上级节点将不再记录这个节点了
	}
}

这个时候我们再来看看Done函数（实际上先看Done函数比较好一点），我们可以发现Done其实就是返回一个channel而已

func (c *cancelCtx) Done() <-chan struct{} {
	d := c.done.Load()
	if d != nil { // 如果done不为空，则直接返回
		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{})
}

这个时候我们结合Done与cancel函数来整体看看，实际上它们是以channel进行通信，如果channel已经关闭，则它们的节点以及子节点都将会接受到channel的通信，从而进行销毁。

这里我们同样的可以了解到context标准库之中依然存在着两个较为关键的函数WithDeadline、WithTimeout，不过我们有了以上的理解，那么这个其实就比较好理解了。这里我们拿WithDeadline来说:

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) { // 如果发现时间已然超出，则退化成WithCancel(parent)函数
		// The current deadline is already sooner than the new one.
		return WithCancel(parent)
	}
	c := &timerCtx{ // 封装context、以及截止时间
		cancelCtx: newCancelCtx(parent),
		deadline:  d,
	}
	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 {
		c.timer = time.AfterFunc(dur, func() { // 添加一个监听器
			c.cancel(true, DeadlineExceeded)
		})
	}
	return c, func() { c.cancel(true, Canceled) }
}

很简单，其实相比于WithCancel(parent)函数仅仅多了一步time.AfterFunc而已。同理WithTimeout也是如此。