support on-new-leader hook for leader elector

Author: yue9944882

extended/src/main/java/io/kubernetes/client/extended/leaderelection/LeaderElector.java

@@ -25,6 +25,11 @@ public class LeaderElector {
   private LeaderElectionRecord observedRecord;
   private long observedTimeMilliSeconds;
   private final Consumer<Throwable> exceptionHandler;
+  // used to implement OnNewLeader(), may lag slightly from the
+  // value observedRecord.HolderIdentity if the transition has
+  // not yet been reported.
+  private String reportedLeader;
+  private Consumer<String> onNewLeaderHook;
 
   private ScheduledExecutorService scheduledWorkers =
       Executors.newSingleThreadScheduledExecutor(
@@ -81,7 +86,27 @@ public LeaderElector(LeaderElectionConfig config, Consumer<Throwable> exceptionH
     this.exceptionHandler = exceptionHandler;
   }
 
+  /**
+   * Runs the leader election in foreground.
+   *
+   * @param startLeadingHook called when a LeaderElector client starts leading
+   * @param stopLeadingHook called when a LeaderElector client stops leading
+   */
   public void run(Runnable startLeadingHook, Runnable stopLeadingHook) {
+    run(startLeadingHook, stopLeadingHook, null);
+  }
+
+  /**
+   * Runs the leader election in foreground.
+   *
+   * @param startLeadingHook called when a LeaderElector client starts leading
+   * @param stopLeadingHook called when a LeaderElector client stops leading
+   * @param onNewLeaderHook called when the client observes a leader that is not the previously
+   *     observed leader. This includes the first observed leader when the client starts.
+   */
+  public void run(
+      Runnable startLeadingHook, Runnable stopLeadingHook, Consumer<String> onNewLeaderHook) {
+    this.onNewLeaderHook = onNewLeaderHook;
     log.info("Start leader election with lock {}", config.getLock().describe());
     try {
       if (!acquire()) {
@@ -122,7 +147,9 @@ private boolean acquire() {
                 log.info("Processing tryAcquireOrRenew successfully canceled");
               } catch (Throwable t) {
                 this.exceptionHandler.accept(t);
-                future.cancel(true);
+                future.cancel(true); // make sure the acquire work doesn't overlap
+              } finally {
+                maybeReportTransition();
               }
             },
             0,
@@ -160,6 +187,7 @@ private void renewLoop() {
                     // retry until success or interrupted
                     while (!tryAcquireOrRenew()) {
                       Thread.sleep(retryPeriodMillis);
+                      maybeReportTransition();
                     }
                   } catch (InterruptedException e) {
                     return false;
@@ -176,7 +204,7 @@ private void renewLoop() {
           }
           renewResult = false;
         } finally {
-          future.cancel(true);
+          future.cancel(true); // make the lease worker doesn't overlap
         }
         if (renewResult) {
           if (log.isDebugEnabled()) {
@@ -280,6 +308,20 @@ private boolean isLeader() {
     return this.config.getLock().identity().equals(this.observedRecord.getHolderIdentity());
   }
 
+  private void maybeReportTransition() {
+    if (this.observedRecord == null) {
+      return;
+    }
+    if (this.observedRecord.getHolderIdentity().equals(this.reportedLeader)) {
+      return;
+    }
+    this.reportedLeader = this.observedRecord.getHolderIdentity();
+
+    if (this.onNewLeaderHook != null) {
+      this.hookWorkers.submit(() -> onNewLeaderHook.accept(this.reportedLeader));
+    }
+  }
+
   private ThreadFactory makeThreadFactory(String nameFormat) {
     return new ThreadFactoryBuilder().setDaemon(true).setNameFormat(nameFormat).build();
   }

extended/src/test/java/io/kubernetes/client/extended/leaderelection/LeaderElectionTest.java

@@ -7,6 +7,7 @@
 import io.kubernetes.client.openapi.ApiException;
 import java.time.Duration;
 import java.util.ArrayList;
+import java.util.Date;
 import java.util.List;
 import java.util.concurrent.CountDownLatch;
 import java.util.concurrent.ExecutorService;
@@ -282,6 +283,86 @@ public void testLeaderElectionCaptureException() throws ApiException, Interrupte
     assertEquals(expectedException, actualException.get().getCause());
   }
 
+  @Test
+  public void testLeaderElectionReportLeaderOnStart() throws ApiException, InterruptedException {
+    when(lock.identity()).thenReturn("foo1");
+    when(lock.get())
+        .thenReturn(
+            new LeaderElectionRecord() {
+              {
+                setHolderIdentity("foo2");
+                setAcquireTime(new Date());
+                setRenewTime(new Date());
+                setLeaderTransitions(1);
+                setLeaseDurationSeconds(60);
+              }
+            });
+    List<String> notifications = new ArrayList<>();
+    LeaderElectionConfig leaderElectionConfig = new LeaderElectionConfig();
+    leaderElectionConfig.setLock(lock);
+    leaderElectionConfig.setLeaseDuration(Duration.ofMillis(1000));
+    leaderElectionConfig.setRetryPeriod(Duration.ofMillis(200));
+    leaderElectionConfig.setRenewDeadline(Duration.ofMillis(700));
+    LeaderElector leaderElector = new LeaderElector(leaderElectionConfig);
+    ExecutorService leaderElectionWorker = Executors.newFixedThreadPool(1);
+    leaderElectionWorker.submit(
+        () -> {
+          leaderElector.run(() -> {}, () -> {}, (id) -> notifications.add(id));
+        });
+
+    Thread.sleep(Duration.ofSeconds(2).toMillis());
+
+    when(lock.get())
+        .thenReturn(
+            new LeaderElectionRecord() {
+              {
+                setHolderIdentity("foo3");
+                setAcquireTime(new Date());
+                setRenewTime(new Date());
+                setLeaderTransitions(1);
+                setLeaseDurationSeconds(60);
+              }
+            });
+    Thread.sleep(Duration.ofSeconds(2).toMillis());
+
+    assertEquals(2, notifications.size());
+    assertEquals("foo2", notifications.get(0));
+    assertEquals("foo3", notifications.get(1));
+  }
+
+  @Test
+  public void testLeaderElectionShouldReportLeaderItAcquiresOnStart()
+      throws ApiException, InterruptedException {
+    when(lock.identity()).thenReturn("foo1");
+    when(lock.get())
+        .thenReturn(
+            new LeaderElectionRecord() {
+              {
+                setHolderIdentity("foo1");
+                setAcquireTime(new Date());
+                setRenewTime(new Date());
+                setLeaderTransitions(1);
+                setLeaseDurationSeconds(60);
+              }
+            });
+    List<String> notifications = new ArrayList<>();
+    LeaderElectionConfig leaderElectionConfig = new LeaderElectionConfig();
+    leaderElectionConfig.setLock(lock);
+    leaderElectionConfig.setLeaseDuration(Duration.ofMillis(1000));
+    leaderElectionConfig.setRetryPeriod(Duration.ofMillis(200));
+    leaderElectionConfig.setRenewDeadline(Duration.ofMillis(700));
+    LeaderElector leaderElector = new LeaderElector(leaderElectionConfig);
+    ExecutorService leaderElectionWorker = Executors.newFixedThreadPool(1);
+    leaderElectionWorker.submit(
+        () -> {
+          leaderElector.run(() -> {}, () -> {}, (id) -> notifications.add(id));
+        });
+
+    Thread.sleep(Duration.ofSeconds(2).toMillis());
+    assertEquals(1, notifications.size());
+    assertEquals("foo1", notifications.get(0));
+  }
+
   public static class MockResourceLock implements Lock {
 
     public static ReentrantLock lock;