实用的context

context是什么

context翻译过来是上下文的意思，在go里面是用来给goroutine传递上下文的，可以传递取消信号、超时、k-v等等，比如说我们经常用它来存储客户端传过来的原子参数，token、COOKIE等等

context的使用

上文说了，context可以传递cv,设置goroutine取消、超时等，我们用一些简单代码来看看context是怎么实现这些功能的

传递k-v

func TestCommon(t *testing.T) {
ctx := setTokenContext(context.Background())
wg := new(sync.WaitGroup)
wg.Add(2)
go run(ctx, wg, 1)
go run(ctx, wg, 2)
wg.Wait()
}
func getToken(ctx context.Context) string {
return ctx.Value("token").(string) // 从context中获取token，因为value方法取出来是interface{},所以需要类型断言成string
}
func setTokenContext(ctx context.Context) context.Context{
return context.WithValue(ctx, "token", "abc-abc-abc") // context中设置token
}
func run(ctx context.Context, wg *sync.WaitGroup, num int) {
//todo 业务
defer wg.Done()
token := getToken(ctx)
fmt.Printf("这里是第%d个goroutian, token为：%s \n", num, token)
}

上面这段代码我们现创建了一个context，然后通过WithValue函数往context中存储了token数据，然后开启两个groutine去获取token,最终的打印结果为下图

关闭goroutine

通常我们可以通过context发送一个信号来关闭我们需要关闭的goroutine

func TestCommon(t *testing.T) {
ctx, cancel := context.WithCancel(context.Background())
go run(ctx, 1) //goroutian1
go run(ctx, 2) //goroutian2
time.Sleep(time.Second)
cancel() // 发送取消信号
time.Sleep(time.Second)
}
func run(ctx context.Context, num int) {
for {
select {
case <- ctx.Done():
fmt.Printf("groutine%d被关闭了\n", num)
return
default:
//执行正常逻辑
fmt.Printf("groutine%d正在执行\n", num)
}
time.Sleep(time.Second) // 每秒钟执行一次
}
}

先调用了一下WithCancel，可以看到函数返回了一个context，还有一个cancel函数，这个cancel函数其实就是发送一个信号给ctx.Done()返回的channel，所以我们通过监听这个channel就可以友好的关闭goroutine，具体的实现逻辑等下文源码解析时会说

打印结果

控制超时

可以使用context来控制groutian的执行时间，超过某个时间就关闭掉，例如http超时，在golang里就是通过context实现的,先看示例代码：

func TestCommon(t *testing.T) {
//ctx, cancel := context.WithDeadline(context.Background(), time.Now().Add(10 * time.Second))
ctx, _ := context.WithTimeout(context.Background(), time.Second * 3)
go run(ctx) //goroutian1
time.Sleep(5 * time.Second)
}
func run(ctx context.Context) {
for {
select {
case <- ctx.Done():
fmt.Printf("groutine超时被关闭了\n")
return
default:
//执行正常逻辑
fmt.Printf("groutine正在执行\n")
}
time.Sleep(time.Second) // 每秒钟执行一次
}
}

先调用了WithTimeout函数，返回两个数据，第一个是context，第二个和上面取消的逻辑一样，也会返回一个cancel函数，所以，也是可以手动cancel的，这里主要说超时，所以先把它过滤掉。WithTimeout函数第二个参数是个时间，time.xx，这里可以传秒、分钟、小时等等等，传多少时间进去，代表我们希望groutine多长时间超时，超时以后和上文提到的手动取消一样，也是发送一个信号给ctx.Done()返回的channel

注意：代码中注释了一行context.WithDeadline其实它和WithTimeout的执行是完全一样的，就是参数不一样而已，这个一会解析源码的时候也会说
打印结果：

原理解析

先看下context源码包，位于$GOPATH/src/context文件夹下，context.go，很简单，总共500多行代码，注释就有200多行

接口

context

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

Context 是一个接口，定义了 4 个方法

Done() 返回一个 channel，可以表示 context 被取消的信号：当这个 channel 被关闭时，说明 context 被取消了。注意，这是一个只读的channel。 我们知道，读一个已关闭的 channel 会读出相应类型的零值。并且源码里没有地方会向这个 channel 里面放入值。换句话说，这是一个 receive-only 的 channel。因此在子协程里读这个 channel，除非被关闭，否则读不出来任何东西。也正是利用了这一点，子协程从 channel 里读出了值（零值）后，就可以做一些收尾工作，尽快退出。

Err() 返回一个错误，表示 channel 被关闭的原因。例如是被取消，还是超时。

Deadline() 返回 context 的截止时间，通过此时间，函数就可以决定是否进行接下来的操作，如果时间太短，就可以不往下做了，否则浪费系统资源。当然，也可以用这个 deadline 来设置一个 I/O 操作的超时时间。

Value() 获取之前设置的 key 对应的 value。

canceler

再来看另一个接口

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

实现了上面定义的两个方法的 Context，就表明该 Context 是可取消的。源码中有两个类型实现了 canceler 接口：cancelCtx 和 timerCtx。

结构体

emptyCtx

源码中给出了一个context接口的空实现，一般用来当源数据

type emptyCtx int
func (*emptyCtx) Deadline() (deadline time.Time, ok bool) {
return
}
func (*emptyCtx) Done() <-chan struct{} {
return nil
}
func (*emptyCtx) Err() error {
return nil
}
func (*emptyCtx) Value(key interface{}) interface{} {
return nil
}

这段代码太简单了，一个emptyCtx实现了Context接口，并且每个函数都返回都是空的这是一个空的Context,不会被cancel,不会用来储存值，源码把它包装成了两个对象

var (
background = new(emptyCtx)
todo = new(emptyCtx)
)
func Background() Context {
return background
}
func TODO() Context {
return todo
}

Background 通常用在 main 函数中，作为所有 context 的根节点。TODO一般是在那种我们不知道要传什么ctx时，作为一个占位符来使用

valueCtx

type valueCtx struct {
Context
key, val interface{}
}
func (c *valueCtx) Value(key interface{}) interface{} {
if c.key == key {
return c.val
}
return c.Context.Value(key)
}
//创建一个带key-value的context
func WithValue(parent Context, key, val interface{}) Context {
if key == nil {
panic("nil key")
}
if !reflectlite.TypeOf(key).Comparable() {
panic("key is not comparable")
}
return &valueCtx{parent, key, val}
}

valueCtx的结构很简单，匿名字段Context，还有key，value为两个interface用来存储值
先看一下获取valueCtx的方法WithValue，函数有三个参数，刚好和结构体对应，比如我们上面示例代码那块写的ctx := context.WithValue(context.Background(), "token", "abc"),这个background就是parent,然后赋值给结构体的第一个参数，最后返回。我们发现ctx是没有其他结构类似SetValue这种可以设置key和值的方法的，那假如说我想再往里面增加一条数据怎么办呢，我们只能像这样：

ctx := context.WithValue(context.Background(), "a", 1)
ctx = context.WithValue(ctx, "b", 2)

这时候两个ctx变量已经不是同一个context了，第一个ctx变成了第二个ctx的parent,那这个时候第二个ctx是个什么结构，根据上面的代码，生成的ctx结构应该是这样的

valueCtx {
parent: {
parent:emptyCtx //background
key: "a",
value: 1,
}
key:"b",
value:2
}

很像一个倒着的链表，用parent来关联父级，现在再来看看Value函数，代码很简单，判断key是否和传过来的key相等，如果相等直接返回，否则调用父级的value方法，这里相当于递归寻找，直到最后找到background根节点。

cancelCtx

我们再开看可取消的context实现

type cancelCtx struct {
Context
mu sync.Mutex // 用来加锁
done chan struct{} // 接收关闭信号的channel
children map[canceler]struct{} // 所有的子节点存为一个map
err error // 保存被取消以后的取消原因，cancel或者超时
}


它实现了canceler接口，是一个可取消的context，由于里面也含有一个Context匿名字段，所以可以被看作是Context类型，先来看看怎么创建它

func WithCancel(parent Context) (ctx Context, cancel CancelFunc) {
c := newCancelCtx(parent)
propagateCancel(parent, &c)
return &c, func() { c.cancel(true, Canceled) }
}
func newCancelCtx(parent Context) cancelCtx {
return cancelCtx{Context: parent}
}


代码很简单，初始化一个cancelCtx,和valueCtx一样，把parent赋值给了匿名属性Context，然后返回一个Context和cancel函数，主要看一下propagateCancel函数做了些什么

func propagateCancel(parent Context, child canceler) {
// 获取父级的done,判断父级是否含有cancelCtx
// 如果父级无cancelCtx，直接返回
// 否则使用map增加关联关系，方便将来cancel父级的时候把子级一起cancel
done := parent.Done()
if done == nil {
return // parent is never canceled
}
// 如果父级已经取消了，取消当前ctx
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 {
// 核心！！！！！父级的children中保存当前ctx
if p.children == nil {
p.children = make(map[canceler]struct{})
}
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():
}
}()
}
}


其实这段代码的核心就两点，判断是否有上级，并且是cancelCtx类型，如果没有，什么也不用做，如果有的话就用上级的children保存当前ctx，方便将来级联取消。

接下来我们看看cancel函数

// cancel closes c.done, cancels each of c's children, and, if
// removeFromParent is true, removes c from its parent's children.
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
// 获取ctx的done（是个channel），然后关闭这个channel
d, _ := c.done.Load().(chan struct{})
if d == nil {
c.done.Store(closedchan)
} else {
close(d)
}
// 获取下级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 {
removeChild(c.Context, c)
}
}

这段代码的核心其实就是close这个ctx中的done，还记得上文中提到我们怎么监听协程是否需要关闭：

case <- ctx.Done():

通过监听ctx.Done()，如果拿到值就代表被cancel了，当调用close时,这边的监听就可以拿到数据了。

最后，让我们看下timeCtx的实现

// WithDeadline returns a copy of the parent context with the deadline adjusted
// to be no later than d. If the parent's deadline is already earlier than d,
// WithDeadline(parent, d) is semantically equivalent to parent. The returned
// context's Done channel is closed when the deadline expires, when the returned
// cancel function is called, or when the parent context's Done channel is
// closed, whichever happens first.
//
// Canceling this context releases resources associated with it, so code should
// call cancel as soon as the operations running in this Context complete.
func WithDeadline(parent Context, d time.Time) (Context, CancelFunc) {
if parent == nil {
panic("cannot create context from nil parent")
}
// 如果父级也是时间类型的ctx,判断超时时间是否大于父级的超时时间，如果大于其实是不需要生成新的ctx，上文有讲，因为父级超时时间到了会直接调用cancel,把子级也一起取消。
if cur, ok := parent.Deadline(); ok && cur.Before(d) {
// The current deadline is already sooner than the new one.
return WithCancel(parent)
}
// 生成时间类型ctx,并设置cancelCtx属性
c := &timerCtx{
cancelCtx: newCancelCtx(parent),
deadline: d,
}
// 上文有专门针对这个函数的讲解
propagateCancel(parent, c)
dur := time.Until(d)
// 如果超时时间小于0，直接取消
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 {
// 重点！！！！！！！开启定时器，并且到时调用cancel函数
c.timer = time.AfterFunc(dur, func() {
c.cancel(true, DeadlineExceeded)
})
}
return c, func() { c.cancel(true, Canceled) }
}

代码很简单，如果理解了上面cancel类型的context，那这个无非就是加了一个定时器，到时自动触发cancel函数而已。