From 07618cb355318f49d4ac40419e3fa151c75c624e Mon Sep 17 00:00:00 2001
From: Gytis Trikleris <gytis@redhat.com>
Date: Tue, 5 Jun 2018 13:03:53 +0200
Subject: [PATCH] LeaderInitiator should revoke leadership on #stop()

---
 .../kubernetes/leader/LeaderInitiator.java     | 11 ++++++-----
 .../kubernetes/leader/LeaderInitiatorTest.java | 18 +++++++++++++++++-
 2 files changed, 23 insertions(+), 6 deletions(-)

diff --git a/spring-cloud-kubernetes-leader/src/main/java/org/springframework/cloud/kubernetes/leader/LeaderInitiator.java b/spring-cloud-kubernetes-leader/src/main/java/org/springframework/cloud/kubernetes/leader/LeaderInitiator.java
index 33acd987..b839f02c 100644
--- a/spring-cloud-kubernetes-leader/src/main/java/org/springframework/cloud/kubernetes/leader/LeaderInitiator.java
+++ b/spring-cloud-kubernetes-leader/src/main/java/org/springframework/cloud/kubernetes/leader/LeaderInitiator.java
@@ -27,21 +27,21 @@ import org.springframework.integration.leader.Candidate;
  */
 public class LeaderInitiator implements SmartLifecycle {
 
+	private final LeaderProperties leaderProperties;
+
 	private final LeadershipController leadershipController;
 
 	private final Candidate candidate;
 
-	private final LeaderProperties leaderProperties;
-
 	private final ScheduledExecutorService scheduledExecutorService;
 
 	private boolean isRunning;
 
-	public LeaderInitiator(LeadershipController leadershipController, Candidate candidate,
-		LeaderProperties leaderProperties, ScheduledExecutorService scheduledExecutorService) {
+	public LeaderInitiator(LeaderProperties leaderProperties, LeadershipController leadershipController,
+		Candidate candidate, ScheduledExecutorService scheduledExecutorService) {
+		this.leaderProperties = leaderProperties;
 		this.leadershipController = leadershipController;
 		this.candidate = candidate;
-		this.leaderProperties = leaderProperties;
 		this.scheduledExecutorService = scheduledExecutorService;
 	}
 
@@ -61,6 +61,7 @@ public class LeaderInitiator implements SmartLifecycle {
 	@Override
 	public void stop() {
 		if (isRunning()) {
+			scheduledExecutorService.execute(() -> leadershipController.revoke(candidate));
 			scheduledExecutorService.shutdown();
 			isRunning = false;
 		}
diff --git a/spring-cloud-kubernetes-leader/src/test/java/org/springframework/cloud/kubernetes/leader/LeaderInitiatorTest.java b/spring-cloud-kubernetes-leader/src/test/java/org/springframework/cloud/kubernetes/leader/LeaderInitiatorTest.java
index 74594c98..cdbe26a0 100644
--- a/spring-cloud-kubernetes-leader/src/test/java/org/springframework/cloud/kubernetes/leader/LeaderInitiatorTest.java
+++ b/spring-cloud-kubernetes-leader/src/test/java/org/springframework/cloud/kubernetes/leader/LeaderInitiatorTest.java
@@ -15,6 +15,7 @@ import static org.assertj.core.api.Assertions.assertThat;
 import static org.mockito.ArgumentMatchers.any;
 import static org.mockito.ArgumentMatchers.eq;
 import static org.mockito.BDDMockito.given;
+import static org.mockito.Mockito.times;
 import static org.mockito.Mockito.verify;
 
 /**
@@ -42,7 +43,7 @@ public class LeaderInitiatorTest {
 
 	@Before
 	public void before() {
-		leaderInitiator = new LeaderInitiator(mockLeadershipController, mockCandidate, mockLeaderProperties,
+		leaderInitiator = new LeaderInitiator(mockLeaderProperties, mockLeadershipController, mockCandidate,
 			mockScheduledExecutorService);
 	}
 
@@ -68,6 +69,7 @@ public class LeaderInitiatorTest {
 		leaderInitiator.stop();
 
 		assertThat(leaderInitiator.isRunning()).isFalse();
+		verify(mockScheduledExecutorService, times(2)).execute(any());
 		verify(mockScheduledExecutorService).shutdown();
 	}
 
@@ -77,10 +79,24 @@ public class LeaderInitiatorTest {
 		leaderInitiator.stop(mockRunnable);
 
 		assertThat(leaderInitiator.isRunning()).isFalse();
+		verify(mockScheduledExecutorService, times(2)).execute(any());
 		verify(mockScheduledExecutorService).shutdown();
 		verify(mockRunnable).run();
 	}
 
+	@Test
+	public void shouldRevokeLeadershipWhenStopping() {
+		leaderInitiator.start();
+		leaderInitiator.stop();
+
+		ArgumentCaptor<Runnable> revokeRunnableCaptor = ArgumentCaptor.forClass(Runnable.class);
+		verify(mockScheduledExecutorService, times(2)).execute(revokeRunnableCaptor.capture());
+
+		Runnable revokeRunnable = revokeRunnableCaptor.getAllValues().get(1);
+		revokeRunnable.run();
+		verify(mockLeadershipController).revoke(mockCandidate);
+	}
+
 	@Test
 	public void shouldScheduleUpdateIfLeadershipIsAcquired() {
 		given(mockLeadershipController.acquire(mockCandidate)).willReturn(true);