Golang Context源码学习

起因

最近学习golang框架的时候发现许多地方都用到了context的概念,比如grpc请求 etcd访问等许多地方。 本着追根溯源搞清楚实现方式的劲头,决定研究下实现原理。

用处

  1. 一般上用在GRpc等框架内,设置超时时间,比如
ctx, cancel := context.WithTimeout(context.Background(), 2 * time.Second)

dial, err := grpc.DialContext(ctx, etcdAddr, grpc.WithInsecure(),grpc.WithBalancer(balancer))    

cancel()

这里通过WithTimeout获得一个超时的Context 给grpc.DialContext 作为参数,这个Context本身内部有个timer定时器,在timer定时器时间到的时候会自动cancel掉Context 并且关闭Context内部的done chan, 一般使用ctx作参数的方法内部会检查done chan一旦发现chan 关闭,那么就应该认为这个操作需要结束了,从而返回错误(这个错误也是context内部的err,是在Context内部cancel时置的)

  1. 自己程序里用到Context时,内部实现方法类似如下
func doSomething(ctx Context) error{
    //go .....doSomethingLong......

    select{
        case <-ctx.Done():
            return ctx.Err()
        case err <- somethingChan:
            return err
    }
}

大致结构

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


//主要暴露的方法
func WithCancel(parent Context) (ctx Context, cancel CancelFunc)

func WithDeadline(parent Context, deadline time.Time) (Context, CancelFunc)

func WithValue(parent Context, key, val interface{}) Context

func WithTimeout(parent Context, timeout time.Duration) (Context, CancelFunc)

以上即Contxt的接口结构和最常用的暴露出的方法,

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

//内部重要struct
type timerCtx struct {
    cancelCtx
    timer *time.Timer // Under cancelCtx.mu.

    deadline time.Time
}

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
}

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

canceler 接口

cancel方法直接取消context,大致实现方法是置cancelCtx的err字段,当err字段不为空时,即意味着这个context已经失效;
Done方法返回是否完成channel, 判断context是否成功完成

timerCtx

withDeadline和WithTimeout返回的实际结构体(parent未失效时),而其中又包含了一个cancelCtx. cancelCtx的context为timerCtx真正的parent.
实现了接口canceler.

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

cancelCtx

对应withCancel 内含Context为其parent Context.
实现了接口canceler

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

主要干了以上三件事

  1. 置了err
  2. 关闭chan done
  3. 递归调用子context的cancel方法

valueCtx

对应withValue。 内含Context为其parent Context. valueCtx逻辑最简单 只是额外加了一对键值对, 主要提供上下文变量保存的作用. Value方法可以递归向上查找key对应的value。见Value方法 如下

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

主要具体实现

func WithDeadline(parent Context, deadline time.Time) (Context, CancelFunc) {
    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)
    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 {
        c.timer = time.AfterFunc(d, func() {
            c.cancel(true, DeadlineExceeded)
        })
    }
    return c, func() { c.cancel(true, Canceled) }
}


func WithCancel(parent Context) (ctx Context, cancel CancelFunc) {
    c := newCancelCtx(parent)
    propagateCancel(parent, &c)
    return &c, func() { c.cancel(true, Canceled) }
}


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

func removeChild(parent Context, child canceler) {
    p, ok := parentCancelCtx(parent)
    if !ok {
        return
    }
    p.mu.Lock()
    if p.children != nil {
        delete(p.children, child)
    }
    p.mu.Unlock()
}

首先判断parent本身是否已经过期了,如果过期, 只返回cancelCtx,因为已经过期,没必要使用timerCtx设置过期时间等等。 否则创建timerCtx。
然后找出最近的一个父cancelCtx,如果存在将此timerCtx置为他的child,不存在就起一个goroutine轮询等待parent完成,一旦parent完成,cancel掉此timerCtx。
根据child的过期时间作判断,如果已经过期,直接cancel掉timerCtx,并从parent中移除,防止资源堆积,如未过期设置timer过期时cancel掉timerCtx.

cancel方法有参数removeFromParent,表示是否从parent context移除本canceler.
因为只有cancelCtx有child字段,所以需要找到最近cancel parent来移除child.

总结

一圈代码看下来,原理和结构其实了解了,但是其实还是有一些实现的小细节让人绕了半天,比如

  1. cancelCtx的child放的全是canceler接口的map, 因为实现canceler接口的结构体才实现cancel方法。

  2. timerCtx的cancel方法里会先调用c.cancelCtx.cancel(false, err), 然后在判断removeFromParent, 再调用removeChild(c.cancelCtx.Context, c). 因为直接调用c.cancelCtx.canel(true, err)显示达不到移除c的目的,因为这里是从c.cancelCtx的child中移除c,然而c是在c.cancelCtx.Context的child里的。

完结感想

原来一直没有写过这种文章,写了一篇才发现好耗时间。而且写的这两也不咋地,不过的确写写能够加深理解,找出浮光掠影式看代码没发现的地方,还是很有好处的。以后要坚持下去。

你可能感兴趣的:(Golang Context源码学习)