From cc87b094398dbab9a6bf89cbe92157e4f5f3ceab Mon Sep 17 00:00:00 2001
From: Manjari Akella <manjari@uber.com>
Date: Wed, 10 Nov 2021 12:22:37 -0800
Subject: [PATCH] Support app shutdown before app.Done() is called (#805)

* Record app shutdown and offer the signal on app.Done()

* Add data race test

* Add lock for shutdown signal

* Pull test into function and block go routines on ready channel

* lint

* Return channel early if shutdown signal received

* Remove comment

* Remove fmt qualifier for assert error message

* Clarify comment

* Assert no error on shutdown for TestDataRace

* Apply suggestions from code review

* go fmt

Co-authored-by: Abhinav Gupta <abg@uber.com>
---
 app.go           | 17 ++++++++++----
 shutdown.go      |  6 +++--
 shutdown_test.go | 61 ++++++++++++++++++++++++++++++++++++++++++++++++
 3 files changed, 78 insertions(+), 6 deletions(-)

diff --git a/app.go b/app.go
index 2909c5fc3..28fa3bf08 100644
--- a/app.go
+++ b/app.go
@@ -379,8 +379,9 @@ type App struct {
 	errorHooks []ErrorHandler
 	validate   bool
 	// Used to signal shutdowns.
-	donesMu sync.RWMutex
-	dones   []chan os.Signal
+	donesMu     sync.Mutex // guards dones and shutdownSig
+	dones       []chan os.Signal
+	shutdownSig os.Signal
 
 	osExit func(code int) // os.Exit override; used for testing only
 }
@@ -787,11 +788,19 @@ func (app *App) Stop(ctx context.Context) (err error) {
 // using the Shutdown functionality (see the Shutdowner documentation for details).
 func (app *App) Done() <-chan os.Signal {
 	c := make(chan os.Signal, 1)
-	signal.Notify(c, _sigINT, _sigTERM)
 
 	app.donesMu.Lock()
+	defer app.donesMu.Unlock()
+	// If shutdown signal has been received already
+	// send it and return. If not, wait for user to send a termination
+	// signal.
+	if app.shutdownSig != nil {
+		c <- app.shutdownSig
+		return c
+	}
+
+	signal.Notify(c, _sigINT, _sigTERM)
 	app.dones = append(app.dones, c)
-	app.donesMu.Unlock()
 	return c
 }
 
diff --git a/shutdown.go b/shutdown.go
index be751c132..d5b8488c0 100644
--- a/shutdown.go
+++ b/shutdown.go
@@ -57,8 +57,10 @@ func (app *App) shutdowner() Shutdowner {
 }
 
 func (app *App) broadcastSignal(signal os.Signal) error {
-	app.donesMu.RLock()
-	defer app.donesMu.RUnlock()
+	app.donesMu.Lock()
+	defer app.donesMu.Unlock()
+
+	app.shutdownSig = signal
 
 	var unsent int
 	for _, done := range app.dones {
diff --git a/shutdown_test.go b/shutdown_test.go
index 63d519500..b6af93f13 100644
--- a/shutdown_test.go
+++ b/shutdown_test.go
@@ -21,9 +21,12 @@
 package fx_test
 
 import (
+	"context"
+	"sync"
 	"testing"
 
 	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
 	"go.uber.org/fx"
 	"go.uber.org/fx/fxtest"
 )
@@ -65,4 +68,62 @@ func TestShutdown(t *testing.T) {
 			"unexpected error returned when shutdown is called with a blocked channel")
 		assert.NotNil(t, <-done, "done channel did not receive signal")
 	})
+
+	t.Run("shutdown app before calling Done()", func(t *testing.T) {
+		t.Parallel()
+
+		var s fx.Shutdowner
+		app := fxtest.New(
+			t,
+			fx.Populate(&s),
+		)
+
+		require.NoError(t, app.Start(context.Background()), "error starting app")
+		assert.NoError(t, s.Shutdown(), "error in app shutdown")
+		done1, done2 := app.Done(), app.Done()
+		defer app.Stop(context.Background())
+		// Receiving on done1 and done2 will deadlock in the event that app.Done()
+		// doesn't work as expected.
+		assert.NotNil(t, <-done1, "done channel 1 did not receive signal")
+		assert.NotNil(t, <-done2, "done channel 2 did not receive signal")
+	})
+}
+
+func TestDataRace(t *testing.T) {
+	t.Parallel()
+
+	var s fx.Shutdowner
+	app := fxtest.New(
+		t,
+		fx.Populate(&s),
+	)
+	require.NoError(t, app.Start(context.Background()), "error starting app")
+
+	const N = 50
+	ready := make(chan struct{}) // used to orchestrate goroutines for Done() and ShutdownOption()
+	var wg sync.WaitGroup        // tracks and waits for all goroutines
+
+	// Spawn N goroutines, each of which call app.Done() and assert
+	// the signal received.
+	wg.Add(N)
+	for i := 0; i < N; i++ {
+		i := i
+		go func() {
+			defer wg.Done()
+			<-ready
+			done := app.Done()
+			assert.NotNil(t, <-done, "done channel %v did not receive signal", i)
+		}()
+	}
+
+	// call Shutdown()
+	wg.Add(1)
+	go func() {
+		<-ready
+		defer wg.Done()
+		assert.NoError(t, s.Shutdown(), "error in app shutdown")
+	}()
+
+	close(ready)
+	wg.Wait()
 }