@@ -136,40 +136,98 @@ func uploadLibraries(ctx context.Context, b *bundle.Bundle, libs map[string][]li
136136}
137137
138138// registerGracefulCleanup sets up signal handlers to release the lock
139- // before the process terminates. Returns a cleanup function for the normal exit path.
139+ // before the process terminates. Returns a new context that will be cancelled
140+ // when a signal is received, and a cleanup function for the exit path.
141+ //
142+ // This follows idiomatic Go patterns for graceful shutdown:
143+ // 1. Use context cancellation to signal shutdown to the main routine
144+ // 2. Use a done channel to wait for the main routine to complete
145+ // 3. Only exit after confirming the main routine has terminated
140146//
141147// Catches SIGINT (Ctrl+C), SIGTERM, SIGHUP, and SIGQUIT.
142148// Note: SIGKILL and SIGSTOP cannot be caught - the kernel terminates the process directly.
143- func registerGracefulCleanup (ctx context.Context , b * bundle.Bundle , goal lock.Goal ) func () {
149+ func registerGracefulCleanup (ctx context.Context , b * bundle.Bundle , goal lock.Goal ) (context.Context , func ()) {
150+ // Create a cancellable context to propagate cancellation to the main routine
151+ ctx , cancel := context .WithCancel (ctx )
152+
153+ // Channel to signal that the main routine has completed
154+ done := make (chan struct {})
155+
156+ // Channel to signal that a signal was received and handled
157+ signalReceived := make (chan struct {})
158+
159+ // Channel to receive OS signals
144160 sigChan := make (chan os.Signal , 1 )
145161 signal .Notify (sigChan , syscall .SIGINT , syscall .SIGTERM , syscall .SIGHUP , syscall .SIGQUIT )
146162
147- // Start goroutine to handle signals
148- go func () {
163+ signalHandler := func () {
164+ // Wait for a signal to be received.
149165 sig := <- sigChan
150- // Stop listening for more signals
166+
167+ // Stop listening for more signals. This allows for multiple interrupts
168+ // to cause the program to force exit.
151169 signal .Stop (sigChan )
152170
171+ // Signal that we received an interrupt
172+ close (signalReceived )
173+
153174 cmdio .LogString (ctx , "Operation interrupted. Gracefully shutting down..." )
154175
155- // Release the lock.
156- bundle .ApplyContext (ctx , b , lock .Release (goal ))
176+ // Cancel the context to signal the main routine to stop
177+ cancel ()
178+
179+ // Wait for the main routine to complete before releasing the lock
180+ // This ensures we don't exit while operations are still in progress
181+ <- done
157182
158- // Exit immediately with standard signal exit code (128 + signal number).
159- // The deferred cleanup function returned below won't run because we exit here.
183+ // Release the lock using a context without cancellation to avoid cancellation issues
184+ // We use context.WithoutCancel to preserve context values (like user agent)
185+ // but remove the cancellation signal so the lock release can complete
186+ releaseCtx := context .WithoutCancel (ctx )
187+ bundle .ApplyContext (releaseCtx , b , lock .Release (goal ))
188+
189+ // Calculate exit code (128 + signal number)
160190 exitCode := 128
161191 if s , ok := sig .(syscall.Signal ); ok {
162192 exitCode += int (s )
163193 }
194+
195+ // Exit with the appropriate signal exit code
164196 os .Exit (exitCode )
165- }()
197+ }
198+
199+ // Start goroutine to handle signals
200+ go signalHandler ()
166201
167- // Return cleanup function for normal exit path
168- return func () {
202+ // Return cleanup function for the exit path
203+ // This should be called via defer to ensure it runs even if there's a panic
204+ cleanup := func () {
205+ // Stop listening for signals
169206 signal .Stop (sigChan )
170- // Don't close the channel - it causes the goroutine to receive nil
171- bundle .ApplyContext (ctx , b , lock .Release (goal ))
207+
208+ // Release the lock (idempotent thanks to sync.Once in lock.Release)
209+ // Use context.WithoutCancel to preserve context values but remove cancellation
210+ releaseCtx := context .WithoutCancel (ctx )
211+ bundle .ApplyContext (releaseCtx , b , lock .Release (goal ))
212+
213+ // Signal that the main routine has completed.
214+ // Once the signal is recieved,
215+ // This must be done AFTER all cleanup is complete
216+ close (done )
217+
218+ // If a signal was received, wait indefinitely for the signal handler to exit
219+ // This prevents the main function from returning and exiting with a different code
220+ // If no signal was received, signalReceived will never be closed, so we just return
221+ select {
222+ case <- signalReceived :
223+ // Signal was received, wait forever for os.Exit() in signal handler
224+ select {}
225+ default :
226+ // No signal received, proceed with normal exit
227+ }
172228 }
229+
230+ return ctx , cleanup
173231}
174232
175233// The deploy phase deploys artifacts and resources.
@@ -189,7 +247,8 @@ func Deploy(ctx context.Context, b *bundle.Bundle, outputHandler sync.OutputHand
189247 }
190248
191249 // lock is acquired here - set up signal handlers and defer cleanup
192- defer registerGracefulCleanup (ctx , b , lock .GoalDeploy )()
250+ ctx , cleanup := registerGracefulCleanup (ctx , b , lock .GoalDeploy )
251+ defer cleanup ()
193252
194253 libs := deployPrepare (ctx , b , false , directDeployment )
195254 if logdiag .HasError (ctx ) {
0 commit comments