KAFKA-16530: Fix high-watermark calculation to not assume the leader is in the voter set

raft/src/main/java/org/apache/kafka/raft/KafkaRaftClient.java

@@ -284,7 +284,7 @@ private void updateLeaderEndOffsetAndTimestamp(
     ) {
         final LogOffsetMetadata endOffsetMetadata = log.endOffset();
 
-        if (state.updateLocalState(endOffsetMetadata)) {
+        if (state.updateLocalState(endOffsetMetadata, partitionState.lastVoterSet().voterIds())) {
             onUpdateLeaderHighWatermark(state, currentTimeMs);
         }
 

raft/src/main/java/org/apache/kafka/raft/LeaderState.java

@@ -32,6 +32,7 @@
 import java.util.Collections;
 import java.util.HashMap;
 import java.util.HashSet;
+import java.util.Iterator;
 import java.util.List;
 import java.util.Map;
 import java.util.Objects;
@@ -243,8 +244,9 @@ private boolean maybeUpdateHighWatermark() {
                         );
                         return true;
                     } else if (highWatermarkUpdateOffset < currentHighWatermarkMetadata.offset) {
-                        log.error("The latest computed high watermark {} is smaller than the current " +
-                                "value {}, which suggests that one of the voters has lost committed data. " +
+                        log.info("The latest computed high watermark {} is smaller than the current " +
+                                "value {}, which should only happen when voter set membership changes. If the voter " +
+                                "set has not changed this suggests that one of the voters has lost committed data. " +
                                 "Full voter replication state: {}", highWatermarkUpdateOffset,
                             currentHighWatermarkMetadata.offset, voterStates.values());
                         return false;
@@ -296,10 +298,12 @@ private void logHighWatermarkUpdate(
      * Update the local replica state.
      *
      * @param endOffsetMetadata updated log end offset of local replica
+     * @param lastVoterSet the up-to-date voter set
      * @return true if the high watermark is updated as a result of this call
      */
     public boolean updateLocalState(
-        LogOffsetMetadata endOffsetMetadata
+        LogOffsetMetadata endOffsetMetadata,
+        Set<Integer> lastVoterSet
     ) {
         ReplicaState state = getOrCreateReplicaState(localId);
         state.endOffset.ifPresent(currentEndOffset -> {
@@ -308,7 +312,8 @@ public boolean updateLocalState(
                     "end offset: " + currentEndOffset.offset + " -> " + endOffsetMetadata.offset);
             }
         });
-        state.updateLeaderState(endOffsetMetadata);
+        state.updateLeaderEndOffset(endOffsetMetadata);
+        updateVoterSet(lastVoterSet);
         return maybeUpdateHighWatermark();
     }
 
@@ -341,9 +346,18 @@ public boolean updateReplicaState(
                     state.nodeId, currentEndOffset.offset, fetchOffsetMetadata.offset);
             }
         });
-
-        Optional<LogOffsetMetadata> leaderEndOffsetOpt =
-            voterStates.get(localId).endOffset;
+        Optional<LogOffsetMetadata> leaderEndOffsetOpt;
+        ReplicaState leaderVoterState = voterStates.get(localId);
+        ReplicaState leaderObserverState = observerStates.get(localId);
+        if (leaderVoterState != null) {
+            leaderEndOffsetOpt = leaderVoterState.endOffset;
+        } else if (leaderObserverState != null) {
+            // The leader is not guaranteed to be in the voter set when in the process of being removed from the quorum.
+            log.info("Updating end offset for leader {} which is also an observer.", localId);
+            leaderEndOffsetOpt = leaderObserverState.endOffset;
+        } else {
+            throw new IllegalStateException("Leader state not found for localId " + localId);
+        }
 
         state.updateFollowerState(
             currentTimeMs,
@@ -387,12 +401,16 @@ public long epochStartOffset() {
     private ReplicaState getOrCreateReplicaState(int remoteNodeId) {
         ReplicaState state = voterStates.get(remoteNodeId);
         if (state == null) {
-            observerStates.putIfAbsent(remoteNodeId, new ReplicaState(remoteNodeId, false));
-            return observerStates.get(remoteNodeId);
+            return createObserverState(remoteNodeId);
         }
         return state;
     }
 
+    private ReplicaState createObserverState(int remoteNodeId) {
+        observerStates.putIfAbsent(remoteNodeId, new ReplicaState(remoteNodeId, false));
+        return observerStates.get(remoteNodeId);
+    }
+
     public DescribeQuorumResponseData.PartitionData describeQuorum(long currentTimeMs) {
         clearInactiveObservers(currentTimeMs);
 
@@ -445,6 +463,27 @@ private boolean isVoter(int remoteNodeId) {
         return voterStates.containsKey(remoteNodeId);
     }
 
+    // with Jose's changes this will probably make more sense as VoterSet
+    private void updateVoterSet(Set<Integer> lastVoterSet) {
+        // Remove any voter state that is not in the last voter set. They become observers.
+        for (Iterator<Map.Entry<Integer, ReplicaState>> iter = voterStates.entrySet().iterator(); iter.hasNext(); ) {
+            Integer nodeId = iter.next().getKey();
+            if (!lastVoterSet.contains(nodeId)) {
+                createObserverState(nodeId);
+                iter.remove();
+            }
+        }
+
+        // Add any voter state that is in the last voter set but not in the current voter set. They are removed from
+        // the observerStates if they exist.
+        for (int voterId : lastVoterSet) {
+            if (!voterStates.containsKey(voterId)) {
+                voterStates.put(voterId, new ReplicaState(voterId, false));
+                observerStates.remove(voterId);
+            }
+        }
+    }
+
     private static class ReplicaState implements Comparable<ReplicaState> {
         final int nodeId;
         Optional<LogOffsetMetadata> endOffset;
@@ -462,7 +501,7 @@ public ReplicaState(int nodeId, boolean hasAcknowledgedLeader) {
             this.hasAcknowledgedLeader = hasAcknowledgedLeader;
         }
 
-        void updateLeaderState(
+        void updateLeaderEndOffset(
             LogOffsetMetadata endOffsetMetadata
         ) {
             // For the leader, we only update the end offset. The remaining fields