如何与许多goroutine协调关闭

Question

Say I have a function

type Foo struct {}

func (a *Foo) Bar() {
    // some expensive work - does some calls to redis
}

which gets executed within a goroutine at some point in my app. Lots of these may be executing at any given point. Prior to application termination, I would like to ensure all remaining goroutines have finished their work.

Can I do something like this:

type Foo struct {
    wg sync.WaitGroup
}

func (a *Foo) Close() {
    a.wg.Wait()
}

func (a *Foo) Bar() {
    a.wg.Add(1)
    defer a.wg.Done()

    // some expensive work - does some calls to redis
}

Assuming here that Bar gets executed within a goroutine and many of these may be running at a given time and that Bar should not be called once Close is called and Close is called upon a sigterm or sigint.

Does this make sense?

Usually I would see the Bar function look like this:

func (a *Foo) Bar() {
    a.wg.Add(1)

    go func() {
        defer a.wg.Done()
        // some expensive work - does some calls to redis
    }()
}

Answer 1

Yes, WaitGroup is the right answer. You can use WaitGroup.Add at anytime that the counter is greater than zero, as per doc.

Note that calls with a positive delta that occur when the counter is zero must happen before a Wait. Calls with a negative delta, or calls with a positive delta that start when the counter is greater than zero, may happen at any time. Typically this means the calls to Add should execute before the statement creating the goroutine or other event to be waited for. If a WaitGroup is reused to wait for several independent sets of events, new Add calls must happen after all previous Wait calls have returned. See the WaitGroup example.

But one trick is that, you should always keep the counter greater than zero, before Close is called. That usually means you should call wg.Add in NewFoo (or something like that) and wg.Done in Close. And to prevent multiple calls to Done ruining the wait group, you should wrap Close into sync.Once. You may also want to prevent new Bar() from being called.