context包使用及原码分析

context包使用及原码分析

---

前言

context 包在 Go 1.7 中引入，它为我们提供了一种在应用程序中处理 context 的方法。

一、context使用场景

1.值传递

代码如下（示例）：

package main

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

func g1(ctx context.Context) {
	// 从context中取begin的值，如果key不存在，则返回nil，可以当做是个key-value的map
	fmt.Println(ctx.Value("from"))
	fmt.Println("自东土大唐而来")
	// 将context值继续包装，增加新值，如果不增加新值，则直接传递ctx
	go g2(context.WithValue(ctx, "to", "你要到哪里去？"))
}

func g2(ctx context.Context) {
	// 支持取上上次包装的context值，具体原因请看下面的源码分析
	// fmt.Println(ctx.Value("from"))
	fmt.Println(ctx.Value("to"))
	fmt.Println("去女儿国")
}

func main() {
    // context.Background()为跟context，本质是空context
	ctx := context.WithValue(context.Background(), "from", "你从哪里来？")
	go g1(ctx)
	time.Sleep(1 * time.Second)
}

2.取消机制

代码如下（示例）：

package main

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

func g1(ctx context.Context) {
	for {
		time.Sleep(1 * time.Second)
		// 监听取消信号
		select {
		// 如果收到取消信号则直接退出
		case <-ctx.Done():
			fmt.Println("done")
			return
		// 如果没有default，会一直阻塞等待退出信号，不方便直观看到该goroutine运行状态
		default:
			fmt.Println("working")
		}
	}
}

func main() {
	ctx, cancel := context.WithCancel(context.Background())
	go g1(ctx)
	time.Sleep(5 * time.Second)
	// 主动执行取消操作
	cancel()
	fmt.Println("canceled")
}

3.timeout机制

代码如下（示例）：

package main

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

func g(ctx context.Context) {
	for {
		time.Sleep(1 * time.Second)
		if deadline, ok := ctx.Deadline(); ok {
			fmt.Println("deadline set")
			if time.Now().After(deadline) {
				fmt.Println(ctx.Err().Error())
				return
			}
		}

		select {
		case <-ctx.Done():
			fmt.Println("done")
			return
		default:
			fmt.Println("work")
		}
	}
}

func main() {
	fmt.Println("hello world!")
	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)

	go g(ctx)

	time.Sleep(10 * time.Second)
	cancel()
	fmt.Println("timeout")
}

4.deadline机制

代码如下（示例）：

package main

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

func g(ctx context.Context) {
	for {
		time.Sleep(1 * time.Second)
		if deadline, ok := ctx.Deadline(); ok {
			fmt.Println("deadline set")
			if time.Now().After(deadline) {
				fmt.Println(ctx.Err().Error())
				return
			}
		}

		select {
		case <-ctx.Done():
			fmt.Println("done")
			return
		default:
			fmt.Println("work")
		}
	}
}

func main() {
	ctx, cancel := context.WithDeadline(context.Background(), time.Now().Add(5*time.Second))

	go g(ctx)

	time.Sleep(10 * time.Second)
	cancel()
	fmt.Println("deadline")
}

---

二、context源码分析

1.抽象接口

context的核心作用是存储键值对和取消机制。存储键值对比较简单，取消机制比较复杂，先来看一下Context抽象出来的接口：

type Context interface {
  // 如果是timerCtx或者自定义的ctx实现了此方法，返回截止时间和true，否则返回false
  Deadline() (deadline time.Time, ok bool)
  // 这里监听取消信号
  Done() <-chan struct{}
  // ctx取消时，返回对应错误，有context canceled和context deadline exceeded
  Err() error 
  // 返回key的val
  Value(key interface{}) interface{}
}

2.根context

根context有两种：background、todo，根据源码注释得知，这两个值都是空context（nil）。

1.background：永远不会被取消，没有值，并且没有截止日期。 它通常被主函数，初始化和测试，并作为传入的顶级上下文请求。

2.todo：作为待办事项或者不清楚该使用哪个上下文时使用。

源码如下（context.context.go）：

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
}

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

func Background() Context {
	return background
}

func TODO() Context {
	return todo
}

3.值传递（valueCtx）

valueCtx携带一个键值对，支持context嵌套（本质是链表）。

源码如下（context.context.go）：

type valueCtx struct {
    // 支持context类型嵌套
	Context
	// 解释了只能存储一对键值对
	key, val interface{}
}

func (c *valueCtx) Value(key interface{}) interface{} {
	// 如果key在当前context里存在，则返回value值
	if c.key == key {
		return c.val
	}
	// 否则从父节点中调用Value方法继续寻找key，此处会发生递归，直到找到该key对应的值。
	// 如果遍历到根节点也没找到该key，则返回根节点的value值nil。
	return c.Context.Value(key)
}

4.取消传递（cancelCtx）

go源码是怎么实现的取消，首先抽象出了一个canceler接口，这个接口里最重要的就是cancel方法，调用这个方法可以发送取消信号，有两个结构体实现了这个接口，分别是 cancelCtx(普通取消) 和 timerCtx(时间取消)。
源码如下（context.context.go）：

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

cancelCtx对应前文说的普通取消机制，它是context取消机制的基石，也是源码中比较难理解的地方，先来看一下它的结构体：

type cancelCtx struct {
	 Context
	 mu       sync.Mutex            // context号称并发安全的基石
	 done     atomic.Value          // 用于接收ctx的取消信号，这个数据的类型做过优化，之前是 chan struct 类型
	 children map[canceler]struct{} // 储存此节点的实现取消接口的子节点，在根节点取消时，遍历它给子节点发送取消信息
	 err      error                 // 调用取消函数时会赋值这个变量
}

若我们要生成一个可取消的ctx，需要调用WithCancel函数，这个函数的内部逻辑是：

func WithCancel(parent Context) (ctx Context, cancel CancelFunc) {
	 if parent == nil {
		  panic("cannot create context from nil parent")
	 }
	 c := newCancelCtx(parent)    // 基于父节点，new一个CancelCtx对象
	 propagateCancel(parent, &c)  // 挂载c到parent上
	 return &c, func() { c.cancel(true, Canceled) } // 返回子ctx，和返回函数
}

这里逻辑比较重的地方是propagateCancel函数和cancel方法，propagateCancel函数主要工作是把子节点挂载父节点上，下面来看看它的源码：

func propagateCancel(parent Context, child canceler) {
	done := parent.Done()
	// 判断父节点的done是否为nil,若为nil则为不可取消的ctx, 直接返回
	if done == nil {
		return
	}
	// 看能否从done里面读到数据，若能说明父节点已取消，取消子节点，返回即可，不能的话继续流转到后续逻辑
	select {
	case <-done:
		child.cancel(false, parent.Err())
		return
	default:
	}

	// 调用parentCancelCtx函数，看是否能找到ctx上层最接近的可取消的父节点
	if p, ok := parentCancelCtx(parent); ok {
		//这里是可以找到的情况
		p.mu.Lock()
		// 父节点有err，说明已经取消，直接取消子节点
		if p.err != nil {
			child.cancel(false, p.err)
		} else {
			// 把本节点挂载到父节点的children map中
			if p.children == nil {
				p.children = make(map[canceler]struct{})
			}
			p.children[child] = struct{}{}
		}
		p.mu.Unlock()
	// else判断主要是预防用户自己实现了一个定制的Ctx中，随意提供了一个Done chan的情况的，由于找不到可取消的父节点的，只好新起一个协程做监听
	} else {
		// 若没有可取消的父节点挂载
		atomic.AddInt32(&goroutines, +1)
		// 新起一个协程
		go func() {
			select {
			// 监听到父节点取消时，取消子节点
			case <-parent.Done():
				child.cancel(false, parent.Err())
				// 监听到子节点取消时，什么都不做，退出协程
			case <-child.Done():
			}
		}()
	}
}

cancel函数源码：

func (c *cancelCtx) cancel(removeFromParent bool, err error) {
	// 取消时必须传入err，不然panic
	if err == nil {
		panic("context: internal error: missing cancel error")
	}
	c.mu.Lock()
	// 已经出错了，说明已取消，直接返回
	if c.err != nil {
		c.mu.Unlock()
		return
	}
	// 用户传进来的err赋给c.err
	c.err = err
	d, _ := c.done.Load().(chan struct{})
	if d == nil {
		// 这里其实和关闭chan差不多，因为后续会用closedchan作判断
		c.done.Store(closedchan)
	} else {
		// 关闭chan
		close(d)
	}
	// 这里是向下取消，依次取消此节点所有的子节点
	for child := range c.children {
		child.cancel(false, err)
	}
	// 清空子节点
	c.children = nil
	c.mu.Unlock()
	// 这里是向上取消，取消此节点和父节点的联系
	if removeFromParent {
		removeChild(c.Context, c)
	}
}

5.时间取消（cancelCtx）

时间取消ctx可传入两种时间，第一种是传入超时时间戳；第二种是传入ctx持续时间，比如2s之后ctx取消，持续时间在实现上是在time.Now的基础上加了个timeout凑个超时时间戳，本质上都是调用的WithDeadline函数。

WithDeadline 函数内部new了一个timerCtx，先来看一下这个结构体的内容：

type timerCtx struct {
	 cancelCtx
	 timer *time.Timer  // 一个统一的计时器，后续通过 time.AfterFunc 使用
	 deadline time.Time // 过期时间戳
}

可以看到 timerCtx 内嵌了cancelCtx，实际上在超时取消这件事上，timerCtx更多负责的是超时相关的逻辑，而取消主要调用的cancelCtx的cancel方法。先来看一下WithDeadline函数的逻辑，看如何返回一个时间Ctx：

func WithDeadline(parent Context, d time.Time) (Context, CancelFunc) {
	// 若父节点为nil，panic
	if parent == nil {
		panic("cannot create context from nil parent")
	}
	// 如果parent有超时时间、且过期时间早于参数d，那parent取消时，child 一定需要取消，直接通过WithCancel走起
	if cur, ok := parent.Deadline(); ok && cur.Before(d) {
		// The current deadline is already sooner than the new one.
		return WithCancel(parent)
	}
	// 构造一个timerCtx, 主要传入一个过期时间
	c := &timerCtx{
		cancelCtx: newCancelCtx(parent),
		deadline:  d,
	}
	// 把这个节点挂载到父节点上
	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)
		})
	}
	// 返回一个ctx和一个取消函数
	return c, func() { c.cancel(true, Canceled) }
}

看完源码可以知道，除了基于时间的取消，当调用CancelFunc时，也能取消超时ctx。再来看一下*timerCtx的cancel方法的源码：

func (c *timerCtx) cancel(removeFromParent bool, err error) {
	// 调用cancel的cancel取消掉它下游的ctx
	c.cancelCtx.cancel(false, err)
	// 取消掉它上游的ctx的连接
	if removeFromParent {
		removeChild(c.cancelCtx.Context, c)
	}
	// 把timer停掉
	c.mu.Lock()
	if c.timer != nil {
		c.timer.Stop()
		c.timer = nil
	}
	c.mu.Unlock()
}