如何在TestMain中为测试套件执行延迟功能?

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.

Is there any way I can know when the main goroutine has finished?

Yes: When your program has stopped.

Once the main goroutine stops your program terminates including any other goroutine. So you will have to redesign as nothing ever executes after the main goroutine exits.

If you want to ensure coordinated cleanup DON'T use runtime.Goexit. There is no reliable/deterministic mechanism to hook onto this type of exit method.

Instead use context.Context and pass it to your relevant go-routines. You can explicitly cancel a Context and the go-routines can monitor it's state for such an event and do their cleanup.

If there is critical tasks that need executing prior to mains exit, you can employ a sync.WaitGroup. Adding to the wait group for each go-routine. And performing Done() on the WaitGroup as each go-routine completes. The main go-routine can finally Wait() on the WaitGroup and the program will exit, safe in the knowledge all cleanup tasks having been completed.


Example:

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

Working playground version.

Here's a solution I was able to come up with using a finalizer on a variable on the main goroutine and then triggering garbage collection:

package main

import (
  "fmt"
  "os"
  "runtime"
  "time"
)

type toBeCollected struct {
  dummy bool // need to have a non-empty struct
}

func main() {
  collectMe := new(toBeCollected)
  runtime.SetFinalizer(collectMe, func(*toBeCollected) {
    fmt.Println("now I'm really done")
    os.Exit(0)
  })
  go func() {
    for range time.NewTicker(time.Millisecond).C {
      runtime.GC()
    }
  }()
  runtime.Goexit()
}

Curiously while I can run it locally, it doesn't work on go playground (see https://play.golang.org/p/bwf7fMu2w76).