fix(pd): avoid duplicate drain migration submission

hugegraph-pd/hg-pd-core/src/main/java/org/apache/hugegraph/pd/StoreNodeService.java

@@ -25,6 +25,7 @@
 import java.util.Map;
 import java.util.Objects;
 import java.util.Random;
+import java.util.Set;
 import java.util.concurrent.atomic.AtomicBoolean;
 
 import org.apache.commons.lang3.StringUtils;
@@ -273,6 +274,11 @@ public Metapb.Store getStore(long id) throws PDException {
      * Update the store information, detect the change of store status, and notify Hugestore
      */
     public Metapb.Store updateStore(Metapb.Store store) throws PDException {
+        return updateStore(store, Collections.emptySet());
+    }
+
+    public Metapb.Store updateStore(Metapb.Store store,
+                                    Set<Integer> submittedGroups) throws PDException {
         Metapb.Store lastStore;
         boolean stateChanged;
         synchronized (this) {
@@ -322,7 +328,7 @@ public Metapb.Store updateStore(Metapb.Store store) throws PDException {
                     log.info("updateStore removeActiveStores store {}", store.getId());
                     storeInfoMeta.removeActiveStore(store);
                 }
-                storeTurnoff(store);
+                storeTurnoff(store, submittedGroups);
             } else if (stateChanged) {
                 if (store.getState() == Metapb.StoreState.Offline) {
                     // Monitor that the store has gone offline and is removed from the active
@@ -355,8 +361,22 @@ public Metapb.Store updateStore(Metapb.Store store) throws PDException {
      * @throws PDException
      */
     public synchronized void storeTurnoff(Metapb.Store store) throws PDException {
+        storeTurnoff(store, Collections.emptySet());
+    }
+
+    public synchronized void storeTurnoff(Metapb.Store store,
+                                          Set<Integer> submittedGroups)
+            throws PDException {
         // Traverse ShardGroup,redistribution
         for (Metapb.ShardGroup group : getShardGroupsByStore(store.getId())) {
+            // PDService can submit a capacity-aware migration before persisting
+            // Exiting. Keep that accepted destination authoritative instead of
+            // immediately issuing a second, independently planned change.
+            if (submittedGroups.contains(group.getId())) {
+                log.info("Skip duplicate drain submission for shard group {}",
+                         group.getId());
+                continue;
+            }
             Metapb.ShardGroup.Builder builder = Metapb.ShardGroup.newBuilder(group);
             builder.clearShards();
             group.getShardsList().forEach(shard -> {

hugegraph-pd/hg-pd-dist/pom.xml

@@ -42,6 +42,7 @@
             <groupId>org.apache.hugegraph</groupId>
             <artifactId>hg-pd-service</artifactId>
             <version>${revision}</version>
+            <classifier>exec</classifier>
         </dependency>
     </dependencies>
     <build>

hugegraph-pd/hg-pd-dist/src/assembly/descriptor/server-assembly.xml

@@ -66,8 +66,9 @@
             <unpack>false</unpack>
             <scope>runtime</scope>
             <useProjectArtifact>false</useProjectArtifact>
+            <outputFileNameMapping>${artifact.artifactId}-${artifact.version}.${artifact.extension}</outputFileNameMapping>
             <includes>
-                <include>org.apache.hugegraph:${executable.jar.name}:jar:*</include>
+                <include>org.apache.hugegraph:${executable.jar.name}:jar:exec</include>
             </includes>
         </dependencySet>
     </dependencySets>

hugegraph-pd/hg-pd-service/pom.xml

@@ -191,6 +191,7 @@
                 <executions>
                     <execution>
                         <configuration>
+                            <classifier>exec</classifier>
                             <mainClass>
                                 org.apache.hugegraph.pd.boot.HugePDServer
                             </mainClass>

hugegraph-pd/hg-pd-service/src/main/java/org/apache/hugegraph/pd/service/PDService.java

@@ -511,6 +511,7 @@ public void setStore(Pdpb.SetStoreRequest request,
                 }
             }
             Metapb.Store store = request.getStore();
+            Set<Integer> submittedGroups = new HashSet<>();
             // Before going offline, check whether the number of active machines is greater than
             // the minimum threshold
             if (state.equals(Metapb.StoreState.Tombstone)) {
@@ -539,9 +540,16 @@ public void setStore(Pdpb.SetStoreRequest request,
                         // to be migrated
                         store = Metapb.Store.newBuilder(lastStore)
                                             .setState(Metapb.StoreState.Exiting).build();
-                        // Perform partition migration operations
-                        taskService.movePartitions((Map<Integer, KVPair<Long, Long>>) resultMap.get(
-                                "movedPartitions"));
+                        // Perform partition migration operations. Preserve the groups whose
+                        // final placement was accepted so updateStore does not immediately
+                        // submit a second independently planned membership change.
+                        Map<Integer, KVPair<Long, Long>> submitted =
+                                taskService.movePartitions(
+                                        (Map<Integer, KVPair<Long, Long>>) resultMap.get(
+                                                "movedPartitions"));
+                        if (submitted != null) {
+                            submittedGroups.addAll(submitted.keySet());
+                        }
                     } else {
                         // If the store is offline, the replica is not migrated
                         // Change the status to Tombstone
@@ -554,7 +562,7 @@ public void setStore(Pdpb.SetStoreRequest request,
                 }
             }
             // Replace license using grpc
-            store = storeNodeService.updateStore(store);
+            store = storeNodeService.updateStore(store, submittedGroups);
             response =
                     Pdpb.SetStoreResponse.newBuilder().setHeader(okHeader).setStore(store).build();
         } catch (PDException e) {

hugegraph-pd/hg-pd-test/src/main/java/org/apache/hugegraph/pd/core/PDCoreSuiteTest.java

@@ -33,6 +33,7 @@
 import org.apache.hugegraph.pd.raft.RaftReflectionUtilTest;
 import org.apache.hugegraph.pd.raft.RaftStateMachineSnapshotTest;
 import org.apache.hugegraph.pd.raft.ZipUtilsTest;
+import org.apache.hugegraph.pd.service.PDServiceDrainSubmissionTest;
 import org.junit.runner.RunWith;
 import org.junit.runners.Suite;
 
@@ -58,6 +59,7 @@
         StoreServiceTest.class,
         StoreDrainTest.class,
         StorePlacementSubmissionTest.class,
+        PDServiceDrainSubmissionTest.class,
         TaskScheduleServiceTest.class,
         IpAuthHandlerTest.class,
         RaftEngineCliProcessorsTest.class,

hugegraph-pd/hg-pd-test/src/main/java/org/apache/hugegraph/pd/core/StoreDrainTest.java

@@ -22,6 +22,7 @@
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
+import java.util.Set;
 import java.util.concurrent.CountDownLatch;
 import java.util.concurrent.ExecutorService;
 import java.util.concurrent.Executors;
@@ -133,6 +134,36 @@ public void testDrainCannotPlaceReplicaBackOnSourceAndNotifiesSharedGroupOnce()
         Assert.assertEquals(ConfChangeType.CONF_CHANGE_TYPE_ADJUST, changes.get(0).getArgument(2));
     }
 
+    @Test
+    public void testPreSubmittedDrainGroupIsNotDispatchedAgain() throws Exception {
+        Fixture fixture = new Fixture();
+        for (long id = 1; id <= 4; id++) {
+            fixture.addStore(id, Metapb.StoreState.Up);
+        }
+        Metapb.Store source = fixture.stores.get(4L);
+        Mockito.when(fixture.meta.getActiveStores())
+               .thenAnswer(call -> new ArrayList<>(fixture.stores.values()));
+        Metapb.ShardGroup submitted = group(0, 1, 2, 4);
+        fixture.groups.put(0, submitted);
+        fixture.groups.put(1, group(1, 1, 2, 4));
+        Metapb.Partition fallback = Metapb.Partition.newBuilder()
+                                                   .setId(1)
+                                                   .setGraphName("graph").build();
+        Mockito.when(fixture.partitions.getPartitionById(1))
+               .thenReturn(List.of(fallback));
+        Metapb.Store drain = source.toBuilder()
+                                    .setState(Metapb.StoreState.Exiting).build();
+
+        fixture.service.updateStore(drain, Set.of(0));
+
+        Assert.assertEquals(submitted, fixture.groups.get(0));
+        Assert.assertTrue(fixture.groups.get(1).getShardsList().stream()
+                                 .noneMatch(shard -> shard.getStoreId() == 4));
+        List<Invocation> changes = changeShardInvocations(fixture.partitions);
+        Assert.assertEquals(1, changes.size());
+        Assert.assertEquals(fallback, changes.get(0).getArgument(0));
+    }
+
     @Test
     public void testInitialAllocationExcludesExitingStoreWithStaleUpLease() throws Exception {
         Fixture fixture = new Fixture();

hugegraph-pd/hg-pd-test/src/main/java/org/apache/hugegraph/pd/service/PDServiceDrainSubmissionTest.java

@@ -0,0 +1,114 @@
+/*
+ * 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.service;
+
+import java.lang.reflect.Field;
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+import java.util.Set;
+
+import org.apache.hugegraph.pd.LogService;
+import org.apache.hugegraph.pd.StoreNodeService;
+import org.apache.hugegraph.pd.TaskScheduleService;
+import org.apache.hugegraph.pd.common.KVPair;
+import org.apache.hugegraph.pd.config.PDConfig;
+import org.apache.hugegraph.pd.grpc.Metapb;
+import org.apache.hugegraph.pd.grpc.Pdpb;
+import org.junit.Assert;
+import org.junit.Test;
+import org.mockito.ArgumentCaptor;
+import org.mockito.Mockito;
+
+import io.grpc.stub.StreamObserver;
+
+public class PDServiceDrainSubmissionTest {
+
+    @SuppressWarnings({"unchecked", "rawtypes"})
+    @Test
+    public void testSetStorePassesAcceptedDrainGroupsToStateTransition() throws Exception {
+        PDService service = Mockito.spy(new PDService());
+        Mockito.doReturn(true).when(service).isLeader();
+
+        StoreNodeService stores = Mockito.mock(StoreNodeService.class);
+        TaskScheduleService tasks = Mockito.mock(TaskScheduleService.class);
+        LogService logs = Mockito.mock(LogService.class);
+        PDConfig config = new PDConfig();
+        config.setMinStoreCount(1);
+        setField(service, "storeNodeService", stores);
+        setField(service, "taskService", tasks);
+        setField(service, "logService", logs);
+        setField(service, "pdConfig", config);
+
+        long sourceId = 4L;
+        Metapb.Store source = store(sourceId);
+        Mockito.when(stores.getStore(sourceId)).thenReturn(source);
+        Mockito.when(stores.getActiveStores())
+               .thenReturn(List.of(store(1L), store(2L), store(3L), source));
+        Mockito.when(stores.checkStoreCanOffline(Mockito.any())).thenReturn(true);
+
+        Map<Integer, KVPair<Long, Long>> plan = new HashMap<>();
+        plan.put(7, new KVPair<>(sourceId, 1L));
+        Map<String, Object> result = new HashMap<>();
+        result.put("flag", true);
+        result.put("current_store_is_online", true);
+        result.put("movedPartitions", plan);
+        Mockito.when(tasks.canAllPartitionsMovedOut(source)).thenReturn(result);
+        Mockito.when(tasks.movePartitions(plan)).thenReturn(plan);
+
+        Metapb.Store exiting = source.toBuilder()
+                                      .setState(Metapb.StoreState.Exiting).build();
+        Mockito.when(stores.updateStore(Mockito.any())).thenReturn(exiting);
+        Mockito.when(stores.updateStore(Mockito.any(), Mockito.anySet()))
+               .thenReturn(exiting);
+        StreamObserver<Pdpb.SetStoreResponse> observer = Mockito.mock(StreamObserver.class);
+        Pdpb.SetStoreRequest request = Pdpb.SetStoreRequest.newBuilder()
+                                                          .setStore(source.toBuilder()
+                                                                          .setState(
+                                                                                  Metapb.StoreState.Tombstone))
+                                                          .build();
+
+        service.setStore(request, observer);
+
+        ArgumentCaptor<Metapb.Store> storeCaptor =
+                ArgumentCaptor.forClass(Metapb.Store.class);
+        ArgumentCaptor<Set<Integer>> groupsCaptor =
+                ArgumentCaptor.forClass((Class) Set.class);
+        Mockito.verify(stores).updateStore(storeCaptor.capture(), groupsCaptor.capture());
+        Assert.assertEquals(Metapb.StoreState.Exiting, storeCaptor.getValue().getState());
+        Assert.assertEquals(Set.of(7), groupsCaptor.getValue());
+        Mockito.verify(tasks).movePartitions(plan);
+        Mockito.verify(observer).onNext(Mockito.argThat(response ->
+                response.getHeader().getError().getType() == Pdpb.ErrorType.OK));
+        Mockito.verify(observer).onCompleted();
+    }
+
+    private static Metapb.Store store(long id) {
+        return Metapb.Store.newBuilder()
+                           .setId(id)
+                           .setState(Metapb.StoreState.Up)
+                           .setStats(Metapb.StoreStats.newBuilder().setStoreId(id))
+                           .build();
+    }
+
+    private static void setField(Object target, String name, Object value) throws Exception {
+        Field field = PDService.class.getDeclaredField(name);
+        field.setAccessible(true);
+        field.set(target, value);
+    }
+}