Go 语言的 Context 源码分析

Go 语言的 Context 源码分析

Nov 24, 2020
源码分析, Go

研究 Context 的源码,有助于对结合运用 interface 和 struct 的理解,以及对其他三方框架对 Context 接口的重新实现也能有一定认识,在实际开发中更是非常有帮助。Context 的源码非常短小,加上大概一半的注释整个文件也才 500 行,500 行就可以和面试官扯一个小时,ROI 巨大,确定不研究下么。

必备知识 #

WithTimeout 使用 #

func main() {
	ctx, cancel := context.WithTimeout(context.Background(), time.Second*2)
	defer cancel()
	work(ctx)
}

func work(ctx context.Context) {
	done := make(chan bool)
	go func(ctx context.Context) {
		println("工作中...")
		time.Sleep(time.Second * 3)
		done <- true
	}(ctx)
	select {
	case <-ctx.Done():
		fmt.Println("超时了", ctx.Err())
	case <-done:
		fmt.Println("执行完成")
	}
}
//工作中...
//超时了 context deadline exceeded

通过控制 WithTimeout 的时间,可以设置在多久后收到 ctx.Done() 的信号,调用 cancel 也能达到同样的效果。

WithValue 使用 #

package main

import (
	"context"
	"fmt"
)

func main() {
	ctx := context.WithValue(context.Background(), "hello", "world")
	ctx1 := context.WithValue(ctx, "foo", "bar")
	fmt.Println(ctx.Value("hello"))
	fmt.Println(ctx1.Value("hello"))
	fmt.Println(ctx1.Value("foo"))
	fmt.Println(ctx.Value("foo"))
}

func Done() <-chan struct{} {
	return make(chan struct{})
}
// output:
// world
// world
// bar
// <nil>

如上面例子,WithValue 函数可以在上下文中定义变量并返回创建的新子 ctx,通过 ctx.Value() 可以获取到你之前定义过的值,需要注意的是,代码中 ctx 并不能获取到 ctx1 中的 value,而 ctx1 可以获取到 ctx 中的 value,说明父子关系的 context 之间“子对父”是有隔离的。

close 使用 #

这里需要再额外加一个 close 队列的例子,先看懂,下面的源码会用到。

package main

import (
	"context"
	"fmt"
	"time"
)

func main() {
	ctx, _ := context.WithTimeout(context.Background(), time.Second*3)

	done := make(chan struct{})

	go func(c context.Context) {
		time.Sleep(time.Second)
		close(done)
	}(ctx)

	select {
	case <-ctx.Done():
		fmt.Println("timeout")
	case <-done:
		fmt.Println("work done!")
	}
}
//work done!

这里要注意的是,close 一个队列后,select 也能收到信号并进入到了相应的 case,可以说这是 cancel 方法的核心原理了。

结构体嵌套接口 #

interface 可以嵌套在 struct 里,这个知识点我是在看源码的时候才发现的,以前真不知道。源码中的 context.Context 是一个 interface,它定义了四个方法,通过上面的例子多少已经知道了一些使用方法,重点不是 Context,而是 cancelCtx 的实现如下:

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

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 接口,那么 cancelCtx 的实例就可以用所有实现了 Context 接口的其他结构体来初始化了,源码里这么用的:

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

这样 newCancelCtx 返回的新的 cancelCtx 实例就初始化好了:不需要再实现这些接口里定义的方法,就自动获得了调用接口方法的能力。

var cancelCtxKey int
func (c *cancelCtx) Value(key interface{}) interface{} {
	if key == &cancelCtxKey {
		return c
	}
	return c.Context.Value(key)
}

这里先有个理解,下面的源码可以细看。

结构体嵌套结构体 #

说完接口嵌套在结构体里,这里就再拓展说一下结构体嵌套在结构体里的用法。在 Go 中结构体 A 嵌套另一个结构体 B,这个平成经常会遇到,通过嵌套,可以扩展结构体,而且减少重复代码。

package main

import "fmt"

type A struct {
	Name string
}

type B struct {
	A
	Age int
}

func main() {
	b := B{
		A: A{
			Name: "A_name",
		},
		Age: 18,
	}

	fmt.Println(b.Name)
}
// A_name

字段提升 #

在上面代码中,可以通过 b.Name 直接访问到 A 的 Name 属性,但是却不能在初始化 b 的时候直接定义,否则编译器会抛出错误。

源码分析 #

通过上面的例子,已经知道怎么使用了,现在再来分析源码。看 context 的源码时和 channel 一对比就发现了一点好处,就是没有关键字带来的编译转换的复杂度,也就是说如果要实现更复杂功能的 context,你行你也能上。

Background 和 TODO #

var (
	background = new(emptyCtx)
	todo       = new(emptyCtx)
)

func Background() Context {
	return background
}

func TODO() Context {
	return todo
}

👆 先看 Backgroud 和 TODO 的定义,他们都是通过 new 初始化了一个 emptyCtx,看起来 Background 和 TODO 类型没有什么区别,根据官方文档的描述,TODO 可以用在还不知道未来如何应用哪种上下文的情况下使用。

emptyCtx #

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
}

可见 emptyCtx 实现了 Context 接口,但是都是空方法,没有任何功能。这种 ctx 通常作为起始的 context 向下传递。

WithCancel #

WithCancel 应该是 Context 里最核心的实现了,后面几个类型的 Context 都是以此为基础进行拓展的,下面开始详细看下它的源码。

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

// newCancelCtx returns an initialized cancelCtx.
func newCancelCtx(parent Context) cancelCtx {
	return cancelCtx{Context: parent}
}

👆 WithCancel 函数创建并返回了一个 cancelCtx 类型的上下文和一个 cancel 闭包。

cancelCtx #

// A cancelCtx can be canceled. When canceled, it also cancels any children
// that implement canceler.
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
}

type canceler interface {
	cancel(removeFromParent bool, err error)
	Done() <-chan struct{}
}

func (c *cancelCtx) Value(key interface{}) interface{} {
	if key == &cancelCtxKey {
		return c
	}
	return c.Context.Value(key)
}

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
}

func (c *cancelCtx) Err() error {
	c.mu.Lock()
	err := c.err
	c.mu.Unlock()
	return err
}

cancelCtx 包含了一个 Context 的 interface,和多出来的几个字段:

  • mu 互斥锁。
  • done 一个 struct{} 类型的队列,用于接收 cancel 方法发送的取消信号。
  • children 是一个map,key 为 canceler 的接口(定义了 cancel 方法和 Done() 方法),而 value 是一个空 struct,实际上并没有使用。
  • err 调用 cancel 时附带的错误信息。

propagateCancel #

func propagateCancel(parent Context, child canceler) {
	done := parent.Done()
	if done == nil {
		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():
			}
		}()
	}
}

propagateCancel 函数也是非常核心的函数之一,在上下文初始化时通过 children 绑定了父子关系,并在 parentchild 之间同步取消和结束的信号,保证在 parent 被取消时,child 也会收到对应的信号,不会出现状态不一致的情况。

cancel 方法 #

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

👆 cancel 会递归遍历自己的 children 的 cancel 方法,并使用 close 关闭自己 done 队列。

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

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

	deadline time.Time
}

timerCtx #

timerCtx 继承了 cancelCtx 相关的变量和方法,同时又比 cancelCtx 多了 timerdeadlinec.timer 通过 time.AfterFunc 设置了一个定时器,在超过这个时间之后就会调用 cancel 方法。

WithTimeout #

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

WithTimeout 只是对 WithDeadline 的参数的封装,后面就都一样了。

WithValue #

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() {
		panic("key is not comparable")
	}
	return &valueCtx{parent, key, val}
}

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

valueCtx #

WithValue 使用 valueCtx 从 Context 类型的 parent 创建新的上下文,除了多出 keyvalue 字段,没有新的内容了,值得多说一句,没看源码之前我还以为一个 context 里的 value 是一个 map,可以设置多个 key,看了源码才知道只能设置一个 key 值,那些可以设置多个 key 值的是三方库的实现。

func (c *valueCtx) Value(key interface{}) interface{} {
	if c.key == key {
		return c.val
	}
	return c.Context.Value(key)
}

小结 #

这里几乎按照源码顺序,每一段代码都做了分析,基本上可以从底层了解了 Context 的原理。一个 Context 可以无限生成子 ctx 使用,并可以被根 ctx 取消,事实上,这可能就是 context 被设计出来的意义,不管是从 interface 和 struct 运用的语法上,还是 Context 的实现上,相信你肯定也有收获。

参考 #

https://pkg.go.dev/context

https://draveness.me/golang/docs/part3-runtime/ch06-concurrency/golang-context/

本文以 CC 署名-非商业性使用-禁止演绎 4.0 国际 协议进行许可。

上一篇:Go 语言中 Channel 的实现

下一篇:Go 语言的 MPG 并发调度模型