gogo-uber-fx

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


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:

Concerns and questions:


Solution

  • 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()
    }