git.iko.soy/crush
on main help sign in
chore: handle cancel when panic
2f9aa8dKujtim Hoxha
cmd/root.go+22−8
compare with main
unifiedsplit
@@ -6,6 +6,7 @@ import (
6 "io"
6 "io"
7 "log/slog"
7 "log/slog"
8 "os"
8 "os"
9+ "sync"
9 "time"
10 "time"
10
11
11 tea "github.com/charmbracelet/bubbletea/v2"
12 tea "github.com/charmbracelet/bubbletea/v2"
@@ -72,9 +73,8 @@ to assist developers in writing, debugging, and understanding code directly from
72 return err
73 return err
73 }
74 }
74
75
75− // Create main context for the application
76+ // Use the context from the command which includes signal handling
76− ctx, cancel := context.WithCancel(context.Background())
77+ ctx := cmd.Context()
77− defer cancel()
78
78
79 // Connect DB, this will also run migrations
79 // Connect DB, this will also run migrations
80 conn, err := db.Connect(ctx, cfg.Options.DataDirectory)
80 conn, err := db.Connect(ctx, cfg.Options.DataDirectory)
@@ -87,8 +87,23 @@ to assist developers in writing, debugging, and understanding code directly from
87 slog.Error(fmt.Sprintf("Failed to create app instance: %v", err))
87 slog.Error(fmt.Sprintf("Failed to create app instance: %v", err))
88 return err
88 return err
89 }
89 }
90− // Defer shutdown here so it runs for both interactive and non-interactive modes
90+
91+ // Set up shutdown handling that works for both normal exit and signal interruption
92+ var shutdownOnce sync.Once
93+ shutdown := func() {
94+ shutdownOnce.Do(func() {
95+ slog.Info("Shutting down application")
91− defer app.Shutdown()
96+ app.Shutdown()
97+ })
98+ }
99+ defer shutdown()
100+
101+ // Handle context cancellation (from signals) in a goroutine
102+ go func() {
103+ <-ctx.Done()
104+ slog.Info("Context cancelled, initiating shutdown")
105+ shutdown()
106+ }()
92
107
93 // Initialize MCP tools early for both modes
108 // Initialize MCP tools early for both modes
94 initMCPTools(ctx, app, cfg)
109 initMCPTools(ctx, app, cfg)
@@ -121,7 +136,6 @@ to assist developers in writing, debugging, and understanding code directly from
121 slog.Error(fmt.Sprintf("TUI run error: %v", err))
136 slog.Error(fmt.Sprintf("TUI run error: %v", err))
122 return fmt.Errorf("TUI error: %v", err)
137 return fmt.Errorf("TUI error: %v", err)
123 }
138 }
124− app.Shutdown()
125 return nil
139 return nil
126 },
140 },
127}
141}
@@ -140,9 +154,9 @@ func initMCPTools(ctx context.Context, app *app.App, cfg *config.Config) {
140 }()
154 }()
141}
155}
142
156
143−func Execute() {
157+func Execute(ctx context.Context) {
144 if err := fang.Execute(
158 if err := fang.Execute(
145− context.Background(),
159+ ctx,
146 rootCmd,
160 rootCmd,
147 fang.WithVersion(version.Version),
161 fang.WithVersion(version.Version),
148 ); err != nil {
162 ); err != nil {
main.go+17−1
@@ -1,10 +1,13 @@
1package main
1package main
2
2
3import (
3import (
4+ "context"
4 "fmt"
5 "fmt"
5 "log/slog"
6 "log/slog"
6 "net/http"
7 "net/http"
7 "os"
8 "os"
9+ "os/signal"
10+ "syscall"
8
11
9 _ "net/http/pprof" // profiling
12 _ "net/http/pprof" // profiling
10
13
@@ -19,6 +22,19 @@ func main() {
19 slog.Error("Application terminated due to unhandled panic")
22 slog.Error("Application terminated due to unhandled panic")
20 })
23 })
21
24
25+ ctx, cancel := context.WithCancel(context.Background())
26+ defer cancel()
27+
28+ sigChan := make(chan os.Signal, 1)
29+ signal.Notify(sigChan, os.Interrupt, syscall.SIGTERM, syscall.SIGQUIT)
30+
31+ // Start signal handler in a goroutine
32+ go func() {
33+ sig := <-sigChan
34+ slog.Info("Received signal, initiating graceful shutdown", "signal", sig)
35+ cancel()
36+ }()
37+
22 if os.Getenv("CRUSH_PROFILE") != "" {
38 if os.Getenv("CRUSH_PROFILE") != "" {
23 go func() {
39 go func() {
24 slog.Info("Serving pprof at localhost:6060")
40 slog.Info("Serving pprof at localhost:6060")
@@ -28,5 +44,5 @@ func main() {
28 }()
44 }()
29 }
45 }
30
46
31− cmd.Execute()
47+ cmd.Execute(ctx)
32}
48}