package util import ( "os" "os/signal" "syscall" ) var onlyOneSignalHandler = make(chan struct{}) var errorExitCodeChannel = make(chan int, 1) // ExitBehavior controls how the program should be terminated // in response to a shutdown signal. type ExitBehavior int const ( // AlwaysErrCode indicates the exit code of the program should always be nonzero // and should correspond to the numeric value of the signal that was received. AlwaysErrCode ExitBehavior = iota // GracefulShutdown treats a shutdown signal as a request to exit gracefully, terminating // goroutines and returning an exit code of 0 if there are no errors during shutdown. GracefulShutdown ExitBehavior = iota ) // SetupExitHandler: // A stop channel is returned which is closed on receiving a shutdown signal (SIGTERM // or SIGINT). If a second signal is caught, the program is terminated directly with // exit code 130. // SetupExitHandler also returns an exit function, this exit function calls os.Exit(...) // if there is a exit code in the errorExitCodeChannel. // The errorExitCodeChannel receives exit codes when SetExitCode is called or when // a shutdown signal is received (only if exitBehavior is AlwaysErrCode). func SetupExitHandler(exitBehavior ExitBehavior) (<-chan struct{}, func()) { close(onlyOneSignalHandler) // panics when called twice stop := make(chan struct{}) c := make(chan os.Signal, 2) signal.Notify(c, shutdownSignals...) go func() { // first signal. Close stop chan and pass exit code to exitCodeChannel. exitCode := 128 + int((<-c).(syscall.Signal)) if exitBehavior == AlwaysErrCode { errorExitCodeChannel <- exitCode } close(stop) // second signal. Exit directly. <-c os.Exit(130) }() return stop, func() { select { case signal := <-errorExitCodeChannel: os.Exit(signal) default: // Do not exit, there are no exit codes in the channel, // so just continue and let the main function go out of // scope instead. } } }