From 978edc690b57b5591c5cf5a14b741f64d6a80d4d Mon Sep 17 00:00:00 2001 From: shaoqi Date: Mon, 6 Apr 2026 12:45:54 +0800 Subject: [PATCH 1/2] fix(mail): replace os.Exit with graceful shutdown in mail watch Replace os.Exit(0) in the signal handler of `mail +watch` with context-based graceful shutdown using context.WithCancel and a shutdown channel. The os.Exit(0) call had several problems: - Bypasses deferred cleanup functions in the call stack - Makes the code untestable (process terminates immediately) - Not idiomatic for Go CLI tools Changes: - Use context.WithCancel to create a cancellable watch context - Signal handler now cancels the context and closes shutdownCh - cli.Start uses watchCtx instead of the parent context - Distinguish between signal-triggered shutdown and real errors via select on shutdownCh Closes #269 Co-Authored-By: Claude Opus 4.6 --- shortcuts/mail/mail_watch.go | 22 +++++++++++++++++----- 1 file changed, 17 insertions(+), 5 deletions(-) diff --git a/shortcuts/mail/mail_watch.go b/shortcuts/mail/mail_watch.go index c5699427..048f7860 100644 --- a/shortcuts/mail/mail_watch.go +++ b/shortcuts/mail/mail_watch.go @@ -424,6 +424,10 @@ var MailWatch = common.Shortcut{ larkws.WithLogger(sdkLogger), ) + watchCtx, cancelWatch := context.WithCancel(ctx) + defer cancelWatch() + + shutdownCh := make(chan struct{}) sigCh := make(chan os.Signal, 1) signal.Notify(sigCh, syscall.SIGINT, syscall.SIGTERM) go func() { @@ -433,6 +437,7 @@ var MailWatch = common.Shortcut{ } }() <-sigCh + signal.Stop(sigCh) info(fmt.Sprintf("\nShutting down... (received %d events)", eventCount)) info("Unsubscribing mailbox events...") if unsubErr := unsubscribe(); unsubErr != nil { @@ -440,14 +445,21 @@ var MailWatch = common.Shortcut{ } else { info("Mailbox unsubscribed.") } - signal.Stop(sigCh) - os.Exit(0) + cancelWatch() + close(shutdownCh) }() info("Connected. Waiting for mail events... (Ctrl+C to stop)") - if err := cli.Start(ctx); err != nil { - unsubscribe() //nolint:errcheck // best-effort cleanup - return output.ErrNetwork("WebSocket connection failed: %v", err) + if err := cli.Start(watchCtx); err != nil { + // Distinguish between signal-triggered shutdown and real errors. + select { + case <-shutdownCh: + // Graceful shutdown via signal; not an error. + return nil + default: + unsubscribe() //nolint:errcheck // best-effort cleanup + return output.ErrNetwork("WebSocket connection failed: %v", err) + } } return nil }, From 8ad9709725d0de800420aa3bb66f9d66987c656f Mon Sep 17 00:00:00 2001 From: shaoqi Date: Mon, 6 Apr 2026 15:03:56 +0800 Subject: [PATCH 2/2] fix(mail): avoid shutdown race in mail watch graceful exit --- shortcuts/mail/mail_watch.go | 21 +++++++++---------- shortcuts/mail/mail_watch_test.go | 34 +++++++++++++++++++++++++++++++ 2 files changed, 44 insertions(+), 11 deletions(-) diff --git a/shortcuts/mail/mail_watch.go b/shortcuts/mail/mail_watch.go index 048f7860..bf368691 100644 --- a/shortcuts/mail/mail_watch.go +++ b/shortcuts/mail/mail_watch.go @@ -427,7 +427,6 @@ var MailWatch = common.Shortcut{ watchCtx, cancelWatch := context.WithCancel(ctx) defer cancelWatch() - shutdownCh := make(chan struct{}) sigCh := make(chan os.Signal, 1) signal.Notify(sigCh, syscall.SIGINT, syscall.SIGTERM) go func() { @@ -446,25 +445,25 @@ var MailWatch = common.Shortcut{ info("Mailbox unsubscribed.") } cancelWatch() - close(shutdownCh) }() info("Connected. Waiting for mail events... (Ctrl+C to stop)") if err := cli.Start(watchCtx); err != nil { - // Distinguish between signal-triggered shutdown and real errors. - select { - case <-shutdownCh: - // Graceful shutdown via signal; not an error. - return nil - default: - unsubscribe() //nolint:errcheck // best-effort cleanup - return output.ErrNetwork("WebSocket connection failed: %v", err) - } + return handleMailWatchStartError(err, watchCtx, unsubscribe) } return nil }, } +func handleMailWatchStartError(err error, watchCtx context.Context, unsubscribe func() error) error { + if watchCtx.Err() != nil { + // Graceful shutdown via signal cancellation; not an error. + return nil + } + unsubscribe() //nolint:errcheck // best-effort cleanup + return output.ErrNetwork("WebSocket connection failed: %v", err) +} + // extractMailEventBody extracts the event body from the Lark event envelope. func extractMailEventBody(data map[string]interface{}) map[string]interface{} { // V2 envelope: { header: {...}, event: { mail_address, message_id, ... } } diff --git a/shortcuts/mail/mail_watch_test.go b/shortcuts/mail/mail_watch_test.go index 02476fbd..e27c9d14 100644 --- a/shortcuts/mail/mail_watch_test.go +++ b/shortcuts/mail/mail_watch_test.go @@ -539,6 +539,40 @@ func TestWrapWatchSubscribeErrorExitError(t *testing.T) { } } +func TestHandleMailWatchStartErrorGracefulShutdownSkipsCleanup(t *testing.T) { + watchCtx, cancel := context.WithCancel(context.Background()) + cancel() + + unsubscribeCalled := false + err := handleMailWatchStartError(assertErr("context canceled"), watchCtx, func() error { + unsubscribeCalled = true + return nil + }) + if err != nil { + t.Fatalf("expected nil error, got %v", err) + } + if unsubscribeCalled { + t.Fatal("unsubscribe should not be called for graceful shutdown") + } +} + +func TestHandleMailWatchStartErrorNetworkFailureCleansUp(t *testing.T) { + unsubscribeCalled := false + err := handleMailWatchStartError(assertErr("boom"), context.Background(), func() error { + unsubscribeCalled = true + return nil + }) + if err == nil { + t.Fatal("expected error") + } + if !unsubscribeCalled { + t.Fatal("expected unsubscribe to be called for startup failure") + } + if !strings.Contains(err.Error(), "WebSocket connection failed: boom") { + t.Fatalf("unexpected error: %v", err) + } +} + // --- watchFetchFormat --- func TestWatchFetchFormat(t *testing.T) {