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 6acaad24e0700a..0de2ed5c7eb31a 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 @@ -339,7 +339,7 @@ public int hashCode() { } private class InfightTask { - public long pickedTabletId; + public Long pickedTabletId; public long srcBe; public long destBe; public long startTimestamp; @@ -375,12 +375,12 @@ public int hashCode() { } private static class WarmupTabletTask { - private final long pickedTabletId; + private final Long pickedTabletId; private final long srcBe; private final long destBe; private final String clusterId; - WarmupTabletTask(long pickedTabletId, long srcBe, long destBe, String clusterId) { + WarmupTabletTask(Long pickedTabletId, long srcBe, long destBe, String clusterId) { this.pickedTabletId = pickedTabletId; this.srcBe = srcBe; this.destBe = destBe; @@ -1058,6 +1058,15 @@ public void fillBeToTablets(long be, long tableId, long partId, long indexId, lo ConcurrentHashMap>> beToTabletsInTable, ConcurrentHashMap>>> partToTablets) { + fillBeToTablets(Long.valueOf(be), Long.valueOf(tableId), Long.valueOf(partId), Long.valueOf(indexId), + Long.valueOf(tabletId), globalBeToTablets, beToTabletsInTable, partToTablets); + } + + void fillBeToTablets(Long be, Long tableId, Long partId, Long indexId, Long tabletId, + ConcurrentHashMap> globalBeToTablets, + ConcurrentHashMap>> beToTabletsInTable, + ConcurrentHashMap>>> + partToTablets) { // global globalBeToTablets.putIfAbsent(be, ConcurrentHashMap.newKeySet()); globalBeToTablets.get(be).add(tabletId); @@ -1176,25 +1185,29 @@ public void statRouteInfo() { Map tmpDbInternal = new HashMap<>(); loopCloudReplica((Database db, Table table, Partition partition, MaterializedIndex index, String cluster) -> { - boolean isColocated = Env.getCurrentColocateIndex().isColocateTable(table.getId()); - tmpTableToDb.put(table.getId(), db.getId()); - tmpPartitionToDb.put(partition.getId(), db.getId()); - tmpDbInternal.computeIfAbsent(db.getId(), k -> { + Long dbId = db.getId(); + Long tableId = table.getId(); + Long partitionId = partition.getId(); + Long indexId = index.getId(); + boolean isColocated = Env.getCurrentColocateIndex().isColocateTable(tableId); + tmpTableToDb.put(tableId, dbId); + tmpPartitionToDb.put(partitionId, dbId); + tmpDbInternal.computeIfAbsent(dbId, k -> { String name = db.getFullName(); return name != null && INTERNAL_DB_NAMES.contains(name); }); for (Tablet tablet : index.getTablets()) { - long tabletId = tablet.getId(); + Long tabletId = tablet.getId(); // active tablet scoring (used for scheduling order) if (activeTabletIds != null && !activeTabletIds.isEmpty() && activeTabletIds.contains(tabletId)) { - tmpTableActive.merge(table.getId(), 1L, Long::sum); - tmpPartitionActive.merge(partition.getId(), 1L, Long::sum); - tmpDbActive.merge(db.getId(), 1L, Long::sum); + tmpTableActive.merge(tableId, 1L, Long::sum); + tmpPartitionActive.merge(partitionId, 1L, Long::sum); + tmpDbActive.merge(dbId, 1L, Long::sum); } for (Replica r : tablet.getReplicas()) { CloudReplica replica = (CloudReplica) r; if (isColocated) { - long beId = -1L; + Long beId = -1L; try { beId = replica.getColocatedBeId(cluster); } catch (ComputeGroupException e) { @@ -1209,13 +1222,13 @@ public void statRouteInfo() { } Backend be = replica.getPrimaryBackend(cluster, false); - long beId = be == null ? -1L : be.getId(); + Long beId = be == null ? Long.valueOf(-1L) : Long.valueOf(be.getId()); if (!allBes.contains(beId)) { continue; } Backend secondaryBe = replica.getSecondaryBackend(cluster); - long secondaryBeId = secondaryBe == null ? -1L : secondaryBe.getId(); + Long secondaryBeId = secondaryBe == null ? Long.valueOf(-1L) : Long.valueOf(secondaryBe.getId()); if (allBes.contains(secondaryBeId)) { Set tablets = tmpBeToTabletsGlobalInSecondary .computeIfAbsent(secondaryBeId, k -> new HashSet<>()); @@ -1224,11 +1237,12 @@ public void statRouteInfo() { InfightTablet taskKey = new InfightTablet(tabletId, cluster); InfightTask task = tabletToInfightTask.get(taskKey); - long futureBeId = task == null ? beId : task.destBe; - fillBeToTablets(beId, table.getId(), partition.getId(), index.getId(), tabletId, + Long futureBeId = task == null ? beId : Long.valueOf(task.destBe); + Long routeTabletId = task == null ? tabletId : task.pickedTabletId; + fillBeToTablets(beId, tableId, partitionId, indexId, routeTabletId, tmpBeToTabletsGlobal, beToTabletsInTable, this.partitionToTablets); - fillBeToTablets(futureBeId, table.getId(), partition.getId(), index.getId(), tabletId, + fillBeToTablets(futureBeId, tableId, partitionId, indexId, routeTabletId, tmpFutureBeToTabletsGlobal, futureBeToTabletsInTable, futurePartitionToTablets); } } @@ -1618,7 +1632,7 @@ private void handleWarmupCompletion(InfightTask task, String clusterId, boolean } } - private void updateBeToTablets(long tabletId, long srcBe, long destBe, + private void updateBeToTablets(Long tabletId, Long srcBe, Long destBe, ConcurrentHashMap> globalBeToTablets, ConcurrentHashMap>> beToTabletsInTable, ConcurrentHashMap globalSrcTablets = globalBeToTablets.get(srcBe); if (globalSrcTablets == null || !globalSrcTablets.remove(tabletId)) { @@ -1658,8 +1672,8 @@ private void updateBeToTablets(long tabletId, long srcBe, long destBe, } } - fillBeToTablets(destBe, tableId, partId, indexId, tabletId, globalBeToTablets, beToTabletsInTable, - partToTablets); + fillBeToTablets(destBe, tableId, partId, indexId, tabletId, globalBeToTablets, + beToTabletsInTable, partToTablets); } private void updateClusterToBeMap(long tabletId, long destBe, String clusterId, @@ -1915,8 +1929,8 @@ private void balanceImpl(List bes, String clusterId, Map> break; // no need balance } - long srcBe = pairInfo.srcBe; - long destBe = pairInfo.destBe; + Long srcBe = pairInfo.srcBe; + Long destBe = pairInfo.destBe; Long pickedTabletId = pickTabletPreferCold(srcBe, beToTablets.get(srcBe), this.activeTabletIds, pickedTabletIds); @@ -2071,7 +2085,7 @@ private Long reservoirPick(Set tabletIds, Set pickedTabletIds, return chosen; } - private boolean preheatAndUpdateTablet(long pickedTabletId, long srcBe, long destBe, String clusterId, + private boolean preheatAndUpdateTablet(Long pickedTabletId, Long srcBe, Long destBe, String clusterId, BalanceType balanceType) { Backend srcBackend = cloudSystemInfoService.getBackend(srcBe); Backend destBackend = cloudSystemInfoService.getBackend(destBe); @@ -2097,7 +2111,7 @@ private boolean preheatAndUpdateTablet(long pickedTabletId, long srcBe, long des return true; } - private boolean transferTablet(long pickedTabletId, long srcBe, long destBe, String clusterId, + private boolean transferTablet(Long pickedTabletId, Long srcBe, Long destBe, String clusterId, BalanceType balanceType, List infos) { LOG.debug("transfer {} from {} to {}, cluster {}, type {}", pickedTabletId, srcBe, destBe, clusterId, balanceType); 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 8637be3c76a4d1..6ef2286c869a96 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 @@ -17,9 +17,21 @@ package org.apache.doris.cloud.catalog; +import org.apache.doris.catalog.ColocateTableIndex; +import org.apache.doris.catalog.Database; +import org.apache.doris.catalog.Env; +import org.apache.doris.catalog.MaterializedIndex; +import org.apache.doris.catalog.OlapTable; +import org.apache.doris.catalog.Partition; +import org.apache.doris.catalog.Tablet; +import org.apache.doris.catalog.TabletInvertedIndex; +import org.apache.doris.catalog.TabletMeta; +import org.apache.doris.cloud.persist.UpdateCloudReplicaInfo; import org.apache.doris.cloud.system.CloudSystemInfoService; import org.apache.doris.common.Config; +import org.apache.doris.datasource.InternalCatalog; import org.apache.doris.metric.MetricRepo; +import org.apache.doris.system.Backend; import org.junit.jupiter.api.AfterEach; import org.junit.jupiter.api.Assertions; @@ -47,12 +59,14 @@ public class CloudTabletRebalancerTest { private boolean oldEnableActiveScheduling; private long oldActiveTabletIdsRefreshIntervalSecond; private int oldForceInactiveAfterRounds; + private int oldWarmupBatchSize; @BeforeEach public void setUp() { oldEnableActiveScheduling = Config.enable_cloud_active_tablet_priority_scheduling; oldActiveTabletIdsRefreshIntervalSecond = Config.cloud_active_tablet_ids_refresh_interval_second; oldForceInactiveAfterRounds = Config.cloud_active_unbalanced_force_inactive_after_rounds; + oldWarmupBatchSize = Config.cloud_warm_up_batch_size; Config.enable_cloud_active_tablet_priority_scheduling = true; } @@ -61,6 +75,7 @@ public void tearDown() { Config.enable_cloud_active_tablet_priority_scheduling = oldEnableActiveScheduling; Config.cloud_active_tablet_ids_refresh_interval_second = oldActiveTabletIdsRefreshIntervalSecond; Config.cloud_active_unbalanced_force_inactive_after_rounds = oldForceInactiveAfterRounds; + Config.cloud_warm_up_batch_size = oldWarmupBatchSize; } private static class TestRebalancer extends CloudTabletRebalancer { @@ -101,6 +116,321 @@ private static T invokePrivate(Object obj, String method, Class[] types, return (T) m.invoke(obj, args); } + @SuppressWarnings("unchecked") + private static T invokePrivate(Object obj, String method, int parameterCount, Object[] args) + throws Exception { + Method target = null; + for (Method candidate : CloudTabletRebalancer.class.getDeclaredMethods()) { + if (candidate.getName().equals(method) && candidate.getParameterCount() == parameterCount) { + target = candidate; + break; + } + } + Assertions.assertNotNull(target, "Cannot find method " + method); + target.setAccessible(true); + return (T) target.invoke(obj, args); + } + + private static class RouteMaps { + private final ConcurrentHashMap> global = new ConcurrentHashMap<>(); + private final ConcurrentHashMap>> byTable = + new ConcurrentHashMap<>(); + private final ConcurrentHashMap>>> + byPartition = new ConcurrentHashMap<>(); + } + + @Test + public void testFillBeToTabletsReusesBoxedIdsAcrossIndexes() { + TestRebalancer rebalancer = new TestRebalancer(); + Long beId = 10_001L; + Long tableId = 20_001L; + Long partitionId = 30_001L; + Long indexId = 40_001L; + Long tabletId = 50_001L; + + ConcurrentHashMap> currentGlobal = new ConcurrentHashMap<>(); + ConcurrentHashMap>> currentByTable = + new ConcurrentHashMap<>(); + ConcurrentHashMap>>> currentByPartition = + new ConcurrentHashMap<>(); + ConcurrentHashMap> futureGlobal = new ConcurrentHashMap<>(); + ConcurrentHashMap>> futureByTable = + new ConcurrentHashMap<>(); + ConcurrentHashMap>>> futureByPartition = + new ConcurrentHashMap<>(); + + rebalancer.fillBeToTablets(beId, tableId, partitionId, indexId, tabletId, + currentGlobal, currentByTable, currentByPartition); + rebalancer.fillBeToTablets(beId, tableId, partitionId, indexId, tabletId, + futureGlobal, futureByTable, futureByPartition); + + assertSameStoredId(beId, currentGlobal); + assertSameStoredId(beId, currentByTable.get(tableId)); + assertSameStoredId(beId, currentByPartition.get(partitionId).get(indexId)); + assertSameStoredId(beId, futureGlobal); + assertSameStoredId(beId, futureByTable.get(tableId)); + assertSameStoredId(beId, futureByPartition.get(partitionId).get(indexId)); + assertSameStoredId(tableId, currentByTable); + assertSameStoredId(tableId, futureByTable); + assertSameStoredId(partitionId, currentByPartition); + assertSameStoredId(partitionId, futureByPartition); + assertSameStoredId(indexId, currentByPartition.get(partitionId)); + assertSameStoredId(indexId, futureByPartition.get(partitionId)); + assertSameStoredId(tabletId, currentGlobal.get(beId)); + assertSameStoredId(tabletId, currentByTable.get(tableId).get(beId)); + assertSameStoredId(tabletId, currentByPartition.get(partitionId).get(indexId).get(beId)); + assertSameStoredId(tabletId, futureGlobal.get(beId)); + assertSameStoredId(tabletId, futureByTable.get(tableId).get(beId)); + assertSameStoredId(tabletId, futureByPartition.get(partitionId).get(indexId).get(beId)); + } + + @Test + public void testTransferTabletReusesSelectedBoxedIdsAcrossCurrentAndFutureIndexes() throws Exception { + TestRebalancer rebalancer = new TestRebalancer(); + Long srcBe = 10_001L; + Long destBe = 10_002L; + Long tableId = 20_001L; + Long partitionId = 30_001L; + Long indexId = 40_001L; + Long tabletId = 50_001L; + RouteMaps current = new RouteMaps(); + RouteMaps future = new RouteMaps(); + initializeRouteMaps(rebalancer, current, future, srcBe, tableId, partitionId, indexId, tabletId); + + try (MockedStatic ignored = mockTabletMeta(tabletId, tableId, partitionId, indexId)) { + boolean moved = invokePrivate(rebalancer, "transferTablet", 6, + new Object[] {tabletId, srcBe, destBe, "cluster-a", + CloudTabletRebalancer.BalanceType.GLOBAL, new ArrayList()}); + + Assertions.assertTrue(moved); + assertSameRouteIds(destBe, tableId, partitionId, indexId, tabletId, current); + assertSameRouteIds(destBe, tableId, partitionId, indexId, tabletId, future); + } + } + + @Test + public void testPreheatTabletReusesSelectedBoxedIdsInFutureIndexes() throws Exception { + TestRebalancer rebalancer = new TestRebalancer(); + Long srcBe = 10_001L; + Long destBe = 10_002L; + Long tableId = 20_001L; + Long partitionId = 30_001L; + Long indexId = 40_001L; + Long tabletId = 50_001L; + RouteMaps current = new RouteMaps(); + RouteMaps future = new RouteMaps(); + initializeRouteMaps(rebalancer, current, future, srcBe, tableId, partitionId, indexId, tabletId); + setField(rebalancer, "cloudSystemInfoService", mockBackendService(srcBe, destBe)); + Config.cloud_warm_up_batch_size = 10; + + try (MockedStatic ignored = mockTabletMeta(tabletId, tableId, partitionId, indexId)) { + boolean moved = invokePrivate(rebalancer, "preheatAndUpdateTablet", 5, + new Object[] {tabletId, srcBe, destBe, "cluster-a", CloudTabletRebalancer.BalanceType.GLOBAL}); + + Assertions.assertTrue(moved); + assertSameRouteIds(destBe, tableId, partitionId, indexId, tabletId, future); + } + } + + @Test + public void testWarmupRollbackRestoresSelectedBoxedIdsInFutureIndexes() throws Exception { + TestRebalancer rebalancer = new TestRebalancer(); + Long srcBe = 10_001L; + Long destBe = 10_002L; + Long tableId = 20_001L; + Long partitionId = 30_001L; + Long indexId = 40_001L; + Long tabletId = 50_001L; + RouteMaps current = new RouteMaps(); + RouteMaps future = new RouteMaps(); + initializeRouteMaps(rebalancer, current, future, srcBe, tableId, partitionId, indexId, tabletId); + setField(rebalancer, "cloudSystemInfoService", mockBackendService(srcBe, destBe)); + Config.cloud_warm_up_batch_size = 10; + + try (MockedStatic ignored = mockTabletMeta(tabletId, tableId, partitionId, indexId)) { + boolean moved = invokePrivate(rebalancer, "preheatAndUpdateTablet", 5, + new Object[] {tabletId, srcBe, destBe, "cluster-a", CloudTabletRebalancer.BalanceType.GLOBAL}); + Assertions.assertTrue(moved); + + Map warmupBatches = getField(rebalancer, "warmupBatches"); + Object batch = warmupBatches.values().iterator().next(); + Field tasksField = batch.getClass().getDeclaredField("tasks"); + tasksField.setAccessible(true); + Object task = ((List) tasksField.get(batch)).get(0); + invokePrivate(rebalancer, "revertWarmupState", new Class[] {task.getClass()}, new Object[] {task}); + + assertSameRouteIds(srcBe, tableId, partitionId, indexId, tabletId, future); + } + } + + @Test + public void testWarmupRollbackReusesInflightBoxedTabletIdAfterRouteRebuild() throws Exception { + TestRebalancer rebalancer = new TestRebalancer(); + Long srcBe = 10_001L; + Long destBe = 10_002L; + Long dbId = 15_001L; + Long tableId = 20_001L; + Long partitionId = 30_001L; + Long indexId = 40_001L; + Long tabletId = 50_001L; + String clusterId = "cluster-a"; + RouteMaps current = new RouteMaps(); + RouteMaps future = new RouteMaps(); + initializeRouteMaps(rebalancer, current, future, srcBe, tableId, partitionId, indexId, tabletId); + setField(rebalancer, "cloudSystemInfoService", mockBackendService(srcBe, destBe)); + setField(rebalancer, "clusterToBes", Collections.singletonMap(clusterId, List.of(srcBe, destBe))); + setField(rebalancer, "allBes", Set.of(srcBe, destBe)); + Config.cloud_warm_up_batch_size = 10; + + try (MockedStatic ignored = mockRouteEnvironment( + dbId, tableId, partitionId, indexId, tabletId, clusterId, srcBe)) { + boolean moved = invokePrivate(rebalancer, "preheatAndUpdateTablet", 5, + new Object[] {tabletId, srcBe, destBe, clusterId, + CloudTabletRebalancer.BalanceType.GLOBAL}); + Assertions.assertTrue(moved); + + Map warmupBatches = getField(rebalancer, "warmupBatches"); + Object batch = warmupBatches.values().iterator().next(); + Field tasksField = batch.getClass().getDeclaredField("tasks"); + tasksField.setAccessible(true); + Object task = ((List) tasksField.get(batch)).get(0); + + rebalancer.statRouteInfo(); + invokePrivate(rebalancer, "handleWarmupBatchFailure", + new Class[] {List.class, Exception.class}, + new Object[] {Collections.singletonList(task), null}); + invokePrivate(rebalancer, "processFailedWarmupTasks", new Class[] {}, new Object[] {}); + + ConcurrentHashMap> rebuiltCurrentGlobal = getField(rebalancer, "beToTabletsGlobal"); + ConcurrentHashMap> rebuiltFutureGlobal = getField( + rebalancer, "futureBeToTabletsGlobal"); + ConcurrentHashMap>> rebuiltCurrentByTable = getField( + rebalancer, "beToTabletsInTable"); + ConcurrentHashMap>> rebuiltFutureByTable = getField( + rebalancer, "futureBeToTabletsInTable"); + ConcurrentHashMap>>> + rebuiltCurrentByPartition = getField(rebalancer, "partitionToTablets"); + ConcurrentHashMap>>> + rebuiltFutureByPartition = getField(rebalancer, "futurePartitionToTablets"); + Long currentTabletId = getStoredId(tabletId, rebuiltCurrentGlobal.get(srcBe)); + Long futureTabletId = getStoredId(tabletId, rebuiltFutureGlobal.get(srcBe)); + Assertions.assertSame(currentTabletId, futureTabletId); + assertSameStoredId(currentTabletId, rebuiltCurrentByTable.get(tableId).get(srcBe)); + assertSameStoredId(currentTabletId, rebuiltFutureByTable.get(tableId).get(srcBe)); + assertSameStoredId(currentTabletId, + rebuiltCurrentByPartition.get(partitionId).get(indexId).get(srcBe)); + assertSameStoredId(currentTabletId, + rebuiltFutureByPartition.get(partitionId).get(indexId).get(srcBe)); + } + } + + 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, + current.global, current.byTable, current.byPartition); + rebalancer.fillBeToTablets(srcBe, tableId, partitionId, indexId, tabletId, + future.global, future.byTable, future.byPartition); + setField(rebalancer, "beToTabletsGlobal", current.global); + setField(rebalancer, "beToTabletsInTable", current.byTable); + setField(rebalancer, "partitionToTablets", current.byPartition); + setField(rebalancer, "futureBeToTabletsGlobal", future.global); + setField(rebalancer, "futureBeToTabletsInTable", future.byTable); + setField(rebalancer, "futurePartitionToTablets", future.byPartition); + } + + private static MockedStatic mockTabletMeta(Long tabletId, Long tableId, Long partitionId, Long indexId) { + Env env = Mockito.mock(Env.class); + TabletInvertedIndex invertedIndex = Mockito.mock(TabletInvertedIndex.class); + TabletMeta tabletMeta = Mockito.mock(TabletMeta.class); + InternalCatalog catalog = Mockito.mock(InternalCatalog.class); + Mockito.when(env.getTabletInvertedIndex()).thenReturn(invertedIndex); + Mockito.when(invertedIndex.getTabletMeta(tabletId)).thenReturn(tabletMeta); + Mockito.when(tabletMeta.getTableId()).thenReturn(tableId); + Mockito.when(tabletMeta.getPartitionId()).thenReturn(partitionId); + Mockito.when(tabletMeta.getIndexId()).thenReturn(indexId); + MockedStatic mockedEnv = Mockito.mockStatic(Env.class); + mockedEnv.when(Env::getCurrentEnv).thenReturn(env); + mockedEnv.when(Env::getCurrentInternalCatalog).thenReturn(catalog); + return mockedEnv; + } + + private static MockedStatic mockRouteEnvironment(Long dbId, Long tableId, Long partitionId, + Long indexId, Long tabletId, String clusterId, Long srcBe) { + Env env = Mockito.mock(Env.class); + TabletInvertedIndex invertedIndex = Mockito.mock(TabletInvertedIndex.class); + TabletMeta tabletMeta = Mockito.mock(TabletMeta.class); + InternalCatalog catalog = Mockito.mock(InternalCatalog.class); + ColocateTableIndex colocateTableIndex = Mockito.mock(ColocateTableIndex.class); + Database database = Mockito.mock(Database.class); + 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); + Backend primaryBackend = Mockito.mock(Backend.class); + + Mockito.when(env.getTabletInvertedIndex()).thenReturn(invertedIndex); + Mockito.when(invertedIndex.getTabletMeta(tabletId)).thenReturn(tabletMeta); + Mockito.when(tabletMeta.getTableId()).thenReturn(tableId); + Mockito.when(tabletMeta.getPartitionId()).thenReturn(partitionId); + Mockito.when(tabletMeta.getIndexId()).thenReturn(indexId); + Mockito.when(catalog.getDbIds()).thenReturn(Collections.singletonList(dbId)); + Mockito.when(catalog.getDbNullable(dbId)).thenReturn(database); + Mockito.when(database.getTables()).thenReturn(Collections.singletonList(table)); + Mockito.when(database.getId()).thenReturn(dbId); + Mockito.when(table.isManagedTable()).thenReturn(true); + Mockito.when(table.getId()).thenReturn(tableId); + 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(replica.getPrimaryBackend(clusterId, false)).thenReturn(primaryBackend); + Mockito.when(primaryBackend.getId()).thenReturn(srcBe); + + MockedStatic mockedEnv = Mockito.mockStatic(Env.class); + mockedEnv.when(Env::getCurrentEnv).thenReturn(env); + mockedEnv.when(Env::getCurrentInternalCatalog).thenReturn(catalog); + mockedEnv.when(Env::getCurrentColocateIndex).thenReturn(colocateTableIndex); + return mockedEnv; + } + + private static CloudSystemInfoService mockBackendService(Long srcBe, Long destBe) { + CloudSystemInfoService systemInfoService = Mockito.mock(CloudSystemInfoService.class); + Mockito.when(systemInfoService.getBackend(srcBe)).thenReturn(Mockito.mock(Backend.class)); + Mockito.when(systemInfoService.getBackend(destBe)).thenReturn(Mockito.mock(Backend.class)); + return systemInfoService; + } + + private static void assertSameRouteIds(Long beId, Long tableId, Long partitionId, Long indexId, + Long tabletId, RouteMaps routeMaps) { + assertSameStoredId(beId, routeMaps.global); + assertSameStoredId(beId, routeMaps.byTable.get(tableId)); + assertSameStoredId(beId, routeMaps.byPartition.get(partitionId).get(indexId)); + assertSameStoredId(tableId, routeMaps.byTable); + assertSameStoredId(partitionId, routeMaps.byPartition); + assertSameStoredId(indexId, routeMaps.byPartition.get(partitionId)); + assertSameStoredId(tabletId, routeMaps.global.get(beId)); + assertSameStoredId(tabletId, routeMaps.byTable.get(tableId).get(beId)); + assertSameStoredId(tabletId, routeMaps.byPartition.get(partitionId).get(indexId).get(beId)); + } + + private static void assertSameStoredId(Long expected, Map map) { + Long stored = map.keySet().stream().filter(expected::equals).findFirst().orElseThrow(); + Assertions.assertSame(expected, stored); + } + + private static void assertSameStoredId(Long expected, Set ids) { + Assertions.assertSame(expected, getStoredId(expected, ids)); + } + + private static Long getStoredId(Long expected, Set ids) { + return ids.stream().filter(expected::equals).findFirst().orElseThrow(); + } + @Test public void testPickTabletPreferCold_picksColdWhenAvailable() throws Exception { TestRebalancer r = new TestRebalancer();