Repository: uber-go/fx
Pull request #1076 (2023-04-29 19:07:43)

Shutdowner: Support calling from fx.Invoke
Stacked on top of:

- https://github.com/uber-go/fx/pull/1081
- https://github.com/uber-go/fx/pull/1082

However, since I can't push branches directly to this repository,
this PR shows commits from all PRs.

---

App.Start nils out the "last" signal recorded by signalReceivers,
which it otherwise broadcasts to waiters if it was already received.
This is unnecessary especially because there's a discrepancy in behavior
of using App.Start vs App.Run when shutting down from fx.Invoke.

Given a program that calls Shutdown from fx.Invoke,
when we do:

    app := fx.New(...)

The shutdowner has already sent the signal, and signalReceivers has
already recorded it.
At that point, whether we call App.Start or App.Run changes behavior:

- If we call App.Run, that calls App.Done (or App.Wait after #1075),
  which gives it back a channel that already has the signal filled in.
  It then calls App.Start, waits on the channel--which returns
  immediately--and then calls App.Stop.
- If we call App.Start and App.Wait, on the other hand,
  Start will clear the signal recorded in signalReceivers,
  and then App.Wait will build a channel that will block indefinitely
  because Shutdowner.Shutdown will not be called again.

So even though App.Run() and App.Start()+App.Wait() are meant to be
equivalent, this causes a serious discrepancy in behavior.
It makes sense to resolve this by supporting Shutdown from Invoke.

Refs #1074


codecov[bot]: ## [Codecov](https://codecov.io/gh/uber-go/fx/pull/1076?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=uber-go) Report
> Merging [#1076](https://codecov.io/gh/uber-go/fx/pull/1076?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=uber-go) (fd94f3f) into [master](https://codecov.io/gh/uber-go/fx/commit/8f5597f7f53f328e1a9e922c521df590b9b8be2f?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=uber-go) (8f5597f) will **decrease** coverage by `0.92%`.
> The diff coverage is `0.00%`.

```diff
@@            Coverage Diff             @@
##           master    #1076      +/-   ##
==========================================
- Coverage   98.53%   97.62%   -0.92%     
==========================================
  Files          40       35       -5     
  Lines        3003     2862     -141     
==========================================
- Hits         2959     2794     -165     
- Misses         37       61      +24     
  Partials        7        7              
```


| [Impacted Files](https://codecov.io/gh/uber-go/fx/pull/1076?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=uber-go) | Coverage Δ | |
|---|---|---|
| [internal/e2e/shutdowner\_wait\_exitcode/main.go](https://codecov.io/gh/uber-go/fx/pull/1076?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=uber-go#diff-aW50ZXJuYWwvZTJlL3NodXRkb3duZXJfd2FpdF9leGl0Y29kZS9tYWluLmdv) | `0.00% <0.00%> (ø)` | |
| [shutdown.go](https://codecov.io/gh/uber-go/fx/pull/1076?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=uber-go#diff-c2h1dGRvd24uZ28=) | `100.00% <ø> (ø)` | |
| [signal.go](https://codecov.io/gh/uber-go/fx/pull/1076?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=uber-go#diff-c2lnbmFsLmdv) | `100.00% <ø> (ø)` | |

... and [7 files with indirect coverage changes](https://codecov.io/gh/uber-go/fx/pull/1076/indirect-changes?src=pr&el=tree-more&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=uber-go)

:mega: We’re building smart automated test selection to slash your CI/CD build times. [Learn more](https://about.codecov.io/iterative-testing/?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=uber-go)

abhinav: Coverage issue will be resolved by https://github.com/uber-go/fx/pull/1078
JacobOaks: Spent a little bit looking at this and I think it makes sense.

Have we considered switching the order of execution of `Start` vs. `Wait` in the case of calling `Run` to more closely match the `Start` -> `Wait` -> `Stop` flow? For example, `run` can take a channel-generating function rather than a channel itself, and we can pass `Wait`:
```go
func (app *App) Run() {
    if code := app.run(app.Wait); code != 0 {
        // ...
    }
}

func (app *App) run(doneFn func() <-chan ShutdownSignal) (exitCode int) {
// ...
if err := app.Start(); err != nil { /* ... */ }
sig := <-doneFn()
// ...
}
```
This would cause shutdown from invokes to no longer work when using `Run` making it a breaking change by itself, but in combination with this PR it shouldn't be breaking, and it might be a good idea to prevent future changes/features from introducing additional disparities between `app.Run` and `app.Start` -> `app.Wait` -> `app.Stop`. Just a thought.
abhinav: That's a good idea! I'll make the switch.
abhinav: Actually, that's a problem. Looks like receivers.Stop is called from both, App.Stop and Shutdowner.Shutdown. I'll see if that can be fixed.
abhinav: Okay, I fixed it but then the ShutdownTimeout option is meaningless. I'll dig a little.
abhinav: Okay, so @JacobOaks @sywhang, Shutdown has a ShutdownTimeout option that doesn't make sense.
Shutdown can't be a blocking operation.
signalReceivers has the operations: Start, Stop, Broadcast, Wait (and its old version, Done).
Start starts a background goroutine to relay information, and Stop stops it.
It makes no sense for Shutdowner.Shutdown to call Stop if Shutdown is just to unblock a `<-Wait()` and a corresponding Stop.

Let me see if I can wrangle this apart a bit.
abhinav: A good signal that the ShutdownTimeout option makes little sense is this in the documentation:

```
 // [..] As the Shutdown method will block while waiting for a signal receiver relay
 // goroutine to stop.
```

That's an internal implementation detail.
abhinav: Updated. Check PRs stacked below this for more context.

Patch:
diff --git a/app.go b/app.go
index bc3ed09..14e050b 100644
--- a/app.go
+++ b/app.go
@@ -574,12 +574,12 @@ func (app *App) Run() {
 	// Historically, we do not os.Exit(0) even though most applications
 	// cede control to Fx with they call app.Run. To avoid a breaking
 	// change, never os.Exit for success.
-	if code := app.run(app.Wait()); code != 0 {
+	if code := app.run(app.Wait); code != 0 {
 		app.exit(code)
 	}
 }
 
-func (app *App) run(done <-chan ShutdownSignal) (exitCode int) {
+func (app *App) run(done func() <-chan ShutdownSignal) (exitCode int) {
 	startCtx, cancel := app.clock.WithTimeout(context.Background(), app.StartTimeout())
 	defer cancel()
 
@@ -587,7 +587,7 @@ func (app *App) run(done <-chan ShutdownSignal) (exitCode int) {
 		return 1
 	}
 
-	sig := <-done
+	sig := <-done()
 	app.log().LogEvent(&fxevent.Stopping{Signal: sig.Signal})
 	exitCode = sig.ExitCode
 
diff --git a/shutdown.go b/shutdown.go
index eb67caa..b5fda5c 100644
--- a/shutdown.go
+++ b/shutdown.go
@@ -78,8 +78,6 @@ type shutdowner struct {
 // Shutdown broadcasts a signal to all of the application's Done channels
 // and begins the Stop process. Applications can be shut down only after they
 // have finished starting up.
-// In practice this means Shutdowner.Shutdown should not be called from an
-// fx.Invoke, but from a fx.Lifecycle.OnStart hook.
 func (s *shutdowner) Shutdown(opts ...ShutdownOption) error {
 	for _, opt := range opts {
 		opt.apply(s)
diff --git a/signal.go b/signal.go
index b804abe..1dbfd84 100644
--- a/signal.go
+++ b/signal.go
@@ -109,7 +109,6 @@ func (recv *signalReceivers) Start(ctx context.Context) {
 		return
 	}
 
-	recv.last = nil
 	recv.finished = make(chan struct{}, 1)
 	recv.shutdown = make(chan struct{}, 1)
 	recv.notify(recv.signals, os.Interrupt, _sigINT, _sigTERM)
@@ -135,6 +134,7 @@ func (recv *signalReceivers) Stop(ctx context.Context) error {
 		close(recv.finished)
 		recv.shutdown = nil
 		recv.finished = nil
+		recv.last = nil
 		return nil
 	}
 }


Test patch:
diff --git a/app_internal_test.go b/app_internal_test.go
index a377dd7..efdc2f3 100644
--- a/app_internal_test.go
+++ b/app_internal_test.go
@@ -47,7 +47,7 @@ func TestAppRun(t *testing.T) {
 	wg.Add(1)
 	go func() {
 		defer wg.Done()
-		app.run(done)
+		app.run(func() <-chan ShutdownSignal { return done })
 	}()
 
 	done <- ShutdownSignal{Signal: _sigINT}
diff --git a/internal/e2e/shutdowner_wait_exitcode/main.go b/internal/e2e/shutdowner_wait_exitcode/main.go
new file mode 100644
index 0000000..f89bfce
--- /dev/null
+++ b/internal/e2e/shutdowner_wait_exitcode/main.go
@@ -0,0 +1,54 @@
+// Copyright (c) 2023 Uber Technologies, Inc.
+//
+// Permission is hereby granted, free of charge, to any person obtaining a copy
+// of this software and associated documentation files (the "Software"), to deal
+// in the Software without restriction, including without limitation the rights
+// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+// copies of the Software, and to permit persons to whom the Software is
+// furnished to do so, subject to the following conditions:
+//
+// The above copyright notice and this permission notice shall be included in
+// all copies or substantial portions of the Software.
+//
+// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
+// THE SOFTWARE.
+
+package main
+
+import (
+	"context"
+	"log"
+	"os"
+	"time"
+
+	"go.uber.org/fx"
+)
+
+func main() {
+	app := fx.New(
+		fx.Invoke(func(shutdowner fx.Shutdowner) error {
+			shutdowner.Shutdown(fx.ExitCode(20))
+			return nil
+		}),
+	)
+
+	ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
+	defer cancel()
+
+	if err := app.Start(ctx); err != nil {
+		log.Fatal(err)
+	}
+
+	sig := <-app.Wait()
+
+	if err := app.Stop(ctx); err != nil {
+		log.Fatal(err)
+	}
+
+	os.Exit(sig.ExitCode)
+}
diff --git a/internal/e2e/shutdowner_wait_exitcode/main_test.go b/internal/e2e/shutdowner_wait_exitcode/main_test.go
new file mode 100644
index 0000000..849be9f
--- /dev/null
+++ b/internal/e2e/shutdowner_wait_exitcode/main_test.go
@@ -0,0 +1,72 @@
+// Copyright (c) 2023 Uber Technologies, Inc.
+//
+// Permission is hereby granted, free of charge, to any person obtaining a copy
+// of this software and associated documentation files (the "Software"), to deal
+// in the Software without restriction, including without limitation the rights
+// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+// copies of the Software, and to permit persons to whom the Software is
+// furnished to do so, subject to the following conditions:
+//
+// The above copyright notice and this permission notice shall be included in
+// all copies or substantial portions of the Software.
+//
+// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
+// THE SOFTWARE.
+
+package main
+
+import (
+	"os"
+	"os/exec"
+	"path/filepath"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+	"go.uber.org/fx/internal/testutil"
+)
+
+// Hijacks the test binary so that the test can run main() as a subprocess
+// instead of trying to compile the program and run it directly.
+func TestMain(m *testing.M) {
+	// If the test binary is named "app", then we're running as a subprocess.
+	// Otherwise, run the tests.
+	switch filepath.Base(os.Args[0]) {
+	case "app":
+		main()
+		os.Exit(0)
+	default:
+		os.Exit(m.Run())
+	}
+}
+
+// Verifies that an Fx program running with Run
+// exits with the exit code passed to Shutdowner.
+//
+// Regression test for https://github.com/uber-go/fx/issues/1074.
+func TestShutdownExitCode(t *testing.T) {
+	exe, err := os.Executable()
+	require.NoError(t, err)
+
+	out := testutil.WriteSyncer{T: t}
+
+	// Run the test binary with the name 'app' so that it runs main().
+	cmd := exec.Command(exe)
+	cmd.Args[0] = "app"
+	cmd.Stdout = &out
+	cmd.Stderr = &out
+
+	// The program should exit with code 20.
+	err = cmd.Run()
+	require.Error(t, err)
+
+	var exitErr *exec.ExitError
+	require.ErrorAs(t, err, &exitErr)
+
+	assert.Equal(t, 20, exitErr.ExitCode())
+}
diff --git a/shutdown_test.go b/shutdown_test.go
index 7565a11..4995619 100644
--- a/shutdown_test.go
+++ b/shutdown_test.go
@@ -128,6 +128,46 @@ func TestShutdown(t *testing.T) {
 
 		assert.NoError(t, s.Shutdown(fx.ExitCode(2), fx.ShutdownTimeout(time.Second)))
 	})
+
+	t.Run("from invoke", func(t *testing.T) {
+		t.Parallel()
+
+		app := fxtest.New(
+			t,
+			fx.Invoke(func(s fx.Shutdowner) {
+				s.Shutdown()
+			}),
+		)
+
+		ctx, cancel := context.WithTimeout(context.Background(), time.Second)
+		defer cancel()
+
+		require.NoError(t, app.Start(ctx), "error starting app")
+		defer app.Stop(ctx)
+		select {
+		case <-ctx.Done():
+			assert.Fail(t, "app did not shutdown in time")
+		case <-app.Wait():
+			// success
+		}
+	})
+
+	t.Run("many times", func(t *testing.T) {
+		t.Parallel()
+
+		var shutdowner fx.Shutdowner
+		app := fxtest.New(
+			t,
+			fx.Populate(&shutdowner),
+		)
+
+		for i := 0; i < 10; i++ {
+			app.RequireStart()
+			shutdowner.Shutdown(fx.ExitCode(i))
+			assert.Equal(t, i, (<-app.Wait()).ExitCode, "run %d", i)
+			app.RequireStop()
+		}
+	})
 }
 
 func TestDataRace(t *testing.T) {