go语言context包源码解析

go语言context包源码解析

由于需要转go，学了基础和一些服务端框架后决定对go的一些源码进行阅读，并适当去实现。第一份源码就看了context，下面对context进行详细的解读。
源码部分请详细阅读注释，写得很清楚。
注意：源码是我自己看了context官方包造的轮子，基本没太大的差别，注释中写了一些我遇到的问题的详细解释（比如实现过程中go中的锁不可重入带来的问题），代码可以在GitHub中下载，地址，用goland打开阅读体验更佳
强烈建议仔细阅读源码
阅读顺序不一定按照文章顺序，二三可以根据情况交换

一. context用途

context用于停止goroutine，协调多个goroutine的取消，设置超时取消等等。基于channel和select来实现停止，另外还可以用context在不同的goroutine中传递数据。其中停止goroutine是context的核心。
停止goroutine的基本原理：
使用一个channel，在多个goroutine中使用select从channel中取值，context中实际上并不会放任何值到channel中而是关闭channel，这样就channel就不会阻塞，select就能走到case <-ctx.Done()这个分支：
源码使用例子：

//不停的将拿到的返回值v放到out这个channel中，直到发生错误或者收到取消信号
 func Stream(ctx context.Context, out chan<- Value) error {
 	for {
 		v, err := DoSomething(ctx)
 		if err != nil {
 			return err
 		}
 		select {
 		case <-ctx.Done():
 			return ctx.Err()
 		case out <- v:
 		}
 	}
 }

如果doSomething只执行一次还需要用context吗？我觉得不需要，不过要对错误判断，出错则退出goroutine

二.context整体结构

context中结构体和接口之间的关系如图所示，图中只给出了直接关系。没有间接关系，比如timerCtx结构体内部有cancelCtx匿名字段，实际上相当于它也实现了Context和canceler接口。
组合关系：结构体中嵌入了匿名字段（例如valueCtx中嵌入了Context字段）
实现：实现接口

接口和结构体的具体定义：
Context：

//Context 上下文管理接口
type Context interface {
	//返回一个-<chan struct{} 如果context实例是不可取消的，那么返回nil，比如空Context,valueCtx
	Done() <-chan struct{}
	//根据key拿到存储在context中的value
	//递归查找如果当前节点没有找到会往父节点
	Value(key interface{}) interface{}
	//返回任务取消的原因
	Err() error
	//返回任务执行的截止时间
	//非timerCtx类型则返回nil
	Deadline() (deadline time.Time, ok bool)
}

canceler

//canceler 取消context的接口
type canceler interface {
	cancel(removeFromParent bool, err error)
	Done() <-chan struct{}
}

valueCtx

// valueContext 用于存储key value
// 不可取消的context
type valueCtx struct {
	Context
	key   interface{}
	value interface{}
}

cancelCtx

//cancelCtx 实现了context和cancel接口
type cancelCtx struct {
	//保存parent context实例
	Context
	//用于锁下面用到的字段
	lock sync.Mutex
	//用于保存那些实现了canceler接口的子context
	//使用map是为了方便删除
	//方便取消孩子context
	children map[canceler]struct{}
	//实现取消需要用用到的channel
	done chan struct{}
	//取消原因
	err error
}

timerCtx

// 基于timer来实现超时取消的的context
type timerCtx struct {
	*cancelCtx
	// 截止时间
	deadLine time.Time
	// 取消原因
	err error
	// 定时器
	*time.Timer
}

除了上面几种类型还有一种emptyCtx类型，它实现了context接口，但是所有方法都返回nil，一般作为context树的根节点，调用Context.Background返回的就是emptyCtx，它内部就是把int自定义为了一个新的类型emptyCtx

type emptyCtx int

看到这里你只需要大概知道context包中有这些接口和结构体，对接口的方法和功能有个印象。
下面我将介绍多个context如何相互作用的基本概念，以便后面理解。
不同的context之间构成了一颗多叉树，context对外提供的With开头的方法总是需要传入一个parent，比如WithCancel，这个parent就是作为一个父节点。
这颗多叉树你可以构建得很复杂，比如：

理解这个图非常重要：
对于一颗context树我们关心如下几点：

（1）取消信号如何传递

取消信号沿着树的叶子方向进行，比如cancelCtx1调用了cancel方法，那么首先cancelCtx1会取消（即关闭字段done这个channel），然后会遍历离他最近的canceler类型的孩子节点cancelCtx3，（保存在children字段，存储类型canceler接口），valueCtx2被忽略，如果cancelCtx3下面还有canceler类型的节点，那么会继续往下传递取消，cancel这个方法的调用时递归的。
对应的代码实现在cancelCtx实现的cancel方法

（2）key的查找如何进行

key的查找沿着树的根方向进行，比如在valueCtx3中放入了一个key为"key3"，在valueCtx2，cancelCtx1都是查不到，只有valueCtx3中才能查到，在valueCtx2放一个key为"key2"，则valueCtx3和valueCtx2都能查到。

（3）context节点增加删除

使用With就可以添加节点，具体的代码实现在propagateCancel这个方法，这里有一些细节，如果此时parent已经取消，那么会直接调用child的cancel，而不会加入到树中，如果时cancelCtx这种类型的节点还需要往树根方向找最近的一个cancelCtx，把自己添加到这个cancelCtx的children中，在执行cancel后会直接把children这个字段置为nil

（3）冲突如何处理？

比如parent节点已经取消（前面提到过）
比如timerCtx中父节点的截止日期更早，例如图中timerCtx1和timerCtx2，这时候以timerCtx1为准，如果它超时取消，即使timerCtx2还没超时也会被取消

三.context源码详细阅读

通过WithCancel这个方法开始一步步看源码。
注释很详细请仔细看，如果不方便可以下载GitHub源码文件

func demoCancel() {
	var worker = func(ctx context.Context) {
		for {
			//其他逻辑代码
			fmt.Println("working...")
			time.Sleep(time.Second)
			select {
			//取消goroutine代码
			case <-ctx.Done():
				break
			default:

			}
		}
	}
	//开始
	ctx, cancel := context.WithCancel(context.Background())
	go worker(ctx)
	time.Sleep(time.Second * 5)
	cancel()
	fmt.Println("main goroutine done...")
}

WithCancel源码:

// 在parent context基础上添加一个节点cancelCtx
func WithCancel(parent Context) (Context, func()) {
	ctx := newCancelCtx(parent)
	//传播取消行为
	propagateCancel(parent, ctx)
	return ctx, func() {
		ctx.cancel(true, canceled)
	}
}

核心代码propagateCancel:
理解难点：else分支什么情况下才会执行？为什么开启一个goroutine去等待取消信号

// 传播取消行为
// 大致功能：如果parent已经取消，那么将取消信号传递到parent这颗子树
// 否则将child加入到一个根节点方向最近的cancelCtx的children中
func propagateCancel(parent Context, child canceler) {
	if parent.Done() == nil {
		//父Context是不可取消的
		return
	}
	//parentCancelCtx往树的根节点方向找到最近的context是cancelCtx类型的
	if p, ok := parentCancelCtx(parent); ok {
		p.lock.Lock()
		if p.err != nil {
			//祖父cancelCtx已经取消，取消孩子节点
			child.cancel(false, p.err)
		} else {
			if p.children == nil {
				p.children = make(map[canceler]struct{})
			}
			//将child加入到祖父context中
			p.children[child] = struct{}{}
		}
		p.lock.Unlock()
	} else {
		//之前一直不太理解这里为什么出现else分支，找到了一篇知乎文章对这个问题有较好的解释
		//参考：https://zhuanlan.zhihu.com/p/68792989
		// 官方对使用context其中一条建议就是context不要放到结构体内部
		// 如果外部代码这样用了，结构体嵌入了context，类型不再是3大context（cancelCtx，timerCtx，valueCtx），
		// 但它实现了context的接口，可以调用接口的方法
		//导致的问题就是执行parentCancelCtx（在树往根方向查找最近的cancelCtx），因为类型原因是拿不到cancelCtx的
		//那么外部取消parent时，取消信号是无法往树的叶子方向传递，从而可能无法取消孩子context
		//所以这里启动一个goroutine来监控parent的取消信号，如果拿到了那么就可以顺利取消孩子context
		//两个case都是为了测试能够正常取消child
		//parent的cancel还是child的cancel哪个先调用都能退出select
		//如果发生parent的cancel不调用也是可能会出现goroutine泄露的，这里保证不了
		go func() {
			select {
			case <-parent.Done():
				child.cancel(false, parent.Err())
			case <-child.Done():
			}
		}()
	}
}

parentCancel源码：

// 从parent往树的根方向走找到最近的一个cancelCtx类型的Context
func parentCancelCtx(parent Context) (*cancelCtx, bool) {
	//使用类型断言
	for {
		switch c := parent.(type) {
		case *cancelCtx:
			return c, true
		case *valueCtx:
			//遇到valueCtx继续往上查找
			parent = c.Context
		case *timerCtx:
			return c.cancelCtx, true
		default:
			return nil, false
		}
	}
}

核心方法cancel:
理解难点：死锁问题的产生，为什么遍历children调用cancel，removeFromParent必须为false？

func (c *cancelCtx) cancel(removeFromParent bool, err error) {
	if err == nil {
		panic("internal error, Missing cancellation reason")
	}
	c.lock.Lock()
	//已经取消
	if c.err != nil {
		c.lock.Unlock()
		return
	}
	//取消
	c.err = err
	//1.关闭channel
	if c.done == nil {
		c.done = closedChan
	} else {
		close(c.done)
	}
	//2.执行孩子节点的cancel
	//孩子节点会持有父节点的lock
	for child := range c.children {
		//这里如果removeFromParent为true
		//会导致锁不可重入带来的死锁问题
		// 产生死锁的执行流程:
		// cancelCtx1.cancel
		// m1.Lock()
		//取消孩子context
		// cancelCtx2.cancel
		// m2.Lock()
		//取消孩子（假设无）
		//m2.Unlock()
		//removeChild
		//获取到父cancelCtx为cancelCtx1
		//使用m1锁
		// m1.Lock() 产生死锁，因为不可重入
		child.cancel(false, err)
	}
	//手动释放
	c.children = nil
	c.lock.Unlock()
	//从父节点中删除孩子节点
	if removeFromParent {
		//removeChild中使用了cancelCtx的锁
		removeChild(c.Context, c)
	}
}

removeChild：

// 从父context查找最近的cancelCtx并从中删除child
func removeChild(parent Context, child canceler) {
	//这里的不需要加锁，树的增长方向向下
	//parentCancelCtx只会往树的根节点方向走
	//并且只读操作
	ctx, ok := parentCancelCtx(parent)
	if !ok {
		return
	}
	//由于go中的Mutex不可重入。之前的版本中会产生死锁问题
	//cancel中具体解释了原因
	ctx.lock.Lock()
	if ctx.children != nil {
		delete(ctx.children, child)
	}
	ctx.lock.Unlock()
}

以上大概就是cancelCtx的核心流程代码
timerCtx实际是类似的，内部使用了一个定时器定时执行cancel方法
以下是WithDeadline的源码，WithTimeOut内部直接调用了WithDeadline
理解难点：为什么Timer初始化需要加锁

// Deadline超时直接调用timeout即可
func WithDeadline(parent Context, deadline time.Time) (Context, func()) {
	//如果parent的deadline在当前context之前那么不需要开启新的定时器，直接由父context定时取消
	//当前context变为cancelCtx
	if d, ok := parent.Deadline(); ok && d.Before(deadline) {
		return WithCancel(parent)
	}
	ctx := &timerCtx{
		cancelCtx: newCancelCtx(parent),
		deadLine:  deadline,
	}
	propagateCancel(parent, ctx)
	dur := time.Until(deadline)
	if dur <= 0 { //立刻取消
		ctx.cancel(true, timeOutErr)
		return ctx, func() {
			ctx.cancel(false, canceled)
		}
	}
	//产生timer，AfterFunc会自己启动
        //这里为什么加锁？
        //前面的无论是propagate还是cancel内部都是有互斥处理
        // 因为ctx可能已经被加入到某个cancelCtx的children中
        // 如果在未初始化Timer之前被调用了cancel方法，那么就会存在对ctx.err并发访问
	ctx.cancelCtx.lock.Lock()
	defer ctx.cancelCtx.lock.Unlock()
	//确认这期间没有context被取消（比如parent被其他goroutine调用了cancel），不然无需初始化timer
	if ctx.err == nil {
		ctx.Timer = time.AfterFunc(dur, func() {
			ctx.cancel(true, timeOutErr)
		})
	}
	return ctx, func() {
		ctx.cancel(true, canceled)
	}
}

核心代码cancel
理解难点：这里为什么removeFromParent不能使用cancelCtx去remove

func (t *timerCtx) cancel(removeFromParent bool, err error) {
	//timer已经在初始化timerCtx中启动
	//这里的cancel作为timer定时执行的函数。这里直接调用cancel
	t.cancelCtx.cancel(false, err)
	//这里为什么这么写
	//因为在propagate中是将一个timerCtx加入到了parent（timerCtx继承cancelCtx所以实现了canceler接口可以被加入到cancelCtx的children）
	//调用上面的t.cancelCtx.cancel是移除不了的
	//timerCtx的parent应该是t.cancelCtx.Context
	if removeFromParent {
		removeChild(t.cancelCtx.Context, t)
	}
	t.cancelCtx.lock.Lock()
	//如果提前调用了cancel，下面的操作可以停止timer
	if t.Timer != nil {
		t.Timer.Stop()
		t.Timer = nil
	}
	t.cancelCtx.lock.Unlock()
}

四. 使用建议

官方使用context建议

1. Do not store Contexts inside a struct type; instead, pass a Context explicitly to each function that needs it. The Context should be the first parameter, typically named ctx.（不要再结构体内部嵌套一个Context）
2. Do not pass a nil Context, even if a function permits it. Pass context.TODO if you are unsure about which Context to use. (不要将nil Context传递给别的函数，如果不知道传什么，使用内置的TODO)
3. Use context Values only for request-scoped data that transits processes and APIs, not for passing optional parameters to functions.
4. The same Context may be passed to functions running in different goroutines; Contexts are safe for simultaneous use by multiple goroutines. （context被多个goroutine使用并发安全）

个人使用context建议
不要嵌套过多的context，构建非常复杂的context树，这样你可能会疑惑到底该调用哪个cancel，除非你真的非常理解context的内部机制 。

五. 其他

感悟：

1. 使用了大量的懒加载技术，变量用到才初始化
2. 并发安全问题考虑得很周到，哪里需要加锁
   一些建议：
3. 对err字段的读取使用了互斥锁，是否可以考虑用读写锁
4. key，Value使用了类似链表的结构，查询速度不够快