重读 context.WithTimeout

温故而知新

context.WithTimeout 用来创建一个 Context，当指定的超时时间到达时自动取消逻辑的执行。

先来看下代码：

// src/context/context.go

// WithTimeout returns WithDeadline(parent, time.Now().Add(timeout)).
//
// Canceling this context releases resources associated with it, so code should
// call cancel as soon as the operations running in this Context complete:
//
//	func slowOperationWithTimeout(ctx context.Context) (Result, error) {
//		ctx, cancel := context.WithTimeout(ctx, 100*time.Millisecond)
//		defer cancel()  // releases resources if slowOperation completes before timeout elapses
//		return slowOperation(ctx)
//	}
func WithTimeout(parent Context, timeout time.Duration) (Context, CancelFunc) {
	return WithDeadline(parent, time.Now().Add(timeout))
}

本质上调用的是 WithDeadline，它们的功能一样，只不过 WithDeadline 使用的是绝对时间，

// src/context/context.go

// WithDeadline returns a copy of the parent context with the deadline adjusted
// to be no later than d. If the parent's deadline is already earlier than d,
// WithDeadline(parent, d) is semantically equivalent to parent. The returned
// context's Done channel is closed when the deadline expires, when the returned
// cancel function is called, or when the parent context's Done channel is
// closed, whichever happens first.
//
// Canceling this context releases resources associated with it, so code should
// call cancel as soon as the operations running in this Context complete.
func WithDeadline(parent Context, d time.Time) (Context, CancelFunc) {
	if parent == nil {
		panic("cannot create context from nil parent")
	}
    // 如果 parent 是一个 emptyCtx，Deadline() 会是一个空实现，ok=false
    // 如果 parent 是一个 timerCtx，DeadLine() 返回的 deadline 是结束截止时间 ，ok=true
	if cur, ok := parent.Deadline(); ok && cur.Before(d) {
		// The current deadline is already sooner than the new one.
        // 如果父 timerCtx 结束截止时间早于新的截止时间，采用父 timerCtx 截止时间
		return WithCancel(parent)
	}
    // 创建一个 timerCtx
	c := &timerCtx{
		cancelCtx: newCancelCtx(parent),
		deadline:  d,
	}
	propagateCancel(parent, c)
	dur := time.Until(d)
	if dur <= 0 {
        // 已经过了截止时间，直接调用 cancel 结束掉自己，并从父 Context 脱离
		c.cancel(true, DeadlineExceeded) // deadline has already passed
		return c, func() { c.cancel(false, Canceled) }
	}
	c.mu.Lock()
	defer c.mu.Unlock()
	if c.err == nil {
        // 当前 Context 还未结束，创建 timer
		c.timer = time.AfterFunc(dur, func() {
			c.cancel(true, DeadlineExceeded)
		})
	}
	return c, func() { c.cancel(true, Canceled) }
}

WithDeadline 返回的是一个 timerCtx，它继承了 cancelCtx，且维护了一个定时器 timer 和一个截止时间 deadline:

// src/context/context.go

// A timerCtx carries a timer and a deadline. It embeds a cancelCtx to
// implement Done and Err. It implements cancel by stopping its timer then
// delegating to cancelCtx.cancel.
type timerCtx struct {
	cancelCtx
	timer *time.Timer // Under cancelCtx.mu.

	deadline time.Time
}

// A cancelCtx can be canceled. When canceled, it also cancels any children
// that implement canceler.
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
}

从 timerCtx 的构造中，我们知道它维护了一个 canceler 的 children，以便自己在 cancel 时也 cancel 掉 children，同时它是继承 Context 的，作为 Context “树”中的一个节点需要处理父 Context 取消和子 Context 的取消问题。

在 WithDeadline 函数内，有这样的一句

propagateCancel(parent, c)

它的实现如下：

// src/context/context.go

// propagateCancel arranges for child to be canceled when parent is.
func propagateCancel(parent Context, child canceler) {
	done := parent.Done()
	if done == nil {
        // 父 Context 不能被取消，可能是 emptyCtx 或 valueCtx，不用考虑
		return // parent is never canceled
	}

	select {
	case <-done:
		// parent is already canceled
        // 父 Context 已经被取消，同步取消 children
		child.cancel(false, parent.Err())
		return
	default:
	}
    // 尝试向上追溯到可以取消的父 Context 节点
	if p, ok := parentCancelCtx(parent); ok {
        // 找到了可以取消的父 Context 节点
		p.mu.Lock()
		if p.err != nil {
			// parent has already been canceled
            // 父 Context 已经被取消，同步取消 children
			child.cancel(false, p.err)
		} else {
            // 父 Context 没有被取消，将自己挂到父 Context 下
			if p.children == nil {
				p.children = make(map[canceler]struct{})
			}
			p.children[child] = struct{}{}
		}
		p.mu.Unlock()
	} else {
		atomic.AddInt32(&goroutines, +1)
		go func() {
			select {
			case <-parent.Done():
				child.cancel(false, parent.Err())
			case <-child.Done():
			}
		}()
	}
}

如果父 Context 不能被取消，即不是 cancelCtx（只有 cancelCtx 具体实现了 Done() 函数，返回的 channel 不为 nil），则不处理； 如果父 Context 可以被取消，且已经被取消了，则需同步取消 children； 如果父 Context 可以被取消，还没被取消，往上追溯可以被取消的 Context，并将自己挂上去，成为父 Context 的 children；

简而言之，这个函数的作用就是尝试找到可以取消的父 Context，并将自己挂上去，当父 Context 被取消的时候能够“顺藤摸瓜”地取消所有的子 Context。

接下来细想，函数末尾的 else 代码块看上去有些奇怪，既然向上追溯不到可以取消的 Context，为何还要执行看上去无用的 else 语句呢？

回过头来看下函数 parentCancelCtx(parent) ：

// src/context/context.go

// parentCancelCtx returns the underlying *cancelCtx for parent.
// It does this by looking up parent.Value(&cancelCtxKey) to find
// the innermost enclosing *cancelCtx and then checking whether
// parent.Done() matches that *cancelCtx. (If not, the *cancelCtx
// has been wrapped in a custom implementation providing a
// different done channel, in which case we should not bypass it.)
func parentCancelCtx(parent Context) (*cancelCtx, bool) {
	done := parent.Done()
	if done == closedchan || done == nil {
		return nil, false
	}
	p, ok := parent.Value(&cancelCtxKey).(*cancelCtx)
	if !ok {
		return nil, false
	}
	pdone, _ := p.done.Load().(chan struct{})
	if pdone != done {
		return nil, false
	}
	return p, true
}

这个函数其实就是向上追溯可以被取消的 cancelCtx 或 timerCtx（其实也继承了 cancelCtx），找到的话 ok=true。

而函数末尾的 else 代码块其实是处理 if 代码块处理不了的情况，原因在于如果自定义一个结构体继承了 cancelCtx，那么它的类型就不可能是 cancelCtx，parentCancelCtx(parent) 也就找不到这个自定义的结构体。

else 代码块是启动一个 goroutine 监控取消信号，如果收到 parent 的取消信息，则自己取消并退出 select，如果收到自己的取消消息，则直接退出 select。

接下来，我们看下 timerCtx 是怎么取消的：

// src/context/context.go

func (c *timerCtx) cancel(removeFromParent bool, err error) {
	c.cancelCtx.cancel(false, err)
	if removeFromParent {
		// Remove this timerCtx from its parent cancelCtx's children.
		removeChild(c.cancelCtx.Context, c)
	}
	c.mu.Lock()
	if c.timer != nil {
        // 将定时器停止
		c.timer.Stop()
        // 销毁定时器
		c.timer = nil
	}
	c.mu.Unlock()
}

// cancel closes c.done, cancels each of c's children, and, if
// removeFromParent is true, removes c from its parent's children.
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 // already canceled
	}
	c.err = err
	d, _ := c.done.Load().(chan struct{})
	if d == nil {
		c.done.Store(closedchan)
	} else {
		close(d)
	}
    // 一个个取消 child
	for child := range c.children {
		// NOTE: acquiring the child's lock while holding parent's lock.
		child.cancel(false, err)
	}
    // 将 children 销毁
	c.children = nil
	c.mu.Unlock()

	if removeFromParent {
		removeChild(c.Context, c)
	}
}

我们看到 timerCtx 的取消操作做了以下几件事：

• 关闭 channel
• 递归取消 children 并销毁
• 停止定时器并销毁

我们知道在使用 context.WithTimeout 的时候要手动调用下 cancel 函数，基本上是如下的形式：

ctx := context.Background()
...
ctx, cancel := context.WithTimeout(ctx, time.Second)
defer cancel()
doSomething(ctx)

defer cancel() 就是在 doSomethind(ctx) 执行结束后手动调用执行的。

超时的时候是怎么执行的呢？关键是前面讲过的 WithDeadline 中的创建 timer：

c.timer = time.AfterFunc(dur, func() {
			c.cancel(true, DeadlineExceeded)
		})

timer 创建函数 AfterFunc(d Duration, f func()) *Timer 完成的任务就是创建一个 timer，实现当超时时间 dur 到的时候，调用传入的参数函数 f：

// src/time/sleep.go

// AfterFunc waits for the duration to elapse and then calls f
// in its own goroutine. It returns a Timer that can
// be used to cancel the call using its Stop method.
func AfterFunc(d Duration, f func()) *Timer {
	t := &Timer{
		r: runtimeTimer{
			when: when(d),
			f:    goFunc,
			arg:  f,
		},
	}
	startTimer(&t.r)
	return t
}

小结： context.WithTimeout 返回了一个 timerCtx，timerCtx 作为 Context 树的一个节点。当父 Context 已经结束的时候执行自己的 cancel 函数；其拥有的 timer 在当截止时间到来时也会调用自己的 cancel 函数。而 cancel 函数会递归的 cancel 掉其下的子 Context 并销毁，并将自己与父 Context 脱离。