golang - context包使用示例 以及 底层实现

1. context 常用方法，以及各种适用于什么场景

1.1 context含有的方法

	var ctx context.Context
	var cancel context.CancelFunc
	// 1，传递key，value的值
	ctx = context.WithValue(context.Background(), "key", "value")
	// 2，超时控制timeout
	ctx, cancel = context.WithTimeout(context.Background(), time.Second*10)
	// 3，超时控制，deadline
	ctx, cancel = context.WithDeadline(context.Background(), time.Now().Add(time.Second*10))
	// 4，取消控制
	ctx, cancel = context.WithCancel(context.Background())
	// 5，养成好习惯，defer触发cancel函数；ps: cancel函数幂等
	defer cancel()

1.2 方法适用场景和伪代码示例

1.2.1 值传递：比如gin框架中用来传递key,value的值，自己简单示例如下

func readContext(ctx context.Context) {
	traceId, ok := ctx.Value("key").(string)
	if ok {
		fmt.Printf("readContext key=%s\n", traceId)
	} else {
		fmt.Printf("readContext no key\n")
	}
}

func main() {
	ctx := context.Background()
	readContext(ctx) //output: readContext no key
	ctx = context.WithValue(ctx, "key", "beautiful")
	readContext(ctx) //output: readContext key=beautiful

}

func TestWithValueContext(t *testing.T) {
	main()
}

1.2.2 超时控制-timeout: http请求设置超时时间

func httpRequest(ctx context.Context) {
	for {
		// process http requests
		select {
		case <-ctx.Done():
			fmt.Println("http requests cancel")
			return
		case <-time.After(time.Second * 1):
		}
	}

}

func TestTimeoutContext(t *testing.T) {
	fmt.Println("start TestTimeoutContext")
	ctx, cancel := context.WithTimeout(context.Background(), time.Second*2)
	defer cancel()
	httpRequest(ctx)
	time.Sleep(time.Second * 5)
}

1.2.3, 超时控制-deadline: 比如文件io或者网络io等耗时操作，可以查看剩余的时间是否充足，决定是否进行下一步操作

func copyFile(ctx context.Context)  {
	deadline, ok := ctx.Deadline()
	if ok == false {
		return
	}
	isEnough := deadline.Sub(time.Now()) > time.Second*5
	if isEnough {
		// write content to file
		fmt.Println("copy file")
	} else {
		fmt.Println("isEnough is false return")
		return
	}
}

func TestDeadlineContext(t *testing.T) {
	ctx, cancel := context.WithDeadline(context.Background(), time.Now().Add(time.Second*4))
	defer cancel()
	copyFile(ctx)
	time.Sleep(time.Second * 5)
}

1.2.4. 取消控制: goroutine发送取消信号，保证自己这个逻辑中发散出去的goroutine全部成功取消

func gen(ctx context.Context) <-chan int {
	ch := make(chan int)
	go func() {
		var n int
		for {
			select {
			case ch <- n:
				n++
				time.Sleep(time.Second)
			case <-ctx.Done():
				return
			}
		}
	}()
	return ch
}

func TestCancelContext(t *testing.T) {
	// 创建一个Cancel context
	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()
	for n := range gen(ctx) {
		fmt.Println(n)
		if n == 5 {
			// 达到要求后触发 cancel
			cancel()
			break
		}
	}

}

2. context包的底层实现是什么样的？

2.1 key，value 传值底层实现

• 函数底层实现代码（golang v1.16），其核心就是当本context无法获取到key的值的时候，递归父context获取

type valueCtx struct {
	Context // 保存着本节点的父节点
	key, val interface{} // 保存着本节点的key，value值
}

// 新建一个value 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}
}

func (c *valueCtx) Value(key interface{}) interface{} {
	// 如果key，value在本节点就有，那么直接取值
    if c.key == key {
		return c.val
	}
    // 否则递归向父节点获取
	return c.Context.Value(key)
}

2.2 cancel实现

2.2.1 cancelCtx 结构体

// A cancelCtx can be canceled. When canceled, it also cancels any children
// that implement canceler.
// 触发取消以后会递归取消所有子节点
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
}

2.2.2 cancelCtx实现了cancel函数，逻辑如下

• 1， 锁保证并发冲突，避免并发冲突
• 2，关闭c.done这个channel，通过这个传递信号（往后细化分析）
• 3，遍历关闭所有子节点，保证不会内存泄漏
• 4，释放自己的所有子节点后，将自己的子节点map赋值为nil
• 5，将自己从自己的父节点中进行移除，这个只有在调用WithCancel（）方法的时候会触发，也就是说传入参数removeFromParent为true（往后细化分析）

func (c *cancelCtx) cancel(removeFromParent bool, err error) {
	if err == nil {
		panic("context: internal error: missing cancel error")
	}
	c.mu.Lock() // 1，锁保证并发冲突
	if c.err != nil {
		c.mu.Unlock()
		return // already canceled
	}
	c.err = err
	if c.done == nil {
		c.done = closedchan
	} else {
		close(c.done) // 2，关闭done这个channel
	}
    // 3，遍历所有子节点，然后调用cancel函数
	for child := range c.children {
		// NOTE: acquiring the child's lock while holding parent's lock.
		child.cancel(false, err)
	}
    // 4，释放自己的所有子节点后，将自己的子节点map赋值为nil
	c.children = nil
	c.mu.Unlock()
	
    
    // 5，将自己从自己的父节点中进行移除，这个只有在调用withCancel的时候会传入true
	if removeFromParent {
		removeChild(c.Context, c)
	}
}

2.2.3 细化：c.done的信号传递

• 这个是基于所有channel的特性，当监听一个channel，channel为空的时候会阻塞，但是如果channel被关闭，那么将不会阻塞，而会读取到一个空值
• 基于上述特性，实现了关闭这个channel，而其他所有监听此channel的goroutine都收到此信号
• 代码举例

func Done(ch chan struct{}, count int) {
	for {
		select {
		case <-ch:
			fmt.Println(count)
			return
		}
	}
}

func TestCloseChannel(t *testing.T) {
	signalChannel := make(chan struct{})
	go Done(signalChannel, 1)
	go Done(signalChannel, 2)
	go Done(signalChannel, 3)
	time.Sleep(3)
	fmt.Println("close signalChannel")
	close(signalChannel)

	// 阻塞当前goroutine
	select {

	}
}

2.2.3 细化：removeFromParent参数-是否从父节点delete自己

• 先看下removeChild的实现代码

// removeChild removes a context from its parent.
func removeChild(parent Context, child canceler) {
	p, ok := parentCancelCtx(parent)
	if !ok {
		return
	}
	p.mu.Lock()
	if p.children != nil {
		delete(p.children, child)
	}
	p.mu.Unlock()
}

• 为什么调用WithCancel（）的时候，也就是新建一个节点的时候会传入removeFromParent=true然后调用removeChild（）呢？
  • 因为你调用cancel作用的更多的处理的挂靠在你这个context上的子节点，而只有最后一步才是真正的释放自己
  • 举例：
    • 1，第一步：假如你创建的一个cancelContext，挂靠在在根节点上（contextBackgroud）上，那你下面的子节点都会因为你的 c.children = nil 而释放。
    • 2，第二步：然后逻辑上你自己都调用了cancel，那么你自己也要释放了，所以就将自己从从父节点中delete的
• 为什么其他删除子节点的时候不会调用？
  • 1，因为其中有一个操作是 delete(p.children, child) ，这个操作会删除父节点的子节点的map中的自己，而一边遍历和一边删除map是会出问题的
  • 2，同时由于cancel（）函数中有操作为 c.children = nil ，所以也无需说去做这种操作