Go context.WithCancel()源碼剖析
Context 接口
type Context interface {
Deadline() (deadline time.Time, ok bool)
Done() <-chan struct{}
Err() error
Value(key any) any
}
- Deadline() 上下文的截止時間
- Done() 上下文是否已關(guān)閉
- Err() 上下文關(guān)閉的原因
- Value(key any) 上下文存儲的信息
WithCancel 函數(shù)
type CancelFunc func()
func WithCancel(parent Context) (ctx Context, cancel CancelFunc) {
c := withCancel(parent)
return c, func() { c.cancel(true, Canceled, nil) }
}
這是一個入口函數(shù)规揪，gorutine可以通過該函數(shù)生成一個可cancel的context偷厦。cancel的方法就是該函數(shù)的 cancel CancelFunc 返回體。
可以看到 WithCancel 函數(shù)返回2個對象：
-
Context->cancelCtx -
CancelFunc->func() { c.cancel(true, Canceled, nil) }
其中 Context 實(shí)現(xiàn)為 cancelCtx ，CancelFunc 為一個 func, func 底層調(diào)用 cancelCtx.cancel()
因此可以看到代碼核心為 cancelCtx 結(jié)構(gòu)體欠雌。
withCancel 函數(shù)
func withCancel(parent Context) *cancelCtx {
if parent == nil {
panic("cannot create context from nil parent")
}
c := &cancelCtx{}
c.propagateCancel(parent, c)
return c
}
代碼邏輯很清晰也很直接策橘，先是初始化一個 cancelCtx 對象锚扎，然后調(diào)用 cancelCtx.propagateCancel(parent, c) 函數(shù)坎拐。
cancelCtx.propagateCancel(parent, c) 實(shí)際上就是綁定兩個context之間的關(guān)系，這個函數(shù)在后面會有詳細(xì)分析芹缔。
因此可以看到 cancelCtx 結(jié)構(gòu)體的核心代碼為 cancelCtx.cancel() 和 cancelCtx.propagateCancel() 這2個函數(shù)坯癣。
接下來我們先分析 cancelCtx 結(jié)構(gòu)體。
cancelCtx結(jié)構(gòu)體
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結(jié)構(gòu)體: 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() 函數(shù)的實(shí)現(xiàn)僅為返回 itself 或者 value(parentCtx)
其中 value() 函數(shù)實(shí)現(xiàn)比較簡單最欠，就是不斷往上查找 parentCtx 示罗，直至找到第一個匹配key的映射值
cancelCtx結(jié)構(gòu)體: 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() 實(shí)際上是返回一個存儲在 cancelCtx.done 字段內(nèi)的空struct類型的channel通道。
cancelCtx.done 字段的類型為 atomic.Value, 這是一個原子操作存儲對象芝硬。
值得注意的是蚜点，cancelCtx.done 字段初始化使用了"二次校驗(yàn)鎖"，這是懶加載保證線程安全的常用對象初始化方式拌阴。
cancelCtx結(jié)構(gòu)體: func (c *cancelCtx) Err() error
func (c *cancelCtx) Err() error {
c.mu.Lock()
err := c.err
c.mu.Unlock()
return err
}
cancelCtx.Err() 實(shí)現(xiàn)比較簡單绍绘，加鎖拿到上下文關(guān)閉原因，然后返回迟赃。
這里說明陪拘，cancelCtx.err 不為nil則說明 cancelCtx 上下文已關(guān)閉。
cancelCtx結(jié)構(gòu)體: 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() 代碼核心邏輯是將當(dāng)前 cancelCtx 加入到 parentCtx.children 變量中
- 首先
parent.Done()是否為nil?
- 如果為nil纤壁，那就代表著parent之前所有祖先context都是不能主動關(guān)閉的左刽，無需綁定 `parentCtx` 與 當(dāng)前`cancelCtx` 的映射關(guān)系，邏輯結(jié)束
- 如果不為nil酌媒，則代表 `parentCtx` 是可主動關(guān)閉的欠痴，需要綁定映射關(guān)系
- 接著監(jiān)聽
parentCtx是否剛好關(guān)閉上下文?
- 如果是，代表著當(dāng)前 `cancelCtx` 也得關(guān)閉上下文秒咨，于是通過 `cancelCtx.cancel()` 關(guān)閉喇辽，邏輯結(jié)束
- 如果不是，那么代表著需綁定 `parentCtx` 與 當(dāng)前`cancelCtx` 的映射關(guān)系雨席，邏輯繼續(xù)
- 嘗試將
parentCtx強(qiáng)轉(zhuǎn)為cancelCtx類型茵臭，如果強(qiáng)轉(zhuǎn)成功，則加鎖判斷parentCtx.err是否為空?
- 為空，則說明 `parentCtx` 上下文尚未關(guān)閉旦委，需要將當(dāng)前 `cancelCtx` 加入到 `parentCtx.children` 變量中，同時邏輯結(jié)束
- 不為空雏亚，則說明其他gorutine在這期間關(guān)閉了上下文缨硝，需要通過 `cancelCtx.cancel()` 關(guān)閉當(dāng)前 `cancelCtx` 上下文，同時邏輯結(jié)束
嘗試將
parentCtx強(qiáng)轉(zhuǎn)為afterFuncer類型罢低，具體邏輯和 3 類似如果
parentCtx不能強(qiáng)轉(zhuǎn)為cancelCtx和afterFuncer查辩，那么起一個gorutinue監(jiān)聽parentCtx和 當(dāng)前cancelCtx是否關(guān)閉。
- 這里可能會有點(diǎn)費(fèi)解网持，但是在 `cancelCtx` 初始化期間宜岛，`parentCtx` 是有可能被其他goroutine關(guān)閉上下文的，這里就是預(yù)防這種場景功舀。
cancelCtx結(jié)構(gòu)體: 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 結(jié)構(gòu)體的核心代碼, 雖然核心但是卻極為簡潔
-
err為必選入?yún)ⅲ?code>cause 默認(rèn)等于err - 加鎖操作萍倡，判斷
cancelCtx.err是否為空
- 不為空說明當(dāng)前
cancelCtx上下文已關(guān)閉，邏輯結(jié)束 - 為空則說明當(dāng)前
cancelCtx上下文需要關(guān)閉
- 關(guān)閉上下文需要
- 給
cancelCtx.err孝常，cancelCtx.cause賦值 - 關(guān)閉
cancelCtx.done的channel通道 - 遍歷并關(guān)閉
cancelCtx.children的上下文贮竟，最后置空cancelCtx.children
總結(jié)
cancelCtx上下文的關(guān)閉信號存儲在done字段cancelCtx.err、cancelCtx.cause、cancelCtx.done都需要加鎖操作所意，這三個字段都代表著上下文是否已關(guān)閉父context如果關(guān)閉，需要主動負(fù)責(zé)關(guān)閉子context
context關(guān)閉是單向傳導(dǎo)姻檀，并不會導(dǎo)致父context關(guān)閉绣版，只會導(dǎo)致所有的子孫context關(guān)閉
