手摸手Go Context探秘

605 阅读6分钟

使用Go作为服务端开发时,每个请求过来都会分配一个goroutine来处理,请求处理过程中,可能还会创建额外的goroutine访问DB或者RPC服务。这个请求涉及的goroutine可能需要访问一些特定的值比如认证token、用户标识或者请求截止时间。当一个请求取消或者超时,则这个请求涉及的goroutine也应该被终止,这样系统就可以快速回收这部分资源。

基于简化目的,context包定义了Context类型,来传递超时、取消信号以及跨API边界和进程之间request-scope值。当服务器来新的请求应该创建一个Context并且返回请求应该接受一个Context。这其中的函数调用链必须传递Context对象,或者是通过WithCancelWithDeadlineWithTimeWithValue衍生的Context对象。当一个Context取消,所有从这个对象衍生的Context都会被取消。

例如你可以利用context的这种机制,实现一个任务超时保护的方法

func main() {
	cancelJob(time.Second*1, func() error {
		time.Sleep(time.Second * 10)
		return nil
	})
}

func cancelJob(timeout time.Duration, f func() error) error {
	var (
		ctx        context.Context
		cancelFunc context.CancelFunc
	)
	if timeout > 0 {
		ctx, cancelFunc = context.WithTimeout(context.Background(), timeout)
	} else {
		ctx, cancelFunc = context.WithCancel(context.Background())
	}

	defer cancelFunc()
	e := make(chan error, 1)
	go func() {
		e <- f()
	}()
	select {
	case err := <-e:
		return err
	case <-ctx.Done():
		return ctx.Err()
	}
}

Context使用大致步骤:

  • 1 构建一个Context对象,如果你不知道该使用什么Context合适,可以调用context.Backgroundcontext.TODO
  • 2 根据你的需求可以对Context进行包装衍生
    • WithCancel :Context可取消
    • WithDeadline: Context可设置截止时间
    • WithTimeout:实际使用的是WithDeadline
    • WithValue: 需要使用Context传值
  • 3 监听ctx.Done这个channel 一旦Context取消 此channel会被关闭
  • 4 最后在方法处理完毕时请及时调用cancel方法 方便资源回收

数据结构

context包提供了两种创建Context对象的便捷方式

  • context.Background 无法被取消 没有值 没有截止时间,通常用于主函数、初始化、测试或者当新请求来了作为顶层Context
  • context.TODO 当你不知道用啥Context的时候使用

这两种方式都是一个emptyCtx对象 本质上没啥差别

var (
	background = new(emptyCtx)
	todo       = new(emptyCtx)
)
func Background() Context {
	return background
}
func TODO() Context {
	return todo
}

此外context包提供了4个方法:WithCancelWithDeadlineWithTimeoutWithValue,可以对context进行衍生为cancelCtxtimerCtxvalueCtx,他们都实现了context.Context接口

// Context的方法是线程安全的
type Context interface {
  // 返回context何时需要被取消 ok为false表示deadline未设置
	Deadline() (deadline time.Time, ok bool)
	// 当context被取消 Done放回一个关闭的channel 
  // Done返回nil 表示当前context不能被取消
	// Done通常在select语句中使用
	Done() <-chan struct{}
  // 返回context取消的原因
	Err() error
  // 返回context中指定的key关联的value,未指定返回nil
	// 主要用作在进程和API边界间传递request-scoped数据,不要用于可选参数传递
	// key需要支持相等操作,最好定义为不可到处类型 避免混淆
	Value(key interface{}) interface{}
}

这几个对象层次结构

context architecture

衍生contexts

通过WithCancelWithDeadlineWithTimeoutWithValue方法衍生的context为原始context提供了取消、传值、超时取消等功能。

WithCancel

通过WithCancel衍生出新的可取消的Context对象

// WithCancel 返回包含父context拷贝和一个新的channel的context 和一个cancel函数
// 当返回的cancel函数被调用或父context的done channcel关闭 则WitchCancel返回的context的channel也会被关闭
// 当操作完成时应该尽快调用cancel函数 这样就可以释放此context关联的资源
func WithCancel(parent Context) (ctx Context, cancel CancelFunc) {
	c := newCancelCtx(parent)
  // 构建父子上下文之间的关系 保证父上下文取消时子上下文也会被取消
	propagateCancel(parent, &c)
	return &c, func() { c.cancel(true, Canceled) }
}

将Context包装为可取消的Context-->cancelCtx

// cancelCtx可以被取消,当取消时,也会将其实现了canceler的子context也取消
type cancelCtx struct {
	Context

	mu       sync.Mutex            // 保护下面的字段
	done     chan struct{}         // 惰性创建 cancel方法第一次调用时关闭
	children map[canceler]struct{} // cancel第一次调用时置为nil
	err      error                 // cancel第一次调用时设置non-nil
}

其中done这个channel是在cancel调用的时候才会被初始化,cancelCtx子context若可取消需要实现canceler接口

// A canceler is a context type that can be canceled directly. The
// implementations are *cancelCtx and *timerCtx.
type canceler interface {
	cancel(removeFromParent bool, err error)
	Done() <-chan struct{}
}

构建可取消context后,后将取消操作进行传播,如果父context的Done为nil,表示其不可取消直接返回,否则会调用parentCancelCtx直到找到可取消父context ,若找到

  • 若可以找到则

    • 且父context已经取消则会调用子context的cancel方法进行取消;
    • 且父context未取消则将当前子context交给父context管理
  • 若找不到 例如开发者自定义的类型则

    直接启动一个gorountine来监听父子取消事件通知

// propagateCancel arranges for child to be canceled when parent is.
func propagateCancel(parent Context, child canceler) {
	if parent.Done() == nil {
		return // parent is never canceled
	}
	if p, ok := parentCancelCtx(parent); ok {//找到父可取消context
		p.mu.Lock()
		if p.err != nil {
			// 父context已经被取消 取消子context
			child.cancel(false, p.err)
		} else {// 父context未cancel 则将子context交给父context管理,方便父节点取消时将取消事件传播给子context
			if p.children == nil {
				p.children = make(map[canceler]struct{})
			}
			p.children[child] = struct{}{}
		}
		p.mu.Unlock()
	} else {//找不到父可取消context
		go func() {
			select {
			case <-parent.Done():
				child.cancel(false, parent.Err())
			case <-child.Done():
			}
		}()
	}
}

parenCancelCtx循环查找context是否存在可取消父节点

// parentCancelCtx follows a chain of parent references until it finds a
// *cancelCtx. This function understands how each of the concrete types in this
// package represents its parent.
func parentCancelCtx(parent Context) (*cancelCtx, bool) {
	for {
		switch c := parent.(type) {
		case *cancelCtx:
			return c, true
		case *timerCtx:
			return &c.cancelCtx, true
		case *valueCtx:
			parent = c.Context
		default:
			return nil, false
		}
	}
}
取消函数cancel

当你的业务方法执行完毕,你应该尽快调用cancel方法,这样方便快速回收相关资源

//关闭c.done,取消掉c的子context,若removeFromParent为true,则将c从父context的子context集合中删除
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 // c.err不为nil 则表示当前context已经被取消
	}
	c.err = err
	if c.done == nil { //调用cancel方法此时才初始化
		c.done = closedchan//closedchan如其名字 为已经关闭的chan
	} else {//关闭c.done
		close(c.done)
	}
  //对子context进行cancel
	for child := range c.children {
		// 此处在持有父context锁的同时 获取子context的锁
		child.cancel(false, err)
	}
	c.children = nil //cancel完毕 置nil
	c.mu.Unlock()

	if removeFromParent {//将当前context从其父context的子context集合中删除
		removeChild(c.Context, c)
	}
}

WithDeadline

让context具备超时取消功能

// WithDeadline 返回包含父context拷贝和deadline为d的context,如果父deadline早于d
// 则语义上WithDeadline(parent, d) 和父context是相等的
//当deadline过期了 或者返回的cancel函数被调用 或者父context的done channel关闭了则WithDeadline返回的context中的done channel也会关闭
// 当操作完成时应该尽快调用cancel函数 这样就可以释放此context关联的资源
func WithDeadline(parent Context, d time.Time) (Context, CancelFunc) {
	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已经到期
		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) }
}

timerCtx内嵌了cancelCtx,主要的cancel能力进行了代理。额外新增了一个截止时间和一个定时器,初始化此类context时如果未到截止时间且未取消 则会启动一个定时器,超时即会执行cancel操作

// timerCtx包含了一个定时器和一个截止时间 内嵌一个cancelCtx来实现 Done和Err方法
// 取消操作通过停止定时器然后调用cancelCtx.cancel来实现
type timerCtx struct {
	cancelCtx
	timer *time.Timer // Under cancelCtx.mu.

	deadline time.Time
}

timerCtx的cancel操作本身会停掉定时器,然后主要cancel操作代理给了cancelCtx

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()
}

WithTimeut

实际掉用了WithDeadline没啥好说的。

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

总的来说,就是可取消context通过父节点保存子节点集合 若父节点取消则将子节点集合中的context依次调用cancel方法。

may be ugly

WithValue

赋予了Context传值能力,Context的能力代理给了父context,自身新增了一个 Value(key interface{}) interface{}方法,根据指定key获取跟context关联的value,逻辑比较简单 没啥好说的。

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

// A valueCtx carries a key-value pair. It implements Value for that key and
// delegates all other calls to the embedded Context.
type valueCtx struct {
	Context
	key, val interface{}
}

总结

context包核心实现除去注释只有200-300行,整体实现还是短小精悍,为我们提供了跨进程、API边界的数据传递以及并发、超时取消等功能。实际应用过程中也给我们技术实现带来很大便利,比如全链路trace的实现。官方建议我们将context作为函数第一个参数使用,不过实际使用过程中还是会给不少人代理心智负担,所以有人为了尽可能不写context,搞了个Goroutine local storage 有兴趣可以研究下