Go context.AfterFunc 取消回调的幂等设计

在服务端代码里，context.AfterFunc 很适合把“请求取消”转换成一个资源释放动作：回滚事务、解除锁、关闭临时连接，或者唤醒正在等待的 goroutine。问题是，业务正常结束时也要做同一份清理。如果取消回调和主动关闭同时发生，最容易出现两个错误：清理执行两次，或者 Close 已经返回但清理仍在后台运行。

解决思路不是反复判断 ctx.Err()，而是把两个入口收敛到同一个一次性函数，再单独记录“清理已经结束”。本文实现的目标可以量化为三个不变量。

基线问题：stop 只决定能否阻止启动

官方文档明确说明：AfterFunc 会在 Context 被取消后，使用自己的 goroutine 调用回调；返回的 stop 在成功阻止回调启动时返回 true。但当它返回 false 时，可能是回调已经启动，也可能是它先前已经被停止。更关键的是，stop 不会等待已经启动的回调完成。

stop := context.AfterFunc(ctx, func() {
    // 取消路径会在独立 goroutine 中释放资源。
    cleanup()
})

defer func() {
    stop()    // 这里只尝试阻止回调，不能等待已启动的回调。
    cleanup() // 正常路径再次清理，可能与回调并发执行。
}()

这段代码的竞争窗口很直观：Context 刚好先取消，回调 goroutine 已启动，此时 stop() 返回 false；随后正常路径又调用一次 cleanup()。数据库回滚通常能容忍第二次调用，但关闭 channel、归还一次性租约、删除临时文件等动作未必可重复。

图1：stop 只控制回调是否启动，不能同时承担一次性清理和完成等待。

假设：把所有入口收敛为一次副作用

要同时消除重复清理与提前返回，需要把两个问题拆开：

1. sync.Once 回答“清理由谁执行”，无论回调还是主动关闭先到，副作用只能发生一次。
2. 一个只关闭不发送的 done channel 回答“清理是否完成”，所有调用者都可以等待同一完成信号。
3. stop 只保留原本职责：主动关闭时，尽量阻止取消回调启动。

因此，stop() 的布尔值不再被误解为完整状态机。它只是决定主动关闭是否需要接管 finish；无论哪条路径执行清理，Close 最后都等待 done。

改动点：Once 保证一次，done 保证完成

package cancelhook

import (
    "context"
    "sync"
)

type CancelHook struct {
    stop    func() bool
    once    sync.Once
    done    chan struct{}
    cleanup func()
}

func New(ctx context.Context, cleanup func()) *CancelHook {
    h := &CancelHook{
        done:    make(chan struct{}),
        cleanup: cleanup,
    }

    // 回调与主动关闭共享同一个 finish，竞争只发生在 Once 入口。
    h.stop = context.AfterFunc(ctx, h.finish)
    return h
}

func (h *CancelHook) finish() {
    h.once.Do(func() {
        // 即使 cleanup 发生 panic，等待者也不会永远阻塞。
        defer close(h.done)
        h.cleanup()
    })
}

func (h *CancelHook) Close() {
    if h.stop() {
        // 回调尚未启动，由当前调用接管清理。
        h.finish()
    }

    // stop=false 可能表示回调正在执行，必须显式等待完成。
    

这个结构覆盖了三种关键顺序：

• Close 先到：stop() 返回 true，主动路径调用 finish，清理后关闭 done。
• 取消先到：回调调用 finish；Close 即使得到 false，也会等待同一个 done。
• 多个 Close 并发：最多一个调用者接管清理，其他调用者在 sync.Once 外等待同一个完成信号。

图2：sync.Once 守住一次性副作用，done channel 提供可等待的完成事实。

压测方法：先验证竞争不变量

这里不使用吞吐量或纳秒数字来掩盖并发正确性。更有价值的测试是反复制造取消与关闭竞争，然后检查清理次数是否始终为 1。下面的测试同时启动 8 个 Close 调用者，并让取消信号参与竞争。

package cancelhook

import (
    "context"
    "sync"
    "sync/atomic"
    "testing"
)

func TestCloseIsIdempotent(t *testing.T) {
    ctx, cancel := context.WithCancel(context.Background())
    var calls atomic.Int32

    hook := New(ctx, func() {
        // 原子计数只用于验证 cleanup 的真实执行次数。
        calls.Add(1)
    })

    cancel()

    var wg sync.WaitGroup
    for i := 0; i 

还应补一个“主动关闭先于取消”的用例。先调用两次 Close，再调用 cancel，最后仍断言计数为 1。这样能证明 stop=true 的接管路径和回调已启动路径都满足同一契约。测试可以使用 go test -race 运行，命令含义如下。

# 运行单元测试并启用数据竞争检测器。
go test -race ./...

结果对比：修复的是语义，不是一次布尔判断

在这个方案里，性能验收应围绕可观测结果：测试结束时 cleanup 计数必须等于 1，所有 Close 已返回，且没有等待中的回调。若清理动作本身很慢，耗时会如实反映在 Close 上，这是同步关闭语义的成本，而不是额外泄漏。

边界条件：四个容易忽略的约束

1. cleanup 不要重入 Close

cleanup 在 sync.Once.Do 内执行，而 done 要等它返回后才关闭。如果 cleanup 再调用同一个 Hook 的 Close，就会等待自己关闭 done，形成死锁。把这个限制写进类型注释，并避免把 Hook 暴露给 cleanup。

2. panic 策略要明确

示例用 defer close(h.done) 保证等待者不会永久阻塞，但它不会吞掉 panic。取消回调运行在独立 goroutine 中，未恢复的 panic 仍可能终止进程。生产代码应让 cleanup 自身不 panic；若业务确实需要恢复，应在 cleanup 边界记录错误，而不是静默忽略。

3. 多个 AfterFunc 彼此独立

同一个 Context 上注册多次 AfterFunc 不会互相替换。每个资源都应拥有自己的 stop、Once 和 done；不要期待停止一个回调会停止其他回调。

4. Close 的等待应有业务边界

如果 cleanup 可能访问不可靠的远程服务，无限等待并不合适。更稳妥的做法是让 cleanup 只执行本地、可控且有上限的释放动作；远程补偿交给单独的有超时任务。不要简单给本文的 Close 加超时后就丢弃后台状态，否则“返回即完成”的契约会被破坏。

结论

context.AfterFunc 的 stop 是“阻止回调启动”的竞争工具，不是一次性清理器，也不是等待句柄。可靠的取消回调需要三个清晰职责：stop 控制启动，sync.Once 保证副作用只发生一次，关闭 channel 表达清理已经完成。把这三个职责拆开后，Context 先取消、主动关闭先发生以及多个调用者并发关闭，都可以落在同一套可验证语义上。

参考资料

• Go 标准库：context.AfterFunc
• Go 标准库：sync.Once
• Go 标准库：atomic.Int32