@@ -17,6 +17,8 @@ limitations under the License.
1717package util
1818
1919import (
20+ "context"
21+ "fmt"
2022 "os"
2123 "os/signal"
2224 "syscall"
@@ -40,35 +42,35 @@ const (
4042)
4143
4244// SetupExitHandler:
43- // A stop channel is returned which is closed on receiving a shutdown signal (SIGTERM
45+ // A context is returned which is canceled on receiving a shutdown signal (SIGTERM
4446// or SIGINT). If a second signal is caught, the program is terminated directly with
4547// exit code 130.
4648// SetupExitHandler also returns an exit function, this exit function calls os.Exit(...)
4749// if there is a exit code in the errorExitCodeChannel.
4850// The errorExitCodeChannel receives exit codes when SetExitCode is called or when
4951// a shutdown signal is received (only if exitBehavior is AlwaysErrCode).
50- func SetupExitHandler (exitBehavior ExitBehavior ) (<- chan struct {} , func ()) {
52+ func SetupExitHandler (parentCtx context. Context , exitBehavior ExitBehavior ) (context. Context , func ()) {
5153 close (onlyOneSignalHandler ) // panics when called twice
5254
53- stop := make ( chan struct {} )
55+ ctx , cancel := context . WithCancelCause ( parentCtx )
5456 c := make (chan os.Signal , 2 )
5557 signal .Notify (c , shutdownSignals ... )
5658 go func () {
57- // first signal. Close stop chan and pass exit code to exitCodeChannel .
58- exitCode := 128 + int ((<- c ).(syscall.Signal ))
59+ // first signal. Cancel context and pass exit code to errorExitCodeChannel .
60+ signalInt := int ((<- c ).(syscall.Signal ))
5961 if exitBehavior == AlwaysErrCode {
60- errorExitCodeChannel <- exitCode
62+ errorExitCodeChannel <- signalInt
6163 }
62- close ( stop )
64+ cancel ( fmt . Errorf ( "received signal %d" , signalInt ) )
6365 // second signal. Exit directly.
6466 <- c
6567 os .Exit (130 )
6668 }()
6769
68- return stop , func () {
70+ return ctx , func () {
6971 select {
70- case signal := <- errorExitCodeChannel :
71- os .Exit (signal )
72+ case signalInt := <- errorExitCodeChannel :
73+ os .Exit (128 + signalInt )
7274 default :
7375 // Do not exit, there are no exit codes in the channel,
7476 // so just continue and let the main function go out of
0 commit comments