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 d5d7085bfa82..ed4d4d9be8ad 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 b7cf0639f3da..502b018fd59d 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 7459510f4a23..3c990a895787 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 f378ef132182..ea08fbbb6c6d 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 7bbd366be579..91bd3c020ee2 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 b94c96259687..82a71b58b68c 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 000000000000..2b91f3532f10 --- /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(); + } + } +}