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..d33448e0a9b10b 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 @@ -514,16 +514,19 @@ protected void runAfterCatalogReady() { return; } - statRouteInfo(); - boolean migrated = migrateTabletsForSmoothUpgrade(); - if (migrated) { + try { statRouteInfo(); - } - - indexBalanced = true; - tableBalanced = true; + boolean migrated = migrateTabletsForSmoothUpgrade(); + if (migrated) { + statRouteInfo(); + } - performBalancing(); + indexBalanced = true; + tableBalanced = true; + performBalancing(); + } finally { + releaseSchedulingIndexes(); + } checkDecommissionState(clusterToBes); inited = true; @@ -642,6 +645,15 @@ private void performBalancing() { } } + private void releaseSchedulingIndexes() { + // These indexes are no longer used after balancing. Replace their top-level maps instead of clearing + // every entry so the complete tablet membership graphs can become collectible without an O(N) traversal. + partitionToTablets = new ConcurrentHashMap<>(); + futurePartitionToTablets = new ConcurrentHashMap<>(); + beToTabletsInTable = new ConcurrentHashMap<>(); + futureBeToTabletsInTable = new ConcurrentHashMap<>(); + } + private boolean shouldForceInactivePhase(boolean activeBalanced) { if (activeBalanced) { consecutiveActiveUnbalancedRounds = 0; 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..a2b415dd162e19 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 @@ -83,7 +83,11 @@ private static class TestRebalancer extends CloudTabletRebalancer { private final Set internalDbIds = new HashSet<>(); TestRebalancer() { - super(null); + this(null); + } + + TestRebalancer(CloudSystemInfoService cloudSystemInfoService) { + super(cloudSystemInfoService); } void setInternalDbIds(Set ids) { @@ -349,6 +353,107 @@ public void testWarmupRollbackReusesInflightBoxedTabletIdAfterRouteRebuild() thr } } + @Test + public void testReleaseSchedulingIndexesKeepsGlobalRoutesAndAllowsNextRebuild() throws Exception { + TestRebalancer rebalancer = new TestRebalancer(); + Long srcBe = 10_001L; + 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); + + invokePrivate(rebalancer, "releaseSchedulingIndexes", new Class[] {}, new Object[] {}); + + ConcurrentHashMap> currentGlobal = getField(rebalancer, "beToTabletsGlobal"); + ConcurrentHashMap> futureGlobal = getField(rebalancer, "futureBeToTabletsGlobal"); + ConcurrentHashMap>> releasedCurrentByTable = + getField(rebalancer, "beToTabletsInTable"); + ConcurrentHashMap>> releasedFutureByTable = + getField(rebalancer, "futureBeToTabletsInTable"); + ConcurrentHashMap>>> + releasedCurrentByPartition = getField(rebalancer, "partitionToTablets"); + ConcurrentHashMap>>> + releasedFutureByPartition = getField(rebalancer, "futurePartitionToTablets"); + Assertions.assertSame(current.global, currentGlobal); + Assertions.assertSame(future.global, futureGlobal); + Assertions.assertNotSame(current.byTable, releasedCurrentByTable); + Assertions.assertNotSame(future.byTable, releasedFutureByTable); + Assertions.assertNotSame(current.byPartition, releasedCurrentByPartition); + Assertions.assertNotSame(future.byPartition, releasedFutureByPartition); + Assertions.assertTrue(releasedCurrentByTable.isEmpty()); + Assertions.assertTrue(releasedFutureByTable.isEmpty()); + Assertions.assertTrue(releasedCurrentByPartition.isEmpty()); + Assertions.assertTrue(releasedFutureByPartition.isEmpty()); + + setField(rebalancer, "clusterToBes", Collections.singletonMap(clusterId, List.of(srcBe))); + setField(rebalancer, "allBes", Set.of(srcBe)); + try (MockedStatic ignored = mockRouteEnvironment( + dbId, tableId, partitionId, indexId, tabletId, clusterId, srcBe)) { + rebalancer.statRouteInfo(); + } + + ConcurrentHashMap>> rebuiltCurrentByTable = + getField(rebalancer, "beToTabletsInTable"); + ConcurrentHashMap>> rebuiltFutureByTable = + getField(rebalancer, "futureBeToTabletsInTable"); + ConcurrentHashMap>>> + rebuiltCurrentByPartition = getField(rebalancer, "partitionToTablets"); + ConcurrentHashMap>>> + rebuiltFutureByPartition = getField(rebalancer, "futurePartitionToTablets"); + Assertions.assertEquals(Set.of(tabletId), rebuiltCurrentByTable.get(tableId).get(srcBe)); + Assertions.assertEquals(Set.of(tabletId), rebuiltFutureByTable.get(tableId).get(srcBe)); + Assertions.assertEquals(Set.of(tabletId), + rebuiltCurrentByPartition.get(partitionId).get(indexId).get(srcBe)); + Assertions.assertEquals(Set.of(tabletId), + rebuiltFutureByPartition.get(partitionId).get(indexId).get(srcBe)); + } + + @Test + public void testRunAfterCatalogReadyReleasesSchedulingIndexesWhenMigrationFails() throws Exception { + 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"; + CloudSystemInfoService systemInfoService = Mockito.mock(CloudSystemInfoService.class); + Backend srcBackend = Mockito.mock(Backend.class); + Mockito.when(systemInfoService.getAllBackendIds()).thenReturn(List.of(srcBe)); + Mockito.when(systemInfoService.getBackend(srcBe)).thenReturn(srcBackend); + Mockito.when(srcBackend.getCloudClusterId()).thenReturn(clusterId); + TestRebalancer rebalancer = new TestRebalancer(systemInfoService); + rebalancer.addTabletMigrationTask(srcBe, destBe); + + boolean oldEnableCloudMultiReplica = Config.enable_cloud_multi_replica; + Config.enable_cloud_multi_replica = false; + try (MockedStatic ignored = mockRouteEnvironment( + dbId, tableId, partitionId, indexId, tabletId, clusterId, srcBe)) { + TabletInvertedIndex invertedIndex = Env.getCurrentEnv().getTabletInvertedIndex(); + Mockito.when(invertedIndex.getTabletMeta(tabletId)) + .thenThrow(new RuntimeException("injected migration failure")); + + RuntimeException exception = Assertions.assertThrows( + RuntimeException.class, rebalancer::runAfterCatalogReady); + + Assertions.assertEquals("injected migration failure", exception.getMessage()); + ConcurrentHashMap> currentGlobal = getField(rebalancer, "beToTabletsGlobal"); + Assertions.assertEquals(Set.of(tabletId), currentGlobal.get(srcBe)); + Assertions.assertTrue(((Map) getField(rebalancer, "beToTabletsInTable")).isEmpty()); + Assertions.assertTrue(((Map) getField(rebalancer, "futureBeToTabletsInTable")).isEmpty()); + Assertions.assertTrue(((Map) getField(rebalancer, "partitionToTablets")).isEmpty()); + Assertions.assertTrue(((Map) getField(rebalancer, "futurePartitionToTablets")).isEmpty()); + } finally { + Config.enable_cloud_multi_replica = oldEnableCloudMultiReplica; + } + } + 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, @@ -415,6 +520,7 @@ private static MockedStatic mockRouteEnvironment(Long dbId, Long tableId, L Mockito.when(tablet.getReplicas()).thenReturn(Collections.singletonList(replica)); Mockito.when(replica.getPrimaryBackend(clusterId, false)).thenReturn(primaryBackend); Mockito.when(primaryBackend.getId()).thenReturn(srcBe); + Mockito.when(primaryBackend.isQueryAvailable()).thenReturn(true); MockedStatic mockedEnv = Mockito.mockStatic(Env.class); mockedEnv.when(Env::getCurrentEnv).thenReturn(env);