Go语言Context源码解析

溯止
224 阅读4分钟

Context

本文主要介绍Go的Context包的内部设计，Context包如何使用不做介绍。文章主要说明了Context的创建过程和Context是如何结束的。

Context接口定义

//  Context 接口
type Context interface {
    // Deadline 返回截止时间，如果没有截止时间则ok=false。
    Deadline() (deadline time.Time, ok bool)
    // 返回一个channel，这个channel会在context结束时被close，如果是一个永远不会结束的context（emptyCtx）将返回nil。
    Done() <-chan struct{}
    // 如果context未结束，则返回nil，已经结束则返回结束原因。取消结束/超时结束，如果err不为空，重复调用返回相同的错误。
    Err() error
    // 从context中获取键对应的值。
    Value(key interface{}) interface{}
   
}

不同Context的关系

image.png

CancelCtx结构体

type cancelCtx struct {
    // 创建时来自parentCtx，为了调用parentCtx的Value函数
	Context

	mu       sync.Mutex            // 锁
        done     atomic.Value          // 一个chan struct{}，使用了懒加载，并不是创建对象时就创建，显示调用Done方法时创建。
	children map[canceler]struct{} // 保存子Context,第一次cancel后设置为nil
	err      error                 // 保存结束原因，第一cancel时设置为nil
}

创建流程图

image.png

timerCtx结构体

从下面的源代码中，我们看到timerCtx比cancelCtx只是多了一个timer和deadline。 ​

type timerCtx struct {
	cancelCtx
    // 用于定时结束context。
	timer *time.Timer 

    // 用于调用Deadline时返回时间。
	deadline time.Time
}

创建流程图

image.png

源码详细分析

CancelCtx

WithCacnel()

func WithCancel(parent Context) (ctx Context, cancel CancelFunc) {
	if parent == nil {
		panic("cannot create context from nil parent")
	}
    // 将parent赋值给canelCtx中的Context
	c := newCancelCtx(parent)
	propagateCancel(parent, &c)
	return &c, func() { c.cancel(true, Canceled) }
}

newCancelCtx()

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

propagateCancel()

func propagateCancel(parent Context, child canceler) {
	done := parent.Done()
    // 如果是TODO或Background的Context则永远不会canceled
	if done == nil {
		return // parent is never canceled
	}

	select {
	case <-done:
		// parentCtx已经结束则当前ctx也立即结束
        // 返回parentCtx的err
        // 因为未加入到parentCtx的child中所以removeFromParent参数为false
		child.cancel(false, parent.Err())
		return
	default:
	}

    // 详解见下方
	if p, ok := parentCancelCtx(parent); ok {
		p.mu.Lock()
        // 获取锁后再次判断parent是否已经退出
		if p.err != nil {
			// parent已经结束
			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)
        // 开启协程进行监听parent的退出
		go func() {
			select {
			case <-parent.Done():
				child.cancel(false, parent.Err())
			case <-child.Done():
			}
		}()
	}

parentCancelCtx()
// 判断parent是否是cancelCtx，
func parentCancelCtx(parent Context) (*cancelCtx, bool) {
	done := parent.Done()
	if done == closedchan || done == nil {
		return nil, false
	}
        // 这里的设计也很有意思，使用了一个其它任何地方都不可能使用的一个key。使用一个地址来作为key。
	p, ok := parent.Value(&cancelCtxKey).(*cancelCtx)
	if !ok {
		return nil, false
	}
        // 这里应该是判断cancelCtx的done是否被自定义了，如果是则不再被当做一个cancelCtx处理
	pdone, _ := p.done.Load().(chan struct{})
	if pdone != done {
		return nil, false
	}
	return p, true
}

Done()

// 这里使用了懒加载，context中的done并非创建时就初始化，而是当调用Done方法时才会判断是否是nil。
// 主要原因是context的parent可能已经结束，创建一个parent已经结束的context时给done分配一个新chan有些浪费。
// 这里还使用到了go的atomic.Value来保存done。这是go1.17版本中的改变。有兴趣的同学可以自行查看。理论上这样做可以少一次锁操作。
func (c *cancelCtx) Done() <-chan struct{} {
	d := c.done.Load()
	if d != nil {
		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{})
}

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
	d, _ := c.done.Load().(chan struct{})
    // 如果done为nil则将一个已经被关闭的chan赋值给done，否则关闭done。
	if d == nil {
		c.done.Store(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 {
		removeChild(c.Context, c)
	}
}

ContextDeadtime&ContextWithTimeout

ContextWithTimeout内部就是ContextDeadtime。所以这里只介绍ContextDeadtime。从上面的Context的关系，已知ContextDeadtime结构体包含了ContextCancel。ContextCancel这里就不再次展开了。主要介绍如何做到定时结束的。

func WithDeadline(parent Context, d time.Time) (Context, CancelFunc) {
	if parent == nil {
		panic("cannot create context from nil parent")
	}
    // 检查parent是否已经结束，parent的结束时间是否早于设置时间，
    // 如果早则无需设置该context的结束时间，无需开启额外开启定时任务。
	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()
    // 开启定时任务执行cancel
	if c.err == nil {
		c.timer = time.AfterFunc(dur, func() {
			c.cancel(true, DeadlineExceeded)
		})
	}
	return c, func() { c.cancel(true, Canceled) }
}

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()
    // 关闭定时器，并将timer设置为nil，可能是方便GC的操作。
	if c.timer != nil {
		c.timer.Stop()
		c.timer = nil
	}
	c.mu.Unlock()
}

总结

Context包使用起来并不难，也很容易理解。设计却十分精巧，做了很多不容易想到的优化，比如重复使用一个close的chan已降低系统开销，context中的done的懒加载。已经不同的context直接的关系，和parent之间的关系处理等。

题外话

image.png 上图是我在查询context相关设计的时看到的一个问题。大概意思为什么创建超时的ctx时为什么要在函数的参数中传入ctx，而不是直接调用ctx的方法。看到这个问题时我的第一反应也是感觉这样好像也可以。提问下方有优秀的回答。参考回答后我个人的理解。

  • context本身是一个context的interface。interface是对对象进行抽象，得出一组共同的行为。上图中的WithTimeout这个函数并不是所有的context类型都需要这个函数。也就不是所有context的共同行为。Go推崇的是面向接口编程。如果加withTimeout函数加入到接口中，只会让接口越来越庞大。
  • 普通的context与timeout的context的关系我个人觉得更像是组合。如果用一个context对象直接创建出另一个对象。这两个的关系就更像是一个继承。组合一般会比继承来的更好维护。context.WithTimeout方法更像是一个工厂方法。
avatar
文章
阅读
粉丝