diff --git a/cmd/launcher/signal_listener.go b/cmd/launcher/signal_listener.go index be7ea1c2f..2edb4f57b 100644 --- a/cmd/launcher/signal_listener.go +++ b/cmd/launcher/signal_listener.go @@ -14,24 +14,31 @@ type signalListener struct { sigChannel chan os.Signal cancel context.CancelFunc slogger *slog.Logger + interrupt chan struct{} interrupted atomic.Bool } func newSignalListener(sigChannel chan os.Signal, cancel context.CancelFunc, slogger *slog.Logger) *signalListener { + signal.Notify(sigChannel, os.Interrupt, syscall.SIGTERM) return &signalListener{ sigChannel: sigChannel, cancel: cancel, slogger: slogger.With("component", "signal_listener"), + interrupt: make(chan struct{}, 1), } } func (s *signalListener) Execute() error { - signal.Notify(s.sigChannel, os.Interrupt, syscall.SIGTERM) - sig := <-s.sigChannel - s.slogger.Log(context.TODO(), slog.LevelInfo, - "beginning shutdown via signal", - "signal_received", sig, - ) + select { + case sig := <-s.sigChannel: + s.slogger.Log(context.TODO(), slog.LevelInfo, + "beginning shutdown via signal", + "signal_received", sig, + ) + case <-s.interrupt: + // Rungroup shutdown + } + return nil } @@ -44,6 +51,6 @@ func (s *signalListener) Interrupt(_ error) { // tell sender in `os/signal` package to stop sending on `s.sigChannel` // to avoid panics for sending on a closed channel signal.Stop(s.sigChannel) + close(s.interrupt) s.cancel() - close(s.sigChannel) } diff --git a/cmd/launcher/signal_listener_test.go b/cmd/launcher/signal_listener_test.go index f0becf674..9daf2a036 100644 --- a/cmd/launcher/signal_listener_test.go +++ b/cmd/launcher/signal_listener_test.go @@ -5,6 +5,7 @@ import ( "errors" "log/slog" "os" + "syscall" "testing" "time" @@ -19,6 +20,45 @@ func TestMain(m *testing.M) { goleak.VerifyTestMain(m, goleak.IgnoreAnyFunction("github.com/Microsoft/go-winio.ioCompletionProcessor")) } +func TestInterruptBeforeExecute(t *testing.T) { + t.Parallel() + + sigChannel := make(chan os.Signal, 1) + _, cancel := context.WithCancel(t.Context()) + var logBytes threadsafebuffer.ThreadSafeBuffer + slogger := slog.New(slog.NewTextHandler(&logBytes, &slog.HandlerOptions{ + Level: slog.LevelDebug, + })) + sigListener := newSignalListener(sigChannel, cancel, slogger) + + // Call Interrupt and Execute out of order -- Execute should immediately return + sigListener.Interrupt(errors.New("test error")) + err := sigListener.Execute() + require.NoError(t, err) + + // Send a sigterm, confirm no panic + sigChannel <- syscall.SIGTERM +} + +func TestSigtermBeforeExecute(t *testing.T) { + t.Parallel() + + sigChannel := make(chan os.Signal, 1) + _, cancel := context.WithCancel(t.Context()) + var logBytes threadsafebuffer.ThreadSafeBuffer + slogger := slog.New(slog.NewTextHandler(&logBytes, &slog.HandlerOptions{ + Level: slog.LevelDebug, + })) + sigListener := newSignalListener(sigChannel, cancel, slogger) + + // Send a sigterm, confirm no panic + sigChannel <- syscall.SIGTERM + + // Start up the listener, then shut it down + go sigListener.Execute() + sigListener.Interrupt(errors.New("test error")) +} + func TestInterrupt_Multiple(t *testing.T) { t.Parallel()