From 7736228f37e5064bb682ca384178a83b54b1cd4e Mon Sep 17 00:00:00 2001
From: George Shiqi Wu <george.wu@imply.io>
Date: Fri, 6 Dec 2024 16:01:28 -0500
Subject: [PATCH] Separate stop/start logic for LeaderLatch (#17546)

---
 .../discovery/CuratorDruidLeaderSelector.java     | 15 ++++++++++-----
 1 file changed, 10 insertions(+), 5 deletions(-)

diff --git a/server/src/main/java/org/apache/druid/curator/discovery/CuratorDruidLeaderSelector.java b/server/src/main/java/org/apache/druid/curator/discovery/CuratorDruidLeaderSelector.java
index d1129803106..4c3d7de23a0 100644
--- a/server/src/main/java/org/apache/druid/curator/discovery/CuratorDruidLeaderSelector.java
+++ b/server/src/main/java/org/apache/druid/curator/discovery/CuratorDruidLeaderSelector.java
@@ -105,8 +105,8 @@ public class CuratorDruidLeaderSelector implements DruidLeaderSelector
             }
             catch (Exception ex) {
               log.makeAlert(ex, "listener becomeLeader() failed. Unable to become leader").emit();
-
-              recreateLeaderLatch();
+              stopAndCreateNewLeaderLatch();
+              startLeaderLatch();
             }
           }
 
@@ -120,8 +120,10 @@ public class CuratorDruidLeaderSelector implements DruidLeaderSelector
               }
 
               leader = false;
+              // give others a chance to become leader.
+              stopAndCreateNewLeaderLatch();
               listener.stopBeingLeader();
-              recreateLeaderLatch();
+              startLeaderLatch();
             }
             catch (Exception ex) {
               log.makeAlert(ex, "listener.stopBeingLeader() failed. Unable to stopBeingLeader").emit();
@@ -206,15 +208,18 @@ public class CuratorDruidLeaderSelector implements DruidLeaderSelector
     listenerExecutor.shutdownNow();
   }
 
-  private void recreateLeaderLatch()
+  private void stopAndCreateNewLeaderLatch()
   {
-    // give others a chance to become leader.
     CloseableUtils.closeAndSuppressExceptions(
         createNewLeaderLatchWithListener(),
         e -> log.warn("Could not close old leader latch; continuing with new one anyway.")
     );
 
     leader = false;
+  }
+
+  private void startLeaderLatch()
+  {
     try {
       //Small delay before starting the latch so that others waiting are chosen to become leader.
       Thread.sleep(ThreadLocalRandom.current().nextInt(1000, 5000));