What is the right way to implement graceful shutdown of background processes in Uber FX?

Question

Let's say I have a service inside my Uber FX Application that should do some background activities like polling an external API. I can run background tasks by firing goroutines, but what is the right way to stop them?

As a possible implementation let's consider the following example:

package main

import (
    "context"
    "log"
    "sync"
    "time"

    "go.uber.org/fx"
)

type AwesomeService struct {
    // context to use for background processes
    bg context.Context
    // to trigger background processes stopping
    cancel context.CancelFunc
    // to wait for background processes to gracefully finish
    wg *sync.WaitGroup
}

func New(lc fx.Lifecycle) *AwesomeService {
    bg, cancel := context.WithCancel(context.Background())
    service := &AwesomeService{
        bg:     bg,
        cancel: cancel,
        wg:     new(sync.WaitGroup),
    }

    lc.Append(fx.Hook{
        OnStart: service.start,
        OnStop:  service.stop,
    })
    return service
}

func (s *AwesomeService) start(_ context.Context) error {
    s.runBackgroundProcess()
    log.Println("Start done")
    return nil
}

func (s *AwesomeService) stop(_ context.Context) error {
    s.cancel()
    s.wg.Wait()
    log.Println("Stop done")
    return nil
}

// runBackgroundProcess does some work till context is done.
func (s *AwesomeService) runBackgroundProcess() {
    s.wg.Add(1)
    go func() {
        defer s.wg.Done()
        for {
            select {
            case <-s.bg.Done():
                return
            case <-time.After(1 * time.Second):
                log.Println("Working...")
            }
        }
    }()
}

func main() {
    fx.New(
        fx.Provide(New),
        fx.Invoke(func(*AwesomeService) {}),
    ).Run()
}

Some notices:

• The service is wired to the application lifecycle by using fx.Lifecycle hooks.
• I can't rely on and use the context in the OnStart/OnStop methods because they are different contexts and correspond to the start/stop activities, not the app lifecycle context.

Concerns and questions:

• The given example is quite heavy in terms of tracking bg tasks. Also, storing context in a struct is a kind of antipattern. Is there any way to simplify it?
• Should I wait to finish goroutines in case there are no resources to release?

Answer 1

Using context is just fine in my opinion, but you could alternatively communicate a shutdown signal to whatever Go routines you'd like via a channel. See example code below.

And yes, you should also wait for the wait group count to return to zero before completely shutting down the app. So you would first shutdown the channel, and then wait on the wait group.

package main

import (
    "context"
    "log"
    "sync"
    "time"

    "go.uber.org/fx"
)

type AwesomeService struct {
    // channel to shutdown background processes
    shutdown chan struct{}
    // to wait for background processes to gracefully finish
    wg *sync.WaitGroup
}

func New(lc fx.Lifecycle) *AwesomeService {
    service := &AwesomeService{
        shutdown: make(chan struct{}),
        wg:     new(sync.WaitGroup),
    }

    lc.Append(fx.Hook{
        OnStart: service.start,
        OnStop:  service.stop,
    })
    return service
}

func (s *AwesomeService) start(_ context.Context) error {
    s.runBackgroundProcess()
    log.Println("Start done")
    return nil
}

func (s *AwesomeService) stop(_ context.Context) error {
    close(s.shutdown)
    s.wg.Wait()
    log.Println("Stop done")
    return nil
}

// runBackgroundProcess does some work till context is done.
func (s *AwesomeService) runBackgroundProcess() {
    s.wg.Add(1)
    go func() {
        defer s.wg.Done()
        for {
            select {
            case <-s.shutdown:
                return
            case <-time.After(1 * time.Second):
                log.Println("Working...")
            }
        }
    }()
}

func main() {
    fx.New(
        fx.Provide(New),
        fx.Invoke(func(*AwesomeService) {}),
    ).Run()
}