diff --git a/fe/fe-core/src/main/java/org/apache/doris/cloud/catalog/CloudReplica.java b/fe/fe-core/src/main/java/org/apache/doris/cloud/catalog/CloudReplica.java index e507eb49e8dfc4..caec16887ac038 100644 --- a/fe/fe-core/src/main/java/org/apache/doris/cloud/catalog/CloudReplica.java +++ b/fe/fe-core/src/main/java/org/apache/doris/cloud/catalog/CloudReplica.java @@ -283,6 +283,10 @@ public long getClusterPrimaryBackendId(String clusterId) { return primaryClusterToBackend.getOrDefault(clusterId, -1L); } + Long getNonColocatedPrimaryBackendId(String clusterId) { + return primaryClusterToBackend.get(clusterId); + } + // For proc display only. In cloud mode a replica is hashed to a different BE in each // compute group, so expose a clusterId -> backendId mapping; the proc display builds // a separate bucket sequence per compute group from it so each group's sequence is diff --git a/fe/fe-core/src/main/java/org/apache/doris/cloud/catalog/CloudTabletRebalancer.java b/fe/fe-core/src/main/java/org/apache/doris/cloud/catalog/CloudTabletRebalancer.java index fde9c0d4850741..63787f2ccbca14 100644 --- a/fe/fe-core/src/main/java/org/apache/doris/cloud/catalog/CloudTabletRebalancer.java +++ b/fe/fe-core/src/main/java/org/apache/doris/cloud/catalog/CloudTabletRebalancer.java @@ -52,6 +52,7 @@ import org.apache.doris.thrift.TWarmUpCacheAsyncRequest; import org.apache.doris.thrift.TWarmUpCacheAsyncResponse; +import com.google.common.annotations.VisibleForTesting; import com.google.common.base.Preconditions; import com.google.common.base.Strings; import com.google.common.collect.Sets; @@ -77,10 +78,14 @@ import java.util.concurrent.LinkedBlockingQueue; import java.util.concurrent.ScheduledExecutorService; import java.util.concurrent.TimeUnit; +import java.util.function.Function; import java.util.stream.Collectors; public class CloudTabletRebalancer extends MasterDaemon { private static final Logger LOG = LogManager.getLogger(CloudTabletRebalancer.class); + private static final int MAX_GLOBAL_TABLET_SET_INITIAL_CAPACITY = 1 << 16; + private static final Function> DEFAULT_GLOBAL_TABLET_SET_FACTORY = + ignored -> ConcurrentHashMap.newKeySet(); private volatile ConcurrentHashMap> beToTabletsGlobal = new ConcurrentHashMap>(); @@ -311,7 +316,7 @@ public enum StatType { } @Getter - private class InfightTablet { + private static class InfightTablet { private final long tabletId; private final String clusterId; @@ -334,7 +339,9 @@ public boolean equals(Object o) { @Override public int hashCode() { - return Objects.hash(tabletId, clusterId); + int result = 1; + result = 31 * result + Long.hashCode(tabletId); + return 31 * result + clusterId.hashCode(); } } @@ -436,39 +443,50 @@ private class TransferPairInfo { } public Set getSnapshotTabletsInPrimaryByBeId(Long beId) { - Set tabletIds = Sets.newHashSet(); Set tablets = beToTabletsGlobal.get(beId); - if (tablets != null) { - // Create a copy - tabletIds.addAll(new HashSet<>(tablets)); - } - Set colocateTablets = beToColocateTabletsGlobal.get(beId); - if (colocateTablets != null) { - // Create a copy - tabletIds.addAll(new HashSet<>(colocateTablets)); - } + Set tabletIds = newSnapshotTabletSet(tabletSetSize(tablets) + tabletSetSize(colocateTablets)); + addSnapshotTablets(tabletIds, tablets); + addSnapshotTablets(tabletIds, colocateTablets); return tabletIds; } public Set getSnapshotTabletsInSecondaryByBeId(Long beId) { - Set tabletIds = Sets.newHashSet(); Set tablets = beToTabletsGlobalInSecondary.get(beId); - if (tablets != null) { - // Create a copy - tabletIds.addAll(new HashSet<>(tablets)); - } + Set tabletIds = newSnapshotTabletSet(tabletSetSize(tablets)); + addSnapshotTablets(tabletIds, tablets); return tabletIds; } public Set getSnapshotTabletsInPrimaryAndSecondaryByBeId(Long beId) { - Set tabletIds = Sets.newHashSet(); - tabletIds.addAll(getSnapshotTabletsInPrimaryByBeId(beId)); - tabletIds.addAll(getSnapshotTabletsInSecondaryByBeId(beId)); + Set primaryTablets = beToTabletsGlobal.get(beId); + Set colocateTablets = beToColocateTabletsGlobal.get(beId); + Set secondaryTablets = beToTabletsGlobalInSecondary.get(beId); + int expectedSize = tabletSetSize(primaryTablets) + + tabletSetSize(colocateTablets) + tabletSetSize(secondaryTablets); + Set tabletIds = newSnapshotTabletSet(expectedSize); + addSnapshotTablets(tabletIds, primaryTablets); + addSnapshotTablets(tabletIds, colocateTablets); + addSnapshotTablets(tabletIds, secondaryTablets); return tabletIds; } + private static int tabletSetSize(Set tablets) { + return tablets == null ? 0 : tablets.size(); + } + + private static void addSnapshotTablets(Set snapshot, Set tablets) { + if (tablets != null) { + snapshot.addAll(tablets); + } + } + + @VisibleForTesting + protected Set newSnapshotTabletSet(int expectedSize) { + return Sets.newHashSetWithExpectedSize(expectedSize); + } + public int getTabletNumByBackendId(long beId) { Map> sourceMap = beToTabletsGlobal; ConcurrentHashMap> futureMap = futureBeToTabletsGlobal; @@ -964,34 +982,38 @@ private boolean completeRouteInfo() { long needRehashDeadTime = System.currentTimeMillis() - Config.rehash_tablet_after_be_dead_seconds * 1000L; loopCloudReplica((Database db, Table table, Partition partition, MaterializedIndex index, String cluster) -> { boolean assigned = false; - List beIds = new ArrayList(); - List tabletIds = new ArrayList(); + List tablets = index.getTablets(); boolean isColocated = Env.getCurrentColocateIndex().isColocateTable(table.getId()); - for (Tablet tablet : index.getTablets()) { + int routeCount = isColocated ? 0 : tablets.size(); + List beIds = newRouteInfoList(routeCount); + List tabletIds = newRouteInfoList(routeCount); + for (Tablet tablet : tablets) { for (Replica r : tablet.getReplicas()) { CloudReplica replica = (CloudReplica) r; // clean secondary map replica.checkAndClearSecondaryClusterToBe(cluster, needRehashDeadTime); - InfightTablet taskKey = new InfightTablet(tablet.getId(), cluster); // colocate table no need to update primary backends if (isColocated) { replica.clearClusterToBe(cluster); - tabletToInfightTask.remove(taskKey); + tabletToInfightTask.remove(new InfightTablet(tablet.getId(), cluster)); continue; } // primary backend is alive or dead not long - Backend be = replica.getPrimaryBackend(cluster, false); + Long primaryBeId = replica.getNonColocatedPrimaryBackendId(cluster); + Backend be = primaryBeId == null + ? null : Env.getCurrentSystemInfo().getBackendByIdWithBoxedId(primaryBeId); if (be != null && (be.isQueryAvailable() || (!be.isQueryDisabled() // Compatible with older version upgrades, see https://github.com/apache/doris/pull/42986 && (be.getLastUpdateMs() <= 0 || be.getLastUpdateMs() > needRehashDeadTime)))) { - beIds.add(be.getId()); + beIds.add(primaryBeId); tabletIds.add(tablet.getId()); continue; } // primary backend not available too long, change one + InfightTablet taskKey = new InfightTablet(tablet.getId(), cluster); long beId = -1L; be = replica.getSecondaryBackend(cluster); if (be != null && be.isQueryAvailable()) { @@ -1053,6 +1075,11 @@ private boolean completeRouteInfo() { return true; } + @VisibleForTesting + protected List newRouteInfoList(int initialCapacity) { + return new ArrayList<>(initialCapacity); + } + public void fillBeToTablets(long be, long tableId, long partId, long indexId, long tabletId, ConcurrentHashMap> globalBeToTablets, ConcurrentHashMap>> beToTabletsInTable, @@ -1067,8 +1094,18 @@ void fillBeToTablets(Long be, Long tableId, Long partId, Long indexId, Long tabl ConcurrentHashMap>> beToTabletsInTable, ConcurrentHashMap>>> partToTablets) { + fillBeToTablets(be, tableId, partId, indexId, tabletId, DEFAULT_GLOBAL_TABLET_SET_FACTORY, + globalBeToTablets, beToTabletsInTable, partToTablets); + } + + private void fillBeToTablets(Long be, Long tableId, Long partId, Long indexId, Long tabletId, + Function> globalTabletSetFactory, + ConcurrentHashMap> globalBeToTablets, + ConcurrentHashMap>> beToTabletsInTable, + ConcurrentHashMap>>> + partToTablets) { // global - globalBeToTablets.computeIfAbsent(be, ignored -> ConcurrentHashMap.newKeySet()).add(tabletId); + globalBeToTablets.computeIfAbsent(be, globalTabletSetFactory).add(tabletId); // table ConcurrentHashMap> beToTabletsOfTable = @@ -1083,6 +1120,23 @@ void fillBeToTablets(Long be, Long tableId, Long partId, Long indexId, Long tabl beToTabletsOfIndex.computeIfAbsent(be, ignored -> ConcurrentHashMap.newKeySet()).add(tabletId); } + private Function> newGlobalTabletSetFactory(Map> previousBeToTablets) { + Map> previousRoute = previousBeToTablets == null + ? Collections.emptyMap() : previousBeToTablets; + return be -> { + Set previousTablets = previousRoute.get(be); + int initialCapacity = previousTablets == null ? 0 + : Math.min(previousTablets.size(), MAX_GLOBAL_TABLET_SET_INITIAL_CAPACITY); + return newGlobalTabletSet(initialCapacity); + }; + } + + @VisibleForTesting + protected Set newGlobalTabletSet(int initialCapacity) { + return initialCapacity == 0 + ? ConcurrentHashMap.newKeySet() : ConcurrentHashMap.newKeySet(initialCapacity); + } + private void enqueueWarmupTask(WarmupTabletTask task) { WarmupBatchKey key = new WarmupBatchKey(task.srcBe, task.destBe); WarmupBatch batch = warmupBatches.computeIfAbsent(key, WarmupBatch::new); @@ -1158,6 +1212,12 @@ private void flushExpiredWarmupBatches() { } public void statRouteInfo() { + // The previous generation remains live until the temporary global routes are complete, so reuse its + // per-backend cardinalities as allocation hints without extending its lifetime. + Function> currentGlobalTabletSetFactory = + newGlobalTabletSetFactory(beToTabletsGlobal); + Function> futureGlobalTabletSetFactory = + newGlobalTabletSetFactory(futureBeToTabletsGlobal); ConcurrentHashMap> tmpBeToTabletsGlobal = new ConcurrentHashMap>(); ConcurrentHashMap> tmpFutureBeToTabletsGlobal = new ConcurrentHashMap>(); ConcurrentHashMap> tmpBeToTabletsGlobalInSecondary @@ -1201,8 +1261,10 @@ public void statRouteInfo() { tmpPartitionActive.merge(partitionId, 1L, Long::sum); tmpDbActive.merge(dbId, 1L, Long::sum); } - for (Replica r : tablet.getReplicas()) { - CloudReplica replica = (CloudReplica) r; + List replicas = tablet.getReplicas(); + int replicaCount = replicas.size(); + for (int replicaIndex = 0; replicaIndex < replicaCount; replicaIndex++) { + CloudReplica replica = (CloudReplica) replicas.get(replicaIndex); if (isColocated) { Long beId = -1L; try { @@ -1218,8 +1280,10 @@ public void statRouteInfo() { continue; } - Backend be = replica.getPrimaryBackend(cluster, false); - Long beId = be == null ? Long.valueOf(-1L) : Long.valueOf(be.getId()); + Long primaryBeId = replica.getNonColocatedPrimaryBackendId(cluster); + Backend be = primaryBeId == null + ? null : Env.getCurrentSystemInfo().getBackendByIdWithBoxedId(primaryBeId); + Long beId = be == null ? Long.valueOf(-1L) : primaryBeId; if (!allBes.contains(beId)) { continue; } @@ -1232,14 +1296,16 @@ public void statRouteInfo() { tablets.add(tabletId); } - InfightTablet taskKey = new InfightTablet(tabletId, cluster); - InfightTask task = tabletToInfightTask.get(taskKey); + InfightTask task = tabletToInfightTask.isEmpty() ? null + : tabletToInfightTask.get(new InfightTablet(tabletId, cluster)); Long futureBeId = task == null ? beId : Long.valueOf(task.destBe); Long routeTabletId = task == null ? tabletId : task.pickedTabletId; fillBeToTablets(beId, tableId, partitionId, indexId, routeTabletId, + currentGlobalTabletSetFactory, tmpBeToTabletsGlobal, beToTabletsInTable, this.partitionToTablets); fillBeToTablets(futureBeId, tableId, partitionId, indexId, routeTabletId, + futureGlobalTabletSetFactory, tmpFutureBeToTabletsGlobal, futureBeToTabletsInTable, futurePartitionToTablets); } } diff --git a/fe/fe-core/src/main/java/org/apache/doris/system/SystemInfoService.java b/fe/fe-core/src/main/java/org/apache/doris/system/SystemInfoService.java index 8b5a80a978ed8a..5713be5b965395 100644 --- a/fe/fe-core/src/main/java/org/apache/doris/system/SystemInfoService.java +++ b/fe/fe-core/src/main/java/org/apache/doris/system/SystemInfoService.java @@ -335,6 +335,10 @@ public Backend getBackend(long backendId) { return getAllClusterBackendsNoException().get(backendId); } + public Backend getBackendByIdWithBoxedId(Long backendId) { + return getAllClusterBackendsNoException().get(backendId); + } + public List getBackends(List backendIds) { List backends = Lists.newArrayList(); for (long backendId : backendIds) { diff --git a/fe/fe-core/src/test/java/org/apache/doris/cloud/catalog/CloudTabletRebalancerTest.java b/fe/fe-core/src/test/java/org/apache/doris/cloud/catalog/CloudTabletRebalancerTest.java index 520932b5fb6283..8c5f7d8c01c7c8 100644 --- a/fe/fe-core/src/test/java/org/apache/doris/cloud/catalog/CloudTabletRebalancerTest.java +++ b/fe/fe-core/src/test/java/org/apache/doris/cloud/catalog/CloudTabletRebalancerTest.java @@ -23,6 +23,7 @@ import org.apache.doris.catalog.MaterializedIndex; import org.apache.doris.catalog.OlapTable; import org.apache.doris.catalog.Partition; +import org.apache.doris.catalog.Replica; import org.apache.doris.catalog.Tablet; import org.apache.doris.catalog.TabletInvertedIndex; import org.apache.doris.catalog.TabletMeta; @@ -32,6 +33,7 @@ import org.apache.doris.datasource.InternalCatalog; import org.apache.doris.metric.MetricRepo; import org.apache.doris.system.Backend; +import org.apache.doris.system.SystemInfoService; import org.junit.jupiter.api.AfterEach; import org.junit.jupiter.api.Assertions; @@ -40,8 +42,11 @@ import org.mockito.MockedStatic; import org.mockito.Mockito; +import java.lang.reflect.Constructor; import java.lang.reflect.Field; import java.lang.reflect.Method; +import java.lang.reflect.Modifier; +import java.util.AbstractList; import java.util.AbstractMap; import java.util.ArrayList; import java.util.Collections; @@ -97,6 +102,33 @@ protected boolean isInternalDbId(Long dbId) { } } + private static class CapacityTrackingRebalancer extends TestRebalancer { + private final List globalTabletSetInitialCapacities = new ArrayList<>(); + private final List routeInfoListInitialCapacities = new ArrayList<>(); + + @Override + protected Set newGlobalTabletSet(int initialCapacity) { + globalTabletSetInitialCapacities.add(initialCapacity); + return ConcurrentHashMap.newKeySet(); + } + + @Override + protected List newRouteInfoList(int initialCapacity) { + routeInfoListInitialCapacities.add(initialCapacity); + return super.newRouteInfoList(initialCapacity); + } + } + + private static class SnapshotCapacityTrackingRebalancer extends TestRebalancer { + private final List snapshotTabletSetInitialCapacities = new ArrayList<>(); + + @Override + protected Set newSnapshotTabletSet(int expectedSize) { + snapshotTabletSetInitialCapacities.add(expectedSize); + return super.newSnapshotTabletSet(expectedSize); + } + } + private static class CountingConcurrentHashMap extends ConcurrentHashMap { private int computeIfAbsentCalls; private int getCalls; @@ -121,6 +153,39 @@ public V putIfAbsent(K key, V value) { } } + private static class IteratorRejectingList extends AbstractList { + private final E element; + + IteratorRejectingList(E element) { + this.element = element; + } + + @Override + public E get(int index) { + if (index != 0) { + throw new IndexOutOfBoundsException(String.valueOf(index)); + } + return element; + } + + @Override + public int size() { + return 1; + } + + @Override + public java.util.Iterator iterator() { + throw new AssertionError("replica traversal must not allocate an iterator"); + } + } + + private static class EmptyLookupRejectingMap extends ConcurrentHashMap { + @Override + public V get(Object key) { + throw new AssertionError("an empty inflight map must not allocate and probe a composite key"); + } + } + private static void setField(Object obj, String name, Object value) throws Exception { Field f = CloudTabletRebalancer.class.getDeclaredField(name); f.setAccessible(true); @@ -164,6 +229,207 @@ private static class RouteMaps { byPartition = new ConcurrentHashMap<>(); } + @Test + public void testInfightTabletIsStaticAndPreservesHashCode() throws Exception { + long tabletId = 50_001L; + String clusterId = "cluster-a"; + Class infightTabletClass = null; + for (Class nestedClass : CloudTabletRebalancer.class.getDeclaredClasses()) { + if (nestedClass.getSimpleName().equals("InfightTablet")) { + infightTabletClass = nestedClass; + break; + } + } + Assertions.assertNotNull(infightTabletClass); + Assertions.assertTrue(Modifier.isStatic(infightTabletClass.getModifiers())); + Constructor constructor = infightTabletClass.getDeclaredConstructor(long.class, String.class); + constructor.setAccessible(true); + Object infightTablet = constructor.newInstance(tabletId, clusterId); + + int expectedHashCode = 31 * (31 + Long.hashCode(tabletId)) + clusterId.hashCode(); + Assertions.assertEquals(expectedHashCode, infightTablet.hashCode()); + } + + @Test + public void testCloudReplicaReturnsStoredBoxedPrimaryBackendId() throws Exception { + CloudReplica replica = new CloudReplica(); + String clusterId = "cluster-a"; + long backendId = 60_001L; + replica.updateClusterToPrimaryBe(clusterId, backendId); + + Method method = CloudReplica.class.getDeclaredMethod( + "getNonColocatedPrimaryBackendId", String.class); + method.setAccessible(true); + Long first = (Long) method.invoke(replica, clusterId); + Long second = (Long) method.invoke(replica, clusterId); + + Assertions.assertEquals(backendId, first); + Assertions.assertSame(first, second); + } + + @Test + public void testSystemInfoServiceLooksUpBackendWithBoxedId() throws Exception { + SystemInfoService systemInfoService = new SystemInfoService(); + Long backendId = Long.valueOf(60_001L); + Backend backend = new Backend(backendId, "127.0.0.1", 9050); + systemInfoService.addBackend(backend); + + Method method = SystemInfoService.class.getDeclaredMethod("getBackendByIdWithBoxedId", Long.class); + Backend actual = (Backend) method.invoke(systemInfoService, backendId); + + Assertions.assertSame(backend, actual); + } + + @Test + public void testRouteRebuildDoesNotUsePrimitivePrimaryBackendPath() throws Exception { + TestRebalancer rebalancer = new TestRebalancer(); + Long dbId = 10_001L; + Long tableId = 20_001L; + Long partitionId = 30_001L; + Long indexId = 40_001L; + Long tabletId = 50_001L; + Long beId = 60_001L; + String clusterId = "cluster-a"; + Tablet tablet = mockTablet(tabletId); + CloudReplica replica = (CloudReplica) tablet.getReplicas().get(0); + setField(rebalancer, "clusterToBes", Collections.singletonMap(clusterId, List.of(beId))); + setField(rebalancer, "allBes", Set.of(beId)); + + try (MockedStatic ignored = mockRouteEnvironment( + dbId, tableId, partitionId, indexId, tablet, clusterId, beId)) { + boolean completed = invokePrivate(rebalancer, "completeRouteInfo", + new Class[] {}, new Object[] {}); + Assertions.assertTrue(completed); + rebalancer.statRouteInfo(); + } + + Mockito.verify(replica, Mockito.never()).getPrimaryBackend(clusterId, false); + } + + @Test + public void testStatRouteInfoTraversesReplicasWithoutIterator() throws Exception { + TestRebalancer rebalancer = new TestRebalancer(); + Long dbId = 10_001L; + Long tableId = 20_001L; + Long partitionId = 30_001L; + Long indexId = 40_001L; + Long tabletId = 50_001L; + Long beId = 60_001L; + String clusterId = "cluster-a"; + Tablet tablet = mockTablet(tabletId); + CloudReplica replica = (CloudReplica) tablet.getReplicas().get(0); + Mockito.when(tablet.getReplicas()).thenReturn(new IteratorRejectingList(replica)); + setField(rebalancer, "clusterToBes", Collections.singletonMap(clusterId, List.of(beId))); + setField(rebalancer, "allBes", Set.of(beId)); + + try (MockedStatic ignored = mockRouteEnvironment( + dbId, tableId, partitionId, indexId, tablet, clusterId, beId)) { + rebalancer.statRouteInfo(); + } + } + + @Test + public void testStatRouteInfoSkipsCompositeKeyForEmptyInflightMap() throws Exception { + TestRebalancer rebalancer = new TestRebalancer(); + Long dbId = 10_001L; + Long tableId = 20_001L; + Long partitionId = 30_001L; + Long indexId = 40_001L; + Long tabletId = 50_001L; + Long beId = 60_001L; + String clusterId = "cluster-a"; + Tablet tablet = mockTablet(tabletId); + setField(rebalancer, "clusterToBes", Collections.singletonMap(clusterId, List.of(beId))); + setField(rebalancer, "allBes", Set.of(beId)); + setField(rebalancer, "tabletToInfightTask", new EmptyLookupRejectingMap<>()); + + try (MockedStatic ignored = mockRouteEnvironment( + dbId, tableId, partitionId, indexId, tablet, clusterId, beId)) { + rebalancer.statRouteInfo(); + } + } + + @Test + public void testCompleteRouteInfoPresizesRouteListsFromIndexTablets() throws Exception { + CapacityTrackingRebalancer rebalancer = new CapacityTrackingRebalancer(); + Long dbId = 10_001L; + Long tableId = 20_001L; + Long partitionId = 30_001L; + Long indexId = 40_001L; + Long tabletId = 50_001L; + Long beId = 60_001L; + String clusterId = "cluster-a"; + setField(rebalancer, "clusterToBes", Collections.singletonMap(clusterId, List.of(beId))); + + try (MockedStatic ignored = mockRouteEnvironment( + dbId, tableId, partitionId, indexId, tabletId, clusterId, beId, 3)) { + boolean completed = invokePrivate(rebalancer, "completeRouteInfo", + new Class[] {}, new Object[] {}); + + Assertions.assertTrue(completed); + Assertions.assertEquals(List.of(3, 3), rebalancer.routeInfoListInitialCapacities); + } + } + + @Test + public void testCompleteRouteInfoDoesNotPresizeUnusedColocateRouteLists() throws Exception { + CapacityTrackingRebalancer rebalancer = new CapacityTrackingRebalancer(); + Long dbId = 10_001L; + Long tableId = 20_001L; + Long partitionId = 30_001L; + Long indexId = 40_001L; + Long tabletId = 50_001L; + Long beId = 60_001L; + String clusterId = "cluster-a"; + setField(rebalancer, "clusterToBes", Collections.singletonMap(clusterId, List.of(beId))); + + try (MockedStatic ignored = mockRouteEnvironment( + dbId, tableId, partitionId, indexId, tabletId, clusterId, beId, 3, true)) { + boolean completed = invokePrivate(rebalancer, "completeRouteInfo", + new Class[] {}, new Object[] {}); + + Assertions.assertTrue(completed); + Assertions.assertEquals(List.of(0, 0), rebalancer.routeInfoListInitialCapacities); + } + } + + @Test + public void testSnapshotTabletSetsUseOnePresizedResultPerRequest() throws Exception { + SnapshotCapacityTrackingRebalancer rebalancer = new SnapshotCapacityTrackingRebalancer(); + Long beId = 60_001L; + ConcurrentHashMap> primary = new ConcurrentHashMap<>(); + primary.put(beId, Set.of(50_001L, 50_002L)); + ConcurrentHashMap> colocate = new ConcurrentHashMap<>(); + colocate.put(beId, Set.of(50_002L, 50_003L)); + ConcurrentHashMap> secondary = new ConcurrentHashMap<>(); + secondary.put(beId, Set.of(50_003L, 50_004L)); + setField(rebalancer, "beToTabletsGlobal", primary); + setField(rebalancer, "beToColocateTabletsGlobal", colocate); + setField(rebalancer, "beToTabletsGlobalInSecondary", secondary); + + Assertions.assertEquals(Set.of(50_001L, 50_002L, 50_003L), + rebalancer.getSnapshotTabletsInPrimaryByBeId(beId)); + Assertions.assertEquals(Set.of(50_003L, 50_004L), + rebalancer.getSnapshotTabletsInSecondaryByBeId(beId)); + Assertions.assertEquals(Set.of(50_001L, 50_002L, 50_003L, 50_004L), + rebalancer.getSnapshotTabletsInPrimaryAndSecondaryByBeId(beId)); + + Assertions.assertEquals(List.of(4, 2, 6), + rebalancer.snapshotTabletSetInitialCapacities); + } + + @Test + public void testSnapshotTabletSetsRemainEmptyForUnknownBackend() { + SnapshotCapacityTrackingRebalancer rebalancer = new SnapshotCapacityTrackingRebalancer(); + Long unknownBeId = 60_001L; + + Assertions.assertTrue(rebalancer.getSnapshotTabletsInPrimaryByBeId(unknownBeId).isEmpty()); + Assertions.assertTrue(rebalancer.getSnapshotTabletsInSecondaryByBeId(unknownBeId).isEmpty()); + Assertions.assertTrue(rebalancer.getSnapshotTabletsInPrimaryAndSecondaryByBeId(unknownBeId).isEmpty()); + Assertions.assertEquals(List.of(0, 0, 0), + rebalancer.snapshotTabletSetInitialCapacities); + } + @Test public void testFillBeToTabletsReusesBoxedIdsAcrossIndexes() { TestRebalancer rebalancer = new TestRebalancer(); @@ -349,6 +615,96 @@ public void testWarmupRollbackReusesInflightBoxedTabletIdAfterRouteRebuild() thr } } + @Test + public void testStatRouteInfoPresizesGlobalTabletSetsFromPreviousRoute() throws Exception { + CapacityTrackingRebalancer rebalancer = new CapacityTrackingRebalancer(); + Long dbId = 10_001L; + Long tableId = 20_001L; + Long partitionId = 30_001L; + Long indexId = 40_001L; + Long tabletId = 50_001L; + Long beId = 60_001L; + String clusterId = "cluster-a"; + + ConcurrentHashMap> previousCurrent = new ConcurrentHashMap<>(); + previousCurrent.put(beId, Set.of(1L, 2L, 3L)); + ConcurrentHashMap> previousFuture = new ConcurrentHashMap<>(); + previousFuture.put(beId, Set.of(1L, 2L, 3L, 4L, 5L)); + setField(rebalancer, "beToTabletsGlobal", previousCurrent); + setField(rebalancer, "futureBeToTabletsGlobal", previousFuture); + setField(rebalancer, "clusterToBes", Collections.singletonMap(clusterId, List.of(beId))); + setField(rebalancer, "allBes", Set.of(beId)); + + try (MockedStatic ignored = mockRouteEnvironment( + dbId, tableId, partitionId, indexId, tabletId, clusterId, beId)) { + rebalancer.statRouteInfo(); + } + + Assertions.assertEquals(List.of(3, 5), rebalancer.globalTabletSetInitialCapacities); + ConcurrentHashMap> current = getField(rebalancer, "beToTabletsGlobal"); + ConcurrentHashMap> future = getField(rebalancer, "futureBeToTabletsGlobal"); + Assertions.assertEquals(Set.of(tabletId), current.get(beId)); + Assertions.assertEquals(Set.of(tabletId), future.get(beId)); + } + + @Test + @SuppressWarnings("unchecked") + public void testStatRouteInfoBoundsStaleGlobalTabletSetCapacity() throws Exception { + CapacityTrackingRebalancer rebalancer = new CapacityTrackingRebalancer(); + Long dbId = 10_001L; + Long tableId = 20_001L; + Long partitionId = 30_001L; + Long indexId = 40_001L; + Long tabletId = 50_001L; + Long beId = 60_001L; + String clusterId = "cluster-a"; + + Set stalePreviousTablets = Mockito.mock(Set.class); + Mockito.when(stalePreviousTablets.size()).thenReturn(2_000_000); + ConcurrentHashMap> previousCurrent = new ConcurrentHashMap<>(); + previousCurrent.put(beId, stalePreviousTablets); + ConcurrentHashMap> previousFuture = new ConcurrentHashMap<>(); + previousFuture.put(beId, stalePreviousTablets); + setField(rebalancer, "beToTabletsGlobal", previousCurrent); + setField(rebalancer, "futureBeToTabletsGlobal", previousFuture); + setField(rebalancer, "clusterToBes", Collections.singletonMap(clusterId, List.of(beId))); + setField(rebalancer, "allBes", Set.of(beId)); + + try (MockedStatic ignored = mockRouteEnvironment( + dbId, tableId, partitionId, indexId, tabletId, clusterId, beId)) { + rebalancer.statRouteInfo(); + } + + Assertions.assertEquals(List.of(65_536, 65_536), + rebalancer.globalTabletSetInitialCapacities); + ConcurrentHashMap> current = getField(rebalancer, "beToTabletsGlobal"); + ConcurrentHashMap> future = getField(rebalancer, "futureBeToTabletsGlobal"); + Assertions.assertEquals(Set.of(tabletId), current.get(beId)); + Assertions.assertEquals(Set.of(tabletId), future.get(beId)); + } + + @Test + public void testStatRouteInfoUsesZeroCapacityForNewBackend() throws Exception { + CapacityTrackingRebalancer rebalancer = new CapacityTrackingRebalancer(); + Long dbId = 10_001L; + Long tableId = 20_001L; + Long partitionId = 30_001L; + Long indexId = 40_001L; + Long tabletId = 50_001L; + Long beId = 60_001L; + String clusterId = "cluster-a"; + + setField(rebalancer, "clusterToBes", Collections.singletonMap(clusterId, List.of(beId))); + setField(rebalancer, "allBes", Set.of(beId)); + + try (MockedStatic ignored = mockRouteEnvironment( + dbId, tableId, partitionId, indexId, tabletId, clusterId, beId)) { + rebalancer.statRouteInfo(); + } + + Assertions.assertEquals(List.of(0, 0), rebalancer.globalTabletSetInitialCapacities); + } + private static void initializeRouteMaps(TestRebalancer rebalancer, RouteMaps current, RouteMaps future, Long srcBe, Long tableId, Long partitionId, Long indexId, Long tabletId) throws Exception { rebalancer.fillBeToTablets(srcBe, tableId, partitionId, indexId, tabletId, @@ -381,6 +737,36 @@ private static MockedStatic mockTabletMeta(Long tabletId, Long tableId, Lon private static MockedStatic mockRouteEnvironment(Long dbId, Long tableId, Long partitionId, Long indexId, Long tabletId, String clusterId, Long srcBe) { + return mockRouteEnvironment(dbId, tableId, partitionId, indexId, tabletId, clusterId, srcBe, 1, false); + } + + private static MockedStatic mockRouteEnvironment(Long dbId, Long tableId, Long partitionId, + Long indexId, Long tabletId, String clusterId, Long srcBe, int tabletCount) { + return mockRouteEnvironment( + dbId, tableId, partitionId, indexId, tabletId, clusterId, srcBe, tabletCount, false); + } + + private static MockedStatic mockRouteEnvironment(Long dbId, Long tableId, Long partitionId, + Long indexId, Long tabletId, String clusterId, Long srcBe, int tabletCount, boolean colocated) { + return mockRouteEnvironment(dbId, tableId, partitionId, indexId, mockTablet(tabletId), + clusterId, srcBe, tabletCount, colocated); + } + + private static Tablet mockTablet(long tabletId) { + Tablet tablet = Mockito.mock(Tablet.class); + CloudReplica replica = Mockito.mock(CloudReplica.class); + Mockito.when(tablet.getId()).thenReturn(tabletId); + Mockito.when(tablet.getReplicas()).thenReturn(Collections.singletonList(replica)); + return tablet; + } + + private static MockedStatic mockRouteEnvironment(Long dbId, Long tableId, Long partitionId, + Long indexId, Tablet tablet, String clusterId, Long srcBe) { + return mockRouteEnvironment(dbId, tableId, partitionId, indexId, tablet, clusterId, srcBe, 1, false); + } + + private static MockedStatic mockRouteEnvironment(Long dbId, Long tableId, Long partitionId, + Long indexId, Tablet tablet, String clusterId, Long srcBe, int tabletCount, boolean colocated) { Env env = Mockito.mock(Env.class); TabletInvertedIndex invertedIndex = Mockito.mock(TabletInvertedIndex.class); TabletMeta tabletMeta = Mockito.mock(TabletMeta.class); @@ -390,9 +776,10 @@ private static MockedStatic mockRouteEnvironment(Long dbId, Long tableId, L OlapTable table = Mockito.mock(OlapTable.class); Partition partition = Mockito.mock(Partition.class); MaterializedIndex index = Mockito.mock(MaterializedIndex.class); - Tablet tablet = Mockito.mock(Tablet.class); - CloudReplica replica = Mockito.mock(CloudReplica.class); + CloudReplica replica = (CloudReplica) tablet.getReplicas().get(0); Backend primaryBackend = Mockito.mock(Backend.class); + SystemInfoService systemInfoService = Mockito.mock(SystemInfoService.class); + Long tabletId = tablet.getId(); Mockito.when(env.getTabletInvertedIndex()).thenReturn(invertedIndex); Mockito.when(invertedIndex.getTabletMeta(tabletId)).thenReturn(tabletMeta); @@ -405,14 +792,15 @@ private static MockedStatic mockRouteEnvironment(Long dbId, Long tableId, L Mockito.when(database.getId()).thenReturn(dbId); Mockito.when(table.isManagedTable()).thenReturn(true); Mockito.when(table.getId()).thenReturn(tableId); + Mockito.when(colocateTableIndex.isColocateTable(tableId)).thenReturn(colocated); Mockito.when(table.getAllPartitions()).thenReturn(Collections.singletonList(partition)); Mockito.when(partition.getId()).thenReturn(partitionId); Mockito.when(partition.getMaterializedIndices(MaterializedIndex.IndexExtState.VISIBLE)) .thenReturn(Collections.singletonList(index)); Mockito.when(index.getId()).thenReturn(indexId); - Mockito.when(index.getTablets()).thenReturn(Collections.singletonList(tablet)); - Mockito.when(tablet.getId()).thenReturn(tabletId); - Mockito.when(tablet.getReplicas()).thenReturn(Collections.singletonList(replica)); + Mockito.when(index.getTablets()).thenReturn(Collections.nCopies(tabletCount, tablet)); + Mockito.when(replica.getNonColocatedPrimaryBackendId(clusterId)).thenReturn(srcBe); + Mockito.when(systemInfoService.getBackendByIdWithBoxedId(srcBe)).thenReturn(primaryBackend); Mockito.when(replica.getPrimaryBackend(clusterId, false)).thenReturn(primaryBackend); Mockito.when(primaryBackend.getId()).thenReturn(srcBe); @@ -420,6 +808,7 @@ private static MockedStatic mockRouteEnvironment(Long dbId, Long tableId, L mockedEnv.when(Env::getCurrentEnv).thenReturn(env); mockedEnv.when(Env::getCurrentInternalCatalog).thenReturn(catalog); mockedEnv.when(Env::getCurrentColocateIndex).thenReturn(colocateTableIndex); + mockedEnv.when(Env::getCurrentSystemInfo).thenReturn(systemInfoService); return mockedEnv; }