深入理解go-context过程以及应用

背景

对于go中context的使用场景以及原理还不是很清楚，context可以方便地控制上下文的goroutine，是属于go独特的特性。go的gin框架就是结合context来控制每一个请求的上下文，下面结合源码一同深入理解一下go-context。同时分享一下go-context可能带来的问题以及解决办法。

1. Context接口

context接口结构如下：

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

• DeadLine() => 返回Context被cancel的时间，就相当于是Context的有效期

• Done() => 返回一个channel，注意这个channel是一个只读的channel，也就是在context.Context中并不会向该channel发送数据，从这个channel中读出数据的唯一方式是将它close掉，这个特性是groutine之间进行通信的关键，后面对于源码的解读对这块儿的说明。

• Err() => 返回一个错误，当返回的error不为空，说明当前context已经被close。

  • context中定义两个context，一个为Canceled，用于cancel取消。
  • 一个为DeadlineExceeded，用于超时后。

  • var Canceled = errors.New("context canceled")
    var DeadlineExceeded error = deadlineExceededError{}
    

2. context实现类型

context中有4种常用类型，分别为emptyCtx、cancelCtx、timerCtx以及valueCtx，每一种context都有不同的应用场景，下面会结合应用详细地说明。

2.1 emptyCtx

emptyCtx光从名字可以看出是一个空的上下文类型，这个结构很简单，常指context.Background()以及cntext.TODO()

这两个方法通常用于根部上下文，todo在官方的解释中是在你不知道当前处于何种上下文位置时使用。两个方法返回的都是context中内置的context变量，具体定义如下：

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

2.2 valueCtx

valueCtx通常用于传递值时使用的一种上下文，它的结构如下：

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

每一个valueCtx只能存储一个键值对，同时实现了Context接口，因此是其的一个实现。

valueCtx的初始化也是十分简单，context.WithValue(context, key, value)即可完成初始化。返回的是存储其父context以及一对key、value值valueCtx。

对于valueCtx，需要关注Value方法，具体源码如下：

• 可以看到，调用Value时会执行私有函数value自下向上不断找匹配的key对应的value，直到根部结束。

func (c *valueCtx) Value(key interface{}) interface{} {
   if c.key == key {
      return c.val
   }
   return value(c.Context, key) // 实现自下向上查找
}

func value(c Context, key interface{}) interface{} {
   for {
   // 不断循环找到对应的key，直到找到根部context，default走的是自定义context
      switch ctx := c.(type) {
      case *valueCtx:
         if key == ctx.key {
            return ctx.val
         }
         c = ctx.Context
      case *cancelCtx:
         if key == &cancelCtxKey {
            return c
         }
         c = ctx.Context
      case *timerCtx:
         if key == &cancelCtxKey {
            return &ctx.cancelCtx
         }
         c = ctx.Context
      case *emptyCtx:
         return nil
      default:
         return c.Value(key)
      }
   }
}

其具体流程图大致如图所示：

• 比如当前处于valueCtx4中，如果要查找key="a"的值，那么会从下往上找，经过valueCtx4->valueCtx3->valueCtx1。当然，key="b"是无法找到的，因为不在同一个链路中。

从命名上可以看到，其context是一个可以被cancel的上下文功能。在context中定义了一个canceler接口，主要用途是当前goroutine需要结束时，可以向其子context发送信号，子context接收信号后做一系列处理，通常用于上下文的控制。canceler结构如下：

• cancelCtx以及timerCtx都实现了canceler接口

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

cancelCtx结构如下：

• mu是同步锁，用于并发场景下的赋值和传值操作时保持原子性
• children存储子canceler，当此context结束时将会依次close children
• done，懒加载机制的原子变量，当执行Done()方法时才会被加载并创建一个channel变量

type cancelCtx struct {
   Context
   mu       sync.Mutex            
   done     atomic.Value          
   children map[canceler]struct{} 
   err      error                 
}

下面来看下cancelCtx的初始化：cancelCtx可以说是context精髓，下面详细分析其执行原理。

• 传入一个parent context，并返回cancelCtx和一个cancel方法，该方法可以给Done()方法传递一个close的信号，通知本context和其children context。

• propagateCancel(parent, &c)是将当前cancelCtx挂载到父context，如果父context是自定义context那么还会开启协程进行监听。

func WithCancel(parent Context) (ctx Context, cancel CancelFunc) {
   if parent == nil {
      panic("cannot create context from nil parent")
   }
   c := newCancelCtx(parent)
   propagateCancel(parent, &c)
   return &c, func() { c.cancel(true, Canceled) }
}

1. propagateCancel

下面详细分析下该方法的作用

func propagateCancel(parent Context, child canceler) {
    // 1. 获取父context的done，如果为nil说明该父context无法被挂在，无需挂在
   done := parent.Done()
   if done == nil {
      return // parent is never canceled
   }
    // 2.(第一次检查)done变量是否发送close信号，如果发送了说明父context关闭，则该context也发送close信号
   select {
   case <-done:
      // parent is already canceled
      child.cancel(false, parent.Err())
      return
   default:
   }
    // 3.获取父类context，如果不能cancel或者已关闭那么就获取失败，ok为false
   if p, ok := parentCancelCtx(parent); ok {
   // 加锁
      p.mu.Lock()
      // (第四次检查)加锁后再次检查父context是否被close。
      if p.err != nil {
         // parent has already been canceled
         child.cancel(false, p.err)
      } else {
      // 如果还存在，那么久将本context挂载到其父context的children这个map中
         if p.children == nil {
            p.children = make(map[canceler]struct{})
         }
         p.children[child] = struct{}{}
      }
      // 释放锁
      p.mu.Unlock()
   } else {
   // else中表示当前无法挂载或者不是cancelCtx，那么久会开启一个协程监听其状态。这里有一个重点需要注意，为什么它不仅监听parent，也监听当前的context，这是因为如果parent一直未close，而当前context结束了，那么久无法关闭当前创建的协程，从而导致内存的泄漏。
   // 统计多少goroutine被创建，用于测试
      atomic.AddInt32(&goroutines, +1)
      go func() {
         select {
         case <-parent.Done():
            child.cancel(false, parent.Err())
         case <-child.Done():
         }
      }()
   }
}

func parentCancelCtx(parent Context) (*cancelCtx, bool) {
   // (第二次检查)是否发送close信号。done明显不为nil为什么要判断一次呢？因为在removeChild方法的时候需要判断，该方法被多个地方调用。
   done := parent.Done()
   // closedchan是内置的一个变量，用于Done()方法还没执行时就被cancel了，这时候会直接返回closedchan，channel在读时它会直接传递close信号
   if done == closedchan || done == nil {
      return nil, false
   }
   // 判断parent是属于什么类型的context只有cancelCtx才可以使用cancelCtxKey，并返回当前ctx
   p, ok := parent.Value(&cancelCtxKey).(*cancelCtx)
   if !ok {
      return nil, false
   }
   //（第三次检查）判断前后两次获取的done是否一致，不一致说明可能已经被close了。还有对于自定义的context，也是无法通过的，具体后面举例说明
   pdone, _ := p.done.Load().(chan struct{})
   if pdone != done {
      return nil, false
   }
   return p, true
}

2. 自定义的context

前面parent CancelCtx方法提到自定义context来举例说明p.done.Load为什么无法通过

• 由于自定义context会实现Done方法，返回的是由自己控制的channel，而通过Value获取到的为父类的父类，虽然是cancelCtx类型，但是和Done方法得到的是不同的context，因此pdone != done成立导致最后一个条件无法通过。

type MyCtx struct {
   context.Context
}

func (my *MyCtx) Done() <-chan struct{} {
   return make(chan struct{})
}

func (my *MyCtx) Value(key interface{}) interface{} {
 // 获取父context
   return my.Context
}

func main() {
   parent, _ := context.WithCancel(context.Background())
   myCtx := &MyCtx{parent}
   child, _ := context.WithCancel(myCtx)
   println(child)
}

• 具体结构如下，parent和myCtx明显是不同的context，因此done变量肯定也不同

3. cancel方法

cancel就是close当前context的channel通道，然后遍历当前context的所有children并全部close。

代码如下：

func (c *cancelCtx) cancel(removeFromParent bool, err error) {
   if err == nil {
      panic("context: internal error: missing cancel error")
   }
   c.mu.Lock()
   // 如果不为空，则表示当前context已经被close过
   if c.err != nil {
      c.mu.Unlock()
      return // already canceled
   }
   c.err = err
   d, _ := c.done.Load().(chan struct{})
   if d == nil {
     // 如果d为空，则说明还没调用Done，此时将已关闭的通道赋值进去
      c.done.Store(closedchan)
   } else {
   // 不为空，则直接关闭
      close(d)
   }
   for child := range c.children {
      // NOTE: acquiring the child's lock while holding parent's lock.
      // cancel的第一个参数是表示，是否从parent中移除，由于当前parent已关闭，children无需再从parent中移除了。
      child.cancel(false, err)
   }
   c.children = nil
   c.mu.Unlock()

   if removeFromParent {
      removeChild(c.Context, c)
   }
}

2.4 timerCtx

这个context则是在cancelCtx基础上而成的，用于需要时间限制的上下文场景下使用。具体结构如下：

• timer为一个定时器，当时间到时将会自动执行cancel触发channel发送信号
• deadline为截止时间，通过Deadline()方法可以获取到

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

   deadline time.Time
}

1. timerCtx初始化

对于timerCtx初始化有2种方法， 分别为context.WithTimeout和context.WithDeadline。而前者是基于后者实现的，因此这里深入理解下WithDeadline是如何运作的，如下为源码：

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

func WithDeadline(parent Context, d time.Time) (Context, CancelFunc) {
   if parent == nil {
      panic("cannot create context from nil parent")
   }
 // 如果当前结束时间在parent时间之后，那么此时的时间受到parent控制，则不需要再设置，仅需返回cancelCtx即可。比如parent是3s后过期，child是6s后过期，那么child是受到parent控制，则不需要再创建timeCtx，只需要cancelCtx即可。
   if cur, ok := parent.Deadline() ; ok && cur.Before(d) {
      return WithCancel(parent)
   }
   // 初始化一个timerCtx，同时赋值一个cancelCtx
   c := &timerCtx{
      cancelCtx: newCancelCtx(parent),
      deadline:  d,
   }
   // 由于结构上是cancelCtx的实现类，因此可以挂载到timerCtx或cancelCtx之上
   propagateCancel(parent, c)
   dur := time.Until(d)
   if dur <= 0 {
       // 如果已经到期，那么直接执行cancel方法传递信号，并从根部移除该timerCtx，设置error为DeadlineExceeded
      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 {
   // 时间到后会触发函数执行，同时timer.Stop也可以触发
      c.timer = time.AfterFunc(dur, func() {
         c.cancel(true, DeadlineExceeded)
      })
   }
   return c, func() { c.cancel(true, Canceled) }
}

timerCtx和cancelCtx实现了自己的cancel方法，仅在cancelCtx的cancel方法基础上加了一个对timer的控制，即cancel后会执行timer.Stop()触发AfterFunc函数的执行。

3.各个context场景用途

1. 对于仅各个goroutine之间传递数据时，可以使用valueCtx。

2. 对于多级关联的goroutine，同时上级的goroutine失效则代表下级所有gotoutine失效，则可以使用cancelCtx。这个对于复杂的同时又嵌套许多个goroutine时的情况特别好用，例如并发的http请求等。同时可以避免内存泄漏。下面举一个可能发生的场景例子：

   • 在主goroutine结束后，子goroutine还在写入数据，但是由于没有读，因此channel一直处于阻塞状态，此时就会发生goroutine异常增长问题。此时如果使用timerCtx或者cancelCtx就可以避免此问题的发生。

func main() {
   ch := make(chan string)
   go func() {
      data := []string{"aa", "bb", "cc", "dd"}
      for _, item := range data {
         ch <- item
      }
   }()

   name := "bb"
   for {
      data := <-ch
      if data == name {
         fmt.Println("do something")
         break
      }
   }
   time.Sleep(50*time.Second)
}

// 此时代码改为如下即可避免
func main() {
   bg := context.Background()
   ctx, cancel := context.WithTimeout(bg, time.Second*1)
   childCtx, _ := context.WithCancel(ctx)
   ch := make(chan string)
   go func(ctx context.Context) {
      data := []string{"aa", "bb", "cc", "dd"}
      for _, item := range data {
         select {
         case <-ctx.Done():
            fmt.Println("child finish")
            return
         default:
            ch <- item
         }
      }
   }(childCtx)

   name := "bb"
   for {
      select {
      case data := <-ch:
         fmt.Println(data)
         if data == name {
            fmt.Println("do something")
            cancel()  // 如果1s内执行完，关闭父goroutine，子goroutine也会关闭
 }
      case <-ctx.Done():
         fmt.Println("parent finish")
         time.Sleep(time.Second * 5)
         return
      }
   }
}

3. 对于需要限制超时的情况，可以使用timerCtx。例如有一个数据库查询亦或者http请求时间过长，那么在select监听时可以设置一个超时时间，时间到时主context则会收到DeadlineExceeded错误信号，可以立即反馈给上游服务。在数据库查询结束或者http请求结束后子context协程也会自动结束。
4. 对于上述3种情况都有可能使用到的，则需要自定义context。比如gin框架中的context，我们可以在此基础上自定义一个context控件。由于gin.context没有超时机制，这了我们在此基础上封装一层，具体如下：

   
     type WrapperContext struct {
        Context context.Context
        Gin     *gin.Context
     }
   
     type WrapperFunc func(c *WrapperContext)
   
     func WithWrapperContext(wrapperFunc WrapperFunc) gin.HandlerFunc {
        return func(c *gin.Context) {
           c.Set("traceid", "{{traceid}}")
           timeout, _ := context.WithTimeout(c, 250*time.Millisecond)
           ctx := &WrapperContext{Context: timeout, Gin: c}
           wrapperFunc(ctx)
        }
     }
   
     func test(c *WrapperContext) {
        for {
           select {
           case <-c.Context.Done():
              fmt.Println("timeout")
              c.Gin.JSON(0, "")
              return
           default:
              fmt.Println("do something")
              time.Sleep(time.Millisecond * 100)
           }
        }
     }
   
     func main() {
        e := gin.Default()
        e.GET("/test", WithWrapperContext(test))
        e.Run(":8080")
     }
     ```