Go context的使用和源码分析

607 阅读1分钟

什么是context

context是一个可以携带超时,取消信号,或者其他和当前请求相关的数据,用于在api,方法,或goroutine之前进行传递

为什么有context

用go来写后台服务中,通常每个请求都会启动一个goroutine,该goroutine可能会启动多个goroutine相互配合完成本次请求工作,例如一次请求多个模块的数据,可以起多个goroutine同时发起db,rpc请求。

和本次请求有关的基本数据就需要在多个goroutine,不同方法之间共享,例如userId,logId

除此之外,一般服务器会对请求设置超时,避免长时间占用goroutine导致资源耗尽,服务雪崩。当超时时间到了,或被手动取消时,和该次请求相关的goroutine都需要快速结束退出,因为他们的工作成果不再被需要了,同时也能节省资源开销

总结来说,context 是用来解决 goroutine 之间退出通知元数据传递的功能

基本使用

保存,传递kv

// key使用自定义类型

type UserInfo string

const UserId UserInfo = "userId"



func main() {

    ctx := context.Background()

    ctx = context.WithValue(ctx, UserId, "123")

    

    val := ctx.Value(UserId)

    userId, ok := val.(string)

    if ok {

       fmt.Println(userId)

    }

}



// 输出结果

123

设置超时时间:

func main() {

   ctx := context.Background()

   // 创建一个4s后超时的context

   cancelCtx, cancel := context.WithTimeout(ctx, time.Second*4)

   defer cancel()



   go func() {

      for {

         // 模拟耗时业务

 time.Sleep(2 * time.Second)



         select {

         // 接收超时信号,并退出

         case <-cancelCtx.Done():

            fmt.Printf("business %v", cancelCtx.Err())

            return

         case <-time.After(1 * time.Second):

         }

      }

   }()



   // 接收超时信号

   select {

   case <-cancelCtx.Done():

      fmt.Printf("main %v", cancelCtx.Err())

   }

   // 避免go进程退出

   time.Sleep(10 * time.Second)

}



// 结果:

// 4秒后打印

main context deadline exceeded

// 5秒后打印

business context deadline exceeded
  1. 创建一个4s后超时的context

  2. 主goroutine和业务goroutine都监听该cancelCtx的Done信号

  3. 4秒后主goroutine收到超时信号

  4. 5秒后业务goroutine收到超时信号

    1. 经过业务2s + timer 1s + 业务2s后,再次检查Done信号发现已经关闭
    2. 如果这里业务goroutine不监听退出信号,而只是主goroutine超时退出,则可能造成业务goroutine泄露

源码分析

以下源码使用版本:1.16.10

Context

context包中的所有接口及结构体关系如下:

image.png' Context接口定义的方法如下:

type Context interface {

 Deadline() (deadline time.Time, ok bool)

 Done() <-chan struct{}

 Err() error

 Value(key interface{}) interface{}

}

该方法主要给用户使用:

  • Deadline () :获取给该context设置的超时时间

  • Done() :返回一个只读channel,表示该context的是否被取消

    • 若Done() == nil,说明其不可被取消
    • 若Done() 不为 nil,则需要监听该channel,一旦有返回,就表示被取消
  • Err() :当ctx被取消时,该方法返回被取消原因,超时或手动取消

  • Value() :主要用于valueCtx获取设置的kv

canceler接口定义如下:

type canceler interface {

   cancel(removeFromParent bool, err error)

   Done() <-chan struct{}

}

该接口表示一个context是可取消的,主要context包内部使用,cancelCtx和timerCtx实现了 canceler 接口

  • cancel():调用cancel时会发送取消信号,以及将自己从父节点移除
  • Done():和Context接口中的Done一致

cancelCtx

创建可取消ctx:

ctx := context.Background()

cancelCtx, cancel := context.WithCancel(ctx)

context.WithCancel返回了可取消的ctx,cancelCtx和一个取消方法cancel

调用cancel方法,会使得所有其他goroutine中监听cancelCtx.Done()的地方收到退出信号,执行退出操作

func WithCancel(parent Context) (ctx Context, cancel CancelFunc) {

   if parent == nil {

      panic("cannot create context from nil parent")

   }

   // 创建一个新的ctx

   c := newCancelCtx(parent)

   // 检查父ctx是否可被取消,若是则将自己挂载到父ctx

   propagateCancel(parent, &c)

   // 返回新ctx,及取消方法

   return &c, func() { c.cancel(true, Canceled) }

}





func newCancelCtx(parent Context) cancelCtx {

   // 将父ctx放到Context

   return cancelCtx{Context: parent}

}

cancelCtx结构如下:

type cancelCtx struct { 

   Context 

   // 保护下以下字段

   mu       sync.Mutex

   // 懒加载 , 用于保存关闭信号           

 done     chan struct{}      

   // 挂载的子Context,  

 children map[canceler]struct{} 

   // 表示被取消的原因:手动取消或超时取消

 err      error                 // set to non-nil by the first cancel call

}

这是一个可以取消的 Context,实现了 canceler 接口。它直接将接口 Context 作为它的一个匿名字段,这样它就可以被看成一个 Context,同时该字段也保存其父Context

propagateCancel主要作用是向上追溯可取消的Context,若有就将自己注册进Context的children,这样一来当上层调用cancel方法时,就可以往下传递,把挂载的子Content也取消

func propagateCancel(parent Context, child canceler) {

   done := parent.Done()

   // 父节点不可取消

   if done == nil {

      return 

 }



   select {

   //  父节点已经被取消

   case <-done:

 child.cancel(false, parent.Err())

      return

   default:

   }



    // 向上找到最近一个可取消的Context

   if p, ok := parentCancelCtx(parent); ok {

      p.mu.Lock()

       // 若已经被取消

      if p.err != nil {

 child.cancel(false, p.err)

      } else {

         if p.children == nil {

            p.children = make(map[canceler]struct{})

         }

         // 将自己挂载到最近一个可取消的父Context

         p.children[child] = struct{}{}

      }

      p.mu.Unlock()

   // 若找不到可取消的context,但parent实现了done方法,也进行监听   

   } else {

      atomic.AddInt32(&goroutines, +1)

      go func() {

         select {

         case <-parent.Done():

            child.cancel(false, parent.Err())

         case <-child.Done():

         }

      }()

   }

}

propagateCancel具体做了啥?

  1. 若父节点p``arent.Done()为空,直接返回

    1. 只有cancelCtx或自定义的Context才会返回不为空的Done,若parent.Done == nil,说明父节点不可被取消,例如emptyCtx,valueCtx
  2. 如果父节点可以被取消,且已经被取消,则取消当前节点,并返回

  3. 向上找到最近一个可取消的Context,例如以下这种情况,就会找到parent.parent对应的Context

    1. 若已经取消,则取消当前节点
    2. 否则将自己挂在到该节点的children

image.png

  1. 若找不到可取消的context,但parent实现了done方法,也进行监听

    1. 因为找不到可取消的Context,则无法将自己挂在上面,就只能自己另起一个goroutine监听父节点的Done,来完成取消操作

    2. 但该goroutine的退出条件是两个,还有一个是case <-child.Done(),如果子节点自己取消了,就退出,不再管父节点的退出信号。因为如果父节点迟迟不退出,这个goroutine就泄露了,这里保证在子节点退出时,就能终止该监听goroutine

再来看看向上找到最近一个可取消Context的方法:parentCancelCtx

func parentCancelCtx(parent Context) (*cancelCtx, bool) {

   done := parent.Done()

   if done == closedchan || done == nil {

      return nil, false

   }

   p, ok := parent.Value(&cancelCtxKey).(*cancelCtx)

   if !ok {

      return nil, false

   }

   p.mu.Lock()

   ok = p.done == done

   p.mu.Unlock()

   if !ok {

      return nil, false

   }

   return p, true

}
  1. 首先第一个判断

    1. 如果done == closedchan直接返回,后续子节点监听到该done如果已经被关闭,就执行cancel
    2. Done == nil在这个场景不可能成立,因为在上一步parent.Done()中会懒加载done
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

}
  1. 下一步会根据&cancelCtxKey在parent中找最近的一个cancelCtx

    1. &cancelCtxKey是context包中定义的私有变量,如果cancelCtx遇到该key就会返回自己,否则会往上找,看看有没有parent是cancelCtx,若果是就返回该cancelCtx
// context包中的私有变量

var cancelCtxKey int



func (c *cancelCtx) Value(key interface{}) interface{} {

   // cancelCtx遇到cancelCtxKey就返回自己

   if key == &cancelCtxKey {

      return c

   }

   // 否则不断往parent中找

   return c.Context.Value(key)

}
  1. 如果没找到可取消的Context,则返回空,外部就监听parent.Done,而不是挂载到parent上

  2. 如果找到了,判断该可取消的Context的done,是否和parent.Done相等,如果不等还是返回空,如果相等就返回该可取消的Context,这一步该怎么理解?

    1. 如果父节点是调用WithCancel,WithTimeout,WithDeadline生成的context,则parent.Done,一定等于该可取消的Context.Done
    2. 如果父节点是自定义的Context,自己实现了Done方法,并且包装了context包里的cancelCtx,这里就不相等

举个例子,假设自定义了MyContext:

type MyContext struct {

   context.Context

}



// 自定义done方法

func (myc *MyContext) Done() <-chan struct{} {

   return make(chan struct{})

}

用MyContext包装cancelCtx:

ctx := context.Background()



cancelCtx, _ := context.WithCancel(ctx)

// 包装context包的cancelCtx

myContext := &MyContext{

   Context: cancelCtx,

}

此时调context.WithCancel(myContext)时,就会出现两个done不一样的情况

也就是说,parent自己实现了Done接口,其返回的done,和根据parent往上寻找的第一个cancelCtx的done不一样。此时有两个done,当前context监听哪一个呢?这里选择监听自己实现的done,因为这种情况下不应该绕过用户自定义的Done。且按照层级关系来说,也应该监听最解决自已的关闭信号

这样一来,当前Context就和父或祖宗Context产生关联,要么将自己挂到最接近的一个父cancelCtx上,如果没有,且父节点自定义实现了一套产生Done信号的方法,就需要新开goroutine监听该信号

再看看返回的CancelFunc具体执行的操作

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 = err

   if c.done == nil {

      c.done = closedchan

   } else {

       // 发出关闭信号

      close(c.done)

   }

   

   // 关闭子Context

   for child := range c.children {

 child.cancel(false, err)

   }

   c.children = nil

   c.mu.Unlock()

    // 若需要,从父节点删除自己

   if removeFromParent {

      removeChild(c.Context, c)

   }

}

总体来看执行了以下操作

  1. 如果该Context已经被关闭,即err != nil,则直接退出。这也保证了cancel方法的幂等性
  2. 记录err,该值为errors.New("context canceled"),用于其他地方监听该ctx.Done返回时知道关闭原因
  3. 关闭chan,让其他监听了该chan的context知道该 context 已经被取消了
  4. 取消由该Content生成的可取消的子Contet
  5. 若参数中removeFromParent 为true,将自己从父Context中取消挂载

现在问题来了:

  1. 什么时候传true?
  2. 为什么有时传 true,有时传 false?

什么时候会传true?答案是调用WithCancel() 时,其返回的cancelFunc中对cancel的调用会传true

return &c, func() { c.cancel(true, Canceled) }

当调用该cancelFunc,会将自己从父Context中删除,这是因为自己已经被取消了,就没有必要再在父Context的关系里面接收父或祖先节点的取消通知

在cancel方法内部取消子Content时,该参数为false,也就是不需要从父Context中删除。这是因为在取消完子Context后,会执行c.children = nil,将所有子Context和自己断绝关系。这样子Content就不需要把自己从父中移除

cancelCtx的这套设计,让我们可以选择取消一颗子树上的context:

image.png

假设我取消红色的cancelCtx,只会取消这棵子树上的context,对整棵树上的其他context没有影响

timerCtx

timerCtx基于cancelCtx,只是多了个timer,deadline。当deadline到期时,timer会执行cancel方法

type timerCtx struct {

   cancelCtx

   timer *time.Timer .

 deadline time.Time

}

如何创建一个可超时自动取消的context?

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

   return WithDeadline(parent, time.Now().Add(timeout))

}

使用WithTimeout函数底层调用了WithDeadline,将一个timeout相对实现都转化为基于当前时间的绝对时间统一处理

WithDeadline方法:

func WithDeadline(parent Context, d time.Time) (Context, CancelFunc) {

   if parent == nil {

      panic("cannot create context from nil parent")

   }

   // 如果父也是timerCtx,且父的过期时间比当前时间早,就不用新起timer监听

   if cur, ok := parent.Deadline(); ok && cur.Before(d) {

 return WithCancel(parent)

   }

   c := &timerCtx{

      cancelCtx: newCancelCtx(parent),

      deadline:  d,

   }

   // 找到父或祖先的可取消Context进行挂载

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

}
  1. 如果父节点也是timerCtx,且父节点比当前节点早过期,就没必要新起一个timer,因为父节点过期时会通知当前节点一起取消
  2. 创建timerCtx,设置父和到期时间
  3. 和新建cancelCtx一样,找到父或祖先的可取消Context进行挂载,如果找不到,且父节点自定义实现了Done方法,就监听该Done
  4. 如果传进来的deadline已经到期了,或者执行1,2,,3步骤时到期了,就取消当前节点
  5. 新起timer进行定时取消操作,此时设置的错误原因就是超时,而不是被取消
var DeadlineExceeded error = deadlineExceededError{}



type deadlineExceededError struct{}



func (deadlineExceededError) Error() string   { return "context deadline exceeded" }

再看看其返回的cancel方法,除了timerCtx到期取消外,也可通过该方法手动取消:

func (c *timerCtx) cancel(removeFromParent bool, err error) {

    // 调用cancel方法

   c.cancelCtx.cancel(false, err)

   // 若需要,从父节点移除自己

   if removeFromParent {

 removeChild(c.cancelCtx.Context, c)

   }

   c.mu.Lock()

   if c.timer != nil {

       // 取消timer

      c.timer.Stop()

      c.timer = nil

   }

   c.mu.Unlock()

}

除了本身cancelCtx的cancel方法外,还将timer停止并清空

  • 自己被提前手动取消,就没有必要继续用timer到期取消了
  • 这里在停止timer后,还将c.time = nil,保证多次调用cancel的幂等性

valueCtx

type valueCtx struct {

   Context

   key, val interface{}

}

valueCtx有一个k,v对

往valueCtx里塞kv对,以及根据key找value的方法如下:



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

   }

   // 创建一个valueCtx,保存key,val,将父节点作为Context

   return &valueCtx{parent, key, val}

}





func (c *valueCtx) Value(key interface{}) interface{} {

   // 如果当前节点的key就是key,返回当前节点的value

   if c.key == key {

      return c.val

   }

   // 否则向上递归找

   return c.Context.Value(key)

}

这里要求key是可比较的,也就是可以和另一个key比较是否相当,不然没法判断当前节点的key是否是参数中的key

最终会形成一棵树

image.png

可以看到,如果要从C4找key1,需要一直遍历到解决根节点,相比与用map保存kv的方式,时间复杂度较高,那为啥valueCtx这么设计呢?

解决并发修改问题:对于任何拿到valueCtx的人相当于都是只读的,你可以修改,只会往后追加,得到一个更长的链表的指针,而不可能去修改别人已经拿到的context,这显然更安全

但这样也会有一些问题,例如若果将C4.key改为key1,则对于C4来说,C1就没用了,但还占着空间

一般来说WithValue存放的信息为和请求想干的信息,例如userId,logId。key建议用自定义类型,这样即时两个key value一样,但类型不一样,也不会产生冲突

总结

到这里Context中的源码就讲解完了,总的来说其设计比较优雅,解决了goroutine,方法直接传递元数据,及超时控制需求

参考文档

qcrao.com/2019/06/12/…