From 121fbc72cf7741103cc740bb6135f92528e1ac88 Mon Sep 17 00:00:00 2001
From: Luke Addison <lukeaddison785@gmail.com>
Date: Sat, 2 Dec 2023 12:55:59 +0000
Subject: [PATCH] Add support for shutdown delay

---
 pkg/manager/signals/signal.go | 26 +++++++++++++++++++++-----
 1 file changed, 21 insertions(+), 5 deletions(-)

diff --git a/pkg/manager/signals/signal.go b/pkg/manager/signals/signal.go
index a79cfb42df..1a89f634ed 100644
--- a/pkg/manager/signals/signal.go
+++ b/pkg/manager/signals/signal.go
@@ -20,14 +20,19 @@ import (
 	"context"
 	"os"
 	"os/signal"
+	"time"
 )
 
 var onlyOneSignalHandler = make(chan struct{})
 
-// SetupSignalHandler registers for SIGTERM and SIGINT. A context is returned
-// which is canceled on one of these signals. If a second signal is caught, the program
-// is terminated with exit code 1.
-func SetupSignalHandler() context.Context {
+// SetupSignalHandlerWithDelay registers for SIGTERM and SIGINT. A context is
+// returned which is canceled on one of these signals after waiting for the
+// specified delay. In particular, the delay can be used to give external
+// Kubernetes controllers (such as kube-proxy) time to observe the termination
+// of this manager before starting shutdown of any webhook servers to avoid
+// receiving connection attempts after closing webhook listeners. If a second
+// signal is caught, the program is terminated with exit code 1.
+func SetupSignalHandlerWithDelay(delay time.Duration) context.Context {
 	close(onlyOneSignalHandler) // panics when called twice
 
 	ctx, cancel := context.WithCancel(context.Background())
@@ -36,10 +41,21 @@ func SetupSignalHandler() context.Context {
 	signal.Notify(c, shutdownSignals...)
 	go func() {
 		<-c
-		cancel()
+		// Cancel the context after delaying for the specified duration but
+		// avoid blocking if a second signal is caught
+		go func() {
+			<-time.After(delay)
+			cancel()
+		}()
 		<-c
 		os.Exit(1) // second signal. Exit directly.
 	}()
 
 	return ctx
 }
+
+// SetupSignalHandler is a special case of SetupSignalHandlerWithDelay with no
+// delay for backwards compatibility
+func SetupSignalHandler() context.Context {
+	return SetupSignalHandlerWithDelay(time.Duration(0))
+}