fix(store): refresh recovery shard metadata

hugegraph-pd/hg-pd-client/src/main/java/org/apache/hugegraph/pd/client/PDClient.java

@@ -75,7 +75,7 @@
 @Slf4j
 public class PDClient {
 
     private static Map<String, ManagedChannel> channels = new ConcurrentHashMap();
     private static ManagedChannel channel = null;
     private final PDConfig config;
     private final Pdpb.RequestHeader header;
@@ -621,17 +621,7 @@
         KVPair<Metapb.Partition, Metapb.Shard> partShard =
                 cache.getPartitionById(graphName, partId);
         if (partShard == null) {
-            Pdpb.GetPartitionByIDRequest request = Pdpb.GetPartitionByIDRequest.newBuilder()
-                                                                               .setHeader(header)
-                                                                               .setGraphName(
-                                                                                       graphName)
-                                                                               .setPartitionId(
-                                                                                       partId)
-                                                                               .build();
-            GetPartitionResponse response =
-                    blockingUnaryCall(PDGrpc.getGetPartitionByIDMethod(), request);
-            handleResponseError(response.getHeader());
-            partShard = new KVPair<>(response.getPartition(), response.getLeader());
+            partShard = getPartitionByIdDirect(graphName, partId);
             if (config.isEnableCache()) {
                 cache.update(graphName, partShard.getKey().getId(), partShard.getKey());
                 cache.updateShardGroup(getShardGroup(partShard.getKey().getId()));
@@ -653,6 +643,19 @@
         return partShard;
     }
 
+    /**
+     * Query this graph partition without the cache or any additional shard-group lookup.
+     * In particular, a NOT_FOUND error belongs to this partition RPC rather than cache filling.
+     */
+    public KVPair<Metapb.Partition, Metapb.Shard> getPartitionByIdDirect(String graphName,
+                                                                      int partId) throws PDException {
+        Pdpb.GetPartitionByIDRequest request = Pdpb.GetPartitionByIDRequest.newBuilder()
+                .setHeader(header).setGraphName(graphName).setPartitionId(partId).build();
+        GetPartitionResponse response = blockingUnaryCall(PDGrpc.getGetPartitionByIDMethod(), request);
+        handleResponseError(response.getHeader());
+        return new KVPair<>(response.getPartition(), response.getLeader());
+    }
+
     public ShardGroup getShardGroup(int partId) throws PDException {
         ShardGroup group = cache.getShardGroup(partId);
         if (group == null) {
@@ -974,7 +977,7 @@
     private <ReqT, RespT, StubT extends AbstractBlockingStub<StubT>> RespT
     blockingUnaryCall(MethodDescriptor<ReqT, RespT> method, ReqT req, int retry) throws
                                                                                  PDException {
         io.grpc.stub.AbstractBlockingStub<StubT> stub = (AbstractBlockingStub<StubT>) getStub();
         try {
             RespT resp = io.grpc.stub.ClientCalls.blockingUnaryCall(stub.getChannel(), method,
                                                                     stub.getCallOptions(), req);

hugegraph-store/docs/operations-guide.md

@@ -665,6 +665,27 @@ or recovery of a migration task across a Leader change.
 Verify the actual voter and learner lists before treating migration as complete: a target count
 that includes learners is not proof of voting membership.
 
+### Restoring Store Membership after a Restart
+
+For a persisted partition, startup verifies the shard group directly with PD and refreshes
+the derived local `SHARDGROUP` once per group when this Store is still a member and no
+engine is running for the group. This lets JRaft restore its durable configuration even
+when the business database contains older membership metadata. Running engines and
+durable Raft configuration are not replaced by the PD bootstrap metadata.
+
+Only PD's explicit `NOT_FOUND` response for a persisted graph partition skips that graph's
+recovery, for example when it was dropped while the Store was offline. Other graphs can
+still recover. The skipped graph's local data, partition-store location and Raft files
+remain intact; a group with no remaining graph partitions does not start an engine.
+A null/ambiguous result, invalid metadata or a failed PD lookup blocks recovery of the
+whole Store process while retaining local data. This includes members with an unspecified
+`None` role until PD reports usable member roles. Restore PD access or correct the reported
+metadata before retrying startup; do not delete data or reset Raft peers to bypass a failure.
+PD exclusion is not proof that a Store has been safely retired. Check every assigned
+group after a restart before continuing a rolling operation. This recovery correction
+does not establish why older derived metadata was persisted, or provide a complete
+desired/observed membership protocol or protection against partial PD partition queries.
+
 ### Removing Store Nodes
 
 **When to Remove**:

hugegraph-store/hg-store-core/src/main/java/org/apache/hugegraph/store/meta/PartitionManager.java

@@ -36,6 +36,7 @@
 import org.apache.hugegraph.pd.common.PDException;
 import org.apache.hugegraph.pd.grpc.MetaTask;
 import org.apache.hugegraph.pd.grpc.Metapb;
+import org.apache.hugegraph.pd.grpc.Pdpb;
 import org.apache.hugegraph.store.HgStoreEngine;
 import org.apache.hugegraph.store.business.BusinessHandlerImpl;
 import org.apache.hugegraph.store.cmd.HgCmdClient;
@@ -107,6 +108,10 @@ public static final class InvalidShardException extends IllegalStateException {
         private InvalidShardException() {
             super("Local partition does not belong to this store");
         }
+
+        private InvalidShardException(String message, Throwable cause) {
+            super(message, cause);
+        }
     }
 
     public void loadPartition() {
@@ -259,6 +264,7 @@ private void loadPartitions() {
         }
 
         Set<Integer> normalPartitions = new HashSet<>();
+        Set<Integer> retainedPartitions = new HashSet<>();
 
         // Once according to the partition read
         for (int partId : partIds) {
@@ -269,64 +275,108 @@ private void loadPartitions() {
             }
 
             var metaParts = wrapper.scan(partId, Metapb.Partition.parser(), key);
-            int countOfPartition = 0;
-
-            var shards = pdProvider.getShardGroup(partId).getShardsList();
+            if (metaParts.isEmpty()) {
+                wrapper.close(partId);
+                continue;
+            }
 
+            // A migrated replica can have an older derived SHARDGROUP in its data DB than
+            // its durable Raft configuration. Validate recovery against PD before Raft replay.
+            var pdGroup = recoveryShardGroup(partId);
+            var shards = pdGroup.getShardsList();
+            if (shards.stream().noneMatch(s -> s.getStoreId() == storeId)) {
+                log.error("partition {} is illegal. store id {} not in valid shard group:{}",
+                          partId, storeId, shards2Peers(shards));
+                throw new InvalidShardException();
+            }
+            List<Metapb.Partition> recoverableParts = new ArrayList<>();
             for (var metaPart : metaParts) {
                 var graph = metaPart.getGraphName();
-                var pdPartition = pdProvider.getPartitionByID(graph, metaPart.getId());
-                boolean isLegeal = false;
-
-                if (pdPartition != null) {
-                    // Check if it contains this store id
-                    if (shards.stream().anyMatch(s -> s.getStoreId() == storeId)) {
-                        isLegeal = true;
-                    }
-                } else {
-                    continue;
+                if (metaPart.getId() != partId || graph.isEmpty()) {
+                    throw new InvalidShardException("Invalid persisted partition " + graph +
+                                                    "-" + partId, null);
                 }
-
-                if (isLegeal) {
-                    if (!partitions.containsKey(graph)) {
-                        partitions.put(graph, new ConcurrentHashMap<>());
+                Partition pdPartition;
+                try {
+                    pdPartition = pdProvider.getPartitionByIDForRecovery(graph, metaPart.getId());
+                } catch (PDException e) {
+                    if (e.getErrorCode() == Pdpb.ErrorType.NOT_FOUND_VALUE) {
+                        // A graph can be dropped while this Store is offline. Do not serve it or
+                        // infer permission to delete its data, Store path or Raft metadata.
+                        retainedPartitions.add(partId);
+                        log.warn("PD no longer has partition {}-{}; skipping recovery and " +
+                                 "retaining local data", graph, partId);
+                        continue;
                     }
-
-                    countOfPartition += 1;
-
-                    Partition partition = new Partition(metaPart);
-                    partition.setWorkState(
-                            Metapb.PartitionState.PState_Normal);     // Start recovery work state
-                    partitions.get(graph).put(partition.getId(), partition);
-                    log.info("load partition : {} -{}", partition.getGraphName(),
-                             partition.getId());
-                } else {
-                    // Invalid
-                    // removePartitionFromLocalDb(graph, partId);
-                    // var businessHandler = HgStoreEngine.getInstance().getBusinessHandler();
-                    // businessHandler.truncate(graph, partId);
-                    // businessHandler.dbCompaction(graph, partId);
-                    log.error("partition {}-{} is illegal. store id {} not in valid shard group:{}",
-                              graph, partId, getStore().getId(), shards2Peers(shards));
-                    throw new InvalidShardException();
+                    throw new InvalidShardException("Cannot read persisted partition " + graph +
+                                                    "-" + partId + " from PD", e);
+                } catch (RuntimeException e) {
+                    throw new InvalidShardException("Cannot read persisted partition " + graph +
+                                                    "-" + partId + " from PD", e);
                 }
+                if (pdPartition == null || pdPartition.getId() != partId ||
+                    !graph.equals(pdPartition.getGraphName())) {
+                    throw new InvalidShardException("Cannot verify persisted partition " + graph +
+                                                    "-" + partId + " with PD", null);
+                }
+                recoverableParts.add(metaPart);
+            }
+            if (recoverableParts.isEmpty()) {
+                wrapper.close(partId);
+                continue;
             }
 
-            if (countOfPartition > 0) {
-                // Partition data is normal
-                normalPartitions.add(partId);
+            // Only refresh startup metadata, once per group even when several graphs share it.
+            // The running engine and JRaft's durable configuration remain authoritative.
+            if (HgStoreEngine.getInstance().getPartitionEngine(partId) == null) {
+                updateShardGroup(ShardGroup.from(pdGroup));
             }
+            for (var metaPart : recoverableParts) {
+                var graph = metaPart.getGraphName();
+                partitions.computeIfAbsent(graph, g -> new ConcurrentHashMap<>());
+                Partition partition = new Partition(metaPart);
+                partition.setWorkState(Metapb.PartitionState.PState_Normal);
+                partitions.get(graph).put(partition.getId(), partition);
+                log.info("load partition : {} -{}", partition.getGraphName(), partition.getId());
+            }
+            normalPartitions.add(partId);
             wrapper.close(partId);
         }
 
         // Remove redundant partition storage paths, partitions that have been migrated away may migrate back
         for (var location : storeMetadata.getPartitionStores()) {
-            if (!normalPartitions.contains(location.getPartitionId())) {
+            if (!normalPartitions.contains(location.getPartitionId()) &&
+                !retainedPartitions.contains(location.getPartitionId())) {
                 storeMetadata.removePartitionStore(location.getPartitionId());
             }
         }
     }
 
+    private Metapb.ShardGroup recoveryShardGroup(int partId) {
+        Metapb.ShardGroup group;
+        try {
+            group = pdProvider.getShardGroupDirect(partId);
+        } catch (RuntimeException e) {
+            throw new InvalidShardException("Cannot read recovery shard group " + partId +
+                                            " from PD", e);
+        }
+        if (group == null || group.getId() != partId || group.getShardsCount() == 0) {
+            throw new InvalidShardException("Invalid recovery shard group " + partId +
+                                            " from PD", null);
+        }
+        Set<Long> stores = new HashSet<>();
+        for (var shard : group.getShardsList()) {
+            if (shard.getStoreId() == 0 || !stores.add(shard.getStoreId()) ||
+                (shard.getRole() != Metapb.ShardRole.Leader &&
+                 shard.getRole() != Metapb.ShardRole.Follower &&
+                 shard.getRole() != Metapb.ShardRole.Learner)) {
+                throw new InvalidShardException("Invalid recovery shard members for group " +
+                                                partId + " from PD", null);
+            }
+        }
+        return group;
+    }
+
     public List<Metapb.Partition> loadPartitionsFromDb(int partitionId) {
         byte[] key = MetadataKeyHelper.getPartitionPrefixKey();
         return wrapper.scan(partitionId, Metapb.Partition.parser(), key);
@@ -797,7 +847,11 @@ public List<String> getPartitionPeers(ShardGroup shardGroup) {
             for (Shard shard : shardGroup.getShards()) {
                 Store store = getStore(shard.getStoreId());
                 if (store != null && !store.getRaftAddress().isEmpty()) {
-                    peers.add(store.getRaftAddress() + "::" + priority);
+                    String peer = store.getRaftAddress() + "::" + priority;
+                    if (shard.getRole() == Metapb.ShardRole.Learner) {
+                        peer += "/learner";
+                    }
+                    peers.add(peer);
                     final int gap = Math.max(decayPriorityGap, (priority / 5));
                     priority = Math.max(ElectionPriority.MinValue, (priority - gap));
                 }

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

@@ -173,6 +173,12 @@ public Partition getPartitionByID(String graph, int partId) {
         return null;
     }
 
+    @Override
+    public Partition getPartitionByIDForRecovery(String graph, int partId) throws PDException {
+        KVPair<Metapb.Partition, Metapb.Shard> pair = pdClient.getPartitionByIdDirect(graph, partId);
+        return pair == null ? null : new Partition(pair.getKey());
+    }
+
     @Override
     public Metapb.Shard getPartitionLeader(String graph, int partId) {
         try {

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

@@ -43,6 +43,14 @@ public interface PdProvider {
 
     Partition getPartitionByID(String graph, int partId);
 
+    /**
+     * Preserve PD lookup failures for recovery. Only an explicit NOT_FOUND error proves absence;
+     * the legacy lookup's null result remains ambiguous and must not authorize cleanup.
+     */
+    default Partition getPartitionByIDForRecovery(String graph, int partId) throws PDException {
+        return getPartitionByID(graph, partId);
+    }
+
     Metapb.Shard getPartitionLeader(String graph, int partId);
 
     Metapb.Partition getPartitionByCode(String graph, int code);

hugegraph-store/hg-store-test/pom.xml

@@ -247,6 +247,8 @@
                                 <include>**/PartitionStateMachineTest.java</include>
                                 <include>**/PartitionEngineErrorTest.java</include>
                                 <include>**/PartitionEngineMembershipTest.java</include>
+                                <include>**/PartitionRecoveryMembershipTest.java</include>
+                                <include>**/RecoveryPdLookupTest.java</include>
                                 <include>**/BatchGraphIsolationTest.java</include>
                                 <include>**/GraphIDManagerTest.java</include>
                                 <include>**/SchemaDriverTest.java</include>

hugegraph-store/hg-store-test/src/main/java/org/apache/hugegraph/store/HeartbeatShutdownTest.java

@@ -192,10 +192,10 @@ private static void runInvalidShardRecovery(Path marker) throws Exception {
         options.setRaftPath(raft.toString());
         PdProvider provider = Mockito.mock(PdProvider.class);
         Metapb.Partition partition = Metapb.Partition.newBuilder().setId(1).setGraphName("graph").build();
-        Mockito.when(provider.getShardGroup(1)).thenReturn(Metapb.ShardGroup.newBuilder()
+        Mockito.when(provider.getShardGroupDirect(1)).thenReturn(Metapb.ShardGroup.newBuilder()
                 .setId(1).addShards(Metapb.Shard.newBuilder().setStoreId(2)
                                                        .setRole(Metapb.ShardRole.Leader)).build());
-        Mockito.when(provider.getPartitionByID("graph", 1)).thenReturn(new Partition(partition));
+        Mockito.when(provider.getPartitionByIDForRecovery("graph", 1)).thenReturn(new Partition(partition));
         Store remote = new Store();
         remote.setRaftAddress("127.0.0.1:8511");
         Mockito.when(provider.getStoreByID(2L)).thenReturn(remote);