From c651dba6d114086ce3336780aa3877b14b1eb057 Mon Sep 17 00:00:00 2001 From: Caideyipi <87789683+Caideyipi@users.noreply.github.com> Date: Tue, 15 Sep 2026 18:53:49 +0800 Subject: [PATCH] [Subscription] Balance consensus ownership and report WAL backlog --- .../broker/ConsensusSubscriptionBroker.java | 171 ++++++++++++---- .../consensus/ConsensusPrefetchingQueue.java | 18 +- ...sensusSubscriptionBrokerOwnershipTest.java | 192 ++++++++++++++++++ .../ConsensusPrefetchingQueueTest.java | 54 +++++ 4 files changed, 380 insertions(+), 55 deletions(-) create mode 100644 iotdb-core/datanode/src/test/java/org/apache/iotdb/db/subscription/broker/ConsensusSubscriptionBrokerOwnershipTest.java diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/subscription/broker/ConsensusSubscriptionBroker.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/subscription/broker/ConsensusSubscriptionBroker.java index b911478c2b288..4df735a7b76d6 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/subscription/broker/ConsensusSubscriptionBroker.java +++ b/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 queues, final String consumerId) { - final ConcurrentHashMap consumerTimestamps = - topicConsumerLastPollMs.computeIfAbsent(topicName, ignored -> new ConcurrentHashMap<>()); - consumerTimestamps.put(consumerId, System.currentTimeMillis()); - evictInactiveConsumers(consumerTimestamps); - final List sortedConsumers = new ArrayList<>(consumerTimestamps.keySet()); - Collections.sort(sortedConsumers); - - final List 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 consumerTimestamps = + topicConsumerLastPollMs.computeIfAbsent(topicName, ignored -> new ConcurrentHashMap<>()); + consumerTimestamps.put(consumerId, System.currentTimeMillis()); + evictInactiveConsumers(consumerTimestamps); + final List sortedConsumers = new ArrayList<>(consumerTimestamps.keySet()); + Collections.sort(sortedConsumers); + + final List 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 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 activeConsumers; private final List activeRegionIds; @@ -862,19 +869,70 @@ private TopicOwnershipSnapshot( this.generation = generation; } - private static TopicOwnershipSnapshot create( - final List activeConsumers, final List activeRegionIds) { + static TopicOwnershipSnapshot create( + final List activeConsumers, + final List 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 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 ownerByRegionId = new LinkedHashMap<>(); + final Map 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 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)), @@ -882,6 +940,29 @@ private static TopicOwnershipSnapshot create( ownerByRegionId.hashCode()); } + private static TopicOwnershipSnapshot empty() { + return new TopicOwnershipSnapshot( + Collections.emptyList(), Collections.emptyList(), Collections.emptyMap(), 0); + } + + private static String findLeastLoadedConsumer( + final List activeConsumers, final Map regionCountByConsumer) { + return activeConsumers.stream() + .min( + Comparator.comparingInt((String consumer) -> regionCountByConsumer.get(consumer)) + .thenComparing(Comparator.naturalOrder())) + .orElseThrow(IllegalStateException::new); + } + + private static String findMostLoadedConsumer( + final List activeConsumers, final Map 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 regionIds) { return activeRegionIds.equals(regionIds); } - private String getOwnerConsumerId(final String regionId) { + String getOwnerConsumerId(final String regionId) { return ownerByRegionId.get(regionId); } diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/subscription/broker/consensus/ConsensusPrefetchingQueue.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/subscription/broker/consensus/ConsensusPrefetchingQueue.java index c6fd52301a428..ecb01d2868e54 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/subscription/broker/consensus/ConsensusPrefetchingQueue.java +++ b/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. * - *

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. + *

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 ======================== diff --git a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/subscription/broker/ConsensusSubscriptionBrokerOwnershipTest.java b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/subscription/broker/ConsensusSubscriptionBrokerOwnershipTest.java new file mode 100644 index 0000000000000..99bdb83fdd934 --- /dev/null +++ b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/subscription/broker/ConsensusSubscriptionBrokerOwnershipTest.java @@ -0,0 +1,192 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you 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 org.apache.iotdb.db.subscription.broker; + +import org.apache.iotdb.db.subscription.broker.ConsensusSubscriptionBroker.TopicOwnershipSnapshot; + +import org.junit.Test; + +import java.util.ArrayList; +import java.util.Arrays; +import java.util.Collections; +import java.util.HashMap; +import java.util.List; +import java.util.Map; +import java.util.stream.Collectors; +import java.util.stream.IntStream; + +import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertFalse; +import static org.junit.Assert.assertNull; +import static org.junit.Assert.assertTrue; + +public class ConsensusSubscriptionBrokerOwnershipTest { + + @Test + public void testEqualNumbersOfConsumersAndRegionsAssignEveryConsumer() { + final List consumers = consumerIds(40); + final List regions = regionIds(3, 42); + + final TopicOwnershipSnapshot snapshot = TopicOwnershipSnapshot.create(consumers, regions, null); + final Map loads = loads(snapshot, consumers, regions); + + assertEquals(40, loads.size()); + assertTrue(loads.values().stream().allMatch(load -> load == 1)); + } + + @Test + public void testJoiningConsumerTriggersOnlyRequiredMoves() { + final List regions = regionIds(3, 42); + final TopicOwnershipSnapshot oneConsumer = + TopicOwnershipSnapshot.create(Collections.singletonList("consumer_1"), regions, null); + final TopicOwnershipSnapshot twoConsumers = + TopicOwnershipSnapshot.create( + Arrays.asList("consumer_1", "consumer_2"), regions, oneConsumer); + + assertEquals(20, movedRegionCount(oneConsumer, twoConsumers, regions)); + assertBalanced(twoConsumers, Arrays.asList("consumer_1", "consumer_2"), regions); + + final TopicOwnershipSnapshot threeConsumers = + TopicOwnershipSnapshot.create( + Arrays.asList("consumer_1", "consumer_2", "consumer_3"), regions, twoConsumers); + assertEquals(13, movedRegionCount(twoConsumers, threeConsumers, regions)); + assertBalanced( + threeConsumers, Arrays.asList("consumer_1", "consumer_2", "consumer_3"), regions); + } + + @Test + public void testLeavingConsumerOnlyReassignsItsRegions() { + final List consumers = + Arrays.asList("consumer_1", "consumer_2", "consumer_3", "consumer_4"); + final List regions = regionIds(3, 42); + final TopicOwnershipSnapshot before = TopicOwnershipSnapshot.create(consumers, regions, null); + final List remainingConsumers = Arrays.asList("consumer_1", "consumer_3", "consumer_4"); + final TopicOwnershipSnapshot after = + TopicOwnershipSnapshot.create(remainingConsumers, regions, before); + + for (final String region : regions) { + if (!"consumer_2".equals(before.getOwnerConsumerId(region))) { + assertEquals(before.getOwnerConsumerId(region), after.getOwnerConsumerId(region)); + } + } + assertEquals(10, movedRegionCount(before, after, regions)); + assertBalanced(after, remainingConsumers, regions); + } + + @Test + public void testRegionChangesPreserveValidOwnership() { + final List consumers = Arrays.asList("consumer_1", "consumer_2", "consumer_3"); + final List initialRegions = regionIds(3, 11); + final TopicOwnershipSnapshot initial = + TopicOwnershipSnapshot.create(consumers, initialRegions, null); + + final List expandedRegions = new ArrayList<>(initialRegions); + expandedRegions.addAll(regionIds(12, 14)); + final TopicOwnershipSnapshot expanded = + TopicOwnershipSnapshot.create(consumers, expandedRegions, initial); + assertEquals(0, movedRegionCount(initial, expanded, initialRegions)); + assertBalanced(expanded, consumers, expandedRegions); + + final List reducedRegions = new ArrayList<>(expandedRegions); + reducedRegions.remove("DataRegion[12]"); + reducedRegions.remove("DataRegion[13]"); + reducedRegions.remove("DataRegion[14]"); + final TopicOwnershipSnapshot reduced = + TopicOwnershipSnapshot.create(consumers, reducedRegions, expanded); + assertEquals(0, movedRegionCount(expanded, reduced, reducedRegions)); + assertBalanced(reduced, consumers, reducedRegions); + } + + @Test + public void testMoreConsumersThanRegionsLeavesOnlyUnavoidableConsumersEmpty() { + final List consumers = consumerIds(5); + final List regions = regionIds(3, 5); + + final TopicOwnershipSnapshot snapshot = TopicOwnershipSnapshot.create(consumers, regions, null); + final Map loads = loads(snapshot, consumers, regions); + + assertEquals(3, loads.values().stream().filter(load -> load == 1).count()); + assertEquals(2, loads.values().stream().filter(load -> load == 0).count()); + assertBalanced(snapshot, consumers, regions); + } + + @Test + public void testEmptyInputsProduceEmptyOwnership() { + final TopicOwnershipSnapshot noConsumers = + TopicOwnershipSnapshot.create( + Collections.emptyList(), Collections.singletonList("DataRegion[3]"), null); + assertNull(noConsumers.getOwnerConsumerId("DataRegion[3]")); + + final TopicOwnershipSnapshot noRegions = + TopicOwnershipSnapshot.create( + Collections.singletonList("consumer_1"), Collections.emptyList(), null); + assertNull(noRegions.getOwnerConsumerId("DataRegion[3]")); + } + + private static List consumerIds(final int count) { + return IntStream.rangeClosed(1, count) + .mapToObj(index -> "consumer_" + index) + .sorted() + .collect(Collectors.toList()); + } + + private static List regionIds(final int startInclusive, final int endInclusive) { + return IntStream.rangeClosed(startInclusive, endInclusive) + .mapToObj(index -> "DataRegion[" + index + "]") + .sorted() + .collect(Collectors.toList()); + } + + private static Map loads( + final TopicOwnershipSnapshot snapshot, + final List consumers, + final List regions) { + final Map result = new HashMap<>(); + consumers.forEach(consumer -> result.put(consumer, 0)); + for (final String region : regions) { + result.computeIfPresent(snapshot.getOwnerConsumerId(region), (ignored, count) -> count + 1); + } + return result; + } + + private static void assertBalanced( + final TopicOwnershipSnapshot snapshot, + final List consumers, + final List regions) { + final Map loads = loads(snapshot, consumers, regions); + assertFalse(loads.isEmpty()); + final int minimumLoad = Collections.min(loads.values()); + final int maximumLoad = Collections.max(loads.values()); + assertTrue(maximumLoad - minimumLoad <= 1); + assertEquals(regions.size(), loads.values().stream().mapToInt(Integer::intValue).sum()); + } + + private static int movedRegionCount( + final TopicOwnershipSnapshot before, + final TopicOwnershipSnapshot after, + final List regions) { + return (int) + regions.stream() + .filter( + region -> + !before.getOwnerConsumerId(region).equals(after.getOwnerConsumerId(region))) + .count(); + } +} diff --git a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/subscription/broker/consensus/ConsensusPrefetchingQueueTest.java b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/subscription/broker/consensus/ConsensusPrefetchingQueueTest.java index a6d76a8d98aea..16182cae92413 100644 --- a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/subscription/broker/consensus/ConsensusPrefetchingQueueTest.java +++ b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/subscription/broker/consensus/ConsensusPrefetchingQueueTest.java @@ -395,6 +395,7 @@ public void testLagIncludesLingeringBatchUntilCommitted() throws Exception { 1L, 1L, true); + queue.setSubscriptionMemoryManager(new SubscriptionMemoryManager(16L * 1024 * 1024)); final IndexedConsensusRequest request = new IndexedConsensusRequest( 1L, Collections.singletonList(StatementTestUtils.genInsertRowNode(1))) @@ -434,6 +435,59 @@ public void testLagIncludesLingeringBatchUntilCommitted() throws Exception { } } + @Test + public void testLagIncludesUnreadWalSearchIndexDistance() throws Exception { + final String originalSystemDir = IoTDBDescriptor.getInstance().getConfig().getSystemDir(); + final File systemDir = temporaryFolder.newFolder("lagWithUnreadWal"); + ConsensusPrefetchingQueue queue = null; + try { + final DataRegionId regionId = new DataRegionId(9); + final FakeConsensusReqReader reader = new FakeConsensusReqReader(); + reader.currentSearchIndex = 300_000L; + final IoTConsensusServerImpl serverImpl = mock(IoTConsensusServerImpl.class); + when(serverImpl.getConsensusReqReader()).thenReturn(reader); + when(serverImpl.getWriterSafeFrontierTracker()).thenReturn(new WriterSafeFrontierTracker()); + final ConsensusLogToTabletConverter converter = mock(ConsensusLogToTabletConverter.class); + when(converter.getDatabaseName()).thenReturn("db"); + when(converter.convert(any())).thenReturn(Collections.singletonList(createTablet())); + queue = + new ConsensusPrefetchingQueue( + "consumerGroup", + "topic", + TopicConstant.ORDER_MODE_LEADER_ONLY_VALUE, + regionId, + serverImpl, + new SubscriptionWalRetentionPolicy( + "topic", + SubscriptionWalRetentionPolicy.UNBOUNDED, + SubscriptionWalRetentionPolicy.UNBOUNDED), + converter, + newCommitManager(systemDir), + new RegionProgress(Collections.emptyMap()), + 1L, + 1L, + true); + queue.setSubscriptionMemoryManager(new SubscriptionMemoryManager(16L * 1024 * 1024)); + + assertEquals(300_000L, queue.getRawWalGap()); + assertEquals(300_000L, queue.getLag()); + + assertNull(queue.poll("consumer")); + assertTrue(pendingEntries(queue).offer(createRequest(1L))); + queue.drivePrefetchOnce(); + + assertEquals(2L, queue.getCurrentReadSearchIndex()); + assertEquals(299_999L, queue.getRawWalGap()); + assertEquals(1L, queue.getRemainingEventCount()); + assertEquals(300_000L, queue.getLag()); + } finally { + if (queue != null) { + queue.close(); + } + IoTDBDescriptor.getInstance().getConfig().setSystemDir(originalSystemDir); + } + } + @Test public void testFilteredEmptyEntryAdvancesProgressWithoutEvent() throws Exception { final String originalSystemDir = IoTDBDescriptor.getInstance().getConfig().getSystemDir();