go源码之context

定义说明

Package context defines the Context type, which carries deadlines, cancelation signals, and other request-scoped values across API boundaries and between processes.
Incoming requests to a server should create a Context, and outgoing calls to servers should accept a Context. The chain of function calls between them must propagate the Context, optionally replacing it with a derived Context created using WithCancel, WithDeadline, WithTimeout, or WithValue. When a Context is canceled, all Contexts derived from it are also canceled.

粗略翻译一下就是：

context定义了上下文类型，它携带跨越API边界和进程之间的deadlines、取消信号和其他请求范围的值。对服务器的传入请求应该创建上下文，对服务器的传出调用应该接受上下文。它们之间的函数调用链必须传播上下文，可以选择用使用WithCancel、WithDeadline、WithTimeout或WithValue创建的派生上下文替换它。当一个上下文被取消时，所有从它派生的上下文也被取消。

源码结构

名称	类型	说明
Context	interface	定义了 Context 接口的四个方法Deadline、Done、Err、Value
emptyCtx	int	注意：emptyCtx永远不会取消，没有值，也没有截止日期。这里使用的是类型等价定义，emptyCtx等价于int类型。并且定义上面的四个方法和String方法。
Background	func	返回new(emptyCtx)
TODO	func	返回new(emptyCtx)
CancelFunc	func	CancelFunc告诉操作放弃其工作，不等待工作停止。多个goroutine可以同时调用CancelFunc。在第一个调用之后，对CancelFunc的后续调用将不执行任何操作。
WithCancel	func	WithCancel返回一个带有新的Done通道的parent副本。当返回的cancel函数被调用时，上下文的Done通道被关闭或者当父上下文的Done通道关闭时，无论哪个先发生。取消此上下文将释放与之相关的资源，代码也应该如此在此上下文中运行的操作一完成，就调用cancel。
newCancelCtx	func	返回一个初始化的cancelCtx
propagateCancel	func	propagateCancel在父元素被取消时取消子元素，这里有用到原子锁
parentCancelCtx	func	找到第一个可取消的父节点
removeChild	func	移除父节点的子节点
canceler	interface	取消者，定义了cancel和Done两个方法
init	func	初始化方法
cancelCtx	struct	一个可以取消的 Context
contextName	func	返回上下文名称
WithDeadline	func	创建一个有 deadline 的 context
timerCtx	struct	timerCtx带有timeout 和deadline 。它将cancelCtx嵌入到实现Done和Err。它通过停止计时器来实现取消，然后委托给cancelCtx.cancel。
WithTimeout	func	创建一个有 timeout 的 context
WithValue	func	创建一个存储 k-v 对的 context
valueCtx	struct	存储k-v，配合WithValue使用
stringify	func	接口类型返回字符串

Context

type Context interface {
	// 获取设置的截止时间的意思，第一个返回式是截止时间，到了这个时间点，Context会自动发起取消请求；第二个返回值ok==false时表示没有设置截止时间，如果需要取消的话，需要调用取消函数进行取消
	Deadline() (deadline time.Time, ok bool)

	// 返回一个只读的chan，类型为struct{}，我们在goroutine中，如果该方法返回的chan可以读取，则意味着parent context已经发起了取消请求，我们通过Done方法收到这个信号后，就应该做清理操作，然后退出goroutine，释放资源。
	Done() <-chan struct{}

	// 在 channel Done 关闭后，返回 context 取消原因
	Err() error

	// 获取该Context上绑定的值，是一个键值对，所以要通过一个Key才可以获取对应的值，这个值一般是并发安全的
	Value(key interface{}) interface{}
}

canceler

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

实现了上面定义的两个方法的 Context，就表明该 Context 是可取消的。源码中有两个类型实现了 canceler 接口：*cancelCtx 和 *timerCtx。注意是加了 * 号的，是这两个结构体的指针实现了 canceler 接口。

接口设计成这个样子的原因：

• “取消”操作应该是建议性，而非强制性
• “取消”操作应该可传递

emptyCtx

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
}

func (e *emptyCtx) String() string {
	switch e {
	case background:
		return "context.Background"
	case todo:
		return "context.TODO"
	}
	return "unknown empty Context"
}
// cancelCtx
type cancelCtx struct {
	Context

	mu       sync.Mutex            // 互斥锁
	done     chan struct{}        
	children map[canceler]struct{} 
	err      error 
}

我们重点看一下 cancel 这个方法

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 // 已经被其他协程取消
	}
	// 给 err 字段赋值
	c.err = err
	// 关闭 channel，通知其他协程
	if c.done == nil {
		c.done = closedchan
	} else {
		close(c.done)
	}
	
	// 遍历它的所有子节点
	for child := range c.children {
                // 递归地取消所有子节点
		child.cancel(false, err)
	}
	// 将子节点置空
	c.children = nil
	c.mu.Unlock()

	if removeFromParent {
	        // 从父节点中移除自己 
		removeChild(c.Context, c)
	}
}

cancel()方法的功能就是关闭 channel：c.done；递归地取消它的所有子节点；从父节点从删除自己。达到的效果是通过关闭 channel，将取消信号传递给了它的所有子节点。

func WithCancel(parent Context) (ctx Context, cancel CancelFunc)
func WithDeadline(parent Context, deadline time.Time) (Context, CancelFunc)
func WithTimeout(parent Context, timeout time.Duration) (Context, CancelFunc)
func WithValue(parent Context, key, val interface{}) Context

上面四个可以理解为Context的继承衍生。

对于我们日常使用来说，学会Context的继承的4个方法和Background、TODO基本就够了。

Context使用场景

1. 超时请求

   package main
   
   import (
       "context"
       "fmt"
       "time"
   )
   
   func main() {
       ctx, cancel := context.WithTimeout(context.Background(), 50*time.Millisecond)
       defer cancel()
   
       select {
       case <-time.After(1 * time.Second):
           fmt.Println("overslept")
       case <-ctx.Done():
           fmt.Println(ctx.Err()) // prints "context deadline exceeded"
       }
   }

2. 共享数据

package main

import (
	"context"
	"fmt"
)

func main() {
	ctx := context.Background()
	process(ctx)

	ctx = context.WithValue(ctx, "traceId", "qcrao-2019")
	process(ctx)
}

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

这个在web开发中很实用，传递session、token等信息。

3. 防止 goroutine 泄漏

   package main
   
   import (
   	"time"
   	"fmt"
   )
   
   func gen() <-chan int {
   	ch := make(chan int)
   	go func() {
   		var n int
   		for {
   			ch <- n
   			n++
   			time.Sleep(time.Second)
   		}
   	}()
   	return ch
   }
   
   func main() {
   	for n := range gen() {
   		fmt.Println(n)
   		if n == 5 {
   			break
   		}
   	}
   }

当 n = 5 的时候，直接 break 。但是 gen 的协程就会执行无限循环，永远不会停下来。发生了 goroutine 泄漏。

package main

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

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

func main() {
	ctx, cancel := context.WithCancel(context.Background())
	defer cancel() // 避免其他地方忘记 cancel，且重复调用不影响

	for n := range gen(ctx) {
		fmt.Println(n)
		if n == 5 {
			cancel()
			break
		}
	}
}

Context 使用原则

最后记住几个主要的使用原则：

• 不要将 Context 塞到结构体里。直接将 Context 类型作为函数的第一参数，而且一般都命名为 ctx。
• 不要向函数传入一个 nil 的 context，如果你实在不知道传什么，标准库给你准备好了一个 context：todo。
• 不要把本应该作为函数参数的类型塞到 context 中，context 存储的应该是一些共同的数据。例如：登陆的 session、cookie 等。
• 同一个 context 可能会被传递到多个 goroutine，别担心，context 是并发安全的。

参考文章：

• 深度解密Go语言之context
• Go语言实战笔记（二十）| Go Context