From 433cc851fb2a4ff2fcc205477c1a3ec4f0a82ce8 Mon Sep 17 00:00:00 2001
From: laverite <laverite@gmail.com>
Date: Tue, 11 May 2021 22:08:05 -0700
Subject: [PATCH] honor LeaderElectionReleaseOnCancel

---
 pkg/manager/manager.go      | 43 +++++++++++++++++++------------------
 pkg/manager/manager_test.go | 33 ++++++++++++++++++++++++++++
 2 files changed, 55 insertions(+), 21 deletions(-)

diff --git a/pkg/manager/manager.go b/pkg/manager/manager.go
index 21750a3185..843919427d 100644
--- a/pkg/manager/manager.go
+++ b/pkg/manager/manager.go
@@ -366,27 +366,28 @@ func New(config *rest.Config, options Options) (Manager, error) {
 	}
 
 	return &controllerManager{
-		cluster:                 cluster,
-		recorderProvider:        recorderProvider,
-		resourceLock:            resourceLock,
-		metricsListener:         metricsListener,
-		metricsExtraHandlers:    metricsExtraHandlers,
-		controllerOptions:       options.Controller,
-		logger:                  options.Logger,
-		elected:                 make(chan struct{}),
-		port:                    options.Port,
-		host:                    options.Host,
-		certDir:                 options.CertDir,
-		webhookServer:           options.WebhookServer,
-		leaseDuration:           *options.LeaseDuration,
-		renewDeadline:           *options.RenewDeadline,
-		retryPeriod:             *options.RetryPeriod,
-		healthProbeListener:     healthProbeListener,
-		readinessEndpointName:   options.ReadinessEndpointName,
-		livenessEndpointName:    options.LivenessEndpointName,
-		gracefulShutdownTimeout: *options.GracefulShutdownTimeout,
-		internalProceduresStop:  make(chan struct{}),
-		leaderElectionStopped:   make(chan struct{}),
+		cluster:                       cluster,
+		recorderProvider:              recorderProvider,
+		resourceLock:                  resourceLock,
+		metricsListener:               metricsListener,
+		metricsExtraHandlers:          metricsExtraHandlers,
+		controllerOptions:             options.Controller,
+		logger:                        options.Logger,
+		elected:                       make(chan struct{}),
+		port:                          options.Port,
+		host:                          options.Host,
+		certDir:                       options.CertDir,
+		webhookServer:                 options.WebhookServer,
+		leaseDuration:                 *options.LeaseDuration,
+		renewDeadline:                 *options.RenewDeadline,
+		retryPeriod:                   *options.RetryPeriod,
+		healthProbeListener:           healthProbeListener,
+		readinessEndpointName:         options.ReadinessEndpointName,
+		livenessEndpointName:          options.LivenessEndpointName,
+		gracefulShutdownTimeout:       *options.GracefulShutdownTimeout,
+		internalProceduresStop:        make(chan struct{}),
+		leaderElectionStopped:         make(chan struct{}),
+		leaderElectionReleaseOnCancel: options.LeaderElectionReleaseOnCancel,
 	}, nil
 }
 
diff --git a/pkg/manager/manager_test.go b/pkg/manager/manager_test.go
index 59255aa54e..d8cb096d21 100644
--- a/pkg/manager/manager_test.go
+++ b/pkg/manager/manager_test.go
@@ -489,6 +489,39 @@ var _ = Describe("manger.Manager", func() {
 				_, isLeaseLock := cm.resourceLock.(*resourcelock.LeaseLock)
 				Expect(isLeaseLock).To(BeTrue())
 			})
+			It("should release lease if ElectionReleaseOnCancel is true", func() {
+				var rl resourcelock.Interface
+				m, err := New(cfg, Options{
+					LeaderElection:                true,
+					LeaderElectionResourceLock:    resourcelock.LeasesResourceLock,
+					LeaderElectionID:              "controller-runtime",
+					LeaderElectionNamespace:       "my-ns",
+					LeaderElectionReleaseOnCancel: true,
+					newResourceLock: func(config *rest.Config, recorderProvider recorder.Provider, options leaderelection.Options) (resourcelock.Interface, error) {
+						var err error
+						rl, err = fakeleaderelection.NewResourceLock(config, recorderProvider, options)
+						return rl, err
+					},
+				})
+				Expect(err).To(BeNil())
+
+				ctx, cancel := context.WithCancel(context.Background())
+				doneCh := make(chan struct{})
+				go func() {
+					defer GinkgoRecover()
+					defer close(doneCh)
+					Expect(m.Start(ctx)).NotTo(HaveOccurred())
+				}()
+				<-m.(*controllerManager).elected
+				cancel()
+				<-doneCh
+
+				ctx, cancel = context.WithCancel(context.Background())
+				defer cancel()
+				record, _, err := rl.Get(ctx)
+				Expect(err).To(BeNil())
+				Expect(record.HolderIdentity).To(BeEmpty())
+			})
 		})
 
 		It("should create a listener for the metrics if a valid address is provided", func() {