[Subscription] Balance consensus ownership and report WAL backlog

iotdb-core/datanode/src/main/java/org/apache/iotdb/db/subscription/broker/ConsensusSubscriptionBroker.java

@@ -42,6 +42,8 @@
 import java.util.ArrayList;
 import java.util.Collections;
 import java.util.Comparator;
+import java.util.HashMap;
+import java.util.LinkedHashMap;
 import java.util.LinkedHashSet;
 import java.util.List;
 import java.util.Map;
@@ -516,39 +518,46 @@ private TopicOwnershipSnapshot refreshAndGetTopicOwnership(
       final String topicName,
       final List<ConsensusPrefetchingQueue> queues,
       final String consumerId) {
-    final ConcurrentHashMap<String, Long> consumerTimestamps =
-        topicConsumerLastPollMs.computeIfAbsent(topicName, ignored -> new ConcurrentHashMap<>());
-    consumerTimestamps.put(consumerId, System.currentTimeMillis());
-    evictInactiveConsumers(consumerTimestamps);
-    final List<String> sortedConsumers = new ArrayList<>(consumerTimestamps.keySet());
-    Collections.sort(sortedConsumers);
-
-    final List<String> activeRegionIds =
-        queues.stream()
-            .filter(q -> !q.isClosed())
-            .map(q -> q.getConsensusGroupId().toString())
-            .sorted()
-            .collect(Collectors.toList());
-
-    final TopicOwnershipSnapshot existingSnapshot = topicOwnershipSnapshots.get(topicName);
-    if (Objects.nonNull(existingSnapshot)
-        && existingSnapshot.hasSameConsumers(sortedConsumers)
-        && existingSnapshot.hasSameRegions(activeRegionIds)) {
-      return existingSnapshot;
-    }
-
-    final TopicOwnershipSnapshot refreshedSnapshot =
-        TopicOwnershipSnapshot.create(sortedConsumers, activeRegionIds);
-    topicOwnershipSnapshots.put(topicName, refreshedSnapshot);
-    LOGGER.debug(
-        DataNodePipeMessages
-            .PIPE_LOG_CONSENSUSSUBSCRIPTIONBROKER_REFRESHED_OWNERSHIP_FOR_TOPIC_EB11CF64,
-        brokerId,
-        topicName,
-        sortedConsumers,
-        activeRegionIds,
-        refreshedSnapshot.getGeneration());
-    return refreshedSnapshot;
+    synchronized (queueLifecycleLock) {
+      // Do not recreate ownership metadata for a topic that was concurrently removed.
+      if (topicNameToConsensusPrefetchingQueues.get(topicName) != queues) {
+        return TopicOwnershipSnapshot.empty();
+      }
+
+      final ConcurrentHashMap<String, Long> consumerTimestamps =
+          topicConsumerLastPollMs.computeIfAbsent(topicName, ignored -> new ConcurrentHashMap<>());
+      consumerTimestamps.put(consumerId, System.currentTimeMillis());
+      evictInactiveConsumers(consumerTimestamps);
+      final List<String> sortedConsumers = new ArrayList<>(consumerTimestamps.keySet());
+      Collections.sort(sortedConsumers);
+
+      final List<String> activeRegionIds =
+          queues.stream()
+              .filter(q -> !q.isClosed())
+              .map(q -> q.getConsensusGroupId().toString())
+              .sorted()
+              .collect(Collectors.toList());
+
+      final TopicOwnershipSnapshot existingSnapshot = topicOwnershipSnapshots.get(topicName);
+      if (Objects.nonNull(existingSnapshot)
+          && existingSnapshot.hasSameConsumers(sortedConsumers)
+          && existingSnapshot.hasSameRegions(activeRegionIds)) {
+        return existingSnapshot;
+      }
+
+      final TopicOwnershipSnapshot refreshedSnapshot =
+          TopicOwnershipSnapshot.create(sortedConsumers, activeRegionIds, existingSnapshot);
+      topicOwnershipSnapshots.put(topicName, refreshedSnapshot);
+      LOGGER.debug(
+          DataNodePipeMessages
+              .PIPE_LOG_CONSENSUSSUBSCRIPTIONBROKER_REFRESHED_OWNERSHIP_FOR_TOPIC_EB11CF64,
+          brokerId,
+          topicName,
+          sortedConsumers,
+          activeRegionIds,
+          refreshedSnapshot.getGeneration());
+      return refreshedSnapshot;
+    }
   }
 
   private List<ConsensusPrefetchingQueue> getAssignedQueues(
@@ -727,8 +736,6 @@ public int unbindByRegion(final ConsensusGroupId regionId) {
           topicNameToConsensusPrefetchingQueues.remove(entry.getKey(), queues);
           topicConsumerLastPollMs.remove(entry.getKey());
           topicOwnershipSnapshots.remove(entry.getKey());
-        } else {
-          topicOwnershipSnapshots.remove(entry.getKey());
         }
       }
     }
@@ -844,7 +851,7 @@ private void closeAndRemoveConsensusPrefetchingQueues(
         brokerId);
   }
 
-  private static final class TopicOwnershipSnapshot {
+  static final class TopicOwnershipSnapshot {
 
     private final List<String> activeConsumers;
     private final List<String> activeRegionIds;
@@ -862,26 +869,100 @@ private TopicOwnershipSnapshot(
       this.generation = generation;
     }
 
-    private static TopicOwnershipSnapshot create(
-        final List<String> activeConsumers, final List<String> activeRegionIds) {
+    static TopicOwnershipSnapshot create(
+        final List<String> activeConsumers,
+        final List<String> activeRegionIds,
+        final TopicOwnershipSnapshot previousSnapshot) {
       if (activeConsumers.isEmpty() || activeRegionIds.isEmpty()) {
         return new TopicOwnershipSnapshot(
-            Collections.emptyList(), Collections.emptyList(), Collections.emptyMap(), 0);
+            Collections.unmodifiableList(new ArrayList<>(activeConsumers)),
+            Collections.unmodifiableList(new ArrayList<>(activeRegionIds)),
+            Collections.emptyMap(),
+            0);
       }
 
-      final Map<String, String> ownerByRegionId = new ConcurrentHashMap<>();
-      final int consumerCount = activeConsumers.size();
-      for (final String regionId : activeRegionIds) {
-        final int ownerIdx = Math.floorMod(regionId.hashCode(), consumerCount);
-        ownerByRegionId.put(regionId, activeConsumers.get(ownerIdx));
+      final Map<String, String> ownerByRegionId = new LinkedHashMap<>();
+      final Map<String, Integer> regionCountByConsumer = new HashMap<>();
+      activeConsumers.forEach(consumer -> regionCountByConsumer.put(consumer, 0));
+
+      // Keep assignments that are still valid. Reassigning every region whenever membership
+      // changes causes consumers to repeatedly lose their WAL queues and makes empty polls likely.
+      if (Objects.nonNull(previousSnapshot)) {
+        for (final String regionId : activeRegionIds) {
+          final String owner = previousSnapshot.getOwnerConsumerId(regionId);
+          if (Objects.nonNull(owner) && regionCountByConsumer.containsKey(owner)) {
+            ownerByRegionId.put(regionId, owner);
+            regionCountByConsumer.computeIfPresent(owner, (ignored, count) -> count + 1);
+          }
+        }
+      }
+
+      final List<String> unassignedRegionIds =
+          activeRegionIds.stream()
+              .filter(regionId -> !ownerByRegionId.containsKey(regionId))
+              .collect(Collectors.toCollection(ArrayList::new));
+
+      // Assign newly created regions, or regions whose owner left, before moving valid ownership.
+      for (final String regionId : unassignedRegionIds) {
+        final String leastLoadedConsumer =
+            findLeastLoadedConsumer(activeConsumers, regionCountByConsumer);
+        ownerByRegionId.put(regionId, leastLoadedConsumer);
+        regionCountByConsumer.computeIfPresent(leastLoadedConsumer, (ignored, count) -> count + 1);
+      }
+
+      // Move only enough valid ownerships to make the distribution balanced. Choosing consumers
+      // and regions deterministically keeps ownership stable across JVMs.
+      while (true) {
+        final String leastLoadedConsumer =
+            findLeastLoadedConsumer(activeConsumers, regionCountByConsumer);
+        final String mostLoadedConsumer =
+            findMostLoadedConsumer(activeConsumers, regionCountByConsumer);
+        if (regionCountByConsumer.get(mostLoadedConsumer)
+                - regionCountByConsumer.get(leastLoadedConsumer)
+            <= 1) {
+          break;
+        }
+
+        final String regionToMove =
+            activeRegionIds.stream()
+                .filter(regionId -> mostLoadedConsumer.equals(ownerByRegionId.get(regionId)))
+                .max(Comparator.naturalOrder())
+                .orElseThrow(IllegalStateException::new);
+        ownerByRegionId.put(regionToMove, leastLoadedConsumer);
+        regionCountByConsumer.computeIfPresent(mostLoadedConsumer, (ignored, count) -> count - 1);
+        regionCountByConsumer.computeIfPresent(leastLoadedConsumer, (ignored, count) -> count + 1);
       }
+
       return new TopicOwnershipSnapshot(
           Collections.unmodifiableList(new ArrayList<>(activeConsumers)),
           Collections.unmodifiableList(new ArrayList<>(activeRegionIds)),
           Collections.unmodifiableMap(ownerByRegionId),
           ownerByRegionId.hashCode());
     }
 
+    private static TopicOwnershipSnapshot empty() {
+      return new TopicOwnershipSnapshot(
+          Collections.emptyList(), Collections.emptyList(), Collections.emptyMap(), 0);
+    }
+
+    private static String findLeastLoadedConsumer(
+        final List<String> activeConsumers, final Map<String, Integer> regionCountByConsumer) {
+      return activeConsumers.stream()
+          .min(
+              Comparator.comparingInt((String consumer) -> regionCountByConsumer.get(consumer))
+                  .thenComparing(Comparator.naturalOrder()))
+          .orElseThrow(IllegalStateException::new);
+    }
+
+    private static String findMostLoadedConsumer(
+        final List<String> activeConsumers, final Map<String, Integer> regionCountByConsumer) {
+      return activeConsumers.stream()
+          .min(
+              Comparator.comparingInt((String consumer) -> -regionCountByConsumer.get(consumer))
+                  .thenComparing(Comparator.naturalOrder()))
+          .orElseThrow(IllegalStateException::new);
+    }
+
     private boolean isEmpty() {
       return activeConsumers.isEmpty() || activeRegionIds.isEmpty();
     }
@@ -894,7 +975,7 @@ private boolean hasSameRegions(final List<String> regionIds) {
       return activeRegionIds.equals(regionIds);
     }
 
-    private String getOwnerConsumerId(final String regionId) {
+    String getOwnerConsumerId(final String regionId) {
       return ownerByRegionId.get(regionId);
     }
 

iotdb-core/datanode/src/main/java/org/apache/iotdb/db/subscription/broker/consensus/ConsensusPrefetchingQueue.java

@@ -4258,21 +4258,19 @@ public ConsensusGroupId getConsensusGroupId() {
   /**
    * Returns the queue-local lag used by metrics.
    *
-   * <p>Events that have already been materialized in memory are counted exactly. For data that is
-   * still only in WAL, the exact number is not tracked by this queue and computing it would require
-   * scanning WAL only for reporting. Therefore, unread WAL data is represented as one extra unit,
-   * so the metric shows that this queue is not caught up without turning lag reporting into another
-   * WAL reader.
+   * <p>Entries in the materialized lifecycle stages have already advanced the WAL cursor. Pending
+   * entries have not, so they overlap with the raw WAL search-index gap. Taking the maximum for the
+   * unmaterialized part avoids double-counting that overlap while still exposing a large unread WAL
+   * backlog instead of collapsing it to one unit.
    */
   public long getLag() {
-    final long queuedLag =
-        prefetchingQueue.size()
+    final long materializedLag =
+        (long) prefetchingQueue.size()
             + inFlightEvents.size()
-            + pendingEntries.size()
             + getRealtimeBufferedEntryCount()
             + lingerBatch.getEntryCount();
-    final boolean hasUnreadWalEntries = hasUnreadWalEntriesBehindCursor();
-    return queuedLag + (hasUnreadWalEntries ? 1 : 0);
+    final long unmaterializedLag = Math.max((long) pendingEntries.size(), getRawWalGap());
+    return materializedLag + unmaterializedLag;
   }
 
   // ======================== Stringify ========================