Small cleanup in App and Node

app/app.go

@@ -11,7 +11,6 @@ import (
 	"syscall"
 
 	"go.uber.org/zap"
-	"golang.org/x/sync/errgroup"
 
 	"github.com/ava-labs/avalanchego/node"
 	"github.com/ava-labs/avalanchego/utils"
@@ -34,16 +33,16 @@ var _ App = (*app)(nil)
 type App interface {
 	// Start kicks off the application and returns immediately.
 	// Start should only be called once.
-	Start() error
+	Start()
 
 	// Stop notifies the application to exit and returns immediately.
 	// Stop should only be called after [Start].
 	// It is safe to call Stop multiple times.
-	Stop() error
+	Stop()
 
 	// ExitCode should only be called after [Start] returns with no error. It
 	// should block until the application finishes
-	ExitCode() (int, error)
+	ExitCode() int
 }
 
 func New(config nodeconfig.Config) (App, error) {
@@ -89,9 +88,7 @@ func New(config nodeconfig.Config) (App, error) {
 
 func Run(app App) int {
 	// start running the application
-	if err := app.Start(); err != nil {
-		return 1
-	}
+	app.Start()
 
 	// register terminationSignals to kill the application
 	terminationSignals := make(chan os.Signal, 1)
@@ -101,13 +98,15 @@ func Run(app App) int {
 	signal.Notify(stackTraceSignal, syscall.SIGABRT)
 
 	// start up a new go routine to handle attempts to kill the application
-	var eg errgroup.Group
-	eg.Go(func() error {
+	var wg sync.WaitGroup
+	wg.Add(1)
+	go func() {
+		defer wg.Done()
 		for range terminationSignals {
-			return app.Stop()
+			app.Stop()
+			return
 		}
-		return nil
-	})
+	}()
 
 	// start a goroutine to listen on SIGABRT signals,
 	// to print the stack trace to standard error.
@@ -118,7 +117,7 @@ func Run(app App) int {
 	}()
 
 	// wait for the app to exit and get the exit code response
-	exitCode, err := app.ExitCode()
+	exitCode := app.ExitCode()
 
 	// shut down the termination signal go routine
 	signal.Stop(terminationSignals)
@@ -129,9 +128,7 @@ func Run(app App) int {
 	close(stackTraceSignal)
 
 	// if there was an error closing or running the application, report that error
-	if eg.Wait() != nil || err != nil {
-		return 1
-	}
+	wg.Wait()
 
 	// return the exit code that the application reported
 	return exitCode
@@ -148,7 +145,7 @@ type app struct {
 // Start the business logic of the node (as opposed to config reading, etc).
 // Does not block until the node is done. Errors returned from this method
 // are not logged.
-func (a *app) Start() error {
+func (a *app) Start() {
 	// [p.ExitCode] will block until [p.exitWG.Done] is called
 	a.exitWG.Add(1)
 	go func() {
@@ -172,19 +169,17 @@ func (a *app) Start() error {
 			zap.Error(err),
 		)
 	}()
-	return nil
 }
 
 // Stop attempts to shutdown the currently running node. This function will
-// return immediately.
-func (a *app) Stop() error {
+// block until shutdown is complete.
+func (a *app) Stop() {
 	a.node.Shutdown(0)
-	return nil
 }
 
 // ExitCode returns the exit code that the node is reporting. This function
 // blocks until the node has been shut down.
-func (a *app) ExitCode() (int, error) {
+func (a *app) ExitCode() int {
 	a.exitWG.Wait()
-	return a.node.ExitCode(), nil
+	return a.node.ExitCode()
 }

node/node.go

@@ -134,8 +134,6 @@ func New(
 		Config:           config,
 	}
 
-	n.DoneShuttingDown.Add(1)
-
 	pop, err := signer.NewProofOfPossession(n.Config.StakingSigningKey)
 	if err != nil {
 		return nil, fmt.Errorf("problem creating proof of possession: %w", err)
@@ -365,10 +363,6 @@ type Node struct {
 	// Sets the exit code
 	shuttingDownExitCode utils.Atomic[int]
 
-	// Incremented only once on initialization.
-	// Decremented when node is done shutting down.
-	DoneShuttingDown sync.WaitGroup
-
 	// Metrics Registerer
 	MetricsGatherer        metrics.MultiGatherer
 	MeterDBMetricsGatherer metrics.MultiGatherer
@@ -682,7 +676,8 @@ func (n *Node) Dispatch() error {
 			)
 		}
 		// If the API server isn't running, shut down the node.
-		// If node is already shutting down, this does nothing.
+		// If node is already shutting down, this does not tigger shutdown again,
+		// and blocks until the shutdown is complete.
 		n.Shutdown(1)
 	})
 
@@ -716,10 +711,11 @@ func (n *Node) Dispatch() error {
 	}
 
 	// Start P2P connections
-	err := n.Net.Dispatch()
+	retErr := n.Net.Dispatch()
 
 	// If the P2P server isn't running, shut down the node.
-	// If node is already shutting down, this does nothing.
+	// If node is already shutting down, this does not tigger shutdown again,
+	// and blocks until the shutdown is complete.
 	n.Shutdown(1)
 
 	if n.tlsKeyLogWriterCloser != nil {
@@ -732,9 +728,6 @@ func (n *Node) Dispatch() error {
 		}
 	}
 
-	// Wait until the node is done shutting down before returning
-	n.DoneShuttingDown.Wait()
-
 	// Remove the process context file to communicate to an orchestrator
 	// that the node is no longer running.
 	if err := os.Remove(n.Config.ProcessContextFilePath); err != nil && !errors.Is(err, fs.ErrNotExist) {
@@ -744,7 +737,7 @@ func (n *Node) Dispatch() error {
 		)
 	}
 
-	return err
+	return retErr
 }
 
 /*
@@ -1629,10 +1622,9 @@ func (n *Node) initDiskTargeter(
 // Shutdown this node
 // May be called multiple times
 func (n *Node) Shutdown(exitCode int) {
-	if !n.shuttingDown.Get() { // only set the exit code once
+	if !n.shuttingDown.Swap(true) { // only set the exit code once
 		n.shuttingDownExitCode.Set(exitCode)
 	}
-	n.shuttingDown.Set(true)
 	n.shutdownOnce.Do(n.shutdown)
 }
 
@@ -1714,7 +1706,6 @@ func (n *Node) shutdown() {
 		)
 	}
 
-	n.DoneShuttingDown.Done()
 	n.Log.Info("finished node shutdown")
 }
 

utils/atomic.go

@@ -38,6 +38,16 @@ func (a *Atomic[T]) Set(value T) {
 	a.value = value
 }
 
+func (a *Atomic[T]) Swap(value T) T {
+	a.lock.Lock()
+	defer a.lock.Unlock()
+
+	old := a.value
+	a.value = value
+
+	return old
+}
+
 func (a *Atomic[T]) MarshalJSON() ([]byte, error) {
 	a.lock.RLock()
 	defer a.lock.RUnlock()