diff --git a/main.go b/main.go index d888f59e4..057dbca57 100644 --- a/main.go +++ b/main.go @@ -6,6 +6,7 @@ import ( "errors" "os" "os/exec" + "os/signal" "runtime/debug" "strings" @@ -41,13 +42,18 @@ func initGotext() { } } +func newSignalContext() (context.Context, context.CancelFunc) { + return signal.NotifyContext(context.Background(), os.Interrupt) +} + func main() { fallbackLog := text.NewLogger(os.Stdout, os.Stderr, os.Stdin, false, "fallback") var ( - err error - ctx = context.Background() - ret = 0 + err error + ctx, stop = newSignalContext() + ret = 0 ) + defer stop() defer func() { if rec := recover(); rec != nil { diff --git a/main_test.go b/main_test.go new file mode 100644 index 000000000..09b91ecb9 --- /dev/null +++ b/main_test.go @@ -0,0 +1,26 @@ +package main + +import ( + "os" + "testing" + "time" +) + +func TestInterruptCancelsContext(t *testing.T) { + ctx, stop := newSignalContext() + defer stop() + + process, err := os.FindProcess(os.Getpid()) + if err != nil { + t.Fatal(err) + } + if err := process.Signal(os.Interrupt); err != nil { + t.Fatal(err) + } + + select { + case <-ctx.Done(): + case <-time.After(time.Second): + t.Fatal("interrupt did not cancel context") + } +}