Go context.WithCancel()源码剖析
Context 接口
type Context interface {
Deadline() (deadline time.Time, ok bool)
Done() <-chan struct{}
Err() error
Value(key any) any
}
- Deadline() 上下文的截止时间
- Done() 上下文是否已关闭
- Err() 上下文关闭的原因
- Value(key any) 上下文存储的信息
WithCancel 函数
type CancelFunc func()
func WithCancel(parent Context) (ctx Context, cancel CancelFunc) {
c := withCancel(parent)
return c, func() { c.cancel(true, Canceled, nil) }
}
这是一个入口函数，gorutine可以通过该函数生成一个可cancel的context。cancel的方法就是该函数的 cancel CancelFunc 返回体。
可以看到 WithCancel 函数返回2个对象：
-
Context->cancelCtx -
CancelFunc->func() { c.cancel(true, Canceled, nil) }
其中 Context 实现为 cancelCtx ，CancelFunc 为一个 func, func 底层调用 cancelCtx.cancel()
因此可以看到代码核心为 cancelCtx 结构体。
withCancel 函数
func withCancel(parent Context) *cancelCtx {
if parent == nil {
panic("cannot create context from nil parent")
}
c := &cancelCtx{}
c.propagateCancel(parent, c)
return c
}
代码逻辑很清晰也很直接，先是初始化一个 cancelCtx 对象，然后调用 cancelCtx.propagateCancel(parent, c) 函数。
cancelCtx.propagateCancel(parent, c) 实际上就是绑定两个context之间的关系，这个函数在后面会有详细分析。
因此可以看到 cancelCtx 结构体的核心代码为 cancelCtx.cancel() 和 cancelCtx.propagateCancel() 这2个函数。
接下来我们先分析 cancelCtx 结构体。
cancelCtx结构体
type cancelCtx struct {
Context
mu sync.Mutex // protects following fields
done atomic.Value // of 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
cause error // set to non-nil by the first cancel call
}
-
cancelCtx.mu为一个操作锁 -
cancelCtx.done为一个atomic.Value类型的原子容器 -
cancelCtx.children为一个仅使用key的map
cancelCtx结构体: func (c *cancelCtx) Value(key any) any
// &cancelCtxKey is the key that a cancelCtx returns itself for.
var cancelCtxKey int
func (c *cancelCtx) Value(key any) any {
//
if key == &cancelCtxKey {
return c
}
return value(c.Context, key)
}
可以看到 cancelCtx.Value() 函数的实现仅为返回 itself 或者 value(parentCtx)
其中 value() 函数实现比较简单，就是不断往上查找 parentCtx ，直至找到第一个匹配key的映射值
cancelCtx结构体: func (c *cancelCtx) Done() <-chan struct{}
func (c *cancelCtx) Done() <-chan struct{} {
d := c.done.Load()
if d != nil {
return d.(chan struct{})
}
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{})
}
可以看到 cancelCtx.Done() 实际上是返回一个存储在 cancelCtx.done 字段内的空struct类型的channel通道。
cancelCtx.done 字段的类型为 atomic.Value, 这是一个原子操作存储对象。
值得注意的是，cancelCtx.done 字段初始化使用了"二次校验锁"，这是懒加载保证线程安全的常用对象初始化方式。
cancelCtx结构体: func (c *cancelCtx) Err() error
func (c *cancelCtx) Err() error {
c.mu.Lock()
err := c.err
c.mu.Unlock()
return err
}
cancelCtx.Err() 实现比较简单，加锁拿到上下文关闭原因，然后返回。
这里说明，cancelCtx.err 不为nil则说明 cancelCtx 上下文已关闭。
cancelCtx结构体: func (c *cancelCtx) propagateCancel(parent Context, child canceler)
func (c *cancelCtx) propagateCancel(parent Context, child canceler) {
c.Context = parent
done := parent.Done()
if done == nil {
return // parent is never canceled
}
select {
case <-done:
// parent is already canceled
child.cancel(false, parent.Err(), Cause(parent))
return
default:
}
if p, ok := parentCancelCtx(parent); ok {
// parent is a *cancelCtx, or derives from one.
p.mu.Lock()
if p.err != nil {
// parent has already been canceled
child.cancel(false, p.err, p.cause)
} else {
if p.children == nil {
p.children = make(map[canceler]struct{})
}
p.children[child] = struct{}{}
}
p.mu.Unlock()
return
}
if a, ok := parent.(afterFuncer); ok {
// parent implements an AfterFunc method.
c.mu.Lock()
stop := a.AfterFunc(func() {
child.cancel(false, parent.Err(), Cause(parent))
})
c.Context = stopCtx{
Context: parent,
stop: stop,
}
c.mu.Unlock()
return
}
goroutines.Add(1)
go func() {
select {
case <-parent.Done():
child.cancel(false, parent.Err(), Cause(parent))
case <-child.Done():
}
}()
}
cancelCtx.propagateCancel() 代码核心逻辑是将当前 cancelCtx 加入到 parentCtx.children 变量中
- 首先
parent.Done()是否为nil?
- 如果为nil，那就代表着parent之前所有祖先context都是不能主动关闭的，无需绑定 `parentCtx` 与 当前`cancelCtx` 的映射关系，逻辑结束
- 如果不为nil，则代表 `parentCtx` 是可主动关闭的，需要绑定映射关系
- 接着监听
parentCtx是否刚好关闭上下文?
- 如果是，代表着当前 `cancelCtx` 也得关闭上下文，于是通过 `cancelCtx.cancel()` 关闭，逻辑结束
- 如果不是，那么代表着需绑定 `parentCtx` 与 当前`cancelCtx` 的映射关系，逻辑继续
- 尝试将
parentCtx强转为cancelCtx类型，如果强转成功，则加锁判断parentCtx.err是否为空?
- 为空，则说明 `parentCtx` 上下文尚未关闭，需要将当前 `cancelCtx` 加入到 `parentCtx.children` 变量中，同时逻辑结束
- 不为空，则说明其他gorutine在这期间关闭了上下文，需要通过 `cancelCtx.cancel()` 关闭当前 `cancelCtx` 上下文，同时逻辑结束
尝试将
parentCtx强转为afterFuncer类型，具体逻辑和 3 类似如果
parentCtx不能强转为cancelCtx和afterFuncer，那么起一个gorutinue监听parentCtx和 当前cancelCtx是否关闭。
- 这里可能会有点费解，但是在 `cancelCtx` 初始化期间，`parentCtx` 是有可能被其他goroutine关闭上下文的，这里就是预防这种场景。
cancelCtx结构体: func (c *cancelCtx) cancel(removeFromParent bool, err, cause error)
func (c *cancelCtx) cancel(removeFromParent bool, err, cause error) {
if err == nil {
panic("context: internal error: missing cancel error")
}
if cause == nil {
cause = err
}
c.mu.Lock()
if c.err != nil {
c.mu.Unlock()
return // already canceled
}
c.err = err
c.cause = cause
d, _ := c.done.Load().(chan struct{})
if d == nil {
c.done.Store(closedchan)
} else {
close(d)
}
for child := range c.children {
// NOTE: acquiring the child's lock while holding parent's lock.
child.cancel(false, err, cause)
}
c.children = nil
c.mu.Unlock()
if removeFromParent {
removeChild(c.Context, c)
}
}
cancelCtx.cancel 为 cancelCtx 结构体的核心代码, 虽然核心但是却极为简洁
-
err为必选入参，cause默认等于err - 加锁操作，判断
cancelCtx.err是否为空
- 不为空说明当前
cancelCtx上下文已关闭，逻辑结束 - 为空则说明当前
cancelCtx上下文需要关闭
- 关闭上下文需要
- 给
cancelCtx.err，cancelCtx.cause赋值 - 关闭
cancelCtx.done的channel通道 - 遍历并关闭
cancelCtx.children的上下文，最后置空cancelCtx.children
总结
cancelCtx上下文的关闭信号存储在done字段cancelCtx.err、cancelCtx.cause、cancelCtx.done都需要加锁操作，这三个字段都代表着上下文是否已关闭父context如果关闭，需要主动负责关闭子context
context关闭是单向传导，并不会导致父context关闭，只会导致所有的子孙context关闭
