From 537379c7dbb263ab5ead06958b890b923cce1794 Mon Sep 17 00:00:00 2001
From: Janez Podhostnik <janez.podhostnik@gmail.com>
Date: Tue, 26 Nov 2024 15:52:35 +0100
Subject: [PATCH] Cleanup cmd

---
 bootstrap/bootstrap.go    |  6 ++++--
 cmd/run/cmd.go            | 38 ++++++++++++++++++++++++++------------
 tests/helpers.go          |  4 +++-
 tests/integration_test.go | 12 +++++++++---
 4 files changed, 42 insertions(+), 18 deletions(-)

diff --git a/bootstrap/bootstrap.go b/bootstrap/bootstrap.go
index 34467fdf8..7f0c6c0bb 100644
--- a/bootstrap/bootstrap.go
+++ b/bootstrap/bootstrap.go
@@ -7,6 +7,8 @@ import (
 	"math"
 	"time"
 
+	"github.com/onflow/flow-go/module/component"
+
 	pebbleDB "github.com/cockroachdb/pebble"
 
 	"github.com/onflow/flow-go-sdk/access"
@@ -557,7 +559,7 @@ func setupStorage(
 // Run will run complete bootstrap of the EVM gateway with all the engines.
 // Run is a blocking call, but it does signal readiness of the service
 // through a channel provided as an argument.
-func Run(ctx context.Context, cfg *config.Config, ready chan struct{}) error {
+func Run(ctx context.Context, cfg *config.Config, ready component.ReadyFunc) error {
 	boot, err := New(cfg)
 	if err != nil {
 		return err
@@ -580,7 +582,7 @@ func Run(ctx context.Context, cfg *config.Config, ready chan struct{}) error {
 	}
 
 	// mark ready
-	close(ready)
+	ready()
 
 	// if context is canceled start shutdown
 	<-ctx.Done()
diff --git a/cmd/run/cmd.go b/cmd/run/cmd.go
index 6b87e04ce..e8cab5b8e 100644
--- a/cmd/run/cmd.go
+++ b/cmd/run/cmd.go
@@ -3,6 +3,7 @@ package run
 import (
 	"context"
 	"encoding/json"
+	"errors"
 	"fmt"
 	"math/big"
 	"os"
@@ -28,27 +29,38 @@ import (
 var Cmd = &cobra.Command{
 	Use:   "run",
 	Short: "Runs the EVM Gateway Node",
-	Run: func(command *cobra.Command, _ []string) {
+	RunE: func(command *cobra.Command, _ []string) error {
+
+		ctx, cancel := context.WithCancel(command.Context())
+		defer cancel()
+
 		// create multi-key account
+		// TODO(JanezP): move to separate command
 		if _, exists := os.LookupEnv("MULTIKEY_MODE"); exists {
 			bootstrap.RunCreateMultiKeyAccount()
-			return
+			return nil
 		}
 
 		if err := parseConfigFromFlags(); err != nil {
-			log.Err(err).Msg("failed to parse flags")
-			os.Exit(1)
+			return fmt.Errorf("failed to parse flags: %w", err)
 		}
 
-		ctx, cancel := context.WithCancel(command.Context())
 		done := make(chan struct{})
 		ready := make(chan struct{})
 		go func() {
 			defer close(done)
-			err := bootstrap.Run(ctx, cfg, ready)
-			if err != nil {
-				log.Err(err).Msg("failed to run bootstrap")
-				cancel()
+			// In case an error happens before ready is called we need to close the ready channel
+			defer close(ready)
+
+			err := bootstrap.Run(
+				ctx,
+				cfg,
+				func() {
+					close(ready)
+				},
+			)
+			if err != nil && !errors.Is(err, context.Canceled) {
+				log.Err(err).Msg("Gateway runtime error")
 			}
 		}()
 
@@ -57,17 +69,19 @@ var Cmd = &cobra.Command{
 		osSig := make(chan os.Signal, 1)
 		signal.Notify(osSig, syscall.SIGINT, syscall.SIGTERM)
 
+		// wait for gateway to exit or for a shutdown signal
 		select {
 		case <-osSig:
 			log.Info().Msg("OS Signal to shutdown received, shutting down")
 			cancel()
 		case <-done:
 			log.Info().Msg("done, shutting down")
-			cancel()
 		}
 
-		log.Info().Msg("OS Signal to shutdown received, shutting down")
-		cancel()
+		// Wait for the gateway to completely stop
+		<-done
+
+		return nil
 	},
 }
 
diff --git a/tests/helpers.go b/tests/helpers.go
index 600922f8d..6378f5c05 100644
--- a/tests/helpers.go
+++ b/tests/helpers.go
@@ -161,7 +161,9 @@ func servicesSetup(t *testing.T) (emulator.Emulator, func()) {
 
 	bootstrapDone := make(chan struct{})
 	go func() {
-		err = bootstrap.Run(ctx, cfg, bootstrapDone)
+		err = bootstrap.Run(ctx, cfg, func() {
+			close(bootstrapDone)
+		})
 		require.NoError(t, err)
 	}()
 
diff --git a/tests/integration_test.go b/tests/integration_test.go
index 14fa40305..af834bf1d 100644
--- a/tests/integration_test.go
+++ b/tests/integration_test.go
@@ -83,7 +83,9 @@ func Test_ConcurrentTransactionSubmission(t *testing.T) {
 
 	ready := make(chan struct{})
 	go func() {
-		err := bootstrap.Run(ctx, cfg, ready)
+		err := bootstrap.Run(ctx, cfg, func() {
+			close(ready)
+		})
 		require.NoError(t, err)
 	}()
 
@@ -181,7 +183,9 @@ func Test_EthClientTest(t *testing.T) {
 
 	ready := make(chan struct{})
 	go func() {
-		err := bootstrap.Run(ctx, cfg, ready)
+		err := bootstrap.Run(ctx, cfg, func() {
+			close(ready)
+		})
 		require.NoError(t, err)
 	}()
 
@@ -288,7 +292,9 @@ func Test_CloudKMSConcurrentTransactionSubmission(t *testing.T) {
 
 	ready := make(chan struct{})
 	go func() {
-		err := bootstrap.Run(ctx, cfg, ready)
+		err := bootstrap.Run(ctx, cfg, func() {
+			close(ready)
+		})
 		require.NoError(t, err)
 	}()