From a326d1c6d12cf08b419aebde82f41a1546a57ea7 Mon Sep 17 00:00:00 2001 From: deardeng Date: Mon, 3 Aug 2026 22:30:17 +0800 Subject: [PATCH 1/2] [improvement](fe) Reuse boxed IDs in cloud tablet indexes ### What problem does this PR solve? Issue Number: None Related PR: #61318 Problem Summary: Cloud tablet route rebuilding repeatedly boxes primitive backend, table, partition, index, and tablet IDs while inserting the same logical IDs into current and future global, table, and partition indexes. Hoist boxing to traversal callers and preserve selected boxed tablet IDs through direct transfers, warmup moves, and warmup rollbacks so all route index families reuse immutable Long references without changing routing or scheduling semantics. A single-threaded JDK 17 allocation model that keeps eager container candidates in both variants estimates that 4 million tablets across 4 clusters reduce cumulative allocation from 25.47 GiB to 19.76 GiB (22.42%) and approximate post-full-GC retained heap from 6.80 GiB to 4.94 GiB (27.37%). For IDs outside the Long cache, each direct incremental move additionally avoids two tablet wrapper allocations and one duplicate retained tablet wrapper; warmup and rollback each avoid one tablet wrapper allocation. These are path-level model and identity-test results, not production RSS measurements. ### Release note None ### Check List (For Author) - Test: Unit Test and Manual test\n - ./run-fe-ut.sh --run org.apache.doris.cloud.catalog.CloudTabletRebalancerTest (13 tests passed)\n - mvn checkstyle:check -pl fe-core (0 violations)\n - Multi-scale JDK 17 eager-container allocation model at 0.2M, 0.4M, 0.8M, and 3.2M tablet-cluster pairs, three runs per scale\n- Behavior changed: No\n- Does this need documentation: No --- .../cloud/catalog/CloudTabletRebalancer.java | 67 +++--- .../catalog/CloudTabletRebalancerTest.java | 216 ++++++++++++++++++ 2 files changed, 256 insertions(+), 27 deletions(-) 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..ecf4a7f68557fa 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,11 @@ 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); + fillBeToTablets(beId, tableId, partitionId, indexId, tabletId, tmpBeToTabletsGlobal, beToTabletsInTable, this.partitionToTablets); - fillBeToTablets(futureBeId, table.getId(), partition.getId(), index.getId(), tabletId, + fillBeToTablets(futureBeId, tableId, partitionId, indexId, tabletId, tmpFutureBeToTabletsGlobal, futureBeToTabletsInTable, futurePartitionToTablets); } } @@ -1618,7 +1631,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 +1671,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 +1928,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 +2084,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 +2110,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..5224b7e7dfe121 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,15 @@ package org.apache.doris.cloud.catalog; +import org.apache.doris.catalog.Env; +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 +53,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 +69,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 +110,213 @@ 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); + } + } + + 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 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) { + Long stored = ids.stream().filter(expected::equals).findFirst().orElseThrow(); + Assertions.assertSame(expected, stored); + } + @Test public void testPickTabletPreferCold_picksColdWhenAvailable() throws Exception { TestRebalancer r = new TestRebalancer(); From dd6909afab54253408756d1408964855e8612b76 Mon Sep 17 00:00:00 2001 From: deardeng Date: Tue, 4 Aug 2026 11:44:23 +0800 Subject: [PATCH 2/2] [fix](fe) Preserve inflight tablet IDs across route rebuilds ### What problem does this PR solve? Issue Number: None Related PR: #66389 Problem Summary: When an asynchronous warmup task survives into a later rebalance round, statRouteInfo rebuilds current and future route maps with a newly boxed tablet ID. If the queued warmup failure is then rolled back, future maps reinsert the task's older boxed ID while current maps retain the rebuilt object. Numeric routing stays correct, but current and future indexes retain duplicate Long wrappers until the next rebuild, defeating boxed-ID sharing. Reuse the inflight task's canonical tablet ID for both map families during rebuild, avoiding one duplicate retained wrapper per affected failed inflight tablet without changing route or scheduling semantics. ### Release note None ### Check List (For Author) - Test: Unit Test - ./run-fe-ut.sh --run org.apache.doris.cloud.catalog.CloudTabletRebalancerTest (14 tests passed) - cd fe && mvn checkstyle:check -pl fe-core (0 violations) - Behavior changed: No - Does this need documentation: No --- .../cloud/catalog/CloudTabletRebalancer.java | 5 +- .../catalog/CloudTabletRebalancerTest.java | 118 +++++++++++++++++- 2 files changed, 119 insertions(+), 4 deletions(-) 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 ecf4a7f68557fa..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 @@ -1238,10 +1238,11 @@ public void statRouteInfo() { InfightTablet taskKey = new InfightTablet(tabletId, cluster); InfightTask task = tabletToInfightTask.get(taskKey); Long futureBeId = task == null ? beId : Long.valueOf(task.destBe); - fillBeToTablets(beId, tableId, partitionId, indexId, tabletId, + Long routeTabletId = task == null ? tabletId : task.pickedTabletId; + fillBeToTablets(beId, tableId, partitionId, indexId, routeTabletId, tmpBeToTabletsGlobal, beToTabletsInTable, this.partitionToTablets); - fillBeToTablets(futureBeId, tableId, partitionId, indexId, tabletId, + fillBeToTablets(futureBeId, tableId, partitionId, indexId, routeTabletId, tmpFutureBeToTabletsGlobal, futureBeToTabletsInTable, futurePartitionToTablets); } } 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 5224b7e7dfe121..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,7 +17,13 @@ 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; @@ -257,6 +263,67 @@ public void testWarmupRollbackRestoresSelectedBoxedIdsInFutureIndexes() throws E } } + @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, @@ -287,6 +354,50 @@ private static MockedStatic mockTabletMeta(Long tabletId, Long tableId, Lon 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)); @@ -313,8 +424,11 @@ private static void assertSameStoredId(Long expected, Map map) { } private static void assertSameStoredId(Long expected, Set ids) { - Long stored = ids.stream().filter(expected::equals).findFirst().orElseThrow(); - Assertions.assertSame(expected, stored); + Assertions.assertSame(expected, getStoredId(expected, ids)); + } + + private static Long getStoredId(Long expected, Set ids) { + return ids.stream().filter(expected::equals).findFirst().orElseThrow(); } @Test