From 6ce30f958b0b0a5db0e75537df5b38f90ed75516 Mon Sep 17 00:00:00 2001 From: Di Wang Date: Fri, 16 Oct 2026 15:58:04 +0800 Subject: [PATCH] Optimize ConfigNode region lookups and reduce metadata copying Look up the target Region and replica node IDs before copying their metadata. Read leaders directly from the cache instead of building the full leader map. --- .../confignode/manager/load/LoadManager.java | 9 + .../manager/load/cache/LoadCache.java | 10 + .../persistence/partition/PartitionInfo.java | 8 +- .../procedure/env/RegionMaintainHandler.java | 36 +++- .../load/cache/ConsensusGroupCacheTest.java | 28 +++ .../env/RegionMaintainHandlerTest.java | 85 ++++++-- .../env/RegionReplicaSetLookupTest.java | 195 ++++++++++++++++++ 7 files changed, 334 insertions(+), 37 deletions(-) create mode 100644 iotdb-core/confignode/src/test/java/org/apache/iotdb/confignode/procedure/env/RegionReplicaSetLookupTest.java diff --git a/iotdb-core/confignode/src/main/java/org/apache/iotdb/confignode/manager/load/LoadManager.java b/iotdb-core/confignode/src/main/java/org/apache/iotdb/confignode/manager/load/LoadManager.java index d5d7085bfa823..ed4d4d9be8ad6 100644 --- a/iotdb-core/confignode/src/main/java/org/apache/iotdb/confignode/manager/load/LoadManager.java +++ b/iotdb-core/confignode/src/main/java/org/apache/iotdb/confignode/manager/load/LoadManager.java @@ -501,6 +501,15 @@ public void removeRegionGroupRelatedCache(TConsensusGroupId consensusGroupId) { eventService.checkAndBroadcastConsensusGroupStatisticsChangeEventIfNecessary(); } + /** + * Get the latest leader of the specified RegionGroup. + * + * @return the leader ID, or -1 if the cache is absent or no leader has been selected + */ + public int getRegionLeaderId(TConsensusGroupId regionGroupId) { + return loadCache.getRegionLeaderId(regionGroupId); + } + /** * Get the latest RegionLeaderMap. * diff --git a/iotdb-core/confignode/src/main/java/org/apache/iotdb/confignode/manager/load/cache/LoadCache.java b/iotdb-core/confignode/src/main/java/org/apache/iotdb/confignode/manager/load/cache/LoadCache.java index b7cf0639f3da7..502b018fd59d8 100644 --- a/iotdb-core/confignode/src/main/java/org/apache/iotdb/confignode/manager/load/cache/LoadCache.java +++ b/iotdb-core/confignode/src/main/java/org/apache/iotdb/confignode/manager/load/cache/LoadCache.java @@ -799,6 +799,16 @@ public void removeRegionGroupCache(TConsensusGroupId consensusGroupId) { consensusGroupCacheMap.remove(consensusGroupId); } + /** + * Get the latest leader of the specified RegionGroup. + * + * @return the leader ID, or -1 if the cache is absent or no leader has been selected + */ + public int getRegionLeaderId(TConsensusGroupId regionGroupId) { + ConsensusGroupCache cache = consensusGroupCacheMap.get(regionGroupId); + return cache == null ? ConsensusGroupCache.UN_READY_LEADER_ID : cache.getLeaderId(); + } + /** * Safely get the latest RegionLeaderMap. * diff --git a/iotdb-core/confignode/src/main/java/org/apache/iotdb/confignode/persistence/partition/PartitionInfo.java b/iotdb-core/confignode/src/main/java/org/apache/iotdb/confignode/persistence/partition/PartitionInfo.java index 7459510f4a23c..3c990a8957879 100644 --- a/iotdb-core/confignode/src/main/java/org/apache/iotdb/confignode/persistence/partition/PartitionInfo.java +++ b/iotdb-core/confignode/src/main/java/org/apache/iotdb/confignode/persistence/partition/PartitionInfo.java @@ -791,11 +791,9 @@ public List getAllReplicaSets(int dataNodeId) { */ public List getReplicaSets( String database, List regionGroupIds) { - if (databasePartitionTables.containsKey(database)) { - return databasePartitionTables.get(database).getReplicaSets(regionGroupIds); - } else { - return Collections.emptyList(); - } + // A database can be deleted while a Region lookup is in progress. + DatabasePartitionTable table = databasePartitionTables.get(database); + return table == null ? Collections.emptyList() : table.getReplicaSets(regionGroupIds); } /** diff --git a/iotdb-core/confignode/src/main/java/org/apache/iotdb/confignode/procedure/env/RegionMaintainHandler.java b/iotdb-core/confignode/src/main/java/org/apache/iotdb/confignode/procedure/env/RegionMaintainHandler.java index f378ef132182d..ea08fbbb6c6d3 100644 --- a/iotdb-core/confignode/src/main/java/org/apache/iotdb/confignode/procedure/env/RegionMaintainHandler.java +++ b/iotdb-core/confignode/src/main/java/org/apache/iotdb/confignode/procedure/env/RegionMaintainHandler.java @@ -53,6 +53,7 @@ import org.apache.iotdb.confignode.manager.ConfigManager; import org.apache.iotdb.confignode.manager.load.LoadManager; import org.apache.iotdb.confignode.manager.load.cache.consensus.ConsensusGroupHeartbeatSample; +import org.apache.iotdb.confignode.manager.partition.PartitionManager; import org.apache.iotdb.confignode.procedure.exception.ProcedureException; import org.apache.iotdb.confignode.rpc.thrift.TCreatePipeReq; import org.apache.iotdb.consensus.pipe.consensuspipe.ConsensusPipeName; @@ -72,6 +73,7 @@ import java.util.List; import java.util.Map; import java.util.Optional; +import java.util.Set; import java.util.concurrent.TimeUnit; import java.util.stream.Collectors; @@ -771,9 +773,17 @@ public List findRegionLocations(TConsensusGroupId regionId) { } public Optional getRegionReplicaSet(TConsensusGroupId regionId) { - return configManager.getPartitionManager().getAllReplicaSets().stream() - .filter(rg -> rg.regionId.equals(regionId)) - .findAny(); + if (regionId == null) { + return Optional.empty(); + } + PartitionManager partitionManager = configManager.getPartitionManager(); + String database = partitionManager.getRegionDatabase(regionId); + if (database == null) { + return Optional.empty(); + } + // Locate by Region ID before copying: unrelated replica sets need not be locked or copied. + return partitionManager.getReplicaSets(database, Collections.singletonList(regionId)).stream() + .findFirst(); } public String getRegionReplicaSetString(TConsensusGroupId regionId) { @@ -931,11 +941,15 @@ public Optional filterDataNodeWithOtherRegionReplica( return Optional.empty(); } - // Choosing the RUNNING DataNodes to execute firstly - // If all DataNodes are not RUNNING, then choose the REMOVING DataNodes secondly + Set regionDataNodeIds = + regionLocations.stream().map(TDataNodeLocation::getDataNodeId).collect(Collectors.toSet()); + // Filter IDs before reading configurations, so unrelated DataNodes are not deep-copied. List aliveDataNodes = - configManager.getNodeManager().filterDataNodeThroughStatus(allowingStatus).stream() - .map(TDataNodeConfiguration::getLocation) + configManager.getLoadManager().filterDataNodeThroughStatus(allowingStatus).stream() + .filter(regionDataNodeIds::contains) + .map( + nodeId -> + configManager.getNodeManager().getRegisteredDataNode(nodeId).getLocation()) .collect(Collectors.toList()); final int leaderId = getRegionLeaderId(regionId); Collections.shuffle(aliveDataNodes); @@ -962,11 +976,11 @@ public Optional filterDataNodeWithOtherRegionReplica( */ private int getRegionLeaderId(TConsensusGroupId regionId) { final LoadManager loadManager = configManager.getLoadManager(); - Integer leaderId = loadManager.getRegionLeaderMap().get(regionId); - if (leaderId == null || leaderId == -1) { + int leaderId = loadManager.getRegionLeaderId(regionId); + if (leaderId == -1) { loadManager.waitForRegionGroupReady(Collections.singletonList(regionId)); - leaderId = loadManager.getRegionLeaderMap().get(regionId); + leaderId = loadManager.getRegionLeaderId(regionId); } - return leaderId == null ? -1 : leaderId; + return leaderId; } } diff --git a/iotdb-core/confignode/src/test/java/org/apache/iotdb/confignode/manager/load/cache/ConsensusGroupCacheTest.java b/iotdb-core/confignode/src/test/java/org/apache/iotdb/confignode/manager/load/cache/ConsensusGroupCacheTest.java index 7bbd366be5799..91bd3c020ee25 100644 --- a/iotdb-core/confignode/src/test/java/org/apache/iotdb/confignode/manager/load/cache/ConsensusGroupCacheTest.java +++ b/iotdb-core/confignode/src/test/java/org/apache/iotdb/confignode/manager/load/cache/ConsensusGroupCacheTest.java @@ -18,14 +18,42 @@ */ package org.apache.iotdb.confignode.manager.load.cache; +import org.apache.iotdb.common.rpc.thrift.TConsensusGroupId; +import org.apache.iotdb.common.rpc.thrift.TConsensusGroupType; import org.apache.iotdb.confignode.manager.load.cache.consensus.ConsensusGroupCache; import org.apache.iotdb.confignode.manager.load.cache.consensus.ConsensusGroupHeartbeatSample; import org.junit.Assert; import org.junit.Test; +import java.util.Collections; + public class ConsensusGroupCacheTest { + @Test + public void regionLeaderLookupTest() { + LoadCache loadCache = new LoadCache(); + TConsensusGroupId dataRegion = new TConsensusGroupId(TConsensusGroupType.DataRegion, 1); + TConsensusGroupId schemaRegion = new TConsensusGroupId(TConsensusGroupType.SchemaRegion, 1); + Assert.assertEquals(-1, loadCache.getRegionLeaderId(dataRegion)); + Assert.assertEquals(-1, loadCache.getRegionLeaderId(schemaRegion)); + + loadCache.createRegionGroupHeartbeatCache("root.test", dataRegion, Collections.singleton(2)); + loadCache.createRegionGroupHeartbeatCache("root.test", schemaRegion, Collections.singleton(3)); + Assert.assertEquals(-1, loadCache.getRegionLeaderId(dataRegion)); + Assert.assertEquals(-1, loadCache.getRegionLeaderId(schemaRegion)); + + loadCache.cacheConsensusSample(dataRegion, new ConsensusGroupHeartbeatSample(1L, 2)); + loadCache.cacheConsensusSample(schemaRegion, new ConsensusGroupHeartbeatSample(1L, 3)); + loadCache.updateConsensusGroupStatistics(); + Assert.assertEquals(2, loadCache.getRegionLeaderId(dataRegion)); + Assert.assertEquals(3, loadCache.getRegionLeaderId(schemaRegion)); + + loadCache.removeRegionGroupCache(dataRegion); + Assert.assertEquals(-1, loadCache.getRegionLeaderId(dataRegion)); + Assert.assertEquals(3, loadCache.getRegionLeaderId(schemaRegion)); + } + @Test public void periodicUpdateTest() { ConsensusGroupCache consensusGroupCache = new ConsensusGroupCache(); diff --git a/iotdb-core/confignode/src/test/java/org/apache/iotdb/confignode/procedure/env/RegionMaintainHandlerTest.java b/iotdb-core/confignode/src/test/java/org/apache/iotdb/confignode/procedure/env/RegionMaintainHandlerTest.java index b94c96259687d..82a71b58b68cd 100644 --- a/iotdb-core/confignode/src/test/java/org/apache/iotdb/confignode/procedure/env/RegionMaintainHandlerTest.java +++ b/iotdb-core/confignode/src/test/java/org/apache/iotdb/confignode/procedure/env/RegionMaintainHandlerTest.java @@ -91,6 +91,7 @@ public static Collection regionTypes() { private TConsensusGroupId regionId; private LoadManager loadManager; + private NodeManager nodeManager; private RegionMaintainHandler handler; private String originalConsensusProtocol; @@ -104,25 +105,26 @@ public void setUp() { regionId = new TConsensusGroupId(regionType, 1); ConfigManager configManager = mock(ConfigManager.class); - NodeManager nodeManager = mock(NodeManager.class); + nodeManager = mock(NodeManager.class); PartitionManager partitionManager = mock(PartitionManager.class); loadManager = mock(LoadManager.class); when(configManager.getNodeManager()).thenReturn(nodeManager); when(configManager.getPartitionManager()).thenReturn(partitionManager); when(configManager.getLoadManager()).thenReturn(loadManager); - when(partitionManager.getAllReplicaSets()) + when(partitionManager.getRegionDatabase(regionId)).thenReturn("root.test"); + when(partitionManager.getReplicaSets("root.test", Collections.singletonList(regionId))) .thenReturn( Collections.singletonList( new TRegionReplicaSet( regionId, Arrays.asList(original, coordinator, otherReplica)))); - when(nodeManager.filterDataNodeThroughStatus(NodeStatus.Running, NodeStatus.ReadOnly)) - .thenReturn( - Arrays.asList( - new TDataNodeConfiguration().setLocation(original), - new TDataNodeConfiguration().setLocation(coordinator), - new TDataNodeConfiguration().setLocation(otherReplica), - new TDataNodeConfiguration().setLocation(nonReplica))); - when(loadManager.getRegionLeaderMap()).thenReturn(Collections.emptyMap()); + when(loadManager.filterDataNodeThroughStatus(NodeStatus.Running, NodeStatus.ReadOnly)) + .thenReturn(Arrays.asList(1, 2, 3, 4)); + for (TDataNodeLocation location : + Arrays.asList(original, coordinator, otherReplica, nonReplica)) { + when(nodeManager.getRegisteredDataNode(location.getDataNodeId())) + .thenReturn(new TDataNodeConfiguration().setLocation(location)); + } + when(loadManager.getRegionLeaderId(regionId)).thenReturn(-1); handler = new RegionMaintainHandler(configManager); } @@ -144,7 +146,7 @@ public void testFilterReplicaWithoutLeaderCache() { @Test public void testFilterReplicaWithUnknownLeader() { - when(loadManager.getRegionLeaderMap()).thenReturn(Collections.singletonMap(regionId, -1)); + when(loadManager.getRegionLeaderId(regionId)).thenReturn(-1); assertEquals( Optional.of(otherReplica), handler.filterDataNodeWithOtherRegionReplica( @@ -154,8 +156,7 @@ public void testFilterReplicaWithUnknownLeader() { @Test public void testFilterReplicaPrefersKnownLeader() { - when(loadManager.getRegionLeaderMap()) - .thenReturn(Collections.singletonMap(regionId, otherReplica.getDataNodeId())); + when(loadManager.getRegionLeaderId(regionId)).thenReturn(otherReplica.getDataNodeId()); assertEquals( Optional.of(otherReplica), handler.filterDataNodeWithOtherRegionReplica(regionId, original)); @@ -178,17 +179,15 @@ public void testTransferLeaderWithoutLeaderCache() throws Exception { @Test public void testTransferLeaderWithUnknownLeader() throws Exception { - when(loadManager.getRegionLeaderMap()).thenReturn(Collections.singletonMap(regionId, -1)); + when(loadManager.getRegionLeaderId(regionId)).thenReturn(-1); assertLeaderTransfer(); } @Test public void testTransferLeaderWaitsForLeaderElection() throws Exception { // The filter has a known leader, while the transfer lookup observes a stale cache first. - when(loadManager.getRegionLeaderMap()) - .thenReturn(Collections.singletonMap(regionId, otherReplica.getDataNodeId())) - .thenReturn(Collections.emptyMap()) - .thenReturn(Collections.singletonMap(regionId, original.getDataNodeId())); + when(loadManager.getRegionLeaderId(regionId)) + .thenReturn(otherReplica.getDataNodeId(), -1, original.getDataNodeId()); assertLeaderTransfer(); verify(loadManager).waitForRegionGroupReady(Collections.singletonList(regionId)); } @@ -196,9 +195,7 @@ public void testTransferLeaderWaitsForLeaderElection() throws Exception { @Test public void testTransferLeaderWhenCacheEntryDisappears() throws Exception { // The cache can change between replica selection and the transfer's leader lookup. - when(loadManager.getRegionLeaderMap()) - .thenReturn(Collections.singletonMap(regionId, -1)) - .thenReturn(Collections.emptyMap()); + when(loadManager.getRegionLeaderId(regionId)).thenReturn(otherReplica.getDataNodeId(), -1); assertLeaderTransfer(); } @@ -239,6 +236,52 @@ private void assertLeaderTransfer() throws Exception { } } + @Test + public void testFilterReplicaDoesNotReadUnrelatedConfigurations() { + assertEquals( + Optional.of(otherReplica), + handler.filterDataNodeWithOtherRegionReplica( + regionId, Arrays.asList(original, coordinator))); + verify(loadManager, never()).getRegionLeaderMap(); + verify(nodeManager, never()).getRegisteredDataNode(nonReplica.getDataNodeId()); + verify(nodeManager, never()) + .filterDataNodeThroughStatus(NodeStatus.Running, NodeStatus.ReadOnly); + } + + @Test + public void testFilterReplicaSkipsUnregisteredNode() { + when(nodeManager.getRegisteredDataNode(otherReplica.getDataNodeId())) + .thenReturn(new TDataNodeConfiguration()); + assertFalse( + handler + .filterDataNodeWithOtherRegionReplica(regionId, Arrays.asList(original, coordinator)) + .isPresent()); + } + + @Test + public void testFilterReplicaRejectsChangedEndpoint() { + TDataNodeLocation changedLocation = otherReplica.deepCopy(); + changedLocation.setInternalEndPoint(new TEndPoint("127.0.0.2", 10730)); + when(nodeManager.getRegisteredDataNode(otherReplica.getDataNodeId())) + .thenReturn(new TDataNodeConfiguration().setLocation(changedLocation)); + assertFalse( + handler + .filterDataNodeWithOtherRegionReplica(regionId, Arrays.asList(original, coordinator)) + .isPresent()); + } + + @Test + public void testFilterReplicaPreservesRequestedStatuses() { + when(loadManager.filterDataNodeThroughStatus(NodeStatus.Removing)) + .thenReturn(Collections.singletonList(otherReplica.getDataNodeId())); + assertEquals( + Optional.of(otherReplica), + handler.filterDataNodeWithOtherRegionReplica(regionId, original, NodeStatus.Removing)); + verify(nodeManager, never()).getRegisteredDataNode(coordinator.getDataNodeId()); + verify(loadManager, never()).getRegionLeaderMap(); + verify(nodeManager, never()).getRegisteredDataNode(nonReplica.getDataNodeId()); + } + private RegionMaintainHandler stubStatusAndSubmit( final NodeStatus status, final TSStatus submitResult) { RegionMaintainHandler retryHandler = spy(handler); diff --git a/iotdb-core/confignode/src/test/java/org/apache/iotdb/confignode/procedure/env/RegionReplicaSetLookupTest.java b/iotdb-core/confignode/src/test/java/org/apache/iotdb/confignode/procedure/env/RegionReplicaSetLookupTest.java new file mode 100644 index 0000000000000..2b91f3532f105 --- /dev/null +++ b/iotdb-core/confignode/src/test/java/org/apache/iotdb/confignode/procedure/env/RegionReplicaSetLookupTest.java @@ -0,0 +1,195 @@ +/* + * 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.procedure.env; + +import org.apache.iotdb.common.rpc.thrift.TConsensusGroupId; +import org.apache.iotdb.common.rpc.thrift.TConsensusGroupType; +import org.apache.iotdb.common.rpc.thrift.TDataNodeLocation; +import org.apache.iotdb.common.rpc.thrift.TEndPoint; +import org.apache.iotdb.common.rpc.thrift.TRegionReplicaSet; +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.database.DeleteDatabasePlan; +import org.apache.iotdb.confignode.consensus.request.write.database.PreDeleteDatabasePlan; +import org.apache.iotdb.confignode.consensus.request.write.region.CreateRegionGroupsPlan; +import org.apache.iotdb.confignode.manager.ConfigManager; +import org.apache.iotdb.confignode.manager.partition.PartitionManager; +import org.apache.iotdb.confignode.persistence.partition.PartitionInfo; +import org.apache.iotdb.confignode.rpc.thrift.TDatabaseSchema; + +import org.junit.Before; +import org.junit.Test; +import org.junit.runner.RunWith; +import org.junit.runners.Parameterized; + +import java.util.ArrayList; +import java.util.Arrays; +import java.util.Collection; +import java.util.List; +import java.util.Optional; + +import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertFalse; +import static org.junit.Assert.assertNotSame; +import static org.junit.Assert.assertTrue; +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.ArgumentMatchers.anyList; +import static org.mockito.ArgumentMatchers.anyString; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.when; + +@RunWith(Parameterized.class) +public class RegionReplicaSetLookupTest { + @Parameterized.Parameters(name = "{0}") + public static Collection regionTypes() { + return Arrays.asList( + new Object[] {TConsensusGroupType.DataRegion}, + new Object[] {TConsensusGroupType.SchemaRegion}); + } + + @Parameterized.Parameter public TConsensusGroupType regionType; + + private final List replicaSets = new ArrayList<>(); + private PartitionInfo partitionInfo; + private PartitionManager partitionManager; + private RegionMaintainHandler handler; + private TConsensusGroupId target; + + @Before + public void setUp() { + partitionInfo = new PartitionInfo(); + partitionManager = mock(PartitionManager.class); + ConfigManager configManager = mock(ConfigManager.class); + when(configManager.getPartitionManager()).thenReturn(partitionManager); + when(partitionManager.getRegionDatabase(any())) + .thenAnswer(invocation -> partitionInfo.getRegionDatabase(invocation.getArgument(0))); + when(partitionManager.getReplicaSets(anyString(), anyList())) + .thenAnswer( + invocation -> + partitionInfo.getReplicaSets(invocation.getArgument(0), invocation.getArgument(1))); + when(partitionManager.getAllReplicaSets()) + .thenAnswer(invocation -> partitionInfo.getAllReplicaSets()); + handler = new RegionMaintainHandler(configManager); + for (int database = 0; database < 3; database++) { + String name = "root.lookup" + database; + partitionInfo.createDatabase( + new DatabaseSchemaPlan(ConfigPhysicalPlanType.CreateDatabase, new TDatabaseSchema(name))); + CreateRegionGroupsPlan plan = new CreateRegionGroupsPlan(); + for (int index = 0; index < 16; index++) { + for (TConsensusGroupType type : + Arrays.asList(TConsensusGroupType.DataRegion, TConsensusGroupType.SchemaRegion)) { + CountingReplicaSet replicaSet = + new CountingReplicaSet(new TConsensusGroupId(type, database * 16 + index)); + replicaSets.add(replicaSet); + plan.addRegionGroup(name, replicaSet); + } + } + partitionInfo.createRegionGroups(plan); + } + target = new TConsensusGroupId(regionType, 20); + } + + @Test + public void testLookupCopiesOnlyRequestedRegion() { + Optional result = handler.getRegionReplicaSet(target); + assertTrue(result.isPresent()); + assertEquals(target, result.get().getRegionId()); + assertEquals(1, copyCount()); + for (CountingReplicaSet replicaSet : replicaSets) { + assertEquals(replicaSet.getRegionId().equals(target) ? 1 : 0, replicaSet.copies); + } + } + + @Test + public void testReturnedReplicaSetIsAnIndependentDeepCopy() { + TRegionReplicaSet first = handler.getRegionReplicaSet(target).get(); + first.getDataNodeLocations().get(0).getInternalEndPoint().setIp("changed"); + first.getDataNodeLocations().remove(1); + first.getRegionId().setId(-1); + TRegionReplicaSet second = handler.getRegionReplicaSet(target).get(); + assertEquals(target, second.getRegionId()); + assertEquals(3, second.getDataNodeLocationsSize()); + assertEquals("127.0.0.1", second.getDataNodeLocations().get(0).getInternalEndPoint().getIp()); + assertNotSame(first.getDataNodeLocations().get(0), second.getDataNodeLocations().get(0)); + assertEquals(2, copyCount()); + } + + @Test + public void testMissingRegionDoesNotCopyReplicaSets() { + TConsensusGroupId missing = new TConsensusGroupId(regionType, 9999); + assertFalse(handler.getRegionReplicaSet(missing).isPresent()); + assertTrue(handler.findRegionLocations(missing).isEmpty()); + assertEquals("UNKNOWN!", handler.getRegionReplicaSetString(missing)); + assertEquals(0, copyCount()); + } + + @Test + public void testNullRegionRemainsAbsent() { + assertFalse(handler.getRegionReplicaSet(null).isPresent()); + assertEquals(0, copyCount()); + } + + @Test + public void testPreDeletedDatabaseRemainsVisible() { + partitionInfo.preDeleteDatabase( + new PreDeleteDatabasePlan("root.lookup1", PreDeleteDatabasePlan.PreDeleteType.EXECUTE)); + assertTrue(handler.getRegionReplicaSet(target).isPresent()); + assertEquals(1, copyCount()); + } + + @Test + public void testDatabaseDeletedBetweenLookups() { + when(partitionManager.getRegionDatabase(target)) + .thenAnswer( + invocation -> { + String database = partitionInfo.getRegionDatabase(target); + partitionInfo.deleteDatabase(new DeleteDatabasePlan(database)); + return database; + }); + assertFalse(handler.getRegionReplicaSet(target).isPresent()); + assertEquals(0, copyCount()); + } + + private int copyCount() { + return replicaSets.stream().mapToInt(replicaSet -> replicaSet.copies).sum(); + } + + private static class CountingReplicaSet extends TRegionReplicaSet { + private int copies; + + private CountingReplicaSet(TConsensusGroupId regionId) { + setRegionId(regionId); + List locations = new ArrayList<>(); + for (int nodeId = 0; nodeId < 3; nodeId++) { + locations.add( + new TDataNodeLocation() + .setDataNodeId(nodeId) + .setInternalEndPoint(new TEndPoint("127.0.0.1", 10730 + nodeId))); + } + setDataNodeLocations(locations); + } + + @Override + public TRegionReplicaSet deepCopy() { + copies++; + return super.deepCopy(); + } + } +}