Go context的使用和源码分析

亚洲第一中锋_哈达迪

已于 2022-04-13 20:14:05 修改

阅读量1.2k

点赞数

文章标签： go

于 2022-04-13 19:11:44 首次发布

本文链接： https://blog.csdn.net/qq_39383767/article/details/124155381

版权

什么是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

创建一个4s后超时的context
主goroutine和业务goroutine都监听该cancelCtx的Done信号
4秒后主goroutine收到超时信号
5秒后业务goroutine收到超时信号
1. 经过业务2s + timer 1s + 业务2s后，再次检查Done信号发现已经关闭
2. 如果这里业务goroutine不监听退出信号，而只是主goroutine超时退出，则可能造成业务goroutine泄露

源码分析

以下源码使用版本：1.16.10

Context

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

’
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具体做了啥？

若父节点p``arent.Done()为空，直接返回
1. 只有cancelCtx或自定义的Context才会返回不为空的Done，若parent.Done == nil，说明父节点不可被取消，例如emptyCtx，valueCtx
如果父节点可以被取消，且已经被取消，则取消当前节点，并返回
向上找到最近一个可取消的Context，例如以下这种情况，就会找到parent.parent对应的Context
1. 若已经取消，则取消当前节点
2. 否则将自己挂在到该节点的children

若找不到可取消的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. 如果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

}

下一步会根据&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)

}

如果没找到可取消的Context，则返回空，外部就监听parent.Done，而不是挂载到parent上
如果找到了，判断该可取消的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)

   }

}

总体来看执行了以下操作

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

现在问题来了：

什么时候传true？
为什么有时传 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：

假设我取消红色的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) }

}

如果父节点也是timerCtx，且父节点比当前节点早过期，就没必要新起一个timer，因为父节点过期时会通知当前节点一起取消
创建timerCtx，设置父和到期时间
和新建cancelCtx一样，找到父或祖先的可取消Context进行挂载，如果找不到，且父节点自定义实现了Done方法，就监听该Done
如果传进来的deadline已经到期了，或者执行1，2，,3步骤时到期了，就取消当前节点
新起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)

}