From 82a7f869022deba6afc3ec9f928a9a02dec02569 Mon Sep 17 00:00:00 2001
From: mprahl <mprahl@users.noreply.github.com>
Date: Fri, 3 Feb 2023 16:34:00 -0500
Subject: [PATCH] Allow PeriodicallyExecConfigPolicies to finish after
 uninstall

Without this commit, as soon as a SIGINT signal was received, mgr.Start
would exit and cause the main function to finish before
PeriodicallyExecConfigPolicies could finish the removal of finalizers
when the config-policy-controller is being uninstalled.

Co-authored-by: Justin Kulikauskas <jkulikau@redhat.com>
Signed-off-by: mprahl <mprahl@users.noreply.github.com>
---
 build/common/config/.golangci.yml             |  3 +++
 controllers/configurationpolicy_controller.go | 27 ++++++++++++++-----
 main.go                                       | 10 +++++--
 3 files changed, 31 insertions(+), 9 deletions(-)

diff --git a/build/common/config/.golangci.yml b/build/common/config/.golangci.yml
index a2e1e9b9..4dba101a 100644
--- a/build/common/config/.golangci.yml
+++ b/build/common/config/.golangci.yml
@@ -29,6 +29,9 @@ linters:
   enable-all: true
   disable:
     - bodyclose
+    # Disable contextcheck since we don't want to pass the context passed to PeriodicallyExecConfigPolicies to functions
+    # called within it or else cleanup won't occur
+    - contextcheck
     - cyclop
     - depguard
     - dupl
diff --git a/controllers/configurationpolicy_controller.go b/controllers/configurationpolicy_controller.go
index 0d2fd6c5..2a60baae 100644
--- a/controllers/configurationpolicy_controller.go
+++ b/controllers/configurationpolicy_controller.go
@@ -151,15 +151,32 @@ func (r *ConfigurationPolicyReconciler) Reconcile(ctx context.Context, request c
 
 // PeriodicallyExecConfigPolicies loops through all configurationpolicies in the target namespace and triggers
 // template handling for each one. This function drives all the work the configuration policy controller does.
-func (r *ConfigurationPolicyReconciler) PeriodicallyExecConfigPolicies(freq uint, elected <-chan struct{}, test bool) {
+func (r *ConfigurationPolicyReconciler) PeriodicallyExecConfigPolicies(
+	ctx context.Context, freq uint, elected <-chan struct{},
+) {
 	log.Info("Waiting for leader election before periodically evaluating configuration policies")
-	<-elected
+
+	select {
+	case <-elected:
+	case <-ctx.Done():
+		return
+	}
 
 	const waiting = 10 * time.Minute
+	var exiting bool
+	// Loop twice after exit condition is received to account for race conditions and retries.
+	loopsAfterExit := 2
 
-	for {
+	for !exiting || (exiting && loopsAfterExit > 0) {
 		start := time.Now()
 
+		select {
+		case <-ctx.Done():
+			exiting = true
+			loopsAfterExit--
+		default:
+		}
+
 		policiesList := policyv1.ConfigurationPolicyList{}
 
 		var skipLoop bool
@@ -264,10 +281,6 @@ func (r *ConfigurationPolicyReconciler) PeriodicallyExecConfigPolicies(freq uint
 			log.V(2).Info("Sleeping before reprocessing the configuration policies", "seconds", sleepTime)
 			time.Sleep(sleepTime)
 		}
-
-		if test {
-			return
-		}
 	}
 }
 
diff --git a/main.go b/main.go
index c57f04ec..c9e25bd8 100644
--- a/main.go
+++ b/main.go
@@ -320,10 +320,16 @@ func main() {
 		os.Exit(1)
 	}
 
+	terminatingCtx := ctrl.SetupSignalHandler()
+	managerCtx, managerCancel := context.WithCancel(context.Background())
+
 	// PeriodicallyExecConfigPolicies is the go-routine that periodically checks the policies
 	log.V(1).Info("Perodically processing Configuration Policies", "frequency", frequency)
 
-	go reconciler.PeriodicallyExecConfigPolicies(frequency, mgr.Elected(), false)
+	go func() {
+		reconciler.PeriodicallyExecConfigPolicies(terminatingCtx, frequency, mgr.Elected())
+		managerCancel()
+	}()
 
 	// This lease is not related to leader election. This is to report the status of the controller
 	// to the addon framework. This can be seen in the "status" section of the ManagedClusterAddOn
@@ -361,7 +367,7 @@ func main() {
 
 	log.Info("Starting manager")
 
-	if err := mgr.Start(ctrl.SetupSignalHandler()); err != nil {
+	if err := mgr.Start(managerCtx); err != nil {
 		log.Error(err, "Problem running manager")
 		os.Exit(1)
 	}