【4-3 Golang】常用标准库—上下文context

  Context顾名思义上下文，可用于在整个请求上下文传值以及控制超时，本篇文章主要介绍Context的设计思路，以及基本使用方式。

Context 使用方式

  设想有一个Go HTTP服务，在请求的整个处理链路，可能随时需要获取一些公共数据，如用户ID等，怎么办呢？通过参数呗，每个函数的第一个输入参数都是用户ID不就行了，如果再加一个公共数据呢？再加一个参数吗？如果将所有这些公共参数封装成一个结构体呢？这样貌似也可以。

  不过Go语言本身就为我们提供了context.Context（context.valueCtx），其可以存储一个key-value键值对，那不行啊，只能存储一个肯定不够啊；怎么办，基于老的context.valueCtx对象再衍生新的context.valueCtx对象，两个context.valueCtx对象各能存储一个key-value键值对，想存储更多的数据，继续衍生就可以了。事例程序如下：

package main

import (
    "context"
    "fmt"
)

func main() {
    ctx := context.Background()
    ctx1 := context.WithValue(ctx, "k1", "v1")
    ctx2 := context.WithValue(ctx1, "k2", "v2")

    fmt.Println(ctx1.Value("k1"))
    fmt.Println(ctx1.Value("k2")) //返回nil

    fmt.Println(ctx2.Value("k1"))
    fmt.Println(ctx2.Value("k2"))
}

// v1 <nil> v1 v2

  基于context.WithValue函数可以衍生新的context.valueCtx对象，同时传递需要存储的key-value键值对；注意第一个参数需要一个context.Context（这是一个接口，context.valueCtx实现了该接口）对象，context.Background函数可返回一个空的context.valueCtx对象。

  仔细观察输出结果，ctx1对象只能获取到k1，ctx2对象可以获取到k1以及k2，因为ctx2对象是基于ctx1对象衍生出来的，也可以说ctx1对象是ctx2对象的父对象，而ctx2.Value既可以获取到自己存储的数据，也能获取到父对象存储的数据。最后值得一提的是，context.valueCtx存储的key-value键值对，类型都是interface{}，所以获取到数据之后一般需要进行类型转换才能使用。

  既然通过context.Context就能实现key-value数据的存储，那就使用它呗，只需要项目中所有函数的第一个参数都是context.Context，就能实现在整个请求链路传值。

  context.Context就这么点作用？当然不是，最常用的还是它的超时控制功能。假设某项任务有时间限制，最多执行3秒，超时后取消任务的继续执行，这不很简单，通过定时器就能实现，那如果任务比较复杂，又分为多个子任务并启动了多个子协程分别执行呢，3秒超时后能同时结束整个任务吗，包括主任务以及多个子任务？这时候定时器能实现吗？比较困难。

  下面程序展示了基于context.Context实现的任务超时控制。

package main

import (
    "context"
    "fmt"
    "sync"
    "time"
)

func main() {
    ctx := context.Background()
    //可取消的context
    ctx1, cancel := context.WithCancel(ctx)

    //WaitGroup控制并发任务，主协程需等待子任务结束才能退出
    wg := sync.WaitGroup{}
    wg.Add(1)
    go task(ctx1, &wg)

    //三秒后取消子任务
    time.AfterFunc(time.Second*3, func() {
        cancel()
    })

    //等待子任务结束
    wg.Wait()
}

func task(ctx context.Context, wg *sync.WaitGroup) {
    defer wg.Done()
    for {
        //context取消时会关闭管道，从而可读，实现任务结束return
        select {
        case <-ctx.Done():
            fmt.Println("context cancel and return")
            return
        default:

        }
        //子任务1秒周期执行一次，死循环
        fmt.Println("exec task", time.Now())
        time.Sleep(time.Second)
    }
}

  context.WithCancel函数返回可取消的上下文（contetx.cancelCtx对象，同样实现了接口context.Context），函数返回两个值，第一个值就是contetx.cancelCtx对象，第二个值是一个函数，可用于结束该上下文，结束之后ctx.Done函数返回的管道变为可读的，所以子任务可以通过其判断上下文是否被结束，是否该结束当前任务。上面程序事例，定时器3秒超时后结束当前上下文，当然你也可以基于任何条件判断是否需要结束上下文。

  超时控制相对也是比较简单的，只需要主协程控制何时结束上下文，子任务只需监听ctx.Done管道，上下文结束后管道可读，从而结束所有子任务。

  contetx.cancelCtx需要你自己基于定时器实现超时控制，Go语言还提供有两个函数，很方便实现context的超时控制：

//timeout之后，自动结束上下文
func WithTimeout(parent Context, timeout time.Duration) (Context, CancelFunc)
func WithDeadline(parent Context, d time.Time) (Context, CancelFunc)

  这两个函数都是返回新的上下文对象（contetx.timerCtx，同样实现了接口context.Context），函数返回两个值，第一个值就是contetx.timerCtx对象，第二个值是一个函数，同样可用于结束该上下文。也就是说，在timeout之后，会自动结束上下文，但是你也可以通过返回的CancelFunc结束该上下文。

实现原理

  在介绍context的使用方式时，我们提到了几个接口或者结构，如context.Context接口，context.valueCtx结构，contetx.cancelCtx结构，contetx.timerCtx结构，context的所有功能都是基于这几个类型实现的。

  context.Context是一个接口，定义了4个基本方法：

type Context interface {
    //返回上下文结束时间，ok=false说明没有定时结束功能
    Deadline() (deadline time.Time, ok bool)

    //  func Stream(ctx context.Context, out chan<- Value) error {
    //      for {
    //          v, err := DoSomething(ctx)
    //          if err != nil {
    //              return err
    //          }
    //          select {
    //          case <-ctx.Done():
    //              return ctx.Err()
    //          case out <- v:
    //          }
    //      }
    //  }
    //返回一个管道，上下文结束时会关闭该管道，从而可读
    Done() <-chan struct{}

    //如果Done管道关闭则Err不为空，用于表示错误原因；否则为nil
    Err() error

    //     func NewContext(ctx context.Context, u *User) context.Context {
    //         return context.WithValue(ctx, userKey, u)
    //     }
    //
    //     // FromContext returns the User value stored in ctx, if any.
    //     func FromContext(ctx context.Context) (*User, bool) {
    //         u, ok := ctx.Value(userKey).(*User)
    //         return u, ok
    //     }
    // 存取键值对数据
    Value(key interface{}) interface{}
}

  从这四个方法也能看出，context.Context天生就是为了上下文传值以及控制超时的。另外通过上面的几个事例，你有没有发现context.WithXXX等函数，都是基于父context衍生出新的context，也就是说这些context存在父子关系（类似一棵树）；父context超时结束后，也会遍历结束其所有的context（如果子context可结束），依次类推。

  下面我们先看看context.valueCtx的实现原理，其定义如下：

type valueCtx struct {
    //父对象
    Context
    //存储key-value键值对
    key, val any
}

func WithValue(parent Context, key, val any) Context {
    return &valueCtx{parent, key, val}
}

func (c *valueCtx) Value(key any) any {
    //直接返回
    if c.key == key {
        return c.val
    }
    //遍历父节点
    return value(c.Context, key)
}

func value(c Context, key any) any {
    //循环遍历父节点
    for {
        switch ctx := c.(type) {
        case *valueCtx:
            if key == ctx.key {
                return ctx.val
            }
            //接续遍历父节点
            c = ctx.Context
        case *emptyCtx:
            return nil

        //省略其他类型context分支

        default:
            return c.Value(key)
        }
    }
}

  context.valueCtx的实现还是特别简单的，只是要注意其存储的key-value键值对类型都是interface{}，所以在获取到值对象后可能需要进行类型转换；另外，如果当前context对象没有获取到键值对，则遍历父对象获取，依次类推。

  contetx.cancelCtx结构用于实现可取消的context对象，其定义如下：

type cancelCtx struct {
    //父对象
    Context
    //实际类型为管道，关闭管道就是结束当前上下文
    done     atomic.Value          // of chan struct{}, created lazily, closed by first cancel call
    //子对象，在父对象结束时结束所有子对象
    children map[canceler]struct{} // set to nil by the first cancel call
    //上下文结束时，赋值err
    err      error                 // set to non-nil by the first cancel call
}

  这才对嘛，基于这些字段才可以实现context.Context接口，不然怎么实现ctx.Done以及ctx.Err方法呢？contetx.cancelCtx的实现如下：

func WithCancel(parent Context) (ctx Context, cancel CancelFunc) {
    c := newCancelCtx(parent)
    propagateCancel(parent, &c)

    //返回一个函数CancelFunc，用于结束当前context对象
    return &c, func() { c.cancel(true, Canceled) }
}


func propagateCancel(parent Context, child canceler) {
    done := parent.Done()
    select {
    case <-done:
        // 如果父对象已结束，结束当前context对象并返回
        child.cancel(false, parent.Err())
        return
    default:
    }

    //判断父对象类型是否为contetx.cancelCtx
    if p, ok := parentCancelCtx(parent); ok {
        
        //关联父子context对象
        p.children[child] = struct{}{}

    } else {

        //子协程处理：如果父对象结束，则关闭当前context对象
        go func() {
            select {
            case <-parent.Done():
                child.cancel(false, parent.Err())
            case <-child.Done():
            }
        }()
    }
}

func (c *cancelCtx) cancel(removeFromParent bool, err error) {
    //关闭管道，用于通知context结束了
    d, _ := c.done.Load().(chan struct{})
    close(d)

    //结束所有子对象
    for child := range c.children {
        child.cancel(false, err)
    }
}

  基于context.WithCancel衍生新的可取消context对象时，注意要维护父子context对象之间的关系，而且在父对象结束后，也需要结束其所有子对象，依次类推。另外，context.WithCancel返回的CancelFunc函数，主逻辑其实也就是关闭管道，以此传递context对象关闭信号。

  contetx.timerCtx与contetx.cancelCtx还是比较类似的，只不过内置了定时器而已，定时器触发时自动关闭context对象，并没有其他区别，这里就不再赘述。

  再次强调，基于context实现任务的超时控制，在Go语言中非常常见，可以说到处都能看到，所以一定要熟练使用并了解其原理。

全链路追踪

  全链路追踪是什么意思呢？设想你维护着一个Go HTTP服务，请求处理过程肯定会记录一些日志，不然遇到问题怎么排查定位？可是，一个HTTP请求过程肯定会记录多条日志，如何将这些日志汇总起来呢？要知道日志文件里这些日志记录可是分散开的。可以基于traceId实现，即所有日志记录都包含traceId字段，并且同一个请求打印的日志traceId相同，这就是全链路追踪了；更进一步，如果还依赖了其他第三方服务，在向他们发起HTTP请求时，也可以携带traceId，第三方服务打印日志也携带该traceId，这样我们甚至能汇总一个用户请求涉及的多个服务间的日志（这只能汇总，还不能直观分析其调用关系，一般还会有其他字段描述调用关系）。

  为什么全链路追踪要放到context这一篇文章介绍呢？因为日志要记录traceId，可是traceId从哪来？上下文呗，也就是context了。所以你的Go服务，所有函数的第一个参数最好都是context.Context。大概应该是这个样子：

func Handler(ctx context.Context, req Req) (resp Resp, err error) {
    //出错，记录error日志
    if err != nil {
        logger.Errorf(ctx, "xxxx error:%v", err)
        return
    }
}

func Errorf(ctx context.Context, tag, template string, args ...interface{}) {
    logger.With("traceId", ctx.Value("traceId")).Errorf(template, args...)
}

  那如果已有的Go项目，参数确实没有context.Context呢？改造所有函数吗？这成本挺高的，当然也可以尝试改造。退而求其次，其实还有一个下下策，日志中记录协程ID，这样同一个协程记录的日志可以根据该ID汇总（子协程执行的任务就无法汇总了）。

  不过，Go语言貌似没有提供方式获取协程ID，确实没有（不建议）。这里提一个开源组件（ https://github.com/petermattis/goid ），封装了协程ID的获取方式，当然底层也是基于线程本地存储获取的，还记得不，Go语言当前调度的协程g会保存在线程本地存储，有了g对象，是不是就能获取到协程ID。这里稍微了解一下就可以。

func getg() *g

func Get() int64 {
    return getg().goid
}

TEXT ·Get(SB),NOSPLIT,$0-8
    //协程本地存储TLS
    MOVQ (TLS), R14
    MOVQ g_goid(R14), R13
    MOVQ R13, ret+0(FP)
    RET

总结

  基于context实现任务的超时控制，在Go语言中非常常见，可以说到处都能看到，所以本篇文章介绍了context是如何在上下文传递key-value键值对，以及控制超时的。最后还简单描述了全链路追踪的含义，在日常Go项目开发中，一定要记得，所有函数的第一个参数最好都是context.Context。