feat: handle ctrl+c in reporting (#2080)

Author: emosbaugh

cmd/installer/cli/install.go

@@ -82,6 +82,8 @@ type InstallCmdFlags struct {
 func InstallCmd(ctx context.Context, name string) *cobra.Command {
 	var flags InstallCmdFlags
 
+	ctx, cancel := context.WithCancel(ctx)
+
 	cmd := &cobra.Command{
 		Use:   "install",
 		Short: fmt.Sprintf("Install %s", name),
@@ -94,13 +96,20 @@ func InstallCmd(ctx context.Context, name string) *cobra.Command {
 		},
 		PostRun: func(cmd *cobra.Command, args []string) {
 			runtimeconfig.Cleanup()
+			cancel() // Cancel context when command completes
 		},
 		RunE: func(cmd *cobra.Command, args []string) error {
 			clusterID := metrics.ClusterID()
 			metricsReporter := NewInstallReporter(
 				replicatedAppURL(), flags.license.Spec.LicenseID, clusterID, cmd.CalledAs(),
 			)
 			metricsReporter.ReportInstallationStarted(ctx)
+
+			// Setup signal handler with the metrics reporter cleanup function
+			signalHandler(ctx, cancel, func(ctx context.Context, err error) {
+				metricsReporter.ReportInstallationFailed(ctx, err)
+			})
+
 			if err := runInstall(cmd.Context(), name, flags, metricsReporter); err != nil {
 				metricsReporter.ReportInstallationFailed(ctx, err)
 				return err

cmd/installer/cli/join.go

@@ -46,6 +46,8 @@ type JoinCmdFlags struct {
 func JoinCmd(ctx context.Context, name string) *cobra.Command {
 	var flags JoinCmdFlags
 
+	ctx, cancel := context.WithCancel(ctx)
+
 	cmd := &cobra.Command{
 		Use:   "join <url> <token>",
 		Short: fmt.Sprintf("Join %s", name),
@@ -61,6 +63,7 @@ func JoinCmd(ctx context.Context, name string) *cobra.Command {
 		},
 		PostRun: func(cmd *cobra.Command, args []string) {
 			runtimeconfig.Cleanup()
+			cancel() // Cancel context when command completes
 		},
 		RunE: func(cmd *cobra.Command, args []string) error {
 			logrus.Debugf("fetching join token remotely")
@@ -70,6 +73,12 @@ func JoinCmd(ctx context.Context, name string) *cobra.Command {
 			}
 			metricsReporter := NewJoinReporter(jcmd.InstallationSpec.MetricsBaseURL, jcmd.ClusterID, cmd.CalledAs())
 			metricsReporter.ReportJoinStarted(ctx)
+
+			// Setup signal handler with the metrics reporter cleanup function
+			signalHandler(ctx, cancel, func(ctx context.Context, err error) {
+				metricsReporter.ReportJoinFailed(ctx, err)
+			})
+
 			if err := runJoin(cmd.Context(), name, flags, jcmd, metricsReporter); err != nil {
 				metricsReporter.ReportJoinFailed(ctx, err)
 				return err

cmd/installer/cli/signal.go

@@ -0,0 +1,41 @@
+package cli
+
+import (
+	"context"
+	"fmt"
+	"os"
+	"os/signal"
+	"syscall"
+
+	"github.com/sirupsen/logrus"
+)
+
+// osExit is a variable to make testing easier
+var osExit = os.Exit
+
+// signalHandler sets up handling for signals to ensure cleanup functions are called.
+func signalHandler(ctx context.Context, cancel context.CancelFunc, cleanupFuncs ...func(context.Context, error)) {
+	sigChan := make(chan os.Signal, 1)
+	signal.Notify(sigChan, syscall.SIGINT, syscall.SIGTERM)
+
+	go func() {
+		select {
+		case sig := <-sigChan:
+			logrus.Debugf("Received signal: %v", sig)
+			err := fmt.Errorf("command interrupted by signal: %v", sig)
+
+			for _, cleanup := range cleanupFuncs {
+				cleanup(ctx, err)
+			}
+
+			// Cancel the context after cleanup functions run
+			cancel()
+
+			// Exit with non-zero status
+			osExit(1)
+		case <-ctx.Done():
+			// Context was canceled elsewhere, do nothing
+			return
+		}
+	}()
+}

cmd/installer/cli/signal_test.go

@@ -0,0 +1,104 @@
+package cli
+
+import (
+	"context"
+	"os"
+	"sync"
+	"syscall"
+	"testing"
+	"time"
+
+	"github.com/stretchr/testify/assert"
+)
+
+func Test_signalHandler_Signal(t *testing.T) {
+	// Create a context with cancel function
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
+
+	// Create a waitgroup to synchronize the test
+	var wg sync.WaitGroup
+	wg.Add(1)
+
+	// Track if cleanup function was called
+	cleanupCalled := false
+	cleanupError := ""
+
+	// Mock cleanup function
+	cleanup := func(ctx context.Context, err error) {
+		cleanupCalled = true
+		if err != nil {
+			cleanupError = err.Error()
+		}
+		wg.Done()
+	}
+
+	// Save original os.Exit and restore after test
+	originalOsExit := osExit
+	defer func() { osExit = originalOsExit }()
+
+	exitCode := 0
+	osExit = func(code int) {
+		exitCode = code
+		// Instead of exiting, just cancel the context
+		cancel()
+	}
+
+	// Set up the signal handler
+	signalHandler(ctx, cancel, cleanup)
+
+	// Send a signal to trigger the handler
+	p, err := os.FindProcess(os.Getpid())
+	if err != nil {
+		t.Fatalf("Failed to find process: %v", err)
+	}
+
+	// Send SIGINT to trigger the handler
+	err = p.Signal(syscall.SIGINT)
+	if err != nil {
+		t.Fatalf("Failed to send signal: %v", err)
+	}
+
+	// Wait for cleanup to be called with a timeout
+	waitCh := make(chan struct{})
+	go func() {
+		wg.Wait()
+		close(waitCh)
+	}()
+
+	select {
+	case <-waitCh:
+		// Success - cleanup was called
+	case <-time.After(1 * time.Second):
+		t.Fatal("Timed out waiting for cleanup function to be called")
+	}
+
+	// Verify cleanup was called with the expected error
+	assert.True(t, cleanupCalled, "Cleanup function should have been called")
+	assert.Contains(t, cleanupError, "command interrupted by signal: interrupt")
+	assert.Equal(t, 1, exitCode, "Exit code should be 1")
+}
+
+func Test_signalHandler_ContextDone(t *testing.T) {
+	// Create a context with cancel function
+	ctx, cancel := context.WithCancel(context.Background())
+
+	// We expect cleanup NOT to be called when context is cancelled
+	cleanupCalled := false
+
+	cleanup := func(ctx context.Context, err error) {
+		cleanupCalled = true
+	}
+
+	// Set up the signal handler
+	signalHandler(ctx, cancel, cleanup)
+
+	// Cancel the context
+	cancel()
+
+	// Give some time for any handlers to run
+	time.Sleep(100 * time.Millisecond)
+
+	// Verify cleanup was NOT called
+	assert.False(t, cleanupCalled, "Cleanup function should not have been called when context is done")
+}