UT for leader graceful shutdown giving up leadership lock

Author: himanshug

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

@@ -0,0 +1,82 @@
+/*
+Copyright 2020 The Kubernetes Authors.
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+http://www.apache.org/licenses/LICENSE-2.0
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+*/
+
+package io.kubernetes.client.extended.leaderelection;
+
+import org.junit.Test;
+
+import java.time.Duration;
+import java.util.concurrent.CountDownLatch;
+import java.util.concurrent.TimeUnit;
+
+/**
+ * Leader Election tests using "simulated" locks created by {@link LockSmith}
+ */
+public class LeaderElectorTest
+{
+  /**
+   * Tests that when a leader candidate is stopped gracefully, second candidate immediately becomes
+   * leader.
+   */
+  @Test(timeout = 20000L)
+  public void testLeaderGracefulShutdown() throws Exception
+  {
+    LockSmith lockSmith = new LockSmith();
+
+    CountDownLatch startBeingLeader1 = new CountDownLatch(1);
+    CountDownLatch stopBeingLeader1 = new CountDownLatch(1);
+
+    LeaderElector leaderElector1 = makeAndRunLeaderElectorAsync(lockSmith, "candidate1", startBeingLeader1, stopBeingLeader1);
+
+    // wait for candidate1 to become leader
+    startBeingLeader1.await();
+
+    CountDownLatch startBeingLeader2 = new CountDownLatch(1);
+    CountDownLatch stopBeingLeader2 = new CountDownLatch(1);
+
+    LeaderElector leaderElector2 = makeAndRunLeaderElectorAsync(lockSmith, "candidate2", startBeingLeader2, stopBeingLeader2);
+
+    leaderElector1.close();
+
+    // ensure stopBeingLeader hook is called
+    stopBeingLeader1.await();
+
+    // wait for candidate2 to become leader
+    startBeingLeader2.await();
+
+    leaderElector2.close();
+  }
+
+  private LeaderElector makeAndRunLeaderElectorAsync(LockSmith lockSmith, String lockIdentity, CountDownLatch startBeingLeader, CountDownLatch stopBeingLeader)
+  {
+    LeaderElectionConfig leaderElectionConfig =
+        new LeaderElectionConfig(
+            lockSmith.makeLock(lockIdentity),
+            Duration.ofMillis(TimeUnit.MINUTES.toMillis(1)),
+            Duration.ofMillis(TimeUnit.SECONDS.toMillis(51)),
+            Duration.ofMillis(TimeUnit.SECONDS.toMillis(3))
+        );
+    LeaderElector leaderElector = new LeaderElector(leaderElectionConfig);
+
+    Thread thread = new Thread(
+        () -> leaderElector.run(
+            () -> startBeingLeader.countDown(), () -> stopBeingLeader.countDown()
+        ),
+        String.format("%s-leader-elector-main", lockIdentity)
+    );
+    thread.setDaemon(true);
+    thread.start();
+
+    return leaderElector;
+  }
+}

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

@@ -0,0 +1,100 @@
+/*
+Copyright 2020 The Kubernetes Authors.
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+http://www.apache.org/licenses/LICENSE-2.0
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+*/
+
+package io.kubernetes.client.extended.leaderelection;
+
+import io.kubernetes.client.openapi.ApiException;
+
+import java.net.HttpURLConnection;
+import java.util.concurrent.atomic.AtomicReference;
+
+/**
+ * Makes simulated {@link Lock} objects that behave as if they were backed by real API server.
+ */
+public class LockSmith
+{
+  private AtomicReference<Resource> lockResourceRef = new AtomicReference<>();
+
+  public Lock makeLock(String identity)
+  {
+    return new SimulatedLock(identity);
+  }
+
+  private class SimulatedLock implements Lock
+  {
+    private final String identity;
+
+    public SimulatedLock(String identity)
+    {
+      this.identity = identity;
+    }
+
+    @Override
+    public LeaderElectionRecord get() throws ApiException
+    {
+      if (lockResourceRef.get() == null) {
+        throw new ApiException("Record Not Found", HttpURLConnection.HTTP_NOT_FOUND, null, null);
+      }
+
+      return lockResourceRef.get().record;
+    }
+
+    @Override
+    public boolean create(LeaderElectionRecord record)
+    {
+      return lockResourceRef.compareAndSet(null, new Resource(record));
+    }
+
+    @Override
+    public boolean update(LeaderElectionRecord record)
+    {
+      Resource res = lockResourceRef.get();
+      if (res == null) {
+        return false;
+      } else {
+        Resource newResource = new Resource(res.version + 1, record);
+        return lockResourceRef.compareAndSet(res, newResource);
+      }
+    }
+
+    @Override
+    public String identity()
+    {
+      return identity;
+    }
+
+    @Override
+    public String describe()
+    {
+      return "simulated/lock";
+    }
+  }
+
+  private static class Resource
+  {
+    final int version;
+    final LeaderElectionRecord record;
+
+    public Resource(LeaderElectionRecord record)
+    {
+      this.version = 0;
+      this.record = record;
+    }
+
+    public Resource(int version, LeaderElectionRecord record)
+    {
+      this.version = version;
+      this.record = record;
+    }
+  }
+}