Repository: uber-go/fx (go)

Issue:
onStop executed many times
### Discussed in https://github.com/uber-go/fx/discussions/929

<div type='discussions-op-text'>

<sup>Originally posted by **coolyrat** August 22, 2022</sup>
```go
package main

import (
	"context"
	"fmt"
	"log"
	"os"
	"os/signal"
	"syscall"
	"time"

	"go.uber.org/fx"
)

func main() {
	server := fx.New(
		fx.Invoke(func(lc fx.Lifecycle) {
			fmt.Println("first invoke...")
			lc.Append(fx.Hook{
				OnStart: func(ctx context.Context) error {
					fmt.Println("#1 start")
					return nil
				},
				OnStop: func(ctx context.Context) error {
					fmt.Println("#1 stop")
					return nil
				},
			})
		}),

		fx.Invoke(func(lc fx.Lifecycle) {
			fmt.Println("last invoke...")
			lc.Append(fx.Hook{
				OnStart: func(ctx context.Context) error {
					fmt.Println("#2 start")
					return nil
				},
				OnStop: func(ctx context.Context) error {
					fmt.Println("#2 stop")
					return nil
				},
			})
		}),
	)

	go func() {
		server.Run()
	}()

	kill := make(chan os.Signal, 1)
	signal.Notify(kill, syscall.SIGINT, syscall.SIGTERM, syscall.SIGKILL, syscall.SIGQUIT)
	<-kill
	fmt.Println("killed")

	ctx, _ := context.WithTimeout(context.Background(), 10*time.Second)
	if err := server.Stop(ctx); err != nil {
		log.Printf("error stopping gracefully")
	}
}

```

When I give the stop signal, the log
```
...
killed
#2 stop
#1 stop
#2 stop
```

And I am expected to be 
```
...
killed
#2 stop
#1 stop
```

Is there anything I'm doing wrong?</div>

Pull request description:
Ensure OnStart/OnStop hooks can only be called once
It is possible for a user to erroneously call app.Run() then
follow up by calling app.Stop(). In such a case, it is possible
for the app.Stop() method to be called by two goroutines concurrently,
resulting in a race.

This adds a state in the App to keep track of whether Stop() has been
invoked so that such a race can be prevented.

Fix #930
Internal Ref: GO-1606

Patch:
diff --git a/app.go b/app.go
index 8d0132c..67244b9 100644
--- a/app.go
+++ b/app.go
@@ -300,6 +300,10 @@ type App struct {
 	dones       []chan os.Signal
 	shutdownSig os.Signal
 
+	// Used to make sure Start/Stop is called only once.
+	runStart sync.Once
+	runStop  sync.Once
+
 	osExit func(code int) // os.Exit override; used for testing only
 }
 
@@ -658,21 +662,25 @@ var (
 // Note that Start short-circuits immediately if the New constructor
 // encountered any errors in application initialization.
 func (app *App) Start(ctx context.Context) (err error) {
-	defer func() {
-		app.log.LogEvent(&fxevent.Started{Err: err})
-	}()
+	app.runStart.Do(func() {
+		defer func() {
+			app.log.LogEvent(&fxevent.Started{Err: err})
+		}()
 
-	if app.err != nil {
-		// Some provides failed, short-circuit immediately.
-		return app.err
-	}
+		if app.err != nil {
+			// Some provides failed, short-circuit immediately.
+			err = app.err
+			return
+		}
 
-	return withTimeout(ctx, &withTimeoutParams{
-		hook:      _onStartHook,
-		callback:  app.start,
-		lifecycle: app.lifecycle,
-		log:       app.log,
+		err = withTimeout(ctx, &withTimeoutParams{
+			hook:      _onStartHook,
+			callback:  app.start,
+			lifecycle: app.lifecycle,
+			log:       app.log,
+		})
 	})
+	return
 }
 
 func (app *App) start(ctx context.Context) error {
@@ -700,16 +708,20 @@ func (app *App) start(ctx context.Context) error {
 // called are executed. However, all those hooks are executed, even if some
 // fail.
 func (app *App) Stop(ctx context.Context) (err error) {
-	defer func() {
-		app.log.LogEvent(&fxevent.Stopped{Err: err})
-	}()
+	app.runStop.Do(func() {
+		// Protect the Stop hooks from being called multiple times.
+		defer func() {
+			app.log.LogEvent(&fxevent.Stopped{Err: err})
+		}()
 
-	return withTimeout(ctx, &withTimeoutParams{
-		hook:      _onStopHook,
-		callback:  app.lifecycle.Stop,
-		lifecycle: app.lifecycle,
-		log:       app.log,
+		err = withTimeout(ctx, &withTimeoutParams{
+			hook:      _onStopHook,
+			callback:  app.lifecycle.Stop,
+			lifecycle: app.lifecycle,
+			log:       app.log,
+		})
 	})
+	return
 }
 
 // Done returns a channel of signals to block on after starting the


Test patch:
diff --git a/app_test.go b/app_test.go
index 0011272..9718041 100644
--- a/app_test.go
+++ b/app_test.go
@@ -31,6 +31,7 @@ import (
 	"reflect"
 	"runtime"
 	"strings"
+	"sync"
 	"testing"
 	"time"
 
@@ -1281,6 +1282,47 @@ func TestAppStart(t *testing.T) {
 		err := app.Start(context.Background()).Error()
 		assert.Contains(t, err, "OnStart hook added by go.uber.org/fx_test.TestAppStart.func10.1 failed: goroutine exited without returning")
 	})
+
+	t.Run("Start/Stop should be called exactly once only.", func(t *testing.T) {
+		t.Parallel()
+		startCalled := 0
+		stopCalled := 0
+		app := fxtest.New(t,
+			Provide(Annotate(func() int { return 0 },
+				OnStart(func(context.Context) error {
+					startCalled += 1
+					return nil
+				}),
+				OnStop(func(context.Context) error {
+					stopCalled += 1
+					return nil
+				})),
+			),
+			Invoke(func(i int) {
+				assert.Equal(t, 0, i)
+			}),
+		)
+		var wg sync.WaitGroup
+		for i := 0; i < 10; i++ {
+			wg.Add(1)
+			go func() {
+				defer wg.Done()
+				app.Start(context.Background())
+			}()
+		}
+		wg.Wait()
+		assert.Equal(t, 1, startCalled)
+		for i := 0; i < 10; i++ {
+			wg.Add(1)
+			go func() {
+				defer wg.Done()
+				app.Stop(context.Background())
+			}()
+		}
+		wg.Wait()
+		assert.Equal(t, 1, stopCalled)
+	})
+
 }
 
 func TestAppStop(t *testing.T) {