diff --git a/hugegraph-pd/docs/api-reference.md b/hugegraph-pd/docs/api-reference.md index 122ae6c8b4..f1c2ad9be4 100644 --- a/hugegraph-pd/docs/api-reference.md +++ b/hugegraph-pd/docs/api-reference.md @@ -90,6 +90,26 @@ enum ShardRole { Represents a Store node in the cluster. +An `Exiting` Store continues renewing its heartbeat lease while draining, but is excluded from +new replica allocation and migration destinations. Expiry of an old heartbeat does not complete +retirement. PD keeps the Store `Exiting` until its latest heartbeat reports no local partition +engines and PD has no shard group referencing its ID. Before deleting its Pod or volume, also +verify the remaining groups' actual Raft membership, full replicas and fresh heartbeats; a +management state or a successful state-update response alone is not a complete maintenance check. +Replica allocation still enforces the configured minimum Store count using eligible destinations. +If the remaining eligible Stores do not satisfy that minimum, draining remains blocked; expiry of +the source lease must not be used to force retirement or reduce the configured replica policy. +If a drain request fails this capacity check, add enough eligible capacity and repeat the same +`Exiting` management request to resume planning for groups that still reference the source in PD. +The persisted intent stays `Exiting`; there is no need to toggle it to `Up`. This retry does not +establish that an already dispatched migration has completed or recover a lost Store-side task. +`Tombstone` is permanent for that Store identity: management updates cannot return it to an active +state, just as registration and heartbeat already reject a retired identity. +Heartbeats, registration and management state changes serialize their Store-state updates. +Re-registering an `Exiting` Store preserves that state even when its address is in the initial +Store list. This preserves the drain intent, but does not replace checking restored engines and +actual Raft configuration after a process restart. + ```protobuf message Store { uint64 id = 1; diff --git a/hugegraph-pd/hg-pd-core/src/main/java/org/apache/hugegraph/pd/StoreNodeService.java b/hugegraph-pd/hg-pd-core/src/main/java/org/apache/hugegraph/pd/StoreNodeService.java index 0f23195cf8..10a0b1692a 100644 --- a/hugegraph-pd/hg-pd-core/src/main/java/org/apache/hugegraph/pd/StoreNodeService.java +++ b/hugegraph-pd/hg-pd-core/src/main/java/org/apache/hugegraph/pd/StoreNodeService.java @@ -131,7 +131,7 @@ public boolean isOK() { * * @param store */ - public Metapb.Store register(Metapb.Store store) throws PDException { + public synchronized Metapb.Store register(Metapb.Store store) throws PDException { if (store.getId() == 0) { // Initial registration, generate a new ID, and ensure that the ID is not duplicated. store = newStoreNode(store); @@ -167,11 +167,15 @@ public Metapb.Store register(Metapb.Store store) throws PDException { // also exposes liveness that tells stale membership apart. // https://github.com/apache/hugegraph/issues/3229 Metapb.StoreState storeState = lastStore.getState(); - if (storeState == Metapb.StoreState.Offline || storeState == Metapb.StoreState.Up - || inInitialStoreList(store)) { - storeState = Metapb.StoreState.Up; - } else { - storeState = Metapb.StoreState.Pending; + // Restarting an exiting Store must not cancel its persisted drain intent, including + // when its address is still present in the bootstrap allowlist. + if (storeState != Metapb.StoreState.Exiting) { + if (storeState == Metapb.StoreState.Offline || storeState == Metapb.StoreState.Up + || inInitialStoreList(store)) { + storeState = Metapb.StoreState.Up; + } else { + storeState = Metapb.StoreState.Pending; + } } store = Metapb.Store.newBuilder(lastStore) @@ -212,6 +216,8 @@ public Metapb.Store register(Metapb.Store store) throws PDException { storeInfoMeta.keepStoreAlive(store); onStoreStatusChanged(store, Metapb.StoreState.Offline, Metapb.StoreState.Up); checkStoreStatus(); + } else if (storeState == Metapb.StoreState.Exiting) { + storeInfoMeta.keepStoreAlive(store); } // Wait for the store information to be saved before sending the changes @@ -266,36 +272,44 @@ public Metapb.Store getStore(long id) throws PDException { /** * Update the store information, detect the change of store status, and notify Hugestore */ - public synchronized Metapb.Store updateStore(Metapb.Store store) throws PDException { - log.info("updateStore storeId: {}, address: {}, state: {}", store.getId(), - store.getAddress(), store.getState()); - Metapb.Store lastStore = storeInfoMeta.getStore(store.getId()); - if (lastStore == null) { - return null; - } - Metapb.Store.Builder builder = - Metapb.Store.newBuilder(lastStore).clearLabels().clearStats(); - store = builder.mergeFrom(store).build(); - if (store.getState() == Metapb.StoreState.Tombstone) { - List activeStores = getStores(); - if (lastStore.getState() == Metapb.StoreState.Up - && activeStores.size() - 1 < pdConfig.getMinStoreCount()) { - throw new PDException(Pdpb.ErrorType.LESS_ACTIVE_STORE_VALUE, - "The number of active stores is less then " + - pdConfig.getMinStoreCount()); + public Metapb.Store updateStore(Metapb.Store store) throws PDException { + Metapb.Store lastStore; + boolean stateChanged; + synchronized (this) { + log.info("updateStore storeId: {}, address: {}, state: {}", store.getId(), + store.getAddress(), store.getState()); + lastStore = storeInfoMeta.getStore(store.getId()); + if (lastStore == null) { + return null; } - } - - storeInfoMeta.updateStore(store); - if (store.getState() != Metapb.StoreState.Unknown && - store.getState() != lastStore.getState()) { - // If you want to take the store offline - if (store.getState() == Metapb.StoreState.Exiting) { - if (lastStore.getState() == Metapb.StoreState.Exiting) { - // If it is already in the offline state, no further processing will be made - return lastStore; + boolean requestedDrain = store.getState() == Metapb.StoreState.Exiting; + Metapb.Store.Builder builder = + Metapb.Store.newBuilder(lastStore).clearLabels().clearStats(); + store = builder.mergeFrom(store).build(); + // Tombstone is terminal for this Store identity, as in registration and + // heartbeat handling. A deferred retirement notification must not race a + // management request that revives the same identity. + if (lastStore.getState() == Metapb.StoreState.Tombstone && + store.getState() != Metapb.StoreState.Tombstone) { + throw new PDException(Pdpb.ErrorType.STORE_HAS_BEEN_REMOVED_VALUE, + String.format("Store id %d has been removed.", store.getId())); + } + if (store.getState() == Metapb.StoreState.Tombstone) { + List activeStores = getStores(); + if (lastStore.getState() == Metapb.StoreState.Up + && activeStores.size() - 1 < pdConfig.getMinStoreCount()) { + throw new PDException(Pdpb.ErrorType.LESS_ACTIVE_STORE_VALUE, + "The number of active stores is less then " + + pdConfig.getMinStoreCount()); } + } + storeInfoMeta.updateStore(store); + stateChanged = store.getState() != Metapb.StoreState.Unknown && + store.getState() != lastStore.getState(); + // A previous attempt can persist Exiting before capacity checks fail. Repeating + // the same intent must resume the remaining groups after capacity is restored. + if (requestedDrain) { List activeStores = this.getActiveStores(); Map storeMap = new HashMap<>(); activeStores.forEach(s -> { @@ -309,20 +323,26 @@ public synchronized Metapb.Store updateStore(Metapb.Store store) throws PDExcept storeInfoMeta.removeActiveStore(store); } storeTurnoff(store); - } else if (store.getState() == Metapb.StoreState.Offline) { - // Monitor that the store has gone offline and is removed from the active - storeInfoMeta.removeActiveStore(store); - } else if (store.getState() == Metapb.StoreState.Tombstone) { - // When the status changes, the store is shut down, the shardGroup is modified, - // and the replica is migrated - log.info("updateStore removeActiveStores store {}", store.getId()); - storeInfoMeta.removeActiveStore(store); - // Storage goes offline - storeTurnoff(store); - } else if (store.getState() == Metapb.StoreState.Up) { - storeInfoMeta.keepStoreAlive(store); - checkStoreStatus(); + } else if (stateChanged) { + if (store.getState() == Metapb.StoreState.Offline) { + // Monitor that the store has gone offline and is removed from the active + storeInfoMeta.removeActiveStore(store); + } else if (store.getState() == Metapb.StoreState.Tombstone) { + // When the status changes, the store is shut down, the shardGroup is modified, + // and the replica is migrated + log.info("updateStore removeActiveStores store {}", store.getId()); + storeInfoMeta.removeActiveStore(store); + // Storage goes offline + storeTurnoff(store); + } else if (store.getState() == Metapb.StoreState.Up) { + storeInfoMeta.keepStoreAlive(store); + checkStoreStatus(); + } } + } + // Tombstone listeners enter synchronized PartitionService methods. Never call + // them with this monitor held: placement takes PartitionService before this. + if (stateChanged) { onStoreStatusChanged(lastStore, lastStore.getState(), store.getState()); } return store; @@ -409,6 +429,21 @@ public List getActiveStores() throws PDException { return storeInfoMeta.getActiveStores(); } + /** + * Return live Stores that may receive new replicas. The TTL entry records liveness and can + * still contain an older Up state after a Store starts exiting, so consult Store metadata. + */ + public List getPlacementStores() throws PDException { + List stores = new ArrayList<>(); + for (Metapb.Store activeStore : storeInfoMeta.getActiveStores()) { + Metapb.Store store = storeInfoMeta.getStore(activeStore.getId()); + if (store != null && store.getState() == Metapb.StoreState.Up) { + stores.add(store); + } + } + return stores; + } + public List getTombStores() throws PDException { List stores = new ArrayList<>(); for (Metapb.Store store : this.getStores()) { @@ -437,7 +472,7 @@ public synchronized List allocShards(Metapb.Graph graph, int partI // number cannot exceed the number of raft groups if (storeInfoMeta.getShardGroup(partId) == null) { // Get active store key - List stores = storeInfoMeta.getActiveStores(); + List stores = getPlacementStores(); if (stores.size() == 0) { throw new PDException(Pdpb.ErrorType.NO_ACTIVE_STORE_VALUE, @@ -497,7 +532,7 @@ public synchronized List allocShards(Metapb.Graph graph, int partI */ public synchronized List reallocShards(Metapb.ShardGroup shardGroup) throws PDException { - List stores = storeInfoMeta.getActiveStores(); + List stores = getPlacementStores(); if (stores.size() == 0) { throw new PDException(Pdpb.ErrorType.NO_ACTIVE_STORE_VALUE, @@ -727,7 +762,7 @@ public synchronized void updateShardGroupState(int groupId, Metapb.PartitionStat * @param storeStats * @throws PDException */ - public Metapb.ClusterStats heartBeat(Metapb.StoreStats storeStats) throws PDException { + public synchronized Metapb.ClusterStats heartBeat(Metapb.StoreStats storeStats) throws PDException { this.storeInfoMeta.updateStoreStats(storeStats); Metapb.Store lastStore = this.getStore(storeStats.getStoreId()); if (lastStore == null) { @@ -745,21 +780,16 @@ public Metapb.ClusterStats heartBeat(Metapb.StoreStats storeStats) throws PDExce Metapb.Store nowStore; // If you are going to take the store offline if (lastStore.getState() == Metapb.StoreState.Exiting) { - List activeStores = this.getActiveStores(); - Map storeMap = new HashMap<>(); - activeStores.forEach(store -> { - storeMap.put(store.getId(), store); - }); - // If the partition of the offline store is 0, it means that the migration has been - // completed and can be taken offline, if it is not 0, the migration is still in - // progress and you need to wait - if (storeStats.getPartitionCount() > 0 && - storeMap.containsKey(storeStats.getStoreId())) { + // TTL expiry only means the last heartbeat is stale, not that migration finished. + // Require both the source and PD's group view to be empty before retiring it. + if (storeStats.getPartitionCount() > 0 || + !getShardGroupsByStore(storeStats.getStoreId()).isEmpty()) { nowStore = Metapb.Store.newBuilder(lastStore) .setStats(storeStats) .setLastHeartbeat(System.currentTimeMillis()) .setState(Metapb.StoreState.Exiting).build(); this.storeInfoMeta.updateStore(nowStore); + this.storeInfoMeta.keepStoreAlive(nowStore); return this.clusterStats; } else { nowStore = Metapb.Store.newBuilder(lastStore) @@ -892,7 +922,13 @@ public void addShardGroupStatusListener(ShardGroupStatusListener listener) { protected void onStoreStatusChanged(Metapb.Store store, Metapb.StoreState old, Metapb.StoreState stats) { log.info("onStoreStatusChanged storeId = {} from {} to {}", store.getId(), old, stats); - statusListeners.forEach(e -> { + List listeners; + synchronized (statusListeners) { + listeners = new ArrayList<>(statusListeners); + } + // Do not hold the listener-list monitor across callbacks either: register() + // dispatches Up while holding this Store service's monitor. + listeners.forEach(e -> { e.onStoreStatusChanged(store, old, stats); }); } diff --git a/hugegraph-pd/hg-pd-core/src/main/java/org/apache/hugegraph/pd/TaskScheduleService.java b/hugegraph-pd/hg-pd-core/src/main/java/org/apache/hugegraph/pd/TaskScheduleService.java index c111a6f85c..bc63149eea 100644 --- a/hugegraph-pd/hg-pd-core/src/main/java/org/apache/hugegraph/pd/TaskScheduleService.java +++ b/hugegraph-pd/hg-pd-core/src/main/java/org/apache/hugegraph/pd/TaskScheduleService.java @@ -298,7 +298,8 @@ public synchronized Map> balancePartitionShard() thr return null; } - int activeStores = storeService.getActiveStores().size(); + List placementStores = storeService.getPlacementStores(); + int activeStores = placementStores.size(); if (activeStores == 0) { log.warn("balancePartitionShard non active stores, skip to balancePartitionShard"); return null; @@ -318,7 +319,7 @@ public synchronized Map> balancePartitionShard() thr // Count the partitions on each store, StoreId -> PartitionID, ShardRole Map> partitionMap = new HashMap<>(); - storeService.getActiveStores().forEach(store -> { + placementStores.forEach(store -> { partitionMap.put(store.getId(), new HashMap<>()); }); @@ -449,22 +450,61 @@ public synchronized Map> balancePartitionShard() thr maxHeap.addAll(tmpList); } - kvService.put(BALANCE_SHARD_KEY, "DOING", 180 * 1000); - - // Start the migration - movedPartitions.forEach((partId, storePair) -> { - // Neither the source nor destination storeID is 0 - if (storePair.getKey() > 0 && storePair.getValue() > 0) { - partitionService.movePartitionsShard(partId, storePair.getKey(), - storePair.getValue()); - } else { + boolean submitted = false; + Iterator>> submission = + movedPartitions.entrySet().iterator(); + while (submission.hasNext()) { + Map.Entry> move = submission.next(); + KVPair stores = move.getValue(); + if (stores.getKey() <= 0 || stores.getValue() <= 0) { log.warn("balancePartitionShard key or value is zero, partId:{} storePair:{}", - partId, storePair); + move.getKey(), stores); + submission.remove(); + continue; } - }); + if (!movePartitionIfPlacementStore(move.getKey(), stores.getKey(), + stores.getValue())) { + submission.remove(); + continue; + } + submitted = true; + } + if (submitted) { + kvService.put(BALANCE_SHARD_KEY, "DOING", 180 * 1000); + } return movedPartitions; } + /** + * Revalidate destination liveness and state while serializing with Store state changes. + * Preserve the existing service lock order: PartitionService then StoreNodeService. + */ + private boolean movePartitionIfPlacementStore(int partitionId, long sourceStoreId, + long destStoreId) { + synchronized (partitionService) { + synchronized (storeService) { + try { + boolean eligible = storeService.getPlacementStores().stream() + .anyMatch(store -> + store.getId() == destStoreId); + if (!eligible) { + log.warn("Skip moving partition {} from store {} to store {}: " + + "destination is no longer an Up placement store", + partitionId, sourceStoreId, destStoreId); + return false; + } + partitionService.movePartitionsShard(partitionId, sourceStoreId, + destStoreId); + return true; + } catch (PDException e) { + log.error("Cannot revalidate destination store {} for partition {}", + destStoreId, partitionId, e); + return false; + } + } + } + } + /** * Balance the number of leaders of partitions between stores */ @@ -751,12 +791,15 @@ public Map canAllPartitionsMovedOut(Metapb.Store sourceStore) th Map partitionDataSize = new HashMap<>(); storeService.getActiveStores().forEach(store -> { + if (store.getId() == sourceStore.getId()) { + resultMap.put("current_store_is_online", true); + } + }); + storeService.getPlacementStores().forEach(store -> { if (store.getId() != sourceStore.getId()) { otherPartitionMap.put(store.getId(), new HashMap<>()); // Records the remaining disk space of other stores, in bytes availableDiskSpace.put(store.getId(), store.getStats().getAvailable()); - } else { - resultMap.put("current_store_is_online", true); } }); // Count the size of the partition to be migrated (from storeStats in KB) @@ -791,10 +834,9 @@ public Map canAllPartitionsMovedOut(Metapb.Store sourceStore) th } // Count the number of partitions of other stores and save them with a small top heap, so // that stores with fewer partitions are always prioritized - PriorityQueue> minHeap = new PriorityQueue<>(otherPartitionMap.size(), - (o1, o2) -> o1.getValue() - .compareTo( - o2.getValue())); + PriorityQueue> minHeap = + new PriorityQueue<>(Math.max(1, otherPartitionMap.size()), + (o1, o2) -> o1.getValue().compareTo(o2.getValue())); otherPartitionMap.forEach((storeId, shards) -> { minHeap.add(new KVPair(storeId, shards.size())); }); @@ -870,15 +912,19 @@ public Map> movePartitions( if (!isLeader()) { return null; } - // Start the migration + // Start the migration, but revalidate every destination because a plan may be stale. log.info("begin move partitions:"); - movedPartitions.forEach((partId, storePair) -> { - // Neither the source nor destination storeID is 0 - if (storePair.getKey() > 0 && storePair.getValue() > 0) { - partitionService.movePartitionsShard(partId, storePair.getKey(), - storePair.getValue()); + Iterator>> iterator = + movedPartitions.entrySet().iterator(); + while (iterator.hasNext()) { + Map.Entry> move = iterator.next(); + KVPair stores = move.getValue(); + if (stores.getKey() <= 0 || stores.getValue() <= 0 || + !movePartitionIfPlacementStore(move.getKey(), stores.getKey(), + stores.getValue())) { + iterator.remove(); } - }); + } return movedPartitions; } diff --git a/hugegraph-pd/hg-pd-test/src/main/java/org/apache/hugegraph/pd/core/PDCoreSuiteTest.java b/hugegraph-pd/hg-pd-test/src/main/java/org/apache/hugegraph/pd/core/PDCoreSuiteTest.java index ad97e37d55..9377d62de8 100644 --- a/hugegraph-pd/hg-pd-test/src/main/java/org/apache/hugegraph/pd/core/PDCoreSuiteTest.java +++ b/hugegraph-pd/hg-pd-test/src/main/java/org/apache/hugegraph/pd/core/PDCoreSuiteTest.java @@ -56,6 +56,8 @@ PartitionServiceTest.class, StoreMonitorDataServiceTest.class, StoreServiceTest.class, + StoreDrainTest.class, + StorePlacementSubmissionTest.class, TaskScheduleServiceTest.class, IpAuthHandlerTest.class, RaftEngineCliProcessorsTest.class, diff --git a/hugegraph-pd/hg-pd-test/src/main/java/org/apache/hugegraph/pd/core/StoreDrainTest.java b/hugegraph-pd/hg-pd-test/src/main/java/org/apache/hugegraph/pd/core/StoreDrainTest.java new file mode 100644 index 0000000000..e607b05926 --- /dev/null +++ b/hugegraph-pd/hg-pd-test/src/main/java/org/apache/hugegraph/pd/core/StoreDrainTest.java @@ -0,0 +1,424 @@ +/* + * 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.hugegraph.pd.core; + +import java.lang.reflect.Field; +import java.util.ArrayList; +import java.util.HashMap; +import java.util.List; +import java.util.Map; +import java.util.concurrent.CountDownLatch; +import java.util.concurrent.ExecutorService; +import java.util.concurrent.Executors; +import java.util.concurrent.Future; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicBoolean; + +import org.apache.hugegraph.pd.ConfigService; +import org.apache.hugegraph.pd.PartitionService; +import org.apache.hugegraph.pd.StoreNodeService; +import org.apache.hugegraph.pd.StoreStatusListener; +import org.apache.hugegraph.pd.TaskScheduleService; +import org.apache.hugegraph.pd.common.PDException; +import org.apache.hugegraph.pd.config.PDConfig; +import org.apache.hugegraph.pd.grpc.Metapb; +import org.apache.hugegraph.pd.grpc.Pdpb; +import org.apache.hugegraph.pd.grpc.pulse.ConfChangeType; +import org.apache.hugegraph.pd.meta.StoreInfoMeta; +import org.apache.hugegraph.pd.raft.RaftEngine; +import org.junit.Assert; +import org.junit.Test; +import org.mockito.Mockito; +import org.mockito.invocation.Invocation; + +import com.alipay.sofa.jraft.Node; + +public class StoreDrainTest { + + @Test + public void testExpiredLeaseDoesNotRetireStoreWithLocalGroups() throws Exception { + Fixture fixture = new Fixture(); + Metapb.Store exiting = fixture.addStore(4, Metapb.StoreState.Exiting); + // No active entry remains, as when a migration outlasts the heartbeat TTL. + Mockito.when(fixture.meta.getActiveStores()).thenReturn(List.of()); + fixture.heartbeat(4, 3); + Assert.assertEquals(Metapb.StoreState.Exiting, fixture.stores.get(4L).getState()); + Mockito.verify(fixture.meta).keepStoreAlive(fixture.stores.get(4L)); + Mockito.verify(fixture.meta, Mockito.never()).removeActiveStore(Mockito.any()); + Assert.assertEquals(exiting.getId(), fixture.stores.get(4L).getStats().getStoreId()); + } + + @Test + public void testZeroLocalGroupsDoesNotRetireReferencedStore() throws Exception { + Fixture fixture = new Fixture(); + fixture.addStore(4, Metapb.StoreState.Exiting); + fixture.groups.put(0, group(0, 1, 2, 4)); + fixture.heartbeat(4, 0); + Assert.assertEquals(Metapb.StoreState.Exiting, fixture.stores.get(4L).getState()); + Mockito.verify(fixture.meta).keepStoreAlive(fixture.stores.get(4L)); + Mockito.verify(fixture.meta, Mockito.never()).removeActiveStore(Mockito.any()); + } + + @Test + public void testRetiresOnlyAfterBothViewsAreEmpty() throws Exception { + Fixture fixture = new Fixture(); + fixture.addStore(4, Metapb.StoreState.Exiting); + fixture.heartbeat(4, 0); + Assert.assertEquals(Metapb.StoreState.Tombstone, fixture.stores.get(4L).getState()); + Mockito.verify(fixture.meta).removeActiveStore(fixture.stores.get(4L)); + Mockito.verify(fixture.meta, Mockito.never()).keepStoreAlive(Mockito.any()); + } + + @Test + public void testPlacementUsesCurrentStateInsteadOfStaleLease() throws Exception { + Fixture fixture = new Fixture(); + Metapb.Store up = fixture.addStore(1, Metapb.StoreState.Up); + fixture.addStore(2, Metapb.StoreState.Exiting); + fixture.addStore(3, Metapb.StoreState.Pending); + fixture.addStore(4, Metapb.StoreState.Tombstone); + Metapb.Store staleUp = Metapb.Store.newBuilder().setId(2) + .setState(Metapb.StoreState.Up).build(); + Mockito.when(fixture.meta.getActiveStores()).thenReturn(List.of( + up, staleUp, fixture.stores.get(3L), fixture.stores.get(4L), + Metapb.Store.newBuilder().setId(5).setState(Metapb.StoreState.Up).build())); + Assert.assertEquals(List.of(up), fixture.service.getPlacementStores()); + // Liveness remains available for an exiting source that still serves its groups. + Assert.assertEquals(5, fixture.service.getActiveStores().size()); + } + + @Test + public void testDrainCannotPlaceReplicaBackOnSourceAndNotifiesSharedGroupOnce() throws Exception { + Fixture fixture = new Fixture(); + for (long id = 1; id <= 3; id++) { + fixture.addStore(id, Metapb.StoreState.Up); + } + Metapb.Store source = fixture.addStore(4, Metapb.StoreState.Exiting); + fixture.groups.put(0, group(0, 1, 2, 4)); + Mockito.when(fixture.meta.getActiveStores()).thenReturn(List.of( + fixture.stores.get(1L), fixture.stores.get(2L), fixture.stores.get(3L), + source.toBuilder().setState(Metapb.StoreState.Up).build())); + List sharedPartitions = List.of( + Metapb.Partition.newBuilder().setId(0).setGraphName("graph-a").build(), + Metapb.Partition.newBuilder().setId(0).setGraphName("graph-b").build()); + Mockito.when(fixture.partitions.getPartitionById(0)).thenReturn(sharedPartitions); + + fixture.service.storeTurnoff(source); + + List result = fixture.groups.get(0).getShardsList(); + Assert.assertEquals(3, result.size()); + Assert.assertTrue(result.stream().noneMatch(shard -> shard.getStoreId() == 4)); + Assert.assertTrue(result.stream().anyMatch(shard -> shard.getStoreId() == 3)); + List changes = Mockito.mockingDetails(fixture.partitions).getInvocations().stream() + .filter(invocation -> invocation.getMethod().getName() + .equals("fireChangeShard")) + .toList(); + Assert.assertEquals(1, changes.size()); + Assert.assertEquals(sharedPartitions.get(0), changes.get(0).getArgument(0)); + Assert.assertEquals(result, changes.get(0).getArgument(1)); + Assert.assertEquals(ConfChangeType.CONF_CHANGE_TYPE_ADJUST, changes.get(0).getArgument(2)); + } + + @Test + public void testInitialAllocationExcludesExitingStoreWithStaleUpLease() throws Exception { + Fixture fixture = new Fixture(); + for (long id = 1; id <= 3; id++) { + fixture.addStore(id, Metapb.StoreState.Up); + } + Metapb.Store source = fixture.addStore(4, Metapb.StoreState.Exiting); + Mockito.when(fixture.meta.getActiveStores()).thenReturn(List.of( + source.toBuilder().setState(Metapb.StoreState.Up).build(), + fixture.stores.get(1L), fixture.stores.get(2L), fixture.stores.get(3L))); + List shards = fixture.service.allocShards( + Metapb.Graph.newBuilder().setGraphName("new-graph").build(), 0); + Assert.assertEquals(3, shards.size()); + Assert.assertTrue(shards.stream().noneMatch(shard -> shard.getStoreId() == 4)); + } + + @Test + public void testDrainingStoreCannotSatisfyPlacementMinimum() throws Exception { + Fixture fixture = new Fixture(); + fixture.addStore(1, Metapb.StoreState.Up); + fixture.addStore(2, Metapb.StoreState.Up); + Metapb.Store source = fixture.addStore(3, Metapb.StoreState.Exiting); + Mockito.when(fixture.meta.getActiveStores()).thenReturn(new ArrayList<>(fixture.stores.values())); + Metapb.ShardGroup original = group(0, 1, 2, 3); + fixture.groups.put(0, original); + + PDException error = Assert.assertThrows(PDException.class, () -> fixture.service.storeTurnoff(source)); + + Assert.assertEquals(Pdpb.ErrorType.LESS_ACTIVE_STORE_VALUE, error.getErrorCode()); + Assert.assertEquals(original, fixture.groups.get(0)); + Assert.assertEquals(3, fixture.service.getActiveStores().size()); + Assert.assertEquals(2, fixture.service.getPlacementStores().size()); + } + + @Test + public void testRepeatedDrainResumesAfterCapacityIsRestored() throws Exception { + Fixture fixture = new Fixture(); + StoreStatusListener listener = Mockito.mock(StoreStatusListener.class); + setField(fixture.service, "statusListeners", List.of(listener)); + fixture.addStore(1, Metapb.StoreState.Up); + fixture.addStore(2, Metapb.StoreState.Up); + Metapb.Store source = fixture.addStore(3, Metapb.StoreState.Up); + Mockito.when(fixture.meta.getActiveStores()) + .thenAnswer(call -> new ArrayList<>(fixture.stores.values())); + Metapb.ShardGroup original = group(0, 1, 2, 3); + fixture.groups.put(0, original); + Metapb.Partition partition = Metapb.Partition.newBuilder() + .setId(0).setGraphName("graph").build(); + Mockito.when(fixture.partitions.getPartitionById(0)).thenReturn(List.of(partition)); + Metapb.Store drain = source.toBuilder().setState(Metapb.StoreState.Exiting).build(); + + PDException error = Assert.assertThrows(PDException.class, + () -> fixture.service.updateStore(drain)); + Assert.assertEquals(Pdpb.ErrorType.LESS_ACTIVE_STORE_VALUE, error.getErrorCode()); + Assert.assertEquals(Metapb.StoreState.Exiting, fixture.stores.get(3L).getState()); + Assert.assertEquals(original, fixture.groups.get(0)); + Assert.assertTrue(changeShardInvocations(fixture.partitions).isEmpty()); + + // A partial metadata update is not an explicit retry of the drain intent. + fixture.service.updateStore(Metapb.Store.newBuilder().setId(3) + .setVersion("metadata-only").build()); + Assert.assertEquals(original, fixture.groups.get(0)); + + fixture.addStore(4, Metapb.StoreState.Up); + fixture.service.updateStore(drain); + + List expected = fixture.groups.get(0).getShardsList(); + Assert.assertEquals(3, expected.size()); + Assert.assertTrue(expected.stream().noneMatch(shard -> shard.getStoreId() == 3)); + Assert.assertTrue(expected.stream().anyMatch(shard -> shard.getStoreId() == 4)); + Assert.assertEquals(Metapb.StoreState.Exiting, fixture.stores.get(3L).getState()); + List changes = changeShardInvocations(fixture.partitions); + Assert.assertEquals(1, changes.size()); + Assert.assertEquals(partition, changes.get(0).getArgument(0)); + Assert.assertEquals(expected, changes.get(0).getArgument(1)); + Assert.assertEquals(ConfChangeType.CONF_CHANGE_TYPE_ADJUST, changes.get(0).getArgument(2)); + + // Planning removed the source from the PD group; another identical request must + // not resubmit that group or repeat a Store-state notification. This is not a + // barrier proving that the actual Raft membership has committed the target. + fixture.service.updateStore(drain); + Assert.assertEquals(1, changeShardInvocations(fixture.partitions).size()); + Mockito.verifyNoInteractions(listener); + } + + @Test + public void testManagementCannotReviveTombstoneButAllowsPartialMetadataUpdate() throws Exception { + Fixture fixture = new Fixture(); + Metapb.Store retired = fixture.addStore(3, Metapb.StoreState.Tombstone); + for (Metapb.StoreState state : List.of(Metapb.StoreState.Up, + Metapb.StoreState.Exiting, + Metapb.StoreState.Offline, + Metapb.StoreState.Pending)) { + PDException error = Assert.assertThrows(PDException.class, () -> + fixture.service.updateStore(retired.toBuilder().setState(state).build())); + Assert.assertEquals(Pdpb.ErrorType.STORE_HAS_BEEN_REMOVED_VALUE, error.getErrorCode()); + Assert.assertEquals(retired, fixture.stores.get(3L)); + } + Metapb.Store partial = Metapb.Store.newBuilder().setId(3) + .setState(Metapb.StoreState.Unknown) + .setVersion("updated").build(); + Metapb.Store updated = fixture.service.updateStore(partial); + Assert.assertEquals(Metapb.StoreState.Tombstone, updated.getState()); + Assert.assertEquals("updated", updated.getVersion()); + Assert.assertEquals(updated, fixture.service.updateStore(updated)); + } + + private static List changeShardInvocations(PartitionService partitions) { + return Mockito.mockingDetails(partitions).getInvocations().stream() + .filter(invocation -> invocation.getMethod().getName().equals("fireChangeShard")) + .toList(); + } + + @Test + public void testDrainIsBlockedWhenOtherLiveStoresCannotReceiveReplicas() throws Exception { + Fixture fixture = new Fixture(); + Metapb.Store source = fixture.addStore(1, Metapb.StoreState.Up); + fixture.addStore(2, Metapb.StoreState.Exiting); + fixture.addStore(3, Metapb.StoreState.Pending); + Mockito.when(fixture.meta.getActiveStores()).thenReturn(new ArrayList<>(fixture.stores.values())); + fixture.groups.put(0, group(0, 1)); + Mockito.when(fixture.partitions.getPartitions()).thenReturn(List.of( + Metapb.Partition.newBuilder().setId(0).setGraphName("graph-a").build())); + TaskScheduleService tasks = Mockito.mock(TaskScheduleService.class, Mockito.CALLS_REAL_METHODS); + Field storesField = TaskScheduleService.class.getDeclaredField("storeService"); + storesField.setAccessible(true); + storesField.set(tasks, fixture.service); + Field partitionsField = TaskScheduleService.class.getDeclaredField("partitionService"); + partitionsField.setAccessible(true); + partitionsField.set(tasks, fixture.partitions); + RaftEngine raft = RaftEngine.getInstance(); + Field nodeField = RaftEngine.class.getDeclaredField("raftNode"); + nodeField.setAccessible(true); + Object previous = nodeField.get(raft); + Node node = Mockito.mock(Node.class); + Mockito.when(node.isLeader(true)).thenReturn(true); + nodeField.set(raft, node); + try { + Map plan = tasks.canAllPartitionsMovedOut(source); + Assert.assertEquals(true, plan.get("current_store_is_online")); + Assert.assertEquals(false, plan.get("flag")); + Assert.assertNull(plan.get("movedPartitions")); + } finally { + nodeField.set(raft, previous); + } + } + + @Test + public void testReregisteringExitingStorePreservesDrainIntent() throws Exception { + Fixture fixture = new Fixture(); + Metapb.Store source = fixture.addStore(4, Metapb.StoreState.Exiting); + fixture.config.setInitialStoreList("source:8500"); + source = source.toBuilder().setAddress("source:8500").build(); + fixture.stores.put(4L, source); + + Metapb.Store registered = fixture.service.register(source); + + Assert.assertEquals(Metapb.StoreState.Exiting, registered.getState()); + Mockito.verify(fixture.meta).keepStoreAlive(registered); + Assert.assertTrue(fixture.service.getPlacementStores().isEmpty()); + } + + @Test + public void testHeartbeatCannotOverwriteConcurrentDrainIntent() throws Exception { + assertStateUpdateSerializesWithDrain(false); + } + + @Test + public void testRegistrationCannotOverwriteConcurrentDrainIntent() throws Exception { + assertStateUpdateSerializesWithDrain(true); + } + + private void assertStateUpdateSerializesWithDrain(boolean registration) throws Exception { + Fixture fixture = new Fixture(); + Metapb.Store source = fixture.addStore(4, Metapb.StoreState.Up); + CountDownLatch readUp = new CountDownLatch(1); + CountDownLatch continueUpdate = new CountDownLatch(1); + CountDownLatch drainRequested = new CountDownLatch(1); + CountDownLatch drainDone = new CountDownLatch(1); + AtomicBoolean firstRead = new AtomicBoolean(true); + Mockito.when(fixture.meta.getStore(4L)).thenAnswer(call -> { + Metapb.Store observed = fixture.stores.get(4L); + if (firstRead.compareAndSet(true, false)) { + readUp.countDown(); + Assert.assertTrue(continueUpdate.await(5, TimeUnit.SECONDS)); + } + return observed; + }); + ExecutorService workers = Executors.newFixedThreadPool(2); + try { + Future update = workers.submit(() -> { + if (registration) { + fixture.service.register(source); + } else { + fixture.heartbeat(4, 3); + } + return null; + }); + Assert.assertTrue(readUp.await(5, TimeUnit.SECONDS)); + Future drain = workers.submit(() -> { + drainRequested.countDown(); + fixture.service.updateStore(source.toBuilder().setState(Metapb.StoreState.Exiting).build()); + drainDone.countDown(); + return null; + }); + Assert.assertTrue(drainRequested.await(5, TimeUnit.SECONDS)); + Assert.assertFalse(drainDone.await(200, TimeUnit.MILLISECONDS)); + continueUpdate.countDown(); + update.get(5, TimeUnit.SECONDS); + drain.get(5, TimeUnit.SECONDS); + Assert.assertEquals(Metapb.StoreState.Exiting, fixture.stores.get(4L).getState()); + } finally { + continueUpdate.countDown(); + workers.shutdownNow(); + } + } + + private static Metapb.ShardGroup group(int id, long... storeIds) { + Metapb.ShardGroup.Builder group = Metapb.ShardGroup.newBuilder().setId(id); + for (int i = 0; i < storeIds.length; i++) { + group.addShards(Metapb.Shard.newBuilder().setStoreId(storeIds[i]) + .setRole(i == 0 ? Metapb.ShardRole.Leader : + Metapb.ShardRole.Follower)); + } + return group.build(); + } + + private static void setField(StoreNodeService service, String name, Object value) + throws Exception { + Field field = StoreNodeService.class.getDeclaredField(name); + field.setAccessible(true); + field.set(service, value); + } + + private static class Fixture { + private final StoreNodeService service = + Mockito.mock(StoreNodeService.class, Mockito.CALLS_REAL_METHODS); + private final StoreInfoMeta meta = Mockito.mock(StoreInfoMeta.class); + private final PartitionService partitions = Mockito.mock(PartitionService.class); + private final Map stores = new HashMap<>(); + private final Map groups = new HashMap<>(); + private final PDConfig config = new PDConfig(); + + private Fixture() throws Exception { + config.setMinStoreCount(3); + PDConfig.Partition partition = config.new Partition(); + partition.setShardCount(3); + config.setPartition(partition); + ConfigService configService = Mockito.mock(ConfigService.class); + Mockito.when(configService.getPartitionCount()).thenReturn(1); + config.setConfigService(configService); + setField(service, "pdConfig", config); + setField(service, "storeInfoMeta", meta); + setField(service, "partitionService", partitions); + setField(service, "shardGroupStatusListeners", List.of()); + setField(service, "statusListeners", List.of()); + setField(service, "clusterStats", Metapb.ClusterStats.getDefaultInstance()); + Mockito.when(meta.getActiveStores()).thenReturn(List.of()); + Mockito.when(meta.getStore(Mockito.anyLong())).thenAnswer(call -> + stores.get((Long) call.getArgument(0))); + Mockito.when(meta.storeExists(Mockito.anyLong())).thenAnswer(call -> + stores.containsKey((Long) call.getArgument(0))); + Mockito.when(meta.getShardGroup(Mockito.anyInt())).thenAnswer(call -> + groups.get((Integer) call.getArgument(0))); + Mockito.when(meta.getShardGroups()).thenAnswer(call -> new ArrayList<>(groups.values())); + Mockito.doAnswer(call -> { + Metapb.Store store = call.getArgument(0); + stores.put(store.getId(), store); + return null; + }).when(meta).updateStore(Mockito.any()); + Mockito.doAnswer(call -> { + Metapb.ShardGroup group = call.getArgument(0); + groups.put(group.getId(), group); + return null; + }).when(meta).updateShardGroup(Mockito.any()); + } + + private Metapb.Store addStore(long id, Metapb.StoreState state) { + Metapb.Store store = Metapb.Store.newBuilder().setId(id).setState(state).build(); + stores.put(id, store); + return store; + } + + private void heartbeat(long id, int groups) throws Exception { + service.heartBeat(Metapb.StoreStats.newBuilder().setStoreId(id) + .setPartitionCount(groups).build()); + } + } +} diff --git a/hugegraph-pd/hg-pd-test/src/main/java/org/apache/hugegraph/pd/core/StorePlacementSubmissionTest.java b/hugegraph-pd/hg-pd-test/src/main/java/org/apache/hugegraph/pd/core/StorePlacementSubmissionTest.java new file mode 100644 index 0000000000..66b08fca79 --- /dev/null +++ b/hugegraph-pd/hg-pd-test/src/main/java/org/apache/hugegraph/pd/core/StorePlacementSubmissionTest.java @@ -0,0 +1,209 @@ +/* + * 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.hugegraph.pd.core; + +import java.lang.reflect.Field; +import java.util.ArrayList; +import java.util.Collections; +import java.util.HashMap; +import java.util.List; +import java.util.Map; +import java.util.concurrent.CountDownLatch; +import java.util.concurrent.ExecutorService; +import java.util.concurrent.Executors; +import java.util.concurrent.Future; +import java.util.concurrent.TimeUnit; + +import org.apache.hugegraph.pd.PartitionService; +import org.apache.hugegraph.pd.StoreNodeService; +import org.apache.hugegraph.pd.StoreStatusListener; +import org.apache.hugegraph.pd.TaskScheduleService; +import org.apache.hugegraph.pd.common.KVPair; +import org.apache.hugegraph.pd.common.PDException; +import org.apache.hugegraph.pd.config.PDConfig; +import org.apache.hugegraph.pd.grpc.Metapb; +import org.apache.hugegraph.pd.grpc.Pdpb; +import org.apache.hugegraph.pd.meta.PartitionMeta; +import org.apache.hugegraph.pd.meta.StoreInfoMeta; +import org.apache.hugegraph.pd.raft.RaftEngine; +import org.junit.Assert; +import org.junit.Test; +import org.mockito.Mockito; + +import com.alipay.sofa.jraft.Node; + +public class StorePlacementSubmissionTest { + + @Test + public void testTombstoneNotificationDoesNotInvertPlacementLocks() throws Exception { + StoreNodeService stores = Mockito.mock(StoreNodeService.class, + Mockito.CALLS_REAL_METHODS); + StoreInfoMeta meta = Mockito.mock(StoreInfoMeta.class); + PartitionService partitions = Mockito.mock(PartitionService.class, + Mockito.CALLS_REAL_METHODS); + PartitionMeta partitionMeta = Mockito.mock(PartitionMeta.class); + setField(PartitionService.class, partitions, "partitionMeta", partitionMeta); + setField(PartitionService.class, partitions, "storeService", stores); + Metapb.Store source = store(1L, Metapb.StoreState.Up); + Metapb.Store destination = store(2L, Metapb.StoreState.Up); + Metapb.Store retiring = store(3L, Metapb.StoreState.Exiting); + Map current = new HashMap<>(); + current.put(source.getId(), source); + current.put(destination.getId(), destination); + current.put(retiring.getId(), retiring); + Mockito.when(meta.getActiveStores()).thenReturn(List.of(source, destination)); + Mockito.when(meta.getStores(null)).thenReturn(List.of(source, destination, retiring)); + Mockito.when(meta.getShardGroups()).thenReturn(List.of()); + Mockito.when(meta.getStore(Mockito.anyLong())) + .thenAnswer(call -> current.get((Long) call.getArgument(0))); + Mockito.doAnswer(call -> { + Metapb.Store store = call.getArgument(0); + current.put(store.getId(), store); + return null; + }).when(meta).updateStore(Mockito.any()); + setField(StoreNodeService.class, stores, "storeInfoMeta", meta); + setField(StoreNodeService.class, stores, "pdConfig", new PDConfig()); + + CountDownLatch partitionHeld = new CountDownLatch(1); + CountDownLatch callbackEntered = new CountDownLatch(1); + List listeners = Collections.synchronizedList(new ArrayList<>()); + setField(StoreNodeService.class, stores, "statusListeners", listeners); + partitions.init(); + Metapb.Partition partition = Metapb.Partition.newBuilder() + .setId(0).setGraphName("graph").build(); + Metapb.Graph graph = Metapb.Graph.newBuilder().setGraphName("graph").build(); + Mockito.when(partitionMeta.getGraph("graph")).thenReturn(graph); + Mockito.doNothing().when(partitions).reallocPartitionShards(graph, partition); + // Keep the installed production Tombstone listener, storeOffline() and the + // synchronized shardOffline() real. Control only its metadata-query boundary. + Mockito.doAnswer(call -> { + callbackEntered.countDown(); + // Fail before blocking if the old lock order reappears, so the test always + // releases both threads instead of leaving deadlocked monitor owners behind. + Assert.assertFalse(Thread.holdsLock(stores)); + Assert.assertFalse(Thread.holdsLock(listeners)); + return List.of(partition); + }).when(partitions).getPartitionByStore(Mockito.any()); + + TaskScheduleService tasks = Mockito.mock(TaskScheduleService.class, + Mockito.CALLS_REAL_METHODS); + setField(TaskScheduleService.class, tasks, "storeService", stores); + setField(TaskScheduleService.class, tasks, "partitionService", partitions); + RaftEngine raft = RaftEngine.getInstance(); + Field nodeField = RaftEngine.class.getDeclaredField("raftNode"); + nodeField.setAccessible(true); + Object previous = nodeField.get(raft); + Node node = Mockito.mock(Node.class); + Mockito.when(node.isLeader(true)).thenReturn(true); + nodeField.set(raft, node); + ExecutorService workers = Executors.newFixedThreadPool(2); + try { + Future placement = workers.submit(() -> { + synchronized (partitions) { + partitionHeld.countDown(); + Assert.assertTrue(callbackEntered.await(5, TimeUnit.SECONDS)); + // The old retirement callback is waiting for this PartitionService + // monitor. A concurrent management update must not revive its ID. + PDException error = Assert.assertThrows(PDException.class, () -> + stores.updateStore(retiring.toBuilder() + .setState(Metapb.StoreState.Up).build())); + Assert.assertEquals(Pdpb.ErrorType.STORE_HAS_BEEN_REMOVED_VALUE, + error.getErrorCode()); + Map> plan = new HashMap<>(); + plan.put(0, new KVPair<>(source.getId(), destination.getId())); + Assert.assertEquals(1, tasks.movePartitions(plan).size()); + } + return null; + }); + Assert.assertTrue(partitionHeld.await(5, TimeUnit.SECONDS)); + Future retirement = workers.submit(() -> stores.updateStore( + retiring.toBuilder().setState(Metapb.StoreState.Tombstone).build())); + retirement.get(5, TimeUnit.SECONDS); + placement.get(5, TimeUnit.SECONDS); + Mockito.verify(partitions).shardOffline(partition, retiring.getId()); + Mockito.verify(partitions).movePartitionsShard(0, source.getId(), destination.getId()); + Assert.assertEquals(Metapb.StoreState.Tombstone, current.get(retiring.getId()).getState()); + } finally { + callbackEntered.countDown(); + workers.shutdownNow(); + try { + Assert.assertTrue(workers.awaitTermination(5, TimeUnit.SECONDS)); + } finally { + nodeField.set(raft, previous); + } + } + } + + @Test + public void testMovePlanRevalidatesDestinationBeforeSubmission() throws Exception { + StoreNodeService stores = Mockito.mock(StoreNodeService.class, + Mockito.CALLS_REAL_METHODS); + StoreInfoMeta meta = Mockito.mock(StoreInfoMeta.class); + PartitionService partitions = Mockito.mock(PartitionService.class); + Map current = new HashMap<>(); + Metapb.Store source = store(1L, Metapb.StoreState.Up); + Metapb.Store destinationLease = store(2L, Metapb.StoreState.Up); + current.put(source.getId(), source); + current.put(destinationLease.getId(), + destinationLease.toBuilder().setState(Metapb.StoreState.Exiting).build()); + Mockito.when(meta.getActiveStores()).thenReturn(List.of(source, destinationLease)); + Mockito.when(meta.getStore(Mockito.anyLong())).thenAnswer(call -> + current.get((Long) call.getArgument(0))); + setField(StoreNodeService.class, stores, "storeInfoMeta", meta); + + TaskScheduleService tasks = Mockito.mock(TaskScheduleService.class, + Mockito.CALLS_REAL_METHODS); + setField(TaskScheduleService.class, tasks, "storeService", stores); + setField(TaskScheduleService.class, tasks, "partitionService", partitions); + + RaftEngine raft = RaftEngine.getInstance(); + Field nodeField = RaftEngine.class.getDeclaredField("raftNode"); + nodeField.setAccessible(true); + Object previous = nodeField.get(raft); + Node node = Mockito.mock(Node.class); + Mockito.when(node.isLeader(true)).thenReturn(true); + nodeField.set(raft, node); + try { + Map> stalePlan = new HashMap<>(); + stalePlan.put(0, new KVPair<>(source.getId(), destinationLease.getId())); + Assert.assertTrue(tasks.movePartitions(stalePlan).isEmpty()); + Mockito.verify(partitions, Mockito.never()) + .movePartitionsShard(Mockito.anyInt(), Mockito.anyLong(), Mockito.anyLong()); + + current.put(destinationLease.getId(), destinationLease); + Map> validPlan = new HashMap<>(); + validPlan.put(0, new KVPair<>(source.getId(), destinationLease.getId())); + Assert.assertEquals(1, tasks.movePartitions(validPlan).size()); + Mockito.verify(partitions).movePartitionsShard(0, source.getId(), + destinationLease.getId()); + } finally { + nodeField.set(raft, previous); + } + } + + private static Metapb.Store store(long id, Metapb.StoreState state) { + return Metapb.Store.newBuilder().setId(id).setState(state).build(); + } + + private static void setField(Class owner, Object target, String name, Object value) + throws Exception { + Field field = owner.getDeclaredField(name); + field.setAccessible(true); + field.set(target, value); + } +}