how do I execute deferred functions in TestMain for a test suite?

Question

I've updated the question to avoid being accused of posting an XY question.

Previously the question was:

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

I'm trying to write a method that finishes all deferred functions on the main goroutine by calling runtime.Goexit() and then calls os.Exit() from a goroutine a spawned before calling this exit method. The problem I have is that I don't know when runtime.Goexit() has completed. Is there any way I can know when the main goroutine has finished?

UPDATED: Let me elaborate on my use case. Consider this pattern for TestMain:

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

In this case, the database is never dropped because os.Exit stops the program. My goal was to come up with an alternate exit function that executes all the deferred functions in TestMain. Now, I could move everything into another function like this:

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

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

However, this pattern would have to be enforced across our test suite and is ugly and difficult to remember. I was exploring whether I could have a helper that makes it possible to write my setup/teardown as follows:

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

With a helper like this, I could write a simple check to ensure that os.Exit never appears in our test suites.

Answer 1

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

而是使用context.Context并将其传递给您相关的例程。您可以显式取消Context，而例程可以监视此类事件的状态并进行清理。

如果有关键任务需要在市电退出之前执行，则可以使用sync.WaitGroup。为每个执行例程添加到等待组。然后Done()在每个go例程完成时在WaitGroup上执行。主例程可以最终 Wait()在WaitGroup上运行，并且程序会退出，这可以确保所有清理任务已完成。

例：

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)!
}

工作操场版本。