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..1a231170b9eba9 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 @@ -1059,22 +1059,19 @@ public void fillBeToTablets(long be, long tableId, long partId, long indexId, lo ConcurrentHashMap>>> partToTablets) { // global - globalBeToTablets.putIfAbsent(be, ConcurrentHashMap.newKeySet()); - globalBeToTablets.get(be).add(tabletId); + globalBeToTablets.computeIfAbsent(be, ignored -> ConcurrentHashMap.newKeySet()).add(tabletId); // table - beToTabletsInTable.putIfAbsent(tableId, new ConcurrentHashMap>()); - ConcurrentHashMap> beToTabletsOfTable = beToTabletsInTable.get(tableId); - beToTabletsOfTable.putIfAbsent(be, ConcurrentHashMap.newKeySet()); - beToTabletsOfTable.get(be).add(tabletId); + ConcurrentHashMap> beToTabletsOfTable = + beToTabletsInTable.computeIfAbsent(tableId, ignored -> new ConcurrentHashMap<>()); + beToTabletsOfTable.computeIfAbsent(be, ignored -> ConcurrentHashMap.newKeySet()).add(tabletId); // partition - partToTablets.putIfAbsent(partId, new ConcurrentHashMap>>()); - ConcurrentHashMap>> indexToTablets = partToTablets.get(partId); - indexToTablets.putIfAbsent(indexId, new ConcurrentHashMap>()); - ConcurrentHashMap> beToTabletsOfIndex = indexToTablets.get(indexId); - beToTabletsOfIndex.putIfAbsent(be, ConcurrentHashMap.newKeySet()); - beToTabletsOfIndex.get(be).add(tabletId); + ConcurrentHashMap>> indexToTablets = + partToTablets.computeIfAbsent(partId, ignored -> new ConcurrentHashMap<>()); + ConcurrentHashMap> beToTabletsOfIndex = + indexToTablets.computeIfAbsent(indexId, ignored -> new ConcurrentHashMap<>()); + beToTabletsOfIndex.computeIfAbsent(be, ignored -> ConcurrentHashMap.newKeySet()).add(tabletId); } private void enqueueWarmupTask(WarmupTabletTask task) { 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..080ee4f5a34cd0 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 @@ -41,6 +41,7 @@ import java.util.Random; import java.util.Set; import java.util.concurrent.ConcurrentHashMap; +import java.util.function.Function; public class CloudTabletRebalancerTest { @@ -81,6 +82,30 @@ protected boolean isInternalDbId(Long dbId) { } } + private static class CountingConcurrentHashMap extends ConcurrentHashMap { + private int computeIfAbsentCalls; + private int getCalls; + private int putIfAbsentCalls; + + @Override + public V computeIfAbsent(K key, Function mappingFunction) { + computeIfAbsentCalls++; + return super.computeIfAbsent(key, mappingFunction); + } + + @Override + public V get(Object key) { + getCalls++; + return super.get(key); + } + + @Override + public V putIfAbsent(K key, V value) { + putIfAbsentCalls++; + return super.putIfAbsent(key, value); + } + } + private static void setField(Object obj, String name, Object value) throws Exception { Field f = CloudTabletRebalancer.class.getDeclaredField(name); f.setAccessible(true); @@ -101,6 +126,50 @@ private static T invokePrivate(Object obj, String method, Class[] types, return (T) m.invoke(obj, args); } + @Test + public void testFillBeToTabletsUsesComputedContainers() { + TestRebalancer rebalancer = new TestRebalancer(); + long beId = 1L; + long tableId = 2L; + long partitionId = 3L; + long indexId = 4L; + + CountingConcurrentHashMap> globalBeToTablets = new CountingConcurrentHashMap<>(); + CountingConcurrentHashMap>> beToTabletsInTable = + new CountingConcurrentHashMap<>(); + CountingConcurrentHashMap> beToTabletsOfTable = new CountingConcurrentHashMap<>(); + beToTabletsInTable.put(tableId, beToTabletsOfTable); + + CountingConcurrentHashMap>>> + partToTablets = new CountingConcurrentHashMap<>(); + CountingConcurrentHashMap>> indexToTablets = + new CountingConcurrentHashMap<>(); + CountingConcurrentHashMap> beToTabletsOfIndex = new CountingConcurrentHashMap<>(); + partToTablets.put(partitionId, indexToTablets); + indexToTablets.put(indexId, beToTabletsOfIndex); + + rebalancer.fillBeToTablets(beId, tableId, partitionId, indexId, 5L, + globalBeToTablets, beToTabletsInTable, partToTablets); + rebalancer.fillBeToTablets(beId, tableId, partitionId, indexId, 6L, + globalBeToTablets, beToTabletsInTable, partToTablets); + + assertComputedContainerUsed(globalBeToTablets); + assertComputedContainerUsed(beToTabletsInTable); + assertComputedContainerUsed(beToTabletsOfTable); + assertComputedContainerUsed(partToTablets); + assertComputedContainerUsed(indexToTablets); + assertComputedContainerUsed(beToTabletsOfIndex); + Assertions.assertEquals(Set.of(5L, 6L), globalBeToTablets.get(beId)); + Assertions.assertEquals(Set.of(5L, 6L), beToTabletsOfTable.get(beId)); + Assertions.assertEquals(Set.of(5L, 6L), beToTabletsOfIndex.get(beId)); + } + + private static void assertComputedContainerUsed(CountingConcurrentHashMap map) { + Assertions.assertEquals(2, map.computeIfAbsentCalls); + Assertions.assertEquals(0, map.putIfAbsentCalls); + Assertions.assertEquals(0, map.getCalls); + } + @Test public void testPickTabletPreferCold_picksColdWhenAvailable() throws Exception { TestRebalancer r = new TestRebalancer();