diff --git a/cmd/crproxy/main.go b/cmd/crproxy/main.go index 0f9104f..d9577be 100644 --- a/cmd/crproxy/main.go +++ b/cmd/crproxy/main.go @@ -6,6 +6,7 @@ import ( "github.com/daocloud/crproxy/cmd/crproxy/cluster" csync "github.com/daocloud/crproxy/cmd/crproxy/sync" + "github.com/daocloud/crproxy/internal/signals" "github.com/spf13/cobra" _ "github.com/daocloud/crproxy/storage/driver/obs" @@ -26,11 +27,11 @@ var ( return cmd.Usage() }, } - pflag = cmd.Flags() ) func main() { - err := cmd.Execute() + ctx := signals.SetupSignalContext() + err := cmd.ExecuteContext(ctx) if err != nil { slog.Error("execute failed", "error", err) os.Exit(1) diff --git a/internal/signals/signals.go b/internal/signals/signals.go new file mode 100644 index 0000000..9e0daf6 --- /dev/null +++ b/internal/signals/signals.go @@ -0,0 +1,32 @@ +package signals + +import ( + "context" + "os" + "os/signal" +) + +var ( + onlyOneSignalHandler = make(chan struct{}) + shutdownHandler chan os.Signal +) + +// SetupSignalContext is same as SetupSignalHandler, but a context.Context is returned. +// Only one of SetupSignalContext and SetupSignalHandler should be called, and only can +// be called once. +func SetupSignalContext() context.Context { + close(onlyOneSignalHandler) // panics when called twice + + shutdownHandler = make(chan os.Signal, 2) + + ctx, cancel := context.WithCancel(context.Background()) + signal.Notify(shutdownHandler, shutdownSignals...) + go func() { + <-shutdownHandler + cancel() + <-shutdownHandler + os.Exit(1) // second signal. Exit directly. + }() + + return ctx +} diff --git a/internal/signals/signals_other.go b/internal/signals/signals_other.go new file mode 100644 index 0000000..51c3151 --- /dev/null +++ b/internal/signals/signals_other.go @@ -0,0 +1,10 @@ +//go:build !windows + +package signals + +import ( + "os" + "syscall" +) + +var shutdownSignals = []os.Signal{os.Interrupt, syscall.SIGTERM} diff --git a/internal/signals/signals_windows.go b/internal/signals/signals_windows.go new file mode 100644 index 0000000..29ab33f --- /dev/null +++ b/internal/signals/signals_windows.go @@ -0,0 +1,9 @@ +//go:build windows + +package signals + +import ( + "os" +) + +var shutdownSignals = []os.Signal{os.Interrupt}