go - 如何在测试套件的 TestMain 中执行延迟函数？

Question

我已经更新了问题以避免被指责发布 XY 问题。

之前的问题是:

how can I tell when runtime.Goexit has been called on the main goroutine?

我正在尝试编写一个方法，通过调用 runtime.Goexit() 然后调用 os.Exit() 来完成主 goroutine 上的所有延迟函数在调用此退出方法之前生成的 goroutine。我遇到的问题是我不知道 runtime.Goexit() 何时完成。有什么方法可以知道主 goroutine 何时完成？

更新:让我详细说明我的用例。考虑 TestMain 的这种模式:

func TestMain(m *testing.M) {
    db.Create()
    defer db.Drop()
    os.Exit(m.Run())
}

在这种情况下，数据库永远不会被删除，因为 os.Exit 会停止程序。我的目标是想出一个替代方案执行 TestMain 中所有延迟函数的退出函数。现在，我可以把所有东西都移到另一个像这样的功能:

func realTestMain(m *testing.M) int {
    db.Create()
    defer db.Drop()
    return m.Run()
}

func TestMain(m *testing.M) {
    os.Exit(realTestMain(m))
}

但是，这种模式必须在我们的测试套件中强制执行，而且很难记住。我在探索我是否可以有一个帮助程序来编写我的设置/拆卸，如下所示:

func TestMain(m *testing.M) {
    db.Create()
    defer db.Drop()
    helpers.SafeExit(m.Run())
}

有了这样的助手，我可以编写一个简单的检查来确保 os.Exit 永远不会出现在我们的测试套件中。

Answer 1

如果您想确保协调清理不要使用runtime.Goexit。没有可靠/确定性的机制来 Hook 这种类型的退出方法。

改为使用 context.Context并将其传递给相关的 go-routines。您可以显式取消上下文，go-routines 可以监视此类事件的状态并进行清理。

如果在主电源退出之前需要执行关键任务，您可以使用 sync.WaitGroup .添加到每个 go-routine 的 WaitGroup 。并在每个 go-routine 完成时对 WaitGroup 执行 Done()。主 go-routine 最终可以在 WaitGroup 上 Wait() 并且程序将退出，安全地知道所有清理任务都已完成。

---

示例:

func main() {
    mainCtx, cancelFn := context.WithCancel(context.Background())
    var wg sync.WaitGroup

    wg.Add(1)
    go func() { taskA(mainCtx, &wg) }()
    wg.Add(1)
    go func() { manager(mainCtx, &wg, cancelFn) }()

    wg.Wait() // waiting for manager/task(s)
}


func taskA(ctx context.Context, wg *sync.WaitGroup) {
    defer func() {
        log.Println("defer(): taskA DONE!")
        wg.Done() // must be called last
    }()

    t := time.NewTicker(300 * time.Millisecond)

    for {
        select {
        case <-t.C:
            log.Println("taskA is working:", time.Now())
        case <-ctx.Done(): // are we done?
            log.Println("taskA was canceled. Reason:", ctx.Err())
            return
        }
    }
}

func manager(ctx context.Context, wg *sync.WaitGroup, cancelFn func()) {
    defer func() {
        log.Println("defer(): Manager DONE!")
        wg.Done() // must be called last
    }()

    time.Sleep(time.Second) // give worker(s) some time to work...

    cancelFn() // time to wrap up worker(s)!
}

---

工作 playground version .