diff --git a/fe/fe-core/src/main/java/org/apache/doris/planner/OlapTableSink.java b/fe/fe-core/src/main/java/org/apache/doris/planner/OlapTableSink.java index 1c934895b5b472..bdf50f0d388db6 100644 --- a/fe/fe-core/src/main/java/org/apache/doris/planner/OlapTableSink.java +++ b/fe/fe-core/src/main/java/org/apache/doris/planner/OlapTableSink.java @@ -1129,17 +1129,21 @@ public List createDummyLocation(OlapTable table) throws final long fakeTabletId = 0; SystemInfoService clusterInfo = Env.getCurrentSystemInfo(); - List aliveBe = clusterInfo.getAllBackendIds(true); - if (aliveBe.isEmpty()) { + List availableBeIds = clusterInfo.getBackendsByCurrentCluster().values().stream() + .filter(Backend::isLoadAvailable) + .filter(backend -> !backend.isDecommissioned() && !backend.isDecommissioning()) + .map(Backend::getId) + .collect(Collectors.toList()); + if (availableBeIds.isEmpty()) { throw new UserException(InternalErrorCode.REPLICA_FEW_ERR, "no available BE in cluster"); } for (int i = 0; i < table.getIndexNumber(); i++) { // only one fake tablet here - Long[] nodes = aliveBe.toArray(new Long[0]); + Long[] nodes = availableBeIds.toArray(new Long[0]); Random random = new SecureRandom(); int nodeIndex = random.nextInt(nodes.length); if (singleReplicaLoad) { - List slaveBe = aliveBe; + List slaveBe = new ArrayList<>(availableBeIds); locationParam.addToTablets(new TTabletLocation(fakeTabletId, Arrays.asList(nodes[nodeIndex]))); diff --git a/fe/fe-core/src/test/java/org/apache/doris/planner/OlapTableSinkTest.java b/fe/fe-core/src/test/java/org/apache/doris/planner/OlapTableSinkTest.java index 3d88e7b1e692dc..2aa2d6ceacd393 100644 --- a/fe/fe-core/src/test/java/org/apache/doris/planner/OlapTableSinkTest.java +++ b/fe/fe-core/src/test/java/org/apache/doris/planner/OlapTableSinkTest.java @@ -17,20 +17,84 @@ package org.apache.doris.planner; +import org.apache.doris.catalog.Env; +import org.apache.doris.catalog.OlapTable; import org.apache.doris.planner.OlapTableSink.AdaptiveBucketAssignment; import org.apache.doris.planner.OlapTableSink.AdaptiveIndexBucketAssignment; +import org.apache.doris.system.Backend; +import org.apache.doris.system.SystemInfoService; import org.apache.doris.thrift.TOlapTableIndexTablets; +import org.apache.doris.thrift.TOlapTableLocationParam; import org.apache.doris.thrift.TOlapTablePartition; import org.apache.doris.thrift.TTabletLocation; +import com.google.common.collect.ImmutableMap; import org.junit.Assert; import org.junit.Test; +import org.mockito.MockedStatic; +import org.mockito.Mockito; import java.util.Arrays; +import java.util.Collections; import java.util.List; import java.util.Map; public class OlapTableSinkTest { + @Test + public void testCreateDummyLocationUsesLoadAvailableBackendInCurrentComputeGroup() throws Exception { + SystemInfoService systemInfoService = Mockito.mock(SystemInfoService.class); + Backend currentComputeGroupBackend = Mockito.mock(Backend.class); + Backend loadDisabledBackend = Mockito.mock(Backend.class); + OlapTable table = Mockito.mock(OlapTable.class); + + Mockito.when(currentComputeGroupBackend.getId()).thenReturn(1L); + Mockito.when(currentComputeGroupBackend.isLoadAvailable()).thenReturn(true); + Mockito.when(loadDisabledBackend.getId()).thenReturn(2L); + Mockito.when(loadDisabledBackend.isLoadAvailable()).thenReturn(false); + Mockito.when(systemInfoService.getBackendsByCurrentCluster()) + .thenReturn(ImmutableMap.of(1L, currentComputeGroupBackend, 2L, loadDisabledBackend)); + Mockito.when(systemInfoService.getAllBackendIds(true)).thenReturn(Collections.singletonList(3L)); + Mockito.when(table.getIndexNumber()).thenReturn(1); + + try (MockedStatic mockedEnv = Mockito.mockStatic(Env.class)) { + mockedEnv.when(Env::getCurrentSystemInfo).thenReturn(systemInfoService); + + OlapTableSink sink = new OlapTableSink(table, null, Collections.emptyList(), false); + List locationParams = sink.createDummyLocation(table); + + Assert.assertEquals(Collections.singletonList(1L), + locationParams.get(0).getTablets().get(0).getNodeIds()); + Mockito.verify(systemInfoService, Mockito.never()).getAllBackendIds(true); + Mockito.verify(systemInfoService).getBackendsByCurrentCluster(); + } + } + + @Test + public void testCreateDummyLocationDoesNotShareBackendCandidatesAcrossIndexes() throws Exception { + SystemInfoService systemInfoService = Mockito.mock(SystemInfoService.class); + Backend currentComputeGroupBackend = Mockito.mock(Backend.class); + OlapTable table = Mockito.mock(OlapTable.class); + + Mockito.when(currentComputeGroupBackend.getId()).thenReturn(1L); + Mockito.when(currentComputeGroupBackend.isLoadAvailable()).thenReturn(true); + Mockito.when(systemInfoService.getBackendsByCurrentCluster()) + .thenReturn(ImmutableMap.of(1L, currentComputeGroupBackend)); + Mockito.when(table.getIndexNumber()).thenReturn(2); + + try (MockedStatic mockedEnv = Mockito.mockStatic(Env.class)) { + mockedEnv.when(Env::getCurrentSystemInfo).thenReturn(systemInfoService); + + OlapTableSink sink = new OlapTableSink(table, null, Collections.emptyList(), true); + List locationParams = sink.createDummyLocation(table); + + Assert.assertEquals(2, locationParams.get(0).getTabletsSize()); + Assert.assertEquals(Collections.singletonList(1L), + locationParams.get(0).getTablets().get(0).getNodeIds()); + Assert.assertEquals(Collections.singletonList(1L), + locationParams.get(0).getTablets().get(1).getNodeIds()); + } + } + @Test public void testAdaptiveRandomBucketAssignmentIsPerIndex() { TOlapTablePartition partition = new TOlapTablePartition();