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 
}

我们重点看一下这个方法

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等信息。

  1. 防止 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
https://github.13sai.com/2021/01/06/context/
作者
13sai
许可协议