golang context 适用场景与源码导读

背景

用于存储数据

• 在业务执行初期，把关键信息写入 ctx, 比如一些主键 id，事件类型等
• 存储调用链数据
• 存储一些简单的 kv
• …

下面看一个简单的例子

package main

import (
  "context"
  "fmt"
)

func main() {
  ctx := context.TODO()
  ctx = context.WithValue(ctx, "a", 1)
  fmt.Println(ctx.Value("a").(int))
}

特别地，对于 context，如果需要切换 goroutine，

用于做控制信号

对于比较长的业务逻辑
以cancelContext 为例
先看使用

package main

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

func main() {
	back := context.Background() // 占位

	child1, cancel1 := context.WithCancel(back)

	child1_1, _ := context.WithCancel(child1)

	child1_2, cancel1_2 := context.WithCancel(child1)

	go func() {
		for range back.Done() {
			fmt.Println("background done")
			break
		}
	}()

	go func() {
		for {
			select {
			case <-back.Done():
				fmt.Println(time.Now().Format(time.RFC3339), "back done")
				return
			}
		}
	}()

	go func() {
		for {
			select {
			case <-child1.Done():
				fmt.Println(time.Now().Format(time.RFC3339), "child1 done")
				return
			}
		}
	}()

	go func() {
		for {
			select {
			case <-child1_1.Done():
				fmt.Println(time.Now().Format(time.RFC3339), "child1_1 done")
				return
			}
		}
	}()

	go func() {
		for {
			select {
			case <-child1_2.Done():
				fmt.Println(time.Now().Format(time.RFC3339), "child1_2 done")
				return
			}
		}
	}()

	go func() {
		time.Sleep(1 * time.Second)
		// 关 cancel1
		cancel1_2()
		time.Sleep(3 * time.Second)
		cancel1()
	}()

	time.Sleep(10 * time.Second)
}

上述代码的功能比较简单，child1是作为parent， child1_1 child1_2是‘children’。程序先sleep 1秒， 调用child1_2关联的cancel方法，再sleep 3s， 调用 调用child1关联的cancel方法

输出

2020-05-29T00:20:31+08:00 child1_2 done
2020-05-29T00:20:34+08:00 child1_1 done
2020-05-29T00:20:34+08:00 child1 done

可以看到，child1_2 先收到了Done信号，再隔3s， child1_1 child1 也分别收到了Done信号
也就是说，父亲可以触发自身与孩子的Done信号，但孩子只会触发自己的Done信号（没有下一级孩子的前提）

源码

go/src/context/context.go

顶层interface

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

Context的设计，主要是解决上下文的数据以及通信问题。尤其是复杂拓扑数据结构的场景，一个‘全局’上下文可以大大简化API的设计。在http包中有大量的应用。

先看

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

Background返回的是一个空的context，为什么需要一个设计一个emptyContext？
主要是因为context的风格是链式调用，emptyContext可以作为初始化的传参使用。

WithCancel 相关

type cancelCtx struct {
	Context

	mu       sync.Mutex            // protects following fields
	done     chan struct{}         // created lazily, closed by first cancel call
	children map[canceler]struct{} // set to nil by the first cancel call
	err      error                 // set to non-nil by the first cancel call
}

可以看到，cancelCtx 有关联的children cancelCtx ，也就是一对多的关系。实际场景可以是很复杂的树结构，这类场景适用context的好处就体现出来了，可以针对自身与子节点做Done信号触发。

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

// newCancelCtx returns an initialized cancelCtx.
func newCancelCtx(parent Context) cancelCtx {
	return cancelCtx{Context: parent}
}

// propagateCancel arranges for child to be canceled when parent is.
func propagateCancel(parent Context, child canceler) {
	if parent.Done() == nil {
		return // parent is never canceled
	}
	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{})
			}
			p.children[child] = struct{}{}
		}
		p.mu.Unlock()
	} else {
		go func() {
			select {
			case <-parent.Done():
				child.cancel(false, parent.Err())
			case <-child.Done():
			}
		}()
	}
}

// parentCancelCtx follows a chain of parent references until it finds a
// *cancelCtx. This function understands how each of the concrete types in this
// package represents its parent.
func parentCancelCtx(parent Context) (*cancelCtx, bool) {
	for {
		switch c := parent.(type) {
		case *cancelCtx:
			return c, true
		case *timerCtx:
			return &c.cancelCtx, true
		case *valueCtx:
			parent = c.Context
		default:
			return nil, false
		}
	}
}

parentCancelCtx 对于context包的cancelCtx timerCtx， 不需要新开goroutine， 对于其它类型的ctx，则需要开一个goroutine来接受Done信号，然后下发children

WithDeadline 相关

timerCtx 
// A timerCtx carries a timer and a deadline. It embeds a cancelCtx to
// implement Done and Err. It implements cancel by stopping its timer then
// delegating to cancelCtx.cancel.
type timerCtx struct {
	cancelCtx
	timer *time.Timer // Under cancelCtx.mu.

	deadline time.Time
}

这个比cancelCtx主要了一个timer

func WithDeadline(parent Context, d time.Time) (Context, CancelFunc) {
	if cur, ok := parent.Deadline(); ok && cur.Before(d) {
		// The current deadline is already sooner than the new one.
		return WithCancel(parent)
	}
	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(true, 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) }
}

AfterFunc 这里会新起一个goroutine，等到超时就会自动调用cancel方法