Golang基础 - Context的使用（源码分析）

go语言中的goroutine机制天然地适合做server的开发，最近在看鹅厂内部某框架代码的时候看到了关于context的操作，虽然用channel已经可以很好的处理不同goroutine之间的通信，但是context十分适合做一些关于取消相关的动作，在很多场景下还是有着一定作用。go源码中context的代码不长，所以今天就简单总结回顾一下。

Context

本文中的context源码来自于go1.15

context顾名思义“上下文”，是用来传递上下文信息的结构，实际上是一个接口，如下：

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

Deadline() 方法返回两个参数，一个是deadline，类型为time.Time，意为该context被取消时的时间线，第二个参数是bool类型，如果一个context没有设置deadline则会返回false。

什么是context被取消？
一个context如果设置了过期时间，那么它会被取消；如果在代码中执行了cancel()，该context也会被取消(详细内容请看后文)

Done() 方法返回个struct{}类型的channel，用来不同goroutine之间传递消息，通常都会结合select来使用，当该channel被close的时候，会返回对应的0值。

Err() 方法返回一个error，分为以下几种情况：

• 当Done中的channel还未被关闭时，返回nil
• 当Done中的channel被关闭时，返回对应的原因，比如是正常被Canceled了还是过期了DeadlineExceeded。

Value() 可以根据输入的key返回context中对应的value值，可以用来传递参数。

默认上下文

go中提供了默认的上下文Background和TODO，它们都返回了一个空的context——emptyCtx。当代码中前后都没有context时但又需要的时候，一般会使用context.Backgroud()作为传递参数。

func Background() Context {
	return background
}
func TODO() Context {
	return todo
}
var (
	background = new(emptyCtx)
	todo       = new(emptyCtx)
)
复制代码

emptyCtx实现了context所有的接口，不过都是空值，不然怎么叫empty呢...

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

cancelCtx

上面提到的emptyCtx没有任何功能，而cancelCtx则可以实现上下文的取消功能，然后通过Done来改变上下文的状态。

type cancelCtx struct {
	Context

	mu       sync.Mutex            // protects following fields
	done     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
}
复制代码

cancelCtx中使用匿名的方式定义了Context字段，done使用“懒汉式”创建，children是一个map，记录了该上下文所拥有的字上下文，其中canceler是一个接口，代码如下：

type canceler interface {
	cancel(removeFromParent bool, err error)  // removeFromParent如果是true，则会将
                                                  // 该context从其父context中移除
	Done() <-chan struct{}
}
复制代码

使用WithCancel可以创建一个上下文，源码如下：

func WithCancel(parent Context) (ctx Context, cancel CancelFunc) {
	if parent == nil {
		panic("cannot create context from nil parent")
	}
	c := newCancelCtx(parent)
	propagateCancel(parent, &c)
	return &c, func() { c.cancel(true, Canceled) }
}
复制代码

有两种情况下被创建的上下文会被取消，一是执行了返回的的cancel()函数，二是如果创建时的parent被取消了，该上下文也会被取消，给大家举一个示例代码吧，

func main() {
	ctxParent, cancelParent := context.WithCancel(context.Background())
	ctxChild, _ := context.WithCancel(ctxParent)
	// 父ctx执行取消
	cancelParent()

	select {
	case <-ctxParent.Done():
		fmt.Println("父ctx被取消")
	}

	select {
	case <-ctxChild.Done():
		fmt.Println("子ctx被取消")
	}

}
复制代码

上面的代码会输出两行

父ctx被取消
子ctx被取消
复制代码

原因就是因为执行了父ctx取消函数之后，子ctx也会随之取消。
关于WithCancel中的newCancelCtx和propagateCancel这两个函数，有兴趣的同学可以自己去看看源码，主要就是调用cancelCtx的cancel函数，cancel中就是执行如何关闭context中的channel，比较简单。

timerCtx

timerCtx包含了一个定时器timer和时间线deadline，当定时器结束时就会调用cancelCtx的cancel方法来实现上下文的取消操作。

type timerCtx struct {
	cancelCtx
	timer *time.Timer // Under cancelCtx.mu.

	deadline time.Time
}
复制代码

使用WithDeadline或者WithTimeout就可以创建一个带定时器的上下文context，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) {
		// 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) // 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) }
}
复制代码

前半部分都是一些初始化相关，主要看c.timer = time.AfterFunc(dur, func() { c.cancel(true, DeadlineExceeded) })，这里使用time.AfterFunc来定义了一个定时器，在dur时间之后执行c.cancel方法，该方法的源码如下：

func (c *timerCtx) cancel(removeFromParent bool, err error) {
	c.cancelCtx.cancel(false, err)
	if removeFromParent {
		// Remove this timerCtx from its parent cancelCtx's children.
		removeChild(c.cancelCtx.Context, c)
	}
	c.mu.Lock()
	if c.timer != nil {
		c.timer.Stop()
		c.timer = nil
	}
	c.mu.Unlock()
}
复制代码

可以看出其实最核心的就是第一行，c.cancelCtx.cancel(false, err)，也就是上面说的调用cancelCtx的cancel方法。

valueCtx

context包中使用了valueCtx来进行key-value对的值传递，结构如下（已经无法再简单了...）：

type valueCtx struct {
	Context
	key, val interface{}
}
复制代码

valueCtx中同样包含了Context这个匿名接口，因此也具有Context的特性。使用WithValue可以设置一个带有key-value的上下文，使用Value则可以递归的查找到key对应的value值，源码如下：

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() { //必须是可比较的，不然Value方法就没法用了
		panic("key is not comparable")
	}
	return &valueCtx{parent, key, val}
}
复制代码

func (c *valueCtx) Value(key interface{}) interface{} {
	if c.key == key {
		return c.val
	}
	return c.Context.Value(key)   // 注意这是go语言的特性，类似于java中的继承特性
                                      // 可以说是valueCtx继承了Context的特性
}
复制代码

总结

最后，context上下文能够很好的传递一些简单的消息、key-value类型的值，但是频繁使用context可能会导致你的代码处处都存在context，因为你总是需要把context作为参数传递进你的函数中...