Go context源码解析

天唯 · · 1230 次点击 · · 开始浏览    
这是一个创建于 的文章,其中的信息可能已经有所发展或是发生改变。

在上一篇文章 golang context初探 中,已经初步了解了context的用法以及应用的场景。那么接下来深入到源码中来学习一下context是怎么实现的。

emptyCtx

context包的代码很少,一个context.go文件,总共才480行代码,其中还包括大量的注释。context包首先定义了一个Context接口:

type Context interface {
    Deadline() (deadline time.Time, ok bool)

    Done() <-chan struct{}

    Err() error

    Value(key interface{}) interface{}
}

接下来定义了一个emptyCtx类型:

// An emptyCtx is never canceled, has no values, and has no deadline. It is not
// struct{}, since vars of this type must have distinct addresses.
type emptyCtx int

为什么叫做emptyCtx呢?注释说了emptyCtx不能被取消,没有值,也没有deadline。同时,emptyCtx也不是一个struct{},因为这种类型地变量需要有不同的地址。
这个emptyCtx实现了Context接口:

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
}

func (e *emptyCtx) String() string {
    switch e {
    case background:
        return "context.Background"
    case todo:
        return "context.TODO"
    }
    return "unknown empty Context"
}

看了上面这段代码就知道为什么emptyCtx不能被取消,没有值,也没有deadline了,因为上面的实现都是直接return。那么,这个emptyCtx有什么用呢?还记得Background()TODO()函数吗?对的,它们的内部就是直接返回emptyCtx类型的指针:

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

func Background() Context {
    return background
}

func TODO() Context {
    return todo
}

所以这两个函数一般是用在main函数、初始化、测试以及顶层请求的Context。OK,继续往下看。
既然emptyCtx类型什么都不做,那么应该有其他的类型来实现相关的功能才对,也就是cancelCtxtimerCtxvalueCtx三种类型。下面来讲一下这三种类型。

Cancel

在之前的文章中我们知道,WithCancelWithTimeoutWithDeadline这三个方法会返回一个CancelFunc类型的函数,在Context内部就定义了canceler接口:

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

canceler是一种可以直接被取消的context类型,待会继续往下看我们会发现,cancelCtxtimerCtx不单单实现了Context接口(通过匿名成员变量),也实现了canceler接口。

cancelCtx

cancelCtx的结构体定义:

type cancelCtx struct {
    Context

    done chan struct{} // closed by the first cancel call.

    mu       sync.Mutex
    children map[canceler]struct{} // set to nil by the first cancel call
    err      error                 // set to non-nil by the first cancel call
}

方法集:

func (c *cancelCtx) Done() <-chan struct{} {
    return c.done
}

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

func (c *cancelCtx) String() string {
    return fmt.Sprintf("%v.WithCancel", c.Context)
}

// cancel closes c.done, cancels each of c's children, and, if
// removeFromParent is true, removes c from its parent's children.
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
    // 关闭c的done channel,所有监听c.Done()的goroutine都会收到消息
    close(c.done)
    // 取消child,由于是map结构,所以取消的顺序是不固定的
    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()
    // 从c.children中移除取消过的child
    if removeFromParent {
        removeChild(c.Context, c)
    }
}

当我们在调用WithCancel的时候,实际上返回的就是一个cancelCtx指针和cancel()方法:

func WithCancel(parent Context) (ctx Context, cancel CancelFunc) {
    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,
        done:    make(chan struct{}),
    }
}

那么,propagateCancel函数又是干什么的呢?

// 向上找到最近的可以被取消的父context,将子context放入parent.Children中
// 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
    }
    // 判断返回的parent是否是cancelCtx
    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 {
        go func() {
            select {
            case <-parent.Done():
                child.cancel(false, parent.Err())
            case <-child.Done():
            }
        }()
    }
}

// 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.
// 不停地向上寻找最近的可取消的父context
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
        }
    }
}

Timer

WithTimeoutWithDeadline其实是差不多的,只是源码内部帮我们封装了一下:

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

func WithDeadline(parent Context, deadline time.Time) (Context, CancelFunc) {
    // 当前的deadline比新的deadline还要早,直接返回
    if cur, ok := parent.Deadline(); ok && cur.Before(deadline) {
        // The current deadline is already sooner than the new one.
        return WithCancel(parent)
    }
    c := &timerCtx{
        cancelCtx: newCancelCtx(parent),
        deadline:  deadline,
    }
    propagateCancel(parent, c)
    d := time.Until(deadline)
    // deadline已经过了,不再设置定时器
    if d <= 0 {
        c.cancel(true, DeadlineExceeded) // deadline has already passed
        return c, func() { c.cancel(true, Canceled) }
    }
    c.mu.Lock()
    defer c.mu.Unlock()
    if c.err == nil {
        // 设定d时间后执行取消方法
        c.timer = time.AfterFunc(d, func() {
            c.cancel(true, DeadlineExceeded)
        })
    }
    return c, func() { c.cancel(true, Canceled) }
}

timerCtx的代码也实现地比较简洁:

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

    deadline time.Time
}

func (c *timerCtx) Deadline() (deadline time.Time, ok bool) {
    return c.deadline, true
}

func (c *timerCtx) String() string {
    return fmt.Sprintf("%v.WithDeadline(%s [%s])", c.cancelCtx.Context, c.deadline, time.Until(c.deadline))
}

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

这里需要注意的是,timerCtx并没有直接实现了canceler接口,而是使用了匿名成员变量,这样就可以不用重头实现一遍Context接口,而是按需只实现了Deadline方法。timerCtxcancel方法先调用了cancelCtxcancel方法,然后再去停止定时器。

Value

先来看看valueCtx的定义:

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

func (c *valueCtx) String() string {
    return fmt.Sprintf("%v.WithValue(%#v, %#v)", c.Context, c.key, c.val)
}

func (c *valueCtx) Value(key interface{}) interface{} {
    // 找到了直接返回
    if c.key == key {
        return c.val
    }
    // 向上继续查找
    return c.Context.Value(key)
}

这是最简单的Context实现了,在匿名变量Context之外,还增加了两个key,value变量来存储值。看看WithValue的实现:

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

需要注意的是,key不能是nil并且必须是可比较的,否则就会导致panic!可比较的意思是key不能为函数类型或者NaN之类的,具体可以看一下reflect包,这里就不细说了。

最后

整个context看下来,对context是怎么实现的也有了清晰的了解,而且context包有大量的测试代码,非常棒!最近在看go的源码,发现真的是一个很好的学习材料,对于如何写出简洁明了的代码是很有帮助的。


有疑问加站长微信联系(非本文作者)

本文来自:简书

感谢作者:天唯

查看原文:Go context源码解析

入群交流(和以上内容无关):加入Go大咖交流群,或添加微信:liuxiaoyan-s 备注:入群;或加QQ群:692541889

1230 次点击  
加入收藏 微博
暂无回复
添加一条新回复 (您需要 登录 后才能回复 没有账号 ?)
  • 请尽量让自己的回复能够对别人有帮助
  • 支持 Markdown 格式, **粗体**、~~删除线~~、`单行代码`
  • 支持 @ 本站用户;支持表情(输入 : 提示),见 Emoji cheat sheet
  • 图片支持拖拽、截图粘贴等方式上传