From d8bdbefbaaf8532df432709c5819a919e34a7c13 Mon Sep 17 00:00:00 2001 From: Di Wang Date: Fri, 23 Oct 2026 19:11:20 +0800 Subject: [PATCH] Fix AUTO RegionGroup replacement checks across region types Count disabled groups only for the requested region type. Add regression tests for spurious and missed replacement and the maximum limit. --- .../manager/partition/PartitionManager.java | 6 +- .../PartitionManagerAutoExtensionTest.java | 181 ++++++++++++++++++ 2 files changed, 185 insertions(+), 2 deletions(-) create mode 100644 iotdb-core/confignode/src/test/java/org/apache/iotdb/confignode/manager/partition/PartitionManagerAutoExtensionTest.java diff --git a/iotdb-core/confignode/src/main/java/org/apache/iotdb/confignode/manager/partition/PartitionManager.java b/iotdb-core/confignode/src/main/java/org/apache/iotdb/confignode/manager/partition/PartitionManager.java index 48ebe498a50db..cacc7ceadcbd9 100644 --- a/iotdb-core/confignode/src/main/java/org/apache/iotdb/confignode/manager/partition/PartitionManager.java +++ b/iotdb-core/confignode/src/main/java/org/apache/iotdb/confignode/manager/partition/PartitionManager.java @@ -781,9 +781,11 @@ private TSStatus autoExtendRegionGroupIfNecessary( allotmentMap.put(database, delta); } else if (allocatedRegionGroupCount - == filterRegionGroupThroughStatus(database, RegionGroupStatus.Disabled).size() + == filterRegionGroupThroughStatus(database, RegionGroupStatus.Disabled).stream() + .filter(replicaSet -> replicaSet.getRegionId().getType() == consensusGroupType) + .count() && allocatedRegionGroupCount < maxRegionGroupNum) { - // 3. All RegionGroups in the specified Database are disabled currently + // 3. All RegionGroups of the requested type in the specified Database are disabled. allotmentMap.put(database, 1); } } diff --git a/iotdb-core/confignode/src/test/java/org/apache/iotdb/confignode/manager/partition/PartitionManagerAutoExtensionTest.java b/iotdb-core/confignode/src/test/java/org/apache/iotdb/confignode/manager/partition/PartitionManagerAutoExtensionTest.java new file mode 100644 index 0000000000000..14a36369a264d --- /dev/null +++ b/iotdb-core/confignode/src/test/java/org/apache/iotdb/confignode/manager/partition/PartitionManagerAutoExtensionTest.java @@ -0,0 +1,181 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ +package org.apache.iotdb.confignode.manager.partition; + +import org.apache.iotdb.common.rpc.thrift.TConsensusGroupId; +import org.apache.iotdb.common.rpc.thrift.TConsensusGroupType; +import org.apache.iotdb.common.rpc.thrift.TRegionReplicaSet; +import org.apache.iotdb.common.rpc.thrift.TSeriesPartitionSlot; +import org.apache.iotdb.confignode.conf.ConfigNodeConfig; +import org.apache.iotdb.confignode.conf.ConfigNodeDescriptor; +import org.apache.iotdb.confignode.consensus.request.ConfigPhysicalPlanType; +import org.apache.iotdb.confignode.consensus.request.write.database.DatabaseSchemaPlan; +import org.apache.iotdb.confignode.consensus.request.write.region.CreateRegionGroupsPlan; +import org.apache.iotdb.confignode.manager.IManager; +import org.apache.iotdb.confignode.manager.ProcedureManager; +import org.apache.iotdb.confignode.manager.load.LoadManager; +import org.apache.iotdb.confignode.manager.schema.ClusterSchemaManager; +import org.apache.iotdb.confignode.persistence.partition.PartitionInfo; +import org.apache.iotdb.confignode.rpc.thrift.TDatabaseSchema; +import org.apache.iotdb.rpc.RpcUtils; + +import org.junit.After; +import org.junit.Before; +import org.junit.Test; +import org.junit.runner.RunWith; +import org.junit.runners.Parameterized; + +import java.util.Arrays; +import java.util.Collection; +import java.util.Collections; + +import static org.junit.Assert.assertEquals; +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.ArgumentMatchers.anyMap; +import static org.mockito.ArgumentMatchers.eq; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.never; +import static org.mockito.Mockito.verify; +import static org.mockito.Mockito.when; + +@RunWith(Parameterized.class) +public class PartitionManagerAutoExtensionTest { + + private static final String DATABASE = "root.auto"; + private static final ConfigNodeConfig CONF = ConfigNodeDescriptor.getInstance().getConf(); + + @Parameterized.Parameters(name = "{0}") + public static Collection parameters() { + return Arrays.asList( + new Object[] {TConsensusGroupType.SchemaRegion}, + new Object[] {TConsensusGroupType.DataRegion}); + } + + private final TConsensusGroupType type; + private TConsensusGroupId requestedGroup; + private TConsensusGroupId otherGroup; + private PartitionManager partitionManager; + private LoadManager loadManager; + private ClusterSchemaManager schemaManager; + private RegionGroupExtensionPolicy originalSchemaPolicy; + private RegionGroupExtensionPolicy originalDataPolicy; + private int originalSeriesSlotNum; + + public PartitionManagerAutoExtensionTest(TConsensusGroupType type) { + this.type = type; + } + + @Before + public void setUp() throws Exception { + originalSchemaPolicy = CONF.getSchemaRegionGroupExtensionPolicy(); + originalDataPolicy = CONF.getDataRegionGroupExtensionPolicy(); + originalSeriesSlotNum = CONF.getSeriesSlotNum(); + CONF.setSchemaRegionGroupExtensionPolicy(RegionGroupExtensionPolicy.AUTO); + CONF.setDataRegionGroupExtensionPolicy(RegionGroupExtensionPolicy.AUTO); + CONF.setSeriesSlotNum(100); + + requestedGroup = new TConsensusGroupId(type, 1); + otherGroup = + new TConsensusGroupId( + type == TConsensusGroupType.SchemaRegion + ? TConsensusGroupType.DataRegion + : TConsensusGroupType.SchemaRegion, + 2); + PartitionInfo partitionInfo = new PartitionInfo(); + partitionInfo.createDatabase( + new DatabaseSchemaPlan( + ConfigPhysicalPlanType.CreateDatabase, new TDatabaseSchema(DATABASE))); + CreateRegionGroupsPlan existingGroups = new CreateRegionGroupsPlan(); + for (TConsensusGroupId group : Arrays.asList(requestedGroup, otherGroup)) { + existingGroups.addRegionGroup( + DATABASE, new TRegionReplicaSet(group, Collections.emptyList())); + } + partitionInfo.createRegionGroups(existingGroups); + + IManager manager = mock(IManager.class); + loadManager = mock(LoadManager.class); + schemaManager = mock(ClusterSchemaManager.class); + ProcedureManager procedureManager = mock(ProcedureManager.class); + when(manager.getLoadManager()).thenReturn(loadManager); + when(manager.getClusterSchemaManager()).thenReturn(schemaManager); + when(manager.getProcedureManager()).thenReturn(procedureManager); + // Keep per-group status lookups consistent with the filtered status lists in each test. + when(loadManager.getRegionGroupStatus(requestedGroup)).thenReturn(RegionGroupStatus.Running); + when(loadManager.getRegionGroupStatus(otherGroup)).thenReturn(RegionGroupStatus.Disabled); + // Keep the minimum and slot-based growth conditions false, so status drives replacement. + when(schemaManager.getMinRegionGroupNum(DATABASE, type)).thenReturn(1); + when(schemaManager.getMaxRegionGroupNum(DATABASE, type)).thenReturn(3); + when(loadManager.allocateRegionGroups(anyMap(), eq(type))) + .thenReturn(new CreateRegionGroupsPlan()); + when(procedureManager.createRegionGroups(eq(type), any())).thenReturn(RpcUtils.SUCCESS_STATUS); + partitionManager = new PartitionManager(manager, partitionInfo); + } + + @After + public void tearDown() { + if (partitionManager != null) { + partitionManager.stopRegionCleaner(); + } + CONF.setSchemaRegionGroupExtensionPolicy(originalSchemaPolicy); + CONF.setDataRegionGroupExtensionPolicy(originalDataPolicy); + CONF.setSeriesSlotNum(originalSeriesSlotNum); + } + + @Test + public void testDisabledOtherTypeDoesNotTriggerReplacement() throws Exception { + when(loadManager.filterRegionGroupThroughStatus(RegionGroupStatus.Disabled)) + .thenReturn(Collections.singletonList(otherGroup)); + + extend(); + + verify(loadManager, never()).allocateRegionGroups(anyMap(), any()); + } + + @Test + public void testDisabledOtherTypeDoesNotHideReplacement() throws Exception { + when(loadManager.getRegionGroupStatus(requestedGroup)).thenReturn(RegionGroupStatus.Disabled); + when(loadManager.filterRegionGroupThroughStatus(RegionGroupStatus.Disabled)) + .thenReturn(Arrays.asList(requestedGroup, otherGroup)); + + extend(); + + verify(loadManager).allocateRegionGroups(Collections.singletonMap(DATABASE, 1), type); + } + + @Test + public void testReplacementRespectsMaximum() throws Exception { + when(schemaManager.getMaxRegionGroupNum(DATABASE, type)).thenReturn(1); + when(loadManager.getRegionGroupStatus(requestedGroup)).thenReturn(RegionGroupStatus.Disabled); + when(loadManager.filterRegionGroupThroughStatus(RegionGroupStatus.Disabled)) + .thenReturn(Arrays.asList(requestedGroup, otherGroup)); + + extend(); + + verify(loadManager, never()).allocateRegionGroups(anyMap(), any()); + } + + private void extend() { + assertEquals( + RpcUtils.SUCCESS_STATUS, + partitionManager.extendRegionGroupIfNecessary( + Collections.singletonMap( + DATABASE, Collections.singletonList(new TSeriesPartitionSlot(0))), + type)); + } +}