Go组件-Context

十三Tech
95 阅读8分钟

Context 上下文，在Go1.7引入标准库

Context 主要用来在 goroutine 之间传递上下文信息，包括：取消信号、超时时间、截止时间、k-v 等

ctx结构体、方法

  • Context 接口：定义了 Context 接口的四个方法

  • emptyCtx 结构体：一个空 Context

  • CancelFunc 函数类型：Context 的取消函数

  • canceler 接口：Context 取消接口

  • cancelCtx 结构体：实现了取消接口的 Context

  • timerCtx 结构体：超时会取消的 Context

  • valueCtx 结构体：可以存储键值对的 Context

  • Background() 函数：返回空 Context，常作为根 Context

  • TODO() 函数：返回一个空 Context，在需要 Context 的地方又没有合适的 Context 就用这个

  • WithCancel() 函数：基于父 Context，创建一个可取消的 Context

  • newCancelCtx() 函数：创建一个可取消的 Context

  • propagateCancel() 函数：将节点挂载到上游第一个 cancelCtx 上，又或者启动协程监听 Context 取消事件

  • parentCancelCtx() 函数：返回上游的第一个 cancelCtx

  • removeChild() 函数：移除 Context 节点

  • init() 函数：包初始化函数，创建了一个关闭的 chan

  • WithDeadline() 函数：创建一个有 deadline 的 Context

  • WithTimeout() 函数：创建一个有 timeout 的 Context

  • WithValue() 函数：创建一个存储键值对的 Context

context 接口

type Context interface {
    // 返回一个 channel，当 context 被取消或者到了 deadline 的时候，
    // 这个 channel 会被 close，从而 <-chan struct{} 会返回。
    // 在没有关闭之前，一直阻塞，因为不会有任何地方往这个 channel 中发送值。
    Done() <-chan struct{}

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

    // 返回 context 是否会被取消以及自动取消时间（即 deadline）
    // ok 为 true，表明设置了 deadline，第一个返回值就是设置的 deadline
    // ok 为 false，表示没有设置 deadline，第一个返回值没意义。
    Deadline() (deadline time.Time, ok bool)

    // 获取 key 对应的 value
    Value(key interface{}) interface{}
}

Context 接口定义了4个方法，它们都是幂等的

canceler 接口

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

cancel 方法的第一个参数 removeFromParent 表示的是，是否从父 Context 移除自身，这是因为 Context 是一个树状结构。 在 Context 取消的时候，它会给所有派生的 Context 也发送取消信号, 所以派生新的 Context 的时候会记录从当前 Context 派生出去的 Context

cancelCtx 结构体

cancel 的操作实际上只会做一次，后续调用 cancel 的时候会返回第一次 cancel 的结果，cancel 是一个幂等操作

// 可以被取消，取消的时候，所有实现了 canceler 接口的派生出来的 Context 也会被取消。
type cancelCtx struct {
	// cancelCtx 也实现了 `Context` 接口
	Context

	// mu 用以保护后面的 done、children、err 字段
	mu       sync.Mutex
	
	// 是一个 chan struct{}，懒汉式创建，
	// 在第一次 cancel 的时候被关闭
	done     atomic.Value
	// 记录所有可以取消的子 Context
	// 在第一次 cancel 的时候会被设置为 nil。
	children map[canceler]struct{}
	// 在第一次 cancel 的时候会被设置为非 nil 的值
	err      error
}

Done()的实现方式

// 返回一个只读的 chan，但没有任何地方会往这个 chan 写入数据，
// cancel 的时候会关闭这个 chan，从而任何 <-ch 的操作都会立即返回。
func (c *cancelCtx) Done() <-chan struct{} {
	// 如果 done 这个 chan 已经初始化了，就直接返回。
	d := c.done.Load()
	if d != nil {
		return d.(chan struct{})
	}
	// 如果 done 还没初始化，则会进行初始化。
	// 也就是上面说的 "懒汉式" 的创建方式，只有在需要的时候才会初始化。
	c.mu.Lock()
	defer c.mu.Unlock()
	d = c.done.Load()
	if d == nil {
		d = make(chan struct{})
		c.done.Store(d)
	}
	return d.(chan struct{})
}

cancel 方法

// 发送取消信号
func (c *cancelCtx) cancel(removeFromParent bool, err error) {
	// 必须传递一个 err
	if err == nil {
		panic("context: internal error: missing cancel error")
	}
	// 如果已经取消，直接返回。（幂等的设计）
	c.mu.Lock()
	if c.err != nil {
		c.mu.Unlock()
		return // already canceled
	}
	// 记录取消原因，在调用 c.Err() 的时候会返回这个原因
	c.err = err
	// 关闭 done 这个通道，通知其他协程
	d, _ := c.done.Load().(chan struct{})
	if d == nil {
		c.done.Store(closedchan)
	} else {
		close(d)
	}
	// 遍历它的所有子结点，并对其子结点进行取消操作
	for child := range c.children {
		child.cancel(false, err)
	}
	// 将子结点置空
	c.children = nil
	c.mu.Unlock()

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

cancel 操作如下：

  • 关闭 c.done
  • 取消 c 的所有孩子 Context
  • 如果 removeFromParent 为 true，会将 c 从其父 Context 的 children 属性中移除

WithCancel 方法

// 返回值里的 Context 的 Done 方法返回的 channel 关闭或者 parent 被 cancel 的时候，
// 返回值的 CancelFunc 会被执行。
func WithCancel(parent Context) (ctx Context, cancel CancelFunc) {
	// 必须从其他 Context 派生
	if parent == nil {
		panic("cannot create context from nil parent")
	}
	c := newCancelCtx(parent)
	// 将 c 挂靠到 parent 的 children 属性中，
	// 从而在 parent 取消的时候，可以感知得到。
	propagateCancel(parent, &c) // 具体实现后面有详细说明 
	return &c, func() { c.cancel(true, Canceled) }
}

// 创建一个 cancelCtx 实例
func newCancelCtx(parent Context) cancelCtx {
    return cancelCtx{Context: parent}
}

propagateCancel 函数

这个函数会在两个地方调用，一个是 WithCancel，另一个是 WithDeadline，它的主要作用是，找到 parent 以及其父级 Context 路径上 第一个 cancelCtx，目的是，将 child 挂载到找到的这个 cancelCtx 的 children 属性上，从而在这个 cancelCtx 取消的时候， 可以通过遍历 cancelCtx.children 对 child 进行通知。

// 由 parent 往根节点搜索第一个 cancelCtx，如果找到则将 child 写入到 cancelCtx.children 中。
// 如果找到的 cancelCtx 自定义了 Done，则启动协程监听 cancelCtx.Done()。 
func propagateCancel(parent Context, child canceler) {
	// 如果 Context 树上完全不存在 cancelCtx，则直接返回
	done := parent.Done()
	if done == nil {
		return // parent is never canceled
	}

	// 如果 parent 已经取消，则直接取消 child
	select {
	case <-done:
		// parent is already canceled
		child.cancel(false, parent.Err())
		return
	default:
	}

	// 往根节点搜索第一个 cancelCtx
	if p, ok := parentCancelCtx(parent); ok {
		// 找到了，但是已经取消了，则取消 child
		p.mu.Lock()
		if p.err != nil {
			// parent has already been canceled
			child.cancel(false, p.err)
		} else {
			// 找到了，尚未取消。
			// 将 child 写入到 p 的 children 属性中。
			// p.children 是懒汉式创建的。
			if p.children == nil {
				p.children = make(map[canceler]struct{})
			}
			p.children[child] = struct{}{}
		}
		p.mu.Unlock()
	} else {
		// 执行到这里的原因是：
		// 用户自定义了 Done 通道（跟 parent 不是同一个 done），
		// 所以不能以父节点路径上的 done 来决定 child 是否取消，
		// 需要通过启动新协程的方式来监听 Done 通道，从而可以正常取消 parent 的孩子节点。
		atomic.AddInt32(&goroutines, +1)
		go func() {
			select {
			case <-parent.Done():
				child.cancel(false, parent.Err())
			case <-child.Done():
			}
		}()
	}
}

child 什么时候从父 Context 移除

需要移除的情况：

  1. WithCancel 派生出新的 Context 的时候，假设叫 root，这个时候派生的这个 root 也是可以继续派生出新的 Context 的，而这个 root 对于它的子孙 Context 它就是根节点，所以当 root 被取消的时候，它和它的子孙 Context 也要被取消了，所以以 root 为根节点的子树需要被移除。
  2. WithDeadline 里面，当给定的 d 其实已经小于当前时间的时候（也就是父 Context 已经超时了），这个时候会将刚挂载到父节点的 timerCtx 移除，同时返回的 CancelFunc 中，cancel 的第一个参数是 false，因为它已经被移除了。

不需要移除的情况：

  1. 在 propagateCancel 中监测到 parent 已经被取消的时候，因为这个时候 child 并没有关联上 parent，所以自然也没有移除的这种操作。
  2. 就是上面提到的第二种情况中，WithDeadline 的时候就监测到 deadline 已经比当前时间小了（超时了）。
  3. 在 cancelCtx 的 cancel 方法里面，遍历 cancelCtx 的孩子节点的时候，不需要做移除的操作，因为 cancelCtx 本身就需要被从 Context 树中移除。
  4. timerCtx 在没有挂载到 parent 上就已经过期了

valueCtx

type valueCtx struct {
	Context
	key, val any
}

value 方法用以从 Context 中获取对应的值，它会从 Context 树自底向上进行递归搜索，具体来说会有以下几种情况：

  • 如果 ctx 是 *valueCtx，则会判断 key 是否等于 ctx 里面的 key，如果相等，返回 ctx.val。否则，再去搜索 ctx 的父 Context。
  • 如果 ctx 是 *cancelCtx，同时 key 是 &cancelCtxKey，则会返回 ctx。否则，会继续搜索 ctx 到根结点这个路径上的第一个 cancelCtx。
  • 如果 ctx 是 *timerCtx，同时 key 是 &cancelCtxKey，则会返回 ctx.cancelCtx。否则，会继续搜索 ctx 到根结点这个路径上的第一个 cancelCtx。
  • 如果 ctx 是 *emptyCtx，则会返回 nil。（因为这时候是最顶层的 Context 了，也找不到对应的值）。
  • 如果都不是以上的几种情况，则有可能是开发者自定义的 Context 实现，则直接返回 c.Value(key)。

它要解决的问题是：

  • 获取父级 Context 中 WithValue 共享的值。
  • 获取父级 Context 中最靠近当前节点的 cancelCtx（非常重要：它的一个很重要的作用是，将当前节点设置为这个 cancelCtx 的 children，从而可以实现在这个父级的 cancelCtx 取消的时候，当前的 Context 可以感知到）。
  • 如果是开发者自己实现的 Context，则直接调用用户自定义的 Value 方法
// 根据 key 从 c 中获取对应的值，会从 Context 树自底向上递归搜索。
func value(c Context, key any) any {
	for {
		switch ctx := c.(type) {
		case *valueCtx:
			if key == ctx.key {
				return ctx.val
			}
			c = ctx.Context
		case *cancelCtx:
			if key == &cancelCtxKey {
				return c
			}
			c = ctx.Context
		case *timerCtx:
			if key == &cancelCtxKey {
				return &ctx.cancelCtx
			}
			c = ctx.Context
		case *emptyCtx:
			return nil
		default:
			return c.Value(key)
		}
	}
}

为什么是通过关闭 chan 的方式取消

如果通过往 chan 写入数据的方式来 通知其他子孙 Context 的话，我们就需要有多少个子孙 Context 就要往 chan 里面发多少次，但是如果选择使用 close 的方式的话， 我们就完全不用管派生出了多少个可以 cancel 的 Context，因为一旦 chan 关闭了，所有的 <-chan 操作立即得以返回

avatar
文章
阅读
粉丝