From b745d7d84b4846e93e0a63cc8df824d9e9c72bc8 Mon Sep 17 00:00:00 2001
From: torredil <torredil@amazon.com>
Date: Thu, 11 Jan 2024 21:32:35 +0000
Subject: [PATCH] Timeout if external-resizer lease is not updated within
 resyncPeriod

Signed-off-by: torredil <torredil@amazon.com>
---
 README.md   |  5 ++++-
 cmd/main.go | 54 ++++++++++++++++++++++++++++++++++++-----------------
 2 files changed, 41 insertions(+), 18 deletions(-)

diff --git a/README.md b/README.md
index 01b866a..25778b4 100644
--- a/README.md
+++ b/README.md
@@ -1,7 +1,10 @@
 ## volume-modifier-for-k8s
 
-volume-modifier-for-k8s is a sidecar deployed alongside CSI drivers to enable volume modification through annotations on the PVC.
+`volume-modifier-for-k8s` is a sidecar deployed alongside CSI drivers to enable volume modification through annotations on the PVC.
 
+## Requirements
+
+Leader election must be enabled in the [external-resizer](https://github.com/kubernetes-csi/external-resizer). This is required in order to efficiently coordinate calls to the EC2 modify-volume API.
 
 ## Security
 
diff --git a/cmd/main.go b/cmd/main.go
index 72cff97..00e3dc1 100644
--- a/cmd/main.go
+++ b/cmd/main.go
@@ -63,6 +63,7 @@ func main() {
 		os.Exit(0)
 	}
 	klog.Infof("Version : %s", version)
+	klog.InfoS("Leader election must be enabled in the external-resizer CSI sidecar")
 
 	podName := os.Getenv("POD_NAME")
 	if podName == "" {
@@ -167,26 +168,45 @@ func main() {
 func leaseHandler(podName string, mc controller.ModifyController, leaseChannel chan *v1.Lease) {
 	var cancel context.CancelFunc = nil
 
-	for lease := range leaseChannel {
-		currentLeader := *lease.Spec.HolderIdentity
+	klog.InfoS("leaseHandler: Looking for external-resizer lease holder")
 
-		klog.V(6).InfoS("leaseHandler: Lease updated", "currentLeader", currentLeader, "podName", podName)
+	timer := time.NewTimer(*resyncPeriod)
+	defer timer.Stop()
 
-		if currentLeader == podName && cancel == nil {
-			var ctx context.Context
-			ctx, cancel = context.WithCancel(context.Background())
-			klog.InfoS("leaseHandler: Starting ModifyController", "podName", podName, "currentLeader", currentLeader)
-			go mc.Run(*workers, ctx)
-		} else if currentLeader != podName && cancel != nil {
-			klog.InfoS("leaseHandler: Stopping ModifyController", "podName", podName, "currentLeader", currentLeader)
-			cancel()
-			cancel = nil
-		}
-	}
+	for {
+		select {
+		case lease, ok := <-leaseChannel:
+			if !ok {
+				if cancel != nil {
+					cancel()
+				}
+				return
+			}
+			currentLeader := *lease.Spec.HolderIdentity
+			klog.V(6).InfoS("leaseHandler: Lease updated", "currentLeader", currentLeader, "podName", podName)
+
+			if currentLeader == podName && cancel == nil {
+				var ctx context.Context
+				ctx, cancel = context.WithCancel(context.Background())
+				klog.InfoS("leaseHandler: Starting ModifyController", "podName", podName, "currentLeader", currentLeader)
+				go mc.Run(*workers, ctx)
+			} else if currentLeader != podName && cancel != nil {
+				klog.InfoS("leaseHandler: Stopping ModifyController", "podName", podName, "currentLeader", currentLeader)
+				cancel()
+				cancel = nil
+			}
+
+			if !timer.Stop() {
+				<-timer.C
+			}
+			timer.Reset(*resyncPeriod)
 
-	// Ensure cancel is called if it's not nil when we exit the function
-	if cancel != nil {
-		cancel()
+		case <-timer.C:
+			if cancel != nil {
+				cancel()
+			}
+			klog.Fatalf("leaseHandler: No external-resizer lease update received within timeout period", "timeout", *resyncPeriod)
+		}
 	}
 }