Newer
Older
AnthosCertManager / cmd / util / signal.go
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.
		}
	}
}