深度解析go context实现原理及其源码

2021-02-26 本文已影响0人迈莫coding

Context 基本使用方法

首先，我们来看一下 Context 接口包含哪些方法，这些方法都是干什么用的。

包 context 定义了 Context 接口，Context 的具体实现包括 4 个方法，分别是Deadline、Done、Err 和 Value，如下所示：

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

Deadline 方法会返回这个 Context 被取消的截止日期。如果没有设置截止日期，ok 的值是 false。后续每次调用这个对象的 Deadline 方法时，都会返回和第一次调用相同的结果。

Done 方法返回一个 Channel 对象。在 Context 被取消时，此 Channel 会被 close，如果没被取消，可能会返回 nil。后续的 Done 调用总是返回相同的结果。当 Done 被 close 的时候，你可以通过 ctx.Err 获取错误信息。Done 这个方法名其实起得并不好，因为名字太过笼统，不能明确反映 Done 被 close 的原因，因为 cancel、timeout、deadline 都可能导致 Done 被 close，不过，目前还没有一个更合适的方法名称。

关于 Done 方法，你必须要记住的知识点就是：如果 Done 没有被 close，Err 方法返回 nil；如果 Done 被 close，Err 方法会返回 Done 被 close 的原因。

Context使用场景

上下文信息传递（request-scoped），比如处理 http 请求、在请求处理链路上传递信息
控制子 goroutine 的运行
超时控制的方法调用
可以取消的方法调用

valueCtx

valueCtx 是基于 parent Context 生成一个新的 Context，保存了一个key-value键值对。它主要用来传递上下文信息。

使用示例

ctx := context.Background()
ctx = context.WithValue(ctx, "key1", "0001")
ctx = context.WithValue(ctx, "key2", "0001")
ctx = context.WithValue(ctx, "key3", "0001")
ctx = context.WithValue(ctx, "key4", "0004")
fmt.Println(ctx.Value("key1")) // 0001

查找过程如图所示：

[图片上传失败...(image-619e51-1614347187242)]

结构体

type valueCtx struct {
   Context  // parent Context
   key, val interface{}  // key-value
}

func (c *valueCtx) Value(key interface{}) interface{} {
   // 若key值 等于 当前valueCtx存储的key值 
   // 则取出其value并返回
   if c.key == key {
      return c.val
   }
   // 否则递归调用valueCtx中Value方法，获取其parent Context中存储的key-value
   return c.Context.Value(key)
}

通过观察 valueCtx 结构体，它利用一个 Context 变量表示其父节点的 context ，这样 valueCtx 也继承了父节点的所有信息；并且它持有一个 key-value 键值对，说明它还可以携带额外的信息。它还覆盖了 Value 方法，优先从自己的存储中检查这个 key，不存在的话会从 parent 中继续检查。

WithValue

WithValue 就是向 context 中添加键值对：

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")
   }
   return &valueCtx{parent, key, val}
}

通过代码可以看出，向 context 中添加键值对并不是在原 context 基础上添加的，而是新建一个 valueCtx 子节点，将原 context 作为父节点。以此类推，就会形成一个 context 链。在查找过程中，如果当前 valueCtx 不存在key值，还会向 parent Context 去查找，如果 parent 还是 valueCtx 的话，还是遵循相同的原则：valueCtx 会嵌入 parent，所以还是会查找 parent 的 Value 方法的。

[图片上传失败...(image-cbad57-1614347187242)]

cancleCtx

在我们开发过程中，我们常常会遇到一些场景，需要主动取消长时间的任务或者中止任务，这个时候就可以使用cancelCtx。通过调用cancel函数就可中止goroutine，进而去释放所占用的资源。

需要注意的是，不是只有中途中止任务时才调用cancel函数，只要任务执行完毕后，就需要调用 cancel，这样，这个 Context 才能释放它的资源（通知它的 children 处理 cancel，从它的 parent 中把自己移除，甚至释放相关的 goroutine）。

使用示例

func main() {
  // gen 在单独的 goroutine 中生成整数 然后将它们发送到返回的管道
  gen := func(ctx context.Context) <-chan int {
     dst := make(chan int)
     n := 1
     go func() {
        for {
           select {
           case <-ctx.Done():
              return // returning not to leak the goroutine
           case dst <- n:
              n++
           }
        }
     }()
     return dst
  }
  ctx, cancel := context.WithCancel(context.Background())
  // 代码完毕后调用cancel函数释放goroutine所占用的资源
  defer cancel() // cancel when we are finished consuming integers
  // 遍历循环获取管道中的值
  for n := range gen(ctx) {
     fmt.Println(n)
     if n == 5 {
        break
     }
  }
}

创建一个 gen函数，在gen函数中创建一个goroutine，专门用来生成整数，然后将他们发送到返回的管道。通过 context.WithCancel 创建可取消的 context ，最后遍历循环获取管道中值，当n的值为5时，退出循环，结束进程。最后调用cancel函数释放goroutine所占用的资源。

结构体

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

cancelCtx和valueCtx类似，结构体中都有一个Context作为其父节点；变量done表示关闭信号传递；变量children表示当前节点所拥有的子节点，err用于存储错误信息表示任务结束的原因。

接下来，看看cancelCtx实现的方法：

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

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
}

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
   // 设置一个关闭的channel或者将done channel关闭，用以发送关闭信号
   if c.done == nil {
      c.done = closedchan
   } else {
      close(c.done)
   }
   // 遍历循环将字节点context取消
   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 {
      // 将当前context节点从父节点上移除
      removeChild(c.Context, c)
   }
}

cancelCtx结构体实现Done和cancel方法，Done方法实现了将done初始化。cancel方法用于将当前节点从父节点上移除以及移除当前节点下的所有子节点。

cancelCtx 被取消时，它的 Err 字段就是下面这个 Canceled 错误：

var Canceled = errors.New("context canceled")

WithCancel

WithCancel函数用来创建一个可取消的context，即cancelCtx类型的context。

WithCancel函数返回值有两个，一个为parent 的副本Context，另一个为触发取消操作的CancelFunc。

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

// newCancelCtx returns an initialized cancelCtx.
func newCancelCtx(parent Context) cancelCtx {
   // 将parent作为父节点context生成一个新的子节点
   return cancelCtx{Context: parent}
}

func propagateCancel(parent Context, child canceler) {
   done := parent.Done()
   if done == nil {
      return // parent is never canceled
   }
   
   select {
   case <-done:
      // parent is already canceled
      child.cancel(false, parent.Err())
      return
   default:
   }
   
   // 获取最近的类型为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{})
         }
         // 将当前子节点加入最近cancelCtx祖先节点的children中
         p.children[child] = struct{}{}
      }
      p.mu.Unlock()
   } else {
      atomic.AddInt32(&goroutines, +1)
      go func() {
         select {
         case <-parent.Done():
            child.cancel(false, parent.Err())
         case <-child.Done():
         }
      }()
   }
}

调用 WithCancel函数时，首先会调用 newCancelCtx函数创建一个以parent作为父节点的context。然后调用propagateCancel函数，用来建立当前context节点与parent节点之间的关系。

在propagateCancel函数中，如果parent节点为nil，说明parent以上的路径没有可取消的cancelCtx，则不需要处理。

否则通过parentCancelCtx函数过去当前节点最近的类型为cancelCtx的祖先节点，首先需要判断该祖先节点是否被取消，若已被取消就取消当前节点；否则将当前节点加入祖先节点的children列表中。

否则的话，则需要新起一个 goroutine，由它来监听 parent 的 Done 是否已关闭。一旦parent.Done()返回的channel关闭，即context链中某个祖先节点context被取消，则将当前context也取消。

WithTimeout

WithTimeout 其实是和 WithDeadline 一样，只不过一个参数是超时时间，一个参数是截止时间。超时时间加上当前时间，其实就是截止时间，因此，WithTimeout 的实现是：

func WithTimeout(parent Context, timeout time.Duration) (Context, CancelFunc) { 
  // 当前时间+timeout就是deadline
  return WithDeadline(parent, time.Now().Add(timeout))
}

WithDeadline

WithDeadline 会返回一个 parent 的副本，并且设置了一个不晚于参数 d 的截止时间，类型为 timerCtx（或者是 cancelCtx）。

如果它的截止时间晚于 parent 的截止时间，那么就以 parent 的截止时间为准，并返回一个类型为 cancelCtx 的 Context，因为 parent 的截止时间到了，就会取消这个 cancelCtx。

如果当前时间已经超过了截止时间，就直接返回一个已经被 cancel 的 timerCtx。否则就会启动一个定时器，到截止时间取消这个 timerCtx。

综合起来，timerCtx 的 Done 被 Close 掉，主要是由下面的某个事件触发的：

截止时间到了
cancel 函数被调用
parent 的 Done 被 close

使用示例

func main() {
  d := time.Now().Add(time.Second * 3)
  ctx, cancel := context.WithDeadline(context.Background(), d)
  defer cancel()
  select {
  case <-time.After(3 * time.Second):
     fmt.Println("overslept")
  case <-ctx.Done():
     fmt.Println(ctx.Err())
  }
}

WithDeadline

func WithDeadline(parent Context, d time.Time) (Context, CancelFunc) {
   if parent == nil {
      panic("cannot create context from nil parent")
   }
   // 如果parent的截止时间更早，直接返回一个cancelCtx即可
   if cur, ok := parent.Deadline(); ok && cur.Before(d) {
      return WithCancel(parent)
   }
   c := &timerCtx{
      cancelCtx: newCancelCtx(parent),
      deadline:  d,
   }
   // 建立新建context与可取消context祖先节点的取消关联关系
   propagateCancel(parent, c)
   dur := time.Until(d)
   if dur <= 0 { //当前时间已经超过了截止时间，直接cancel
      c.cancel(true, DeadlineExceeded) 
      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) }
}

调用 WithDeadline函数，首先判断parent的截止时间是否早于当前timerCtx，若为true的话，直接返回一个cancelCtx即可。否则需要调用propagateCancel函数建议新建context与可取消context祖先节点的取消关联关系，建立关联关系之后，若当前时间已经超过截止时间后，直接cancel。否则的话，需设置一个定时器，到截止时间后取消。

总结

context主要用于父子任务之间的同步取消信号，本质上是一种协程调度的方式。另外在使用context时有两点值得注意：上游任务仅仅使用context通知下游任务不再需要，但不会直接干涉和中断下游任务的执行，由下游任务自行决定后续的处理操作，也就是说context的取消操作是无侵入的；context是线程安全的，因为context本身是不可变的（immutable），因此可以放心地在多个协程中传递使用。

到这里，Context 的源码已解读完毕，希望对您有收获，咱们下期再见。

文章也会持续更新，可以微信搜索「迈莫coding 」第一时间阅读。

深度解析go context实现原理及其源码

目录

Context 基本使用方法

Context使用场景

valueCtx

使用示例

结构体

WithValue

cancleCtx

使用示例

结构体

WithCancel

WithTimeout

WithDeadline

使用示例

WithDeadline

总结

猜你喜欢

热点阅读