chore: handle cancel when panic

Kujtim Hoxha created 3 months ago

Change summary

cmd/root.go | 30 ++++++++++++++++++++++--------
main.go     | 18 +++++++++++++++++-
2 files changed, 39 insertions(+), 9 deletions(-)

Detailed changes

cmd/root.go

@@ -6,6 +6,7 @@ import (
 	"io"
 	"log/slog"
 	"os"
+	"sync"
 	"time"
 
 	tea "github.com/charmbracelet/bubbletea/v2"
@@ -72,9 +73,8 @@ to assist developers in writing, debugging, and understanding code directly from
 			return err
 		}
 
-		// Create main context for the application
-		ctx, cancel := context.WithCancel(context.Background())
-		defer cancel()
+		// Use the context from the command which includes signal handling
+		ctx := cmd.Context()
 
 		// Connect DB, this will also run migrations
 		conn, err := db.Connect(ctx, cfg.Options.DataDirectory)
@@ -87,8 +87,23 @@ to assist developers in writing, debugging, and understanding code directly from
 			slog.Error(fmt.Sprintf("Failed to create app instance: %v", err))
 			return err
 		}
-		// Defer shutdown here so it runs for both interactive and non-interactive modes
-		defer app.Shutdown()
+
+		// Set up shutdown handling that works for both normal exit and signal interruption
+		var shutdownOnce sync.Once
+		shutdown := func() {
+			shutdownOnce.Do(func() {
+				slog.Info("Shutting down application")
+				app.Shutdown()
+			})
+		}
+		defer shutdown()
+
+		// Handle context cancellation (from signals) in a goroutine
+		go func() {
+			<-ctx.Done()
+			slog.Info("Context cancelled, initiating shutdown")
+			shutdown()
+		}()
 
 		// Initialize MCP tools early for both modes
 		initMCPTools(ctx, app, cfg)
@@ -121,7 +136,6 @@ to assist developers in writing, debugging, and understanding code directly from
 			slog.Error(fmt.Sprintf("TUI run error: %v", err))
 			return fmt.Errorf("TUI error: %v", err)
 		}
-		app.Shutdown()
 		return nil
 	},
 }
@@ -140,9 +154,9 @@ func initMCPTools(ctx context.Context, app *app.App, cfg *config.Config) {
 	}()
 }
 
-func Execute() {
+func Execute(ctx context.Context) {
 	if err := fang.Execute(
-		context.Background(),
+		ctx,
 		rootCmd,
 		fang.WithVersion(version.Version),
 	); err != nil {

main.go

@@ -1,10 +1,13 @@
 package main
 
 import (
+	"context"
 	"fmt"
 	"log/slog"
 	"net/http"
 	"os"
+	"os/signal"
+	"syscall"
 
 	_ "net/http/pprof" // profiling
 
@@ -19,6 +22,19 @@ func main() {
 		slog.Error("Application terminated due to unhandled panic")
 	})
 
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
+
+	sigChan := make(chan os.Signal, 1)
+	signal.Notify(sigChan, os.Interrupt, syscall.SIGTERM, syscall.SIGQUIT)
+
+	// Start signal handler in a goroutine
+	go func() {
+		sig := <-sigChan
+		slog.Info("Received signal, initiating graceful shutdown", "signal", sig)
+		cancel()
+	}()
+
 	if os.Getenv("CRUSH_PROFILE") != "" {
 		go func() {
 			slog.Info("Serving pprof at localhost:6060")
@@ -28,5 +44,5 @@ func main() {
 		}()
 	}
 
-	cmd.Execute()
+	cmd.Execute(ctx)
 }