From 4272b5af7a7f4e74657cc0bfcdbbcf85141d1ea4 Mon Sep 17 00:00:00 2001
From: Vince Prignano <vincepri@redhat.com>
Date: Thu, 2 Feb 2023 14:00:51 -0800
Subject: [PATCH] Allow to set GracefulShutdownTimeout to -1, disabling
 timeouts

Signed-off-by: Vince Prignano <vincepri@redhat.com>
---
 pkg/manager/internal.go     |  7 +++++-
 pkg/manager/manager_test.go | 44 +++++++++++++++++++++++++++++++++++++
 2 files changed, 50 insertions(+), 1 deletion(-)

diff --git a/pkg/manager/internal.go b/pkg/manager/internal.go
index 51bb9deea6..86dddf088a 100644
--- a/pkg/manager/internal.go
+++ b/pkg/manager/internal.go
@@ -509,7 +509,12 @@ func (cm *controllerManager) engageStopProcedure(stopComplete <-chan struct{}) e
 	//
 	// The shutdown context immediately expires if the gracefulShutdownTimeout is not set.
 	var shutdownCancel context.CancelFunc
-	cm.shutdownCtx, shutdownCancel = context.WithTimeout(context.Background(), cm.gracefulShutdownTimeout)
+	if cm.gracefulShutdownTimeout < 0 {
+		// We want to wait forever for the runnables to stop.
+		cm.shutdownCtx, shutdownCancel = context.WithCancel(context.Background())
+	} else {
+		cm.shutdownCtx, shutdownCancel = context.WithTimeout(context.Background(), cm.gracefulShutdownTimeout)
+	}
 	defer shutdownCancel()
 
 	// Start draining the errors before acquiring the lock to make sure we don't deadlock
diff --git a/pkg/manager/manager_test.go b/pkg/manager/manager_test.go
index ccade8a9b2..a2beb60ea1 100644
--- a/pkg/manager/manager_test.go
+++ b/pkg/manager/manager_test.go
@@ -936,6 +936,50 @@ var _ = Describe("manger.Manager", func() {
 				<-runnableStopped
 			})
 
+			It("should wait forever for runnables if gracefulShutdownTimeout is <0 (-1)", func() {
+				m, err := New(cfg, options)
+				Expect(err).NotTo(HaveOccurred())
+				for _, cb := range callbacks {
+					cb(m)
+				}
+				m.(*controllerManager).gracefulShutdownTimeout = time.Duration(-1)
+
+				Expect(m.Add(RunnableFunc(func(ctx context.Context) error {
+					<-ctx.Done()
+					time.Sleep(100 * time.Millisecond)
+					return nil
+				}))).ToNot(HaveOccurred())
+				Expect(m.Add(RunnableFunc(func(ctx context.Context) error {
+					<-ctx.Done()
+					time.Sleep(200 * time.Millisecond)
+					return nil
+				}))).ToNot(HaveOccurred())
+				Expect(m.Add(RunnableFunc(func(ctx context.Context) error {
+					<-ctx.Done()
+					time.Sleep(500 * time.Millisecond)
+					return nil
+				}))).ToNot(HaveOccurred())
+				Expect(m.Add(RunnableFunc(func(ctx context.Context) error {
+					<-ctx.Done()
+					time.Sleep(1500 * time.Millisecond)
+					return nil
+				}))).ToNot(HaveOccurred())
+
+				ctx, cancel := context.WithCancel(context.Background())
+				managerStopDone := make(chan struct{})
+				go func() {
+					defer GinkgoRecover()
+					Expect(m.Start(ctx)).NotTo(HaveOccurred())
+					close(managerStopDone)
+				}()
+				<-m.Elected()
+				cancel()
+
+				beforeDone := time.Now()
+				<-managerStopDone
+				Expect(time.Since(beforeDone)).To(BeNumerically(">=", 1500*time.Millisecond))
+			})
+
 		}
 
 		Context("with defaults", func() {