[fix][meta] Fix leader elections settling on a lock deleted while they read or write it

pulsar-metadata/src/main/java/org/apache/pulsar/metadata/coordination/impl/LeaderElectionImpl.java

@@ -74,6 +74,17 @@ private enum InternalState {
 
     private InternalState internalState;
 
+    // Counts the deletions of the lock this election was notified of. A deletion can be notified before the result of
+    // a read or a write of the lock that came before it (Oxia delivers them on different threads), so an election
+    // cycle that sees the count change while it reads or writes the lock reads it again instead of settling on what
+    // may be gone.
+    private long lockDeletions;
+
+    // Whether the election may hold the lock although it doesn't lead: an election cycle took the lock, but a
+    // deletion of the lock was notified meanwhile, which may be of an older lock, so the cycle reads the lock again.
+    // Until a cycle settles the election, closing it releases the lock if it's still the election's, as when it leads.
+    private boolean holdsLockProvisionally;
+
     private static final int LEADER_ELECTION_RETRY_DELAY_SECONDS = 5;
 
     // Upper bound for getLeaderValue() waiting on an election that never settles, aligned with the
@@ -150,21 +161,29 @@ private synchronized CompletableFuture<LeaderElectionState> elect() {
         }
         // First check if there's already a leader elected
         internalState = InternalState.ElectionInProgress;
+        long lockDeletionsBeforeRead = lockDeletions;
         return store.get(path).thenCompose(optLock -> {
             if (optLock.isPresent()) {
-                return handleExistingLeaderValue(optLock.get());
+                return handleExistingLeaderValue(optLock.get(), lockDeletionsBeforeRead);
             } else {
                 return tryToBecomeLeader();
             }
         });
     }
 
-    private synchronized CompletableFuture<LeaderElectionState> handleExistingLeaderValue(GetResult res) {
+    private synchronized CompletableFuture<LeaderElectionState> handleExistingLeaderValue(GetResult res,
+            long lockDeletionsBeforeRead) {
         if (internalState == InternalState.Closed) {
             // Closed while the lock was being read: the value read must not bring the election back, nor make it
             // delete the lock
             return FutureUtil.failedFuture(new AlreadyClosedException("The leader election was already closed"));
         }
+        if (lockDeletions != lockDeletionsBeforeRead) {
+            // The lock may have been deleted after it was read: the leader it names may be gone
+            log.debug().attr("path", path)
+                    .log("A deletion of the lock was notified while reading it. Reading it again.");
+            return elect();
+        }
         T existingValue;
         try {
             existingValue = serde.deserialize(path, res.getValue(), res.getStat());
@@ -211,6 +230,7 @@ private synchronized CompletableFuture<LeaderElectionState> handleExistingLeader
 
     private synchronized void changeState(LeaderElectionState les) {
         internalState = InternalState.LeaderIsPresent;
+        holdsLockProvisionally = false;
         if (this.leaderElectionState != les) {
             this.leaderElectionState = les;
             try {
@@ -236,13 +256,23 @@ private synchronized CompletableFuture<LeaderElectionState> tryToBecomeLeader()
         }
 
         CompletableFuture<LeaderElectionState> result = new CompletableFuture<>();
+        long lockDeletionsBeforeWrite = lockDeletions;
         store.put(path, payload, Optional.of(-1L), EnumSet.of(CreateOption.Ephemeral))
                 .thenAccept(stat -> {
                     synchronized (LeaderElectionImpl.this) {
-                        if (internalState == InternalState.ElectionInProgress) {
+                        if (internalState == InternalState.ElectionInProgress
+                                && lockDeletions != lockDeletionsBeforeWrite) {
+                            // The lock may have been deleted once written: read it again rather than lead without it.
+                            // The lock may also still be there, if the deletion notified was of an older lock.
+                            holdsLockProvisionally = true;
+                            log.debug().attr("path", path).attr("value", value)
+                                    .log("A deletion of the lock was notified while taking it. Reading it again.");
+                            FutureUtil.completeAfter(result, elect());
+                        } else if (internalState == InternalState.ElectionInProgress) {
                             log.info().attr("path", path).attr("value", value)
                                     .log("Acquired leadership");
                             internalState = InternalState.LeaderIsPresent;
+                            holdsLockProvisionally = false;
                             leaderKnown(Optional.of(value));
                             if (leaderElectionState != LeaderElectionState.Leading) {
                                 leaderElectionState = LeaderElectionState.Leading;
@@ -313,9 +343,10 @@ public synchronized CompletableFuture<Void> asyncClose() {
             currentLeaderFuture.complete(Optional.empty());
         }
 
-        if (leaderElectionState != LeaderElectionState.Leading) {
+        if (leaderElectionState != LeaderElectionState.Leading && !holdsLockProvisionally) {
             return CompletableFuture.completedFuture(null);
         }
+        holdsLockProvisionally = false;
 
         // Only delete the leader node if it is still ours: if our session expired, the node we
         // created is already gone and the node at the path, if any, belongs to the participant
@@ -448,8 +479,12 @@ void handleSessionNotification(SessionEvent event) {
     // Called for the notifications on the path of this election
     void handlePathNotification(Notification notification) {
         synchronized (this) {
+            if (notification.getType() == NotificationType.Deleted) {
+                lockDeletions++;
+            }
             if (internalState != InternalState.LeaderIsPresent) {
-                // Ignoring notification since we're not trying to become leader
+                // Not taking part in the election, or an election cycle is in progress: it checks the deletions once
+                // it has read or written the lock
                 return;
             }
 

pulsar-metadata/src/test/java/org/apache/pulsar/metadata/coordination/impl/LeaderElectionImplTest.java

@@ -50,6 +50,7 @@
 import org.apache.pulsar.metadata.api.Option;
 import org.apache.pulsar.metadata.api.Stat;
 import org.apache.pulsar.metadata.api.coordination.CoordinationService;
+import org.apache.pulsar.metadata.api.coordination.LeaderElection;
 import org.apache.pulsar.metadata.api.coordination.LeaderElectionState;
 import org.apache.pulsar.metadata.api.extended.MetadataStoreExtended;
 import org.apache.pulsar.metadata.api.extended.SessionEvent;
@@ -232,4 +233,173 @@ public CompletableFuture<Optional<GetResult>> get(String p, Set<Option> opts) {
         assertThat(le.getState()).isEqualTo(LeaderElectionState.NoLeader);
         assertThat(le.getLeaderValue().join()).isEmpty();
     }
+
+    @Test(dataProvider = "distributedImpl", timeOut = 30000)
+    public void candidateDoesNotFollowALeaderGoneWhileReadingTheLock(String provider, Supplier<String> urlSupplier)
+            throws Exception {
+        String path = newKey();
+        @Cleanup
+        MetadataStoreExtended leaderStore = MetadataStoreExtended.create(urlSupplier.get(),
+                MetadataStoreConfig.builder().build());
+        @Cleanup
+        CoordinationService leaderCs = new CoordinationServiceImpl(leaderStore);
+        LeaderElection<String> leader = leaderCs.getLeaderElection(String.class, path, __ -> {
+        });
+        assertThat(leader.elect("leader").join()).isEqualTo(LeaderElectionState.Leading);
+
+        AtomicBoolean holdNextRead = new AtomicBoolean();
+        CompletableFuture<Optional<GetResult>> heldRead = new CompletableFuture<>();
+        CompletableFuture<Void> heldReadReleased = new CompletableFuture<>();
+        @Cleanup
+        FaultInjectionMetadataStore candidateStore = new FaultInjectionMetadataStore(
+                MetadataStoreExtended.create(urlSupplier.get(), MetadataStoreConfig.builder().build())) {
+            @Override
+            public CompletableFuture<Optional<GetResult>> get(String p, Set<Option> opts) {
+                CompletableFuture<Optional<GetResult>> result = super.get(p, opts);
+                if (p.equals(path) && holdNextRead.compareAndSet(true, false)) {
+                    // The store answers, and the election gets the answer once the test releases it
+                    FutureUtil.completeAfter(heldRead, result);
+                    return result.thenCombine(heldReadReleased, (r, __) -> r);
+                }
+                return result;
+            }
+        };
+        @Cleanup
+        CoordinationService candidateCs = new CoordinationServiceImpl(candidateStore);
+        BlockingQueue<LeaderElectionState> stateChanges = new LinkedBlockingQueue<>();
+        LeaderElection<String> candidate = candidateCs.getLeaderElection(String.class, path, stateChanges::add);
+        CompletableFuture<Void> deletionDispatched = deletionDispatched(candidateStore, path);
+
+        holdNextRead.set(true);
+        CompletableFuture<LeaderElectionState> election = candidate.elect("candidate");
+        assertThat(heldRead.get(10, TimeUnit.SECONDS)).as("lock of the leader").isPresent();
+        // The leader goes away, and the candidate's election is notified of it before it gets the lock it read
+        leader.close();
+        deletionDispatched.get(10, TimeUnit.SECONDS);
+        heldReadReleased.complete(null);
+
+        // The candidate doesn't follow the leader that went away: it takes the leadership
+        assertThat(election.get(10, TimeUnit.SECONDS)).isEqualTo(LeaderElectionState.Leading);
+        assertThat(stateChanges).containsExactly(LeaderElectionState.Leading);
+        assertThat(candidate.getLeaderValue().join()).contains("candidate");
+    }
+
+    @Test(timeOut = 20000)
+    public void electionDoesNotLeadWithALockDeletedWhileTakingIt() throws Exception {
+        String path = "/electionDoesNotLeadWithALockDeletedWhileTakingIt";
+        AtomicBoolean holdNextWrite = new AtomicBoolean();
+        CompletableFuture<Stat> heldWrite = new CompletableFuture<>();
+        CompletableFuture<Void> heldWriteReleased = new CompletableFuture<>();
+        @Cleanup
+        FaultInjectionMetadataStore store = new FaultInjectionMetadataStore(
+                new LocalMemoryMetadataStore("memory:local", MetadataStoreConfig.builder().build())) {
+            @Override
+            public CompletableFuture<Stat> put(String p, byte[] value, Optional<Long> expectedVersion,
+                                               Set<Option> opts) {
+                CompletableFuture<Stat> result = super.put(p, value, expectedVersion, opts);
+                if (p.equals(path) && holdNextWrite.compareAndSet(true, false)) {
+                    // The store writes the lock, and the election learns it once the test releases it
+                    FutureUtil.completeAfter(heldWrite, result);
+                    return result.thenCombine(heldWriteReleased, (r, __) -> r);
+                }
+                return result;
+            }
+        };
+        @Cleanup
+        CoordinationServiceImpl cs = new CoordinationServiceImpl(store);
+        BlockingQueue<LeaderElectionState> stateChanges = new LinkedBlockingQueue<>();
+        LeaderElection<String> le = cs.getLeaderElection(String.class, path, stateChanges::add);
+        CompletableFuture<Void> deletionDispatched = deletionDispatched(store, path);
+
+        holdNextWrite.set(true);
+        CompletableFuture<LeaderElectionState> election = le.elect("test-1");
+        heldWrite.get(5, TimeUnit.SECONDS);
+        // The lock is deleted, and the election is notified of it before it learns that it wrote the lock
+        store.delete(path, Optional.empty()).join();
+        deletionDispatched.get(5, TimeUnit.SECONDS);
+        heldWriteReleased.complete(null);
+
+        // The election doesn't lead without the lock: it takes it again
+        assertThat(election.get(5, TimeUnit.SECONDS)).isEqualTo(LeaderElectionState.Leading);
+        assertThat(stateChanges).containsExactly(LeaderElectionState.Leading);
+        assertThat(store.get(path).join()).as("lock of the leader").isPresent();
+    }
+
+    @Test(timeOut = 20000)
+    public void electionClosedWhileReadingTheLockItTookReleasesIt() throws Exception {
+        String path = "/electionClosedWhileReadingTheLockItTookReleasesIt";
+        AtomicBoolean holdNextWrite = new AtomicBoolean();
+        CompletableFuture<Stat> heldWrite = new CompletableFuture<>();
+        CompletableFuture<Void> heldWriteReleased = new CompletableFuture<>();
+        AtomicBoolean holdNextRead = new AtomicBoolean();
+        CompletableFuture<Optional<GetResult>> heldRead = new CompletableFuture<>();
+        CompletableFuture<Void> heldReadReleased = new CompletableFuture<>();
+        @Cleanup
+        FaultInjectionMetadataStore store = new FaultInjectionMetadataStore(
+                new LocalMemoryMetadataStore("memory:local", MetadataStoreConfig.builder().build())) {
+            @Override
+            public CompletableFuture<Stat> put(String p, byte[] value, Optional<Long> expectedVersion,
+                                               Set<Option> opts) {
+                CompletableFuture<Stat> result = super.put(p, value, expectedVersion, opts);
+                if (p.equals(path) && holdNextWrite.compareAndSet(true, false)) {
+                    // The store writes the lock, and the election learns it once the test releases it
+                    FutureUtil.completeAfter(heldWrite, result);
+                    return result.thenCombine(heldWriteReleased, (r, __) -> r);
+                }
+                return result;
+            }
+
+            @Override
+            public CompletableFuture<Optional<GetResult>> get(String p, Set<Option> opts) {
+                CompletableFuture<Optional<GetResult>> result = super.get(p, opts);
+                if (p.equals(path) && holdNextRead.compareAndSet(true, false)) {
+                    // The store answers, and the election gets the answer once the test releases it
+                    FutureUtil.completeAfter(heldRead, result);
+                    return result.thenCombine(heldReadReleased, (r, __) -> r);
+                }
+                return result;
+            }
+        };
+        @Cleanup
+        CoordinationServiceImpl cs = new CoordinationServiceImpl(store);
+        @Cleanup("shutdownNow")
+        ScheduledExecutorService executor = Executors.newSingleThreadScheduledExecutor();
+        BlockingQueue<LeaderElectionState> stateChanges = new LinkedBlockingQueue<>();
+        LeaderElectionImpl<String> le = new LeaderElectionImpl<>(cs, store, String.class, path, stateChanges::add,
+                executor);
+
+        holdNextWrite.set(true);
+        CompletableFuture<LeaderElectionState> election = le.elect("test-1");
+        heldWrite.get(5, TimeUnit.SECONDS);
+        // The election is notified of an older deletion of the lock before it learns that it wrote the lock, so it
+        // reads the lock again, and it's closed before it gets the answer
+        le.handlePathNotification(new Notification(NotificationType.Deleted, path));
+        holdNextRead.set(true);
+        heldWriteReleased.complete(null);
+        assertThat(heldRead.get(5, TimeUnit.SECONDS)).as("lock read again").isPresent();
+        le.close();
+
+        // close() released the lock the election took, although the election didn't lead
+        assertThat(store.get(path).join()).as("lock of the closed election").isEmpty();
+        heldReadReleased.complete(null);
+        assertThatThrownBy(() -> election.get(5, TimeUnit.SECONDS))
+                .hasCauseInstanceOf(AlreadyClosedException.class);
+        assertThat(stateChanges).as("state changes").isEmpty();
+        assertThat(le.getState()).isEqualTo(LeaderElectionState.NoLeader);
+        assertThat(store.get(path).join()).as("lock of the closed election").isEmpty();
+    }
+
+    /**
+     * Completes once the store has called the listeners registered before this one, such as the one of the leader
+     * elections, with the deletion of the path.
+     */
+    private static CompletableFuture<Void> deletionDispatched(MetadataStoreExtended store, String path) {
+        CompletableFuture<Void> dispatched = new CompletableFuture<>();
+        store.registerCancellableListener(n -> {
+            if (n.getType() == NotificationType.Deleted && n.getPath().equals(path)) {
+                dispatched.complete(null);
+            }
+        });
+        return dispatched;
+    }
 }