From a63be1c7f7d0b6c4c94aa45d00f2263f4affb87e Mon Sep 17 00:00:00 2001 From: pallavithegod Date: Fri, 16 Oct 2026 16:25:52 +0530 Subject: [PATCH 1/4] fix(store): gate up registration on restore and report offline on shutdown --- docker/README.md | 19 +++ helm/hugegraph/docs/operations.md | 7 + hugegraph-pd/docs/api-reference.md | 16 +++ .../apache/hugegraph/pd/StoreNodeService.java | 22 +-- .../apache/hugegraph/pd/rest/StoreAPI.java | 3 + .../pd/rest/StoreHeartbeatAgeTest.java | 60 ++++++++ .../hugegraph/pd/core/PDCoreSuiteTest.java | 1 + .../pd/core/StoreRestoreRegistrationTest.java | 92 ++++++++++++ .../hugegraph/store/HeartbeatService.java | 50 +++++++ .../apache/hugegraph/store/HgStoreEngine.java | 15 +- .../hugegraph/store/pd/DefaultPdProvider.java | 5 +- hugegraph-store/hg-store-test/pom.xml | 1 + .../store/HeartbeatShutdownTest.java | 5 +- .../store/StoreRegistrationTest.java | 135 ++++++++++++++++++ 14 files changed, 409 insertions(+), 22 deletions(-) create mode 100644 hugegraph-pd/hg-pd-service/src/test/java/org/apache/hugegraph/pd/rest/StoreHeartbeatAgeTest.java create mode 100644 hugegraph-pd/hg-pd-test/src/main/java/org/apache/hugegraph/pd/core/StoreRestoreRegistrationTest.java create mode 100644 hugegraph-store/hg-store-test/src/main/java/org/apache/hugegraph/store/StoreRegistrationTest.java diff --git a/docker/README.md b/docker/README.md index 233cd4ef71..559ef8d649 100644 --- a/docker/README.md +++ b/docker/README.md @@ -233,6 +233,25 @@ Open `http://localhost:8088` and sign in as `admin` with the password from `.env +### Store registration during restart + +Stores first register `Offline` to obtain their persisted identity, restore local +partition engines, and then register `Up`. A restore failure keeps the Store +offline and the heartbeat loop retries it. Graceful shutdown stops heartbeat +producers before reporting `Offline` to PD; if PD is unavailable, normal heartbeat +expiry remains the fallback. A temporary stop does not request permanent retirement. + +Upgrade PD before Store to enable this registration handshake. Older Stores keep +their existing registration behavior; older PD versions ignore the restore gate. +`Up` confirms local engine restore, not Raft catch-up or full shard-group health. +PD's authenticated `/v1/stores` and `/v1/store/{id}` responses also expose +`heartbeatAgeMillis`, computed with PD's clock (`-1` when no heartbeat is known). +Check it against the Store heartbeat interval to distinguish a stale `Up` record +from a recent heartbeat without waiting for keep-alive expiry. A future timestamp +caused by a PD clock adjustment reports zero age. +Keep the cluster preset's manual rollout barrier and verify fresh heartbeats, +membership and leaders before replacing another Store. + ### Select image versions
diff --git a/helm/hugegraph/docs/operations.md b/helm/hugegraph/docs/operations.md index c227dd06c8..b74ed7c589 100644 --- a/helm/hugegraph/docs/operations.md +++ b/helm/hugegraph/docs/operations.md @@ -198,6 +198,13 @@ A single Store replacement can still cause request errors or timeouts while Raft leadership and client routes converge. Observe application errors as well as shard recovery before continuing; do not blindly retry writes. +With PD and Store images containing #3229's registration changes, `Up` follows +local engine restore and graceful stops report `Offline`. Use PD's +`heartbeatAgeMillis` to reject stale `Up` records after an abrupt stop. +These signals do not prove Raft catch-up. + + ### PD patrol and leader balance Partition patrol repairs shard membership; it does not balance shard leaders. diff --git a/hugegraph-pd/docs/api-reference.md b/hugegraph-pd/docs/api-reference.md index 79d30ce69e..5dffe754ab 100644 --- a/hugegraph-pd/docs/api-reference.md +++ b/hugegraph-pd/docs/api-reference.md @@ -200,6 +200,12 @@ channel.shutdown(); Register a new Store node with PD. +Updated Stores first register `Offline` to obtain their identity, restore local +partition engines, then register `Up`. Upgrade PD before Store: older PD versions +ignore this restore gate, and older Stores retain their previous registration +behavior. Pending Stores still require activation. `Up` does not guarantee Raft +catch-up; graceful shutdown reports `Offline` without requesting permanent retirement. + **Request**: ```protobuf message RegisterStoreRequest { @@ -914,6 +920,16 @@ curl http://localhost:8620/v1/partitions?graph_name=social_network ### Store API +#### Heartbeat freshness + +The authenticated `/v1/stores` and `/v1/store/{id}` responses expose +`heartbeatAgeMillis` alongside the existing `lastHeartBeat`. Age is computed with +PD's clock, is `-1` if no heartbeat is recorded, and is clamped to zero for future +timestamps after a clock adjustment. Compare it with the configured Store +heartbeat interval to detect stale `Up` records before keep-alive expiry. +Read this state from the PD leader and continue to verify shard-group membership +and leaders during rolling updates; fresh heartbeats do not prove Raft catch-up. + #### List Stores ```bash 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 f8d8b81ed6..b0cb7c6e83 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 @@ -133,6 +133,7 @@ public boolean isOK() { * @param store */ public Metapb.Store register(Metapb.Store store) throws PDException { + boolean restoring = store.getState() == Metapb.StoreState.Offline; if (store.getId() == 0) { // Initial registration, generate a new ID, and ensure that the ID is not duplicated. store = newStoreNode(store); @@ -158,15 +159,6 @@ public Metapb.Store register(Metapb.Store store) throws PDException { } // offline or up, or in the initial activation list, go live automatically - // TODO: do not mark a re-registering Store Up before it has restored its partition engines - // (HgStoreEngine.restoreLocalPartitionEngine); report a restoring state, or expose - // restore-complete per shard group, so a rolling restart can wait on it. The Helm chart - // (helm/hugegraph) production preset (values-cluster.yaml) keeps Store rollouts on - // OnDelete with a manual /v1/shardGroups barrier between Pod deletions; the default - // values use RollingUpdate. Restore-complete alone does not retire the barrier: a stopped - // Store stays Up until its keep-alive entry expires, so keep the rollout checks until PD - // 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)) { @@ -206,6 +198,13 @@ public Metapb.Store register(Metapb.Store store) throws PDException { } } + // Updated Stores register Offline to obtain their identity before restoring partitions. + // Older Stores still request Up, and Pending activation and duplicate checks stay intact. + if (restoring && storeState == Metapb.StoreState.Up) { + storeState = Metapb.StoreState.Offline; + store = store.toBuilder().setState(storeState).build(); + } + // Store information storeInfoMeta.updateStore(store); if (storeState == Metapb.StoreState.Up) { @@ -213,6 +212,11 @@ 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.Offline && + lastStore.getState() == Metapb.StoreState.Up) { + storeInfoMeta.removeActiveStore(store); + onStoreStatusChanged(store, Metapb.StoreState.Up, Metapb.StoreState.Offline); + checkStoreStatus(); } // Wait for the store information to be saved before sending the changes diff --git a/hugegraph-pd/hg-pd-service/src/main/java/org/apache/hugegraph/pd/rest/StoreAPI.java b/hugegraph-pd/hg-pd-service/src/main/java/org/apache/hugegraph/pd/rest/StoreAPI.java index d9104544bb..eb84c4ca17 100644 --- a/hugegraph-pd/hg-pd-service/src/main/java/org/apache/hugegraph/pd/rest/StoreAPI.java +++ b/hugegraph-pd/hg-pd-service/src/main/java/org/apache/hugegraph/pd/rest/StoreAPI.java @@ -298,6 +298,7 @@ class StoreStatistics { // For the time being, the time of the first heartbeat is taken as the registration time long registedTimeStamp; long lastHeartBeat; // Last heartbeat time + long heartbeatAgeMillis; // PD's elapsed time since the last heartbeat; -1 if unknown long capacity; long available; int partitionCount; @@ -341,6 +342,8 @@ class StoreStatistics { } registedTimeStamp = store.getStartTimestamp(); // Time of registration lastHeartBeat = store.getLastHeartbeat(); + heartbeatAgeMillis = lastHeartBeat > 0 ? + Math.max(0L, System.currentTimeMillis() - lastHeartBeat) : -1L; capacity = store.getStats().getCapacity(); available = store.getStats().getAvailable(); partitionCount = store.getStats().getPartitionCount(); diff --git a/hugegraph-pd/hg-pd-service/src/test/java/org/apache/hugegraph/pd/rest/StoreHeartbeatAgeTest.java b/hugegraph-pd/hg-pd-service/src/test/java/org/apache/hugegraph/pd/rest/StoreHeartbeatAgeTest.java new file mode 100644 index 0000000000..2a3638eaba --- /dev/null +++ b/hugegraph-pd/hg-pd-service/src/test/java/org/apache/hugegraph/pd/rest/StoreHeartbeatAgeTest.java @@ -0,0 +1,60 @@ +/* + * 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.rest; + +import org.apache.hugegraph.pd.grpc.Metapb; +import org.apache.hugegraph.pd.service.PDRestService; +import org.junit.Assert; +import org.junit.Test; + +public class StoreHeartbeatAgeTest { + + @Test + public void testStaleHeartbeatIsVisibleEvenWhileStoreIsUp() { + long before = System.currentTimeMillis(); + StoreAPI.StoreStatistics statistics = statistics(before - 60000L); + long after = System.currentTimeMillis(); + Assert.assertEquals("Up", statistics.getState()); + Assert.assertEquals(before - 60000L, statistics.getLastHeartBeat()); + Assert.assertTrue(statistics.getHeartbeatAgeMillis() >= 60000L); + Assert.assertTrue(statistics.getHeartbeatAgeMillis() <= after - before + 60000L); + } + + @Test + public void testMissingHeartbeatHasUnknownAge() { + Assert.assertEquals(-1L, statistics(0L).getHeartbeatAgeMillis()); + } + + @Test + public void testFutureHeartbeatDoesNotHaveNegativeAge() { + Assert.assertEquals(0L, statistics(System.currentTimeMillis() + 3600000L).getHeartbeatAgeMillis()); + } + + private static StoreAPI.StoreStatistics statistics(long lastHeartbeat) { + Metapb.Store store = Metapb.Store.newBuilder().setId(7).setState(Metapb.StoreState.Up) + .setLastHeartbeat(lastHeartbeat).build(); + StoreAPI api = new StoreAPI(); + api.pdRestService = new PDRestService() { + @Override + public Metapb.Store getStore(long storeId) { + return store; + } + }; + return api.new StoreStatistics(store); + } +} 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 23ca938c51..5ec2deede5 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 @@ -47,6 +47,7 @@ PartitionServiceTest.class, StoreMonitorDataServiceTest.class, StoreServiceTest.class, + StoreRestoreRegistrationTest.class, TaskScheduleServiceTest.class, IpAuthHandlerTest.class, RaftEngineIpAuthIntegrationTest.class, diff --git a/hugegraph-pd/hg-pd-test/src/main/java/org/apache/hugegraph/pd/core/StoreRestoreRegistrationTest.java b/hugegraph-pd/hg-pd-test/src/main/java/org/apache/hugegraph/pd/core/StoreRestoreRegistrationTest.java new file mode 100644 index 0000000000..be187ecf22 --- /dev/null +++ b/hugegraph-pd/hg-pd-test/src/main/java/org/apache/hugegraph/pd/core/StoreRestoreRegistrationTest.java @@ -0,0 +1,92 @@ +/* + * 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.util.List; +import java.util.Map; +import java.util.concurrent.atomic.AtomicReference; + +import org.apache.hugegraph.pd.StoreNodeService; +import org.apache.hugegraph.pd.config.PDConfig; +import org.apache.hugegraph.pd.grpc.Metapb; +import org.apache.hugegraph.pd.meta.StoreInfoMeta; +import org.apache.hugegraph.testutil.Whitebox; +import org.junit.Assert; +import org.junit.Before; +import org.junit.Test; +import org.mockito.Mockito; + +public class StoreRestoreRegistrationTest { + + private StoreNodeService service; + private StoreInfoMeta metadata; + private AtomicReference current; + + @Before + public void setUp() throws Exception { + service = Mockito.mock(StoreNodeService.class, Mockito.CALLS_REAL_METHODS); + metadata = Mockito.mock(StoreInfoMeta.class); + current = new AtomicReference<>(store(Metapb.StoreState.Up)); + PDConfig config = Mockito.mock(PDConfig.class); + Mockito.when(config.getInitialStoreMap()).thenReturn(Map.of()); + Whitebox.setInternalState(service, "pdConfig", config); + Whitebox.setInternalState(service, "storeInfoMeta", metadata); + Whitebox.setInternalState(service, "statusListeners", List.of()); + Mockito.doNothing().when(service).checkStoreStatus(); + Mockito.when(metadata.storeExists(7L)).thenReturn(true); + Mockito.when(metadata.getStore(7L)).thenAnswer(call -> current.get()); + Mockito.doAnswer(call -> { + current.set(call.getArgument(0)); + return null; + }).when(metadata).updateStore(Mockito.any()); + } + + @Test + public void testRestoringStoreRemainsOfflineUntilSecondRegistration() throws Exception { + Metapb.Store restoring = service.register(store(Metapb.StoreState.Offline)); + Assert.assertEquals(Metapb.StoreState.Offline, restoring.getState()); + Mockito.verify(metadata).removeActiveStore(restoring); + Mockito.verify(metadata, Mockito.never()).keepStoreAlive(Mockito.any()); + + Metapb.Store restored = service.register(store(Metapb.StoreState.Up)); + Assert.assertEquals(Metapb.StoreState.Up, restored.getState()); + Mockito.verify(metadata).keepStoreAlive(restored); + } + + @Test + public void testPendingStoreStillRequiresActivation() throws Exception { + current.set(store(Metapb.StoreState.Pending)); + Assert.assertEquals(Metapb.StoreState.Pending, + service.register(store(Metapb.StoreState.Offline)).getState()); + Assert.assertEquals(Metapb.StoreState.Pending, + service.register(store(Metapb.StoreState.Up)).getState()); + Mockito.verify(metadata, Mockito.never()).keepStoreAlive(Mockito.any()); + } + + @Test + public void testLegacyUpRegistrationStillActivatesOfflineStore() throws Exception { + current.set(store(Metapb.StoreState.Offline)); + Assert.assertEquals(Metapb.StoreState.Up, + service.register(store(Metapb.StoreState.Up)).getState()); + } + + private static Metapb.Store store(Metapb.StoreState state) { + return Metapb.Store.newBuilder().setId(7).setAddress("127.0.0.1:8500") + .setRaftAddress("127.0.0.1:8510").setState(state).build(); + } +} diff --git a/hugegraph-store/hg-store-core/src/main/java/org/apache/hugegraph/store/HeartbeatService.java b/hugegraph-store/hg-store-core/src/main/java/org/apache/hugegraph/store/HeartbeatService.java index ea6f1b5870..06105f7d7c 100644 --- a/hugegraph-store/hg-store-core/src/main/java/org/apache/hugegraph/store/HeartbeatService.java +++ b/hugegraph-store/hg-store-core/src/main/java/org/apache/hugegraph/store/HeartbeatService.java @@ -200,10 +200,36 @@ protected void registerStore() { this.storeInfo.setStoreAddress(IpUtil.getNearestAddress(options.getGrpcAddress())); this.storeInfo.setRaftAddress(IpUtil.getNearestAddress(options.getRaftAddress())); + // Obtain the identity needed by partition recovery without advertising this Store Up. + this.storeInfo.setState(Metapb.StoreState.Offline); long storeId = pdProvider.registerStore(this.storeInfo); if (storeId != 0) { storeInfo.setId(storeId); storeMetadata.save(storeInfo); + if (storeEngine != null) { + try { + storeEngine.stateChanged(storeInfo, Metapb.StoreState.Offline, Metapb.StoreState.Up); + } catch (RuntimeException e) { + log.error("Store {} remains Offline after partition restore failed", storeId, e); + timerNextDelay = REGISTER_RETRY_INTERVAL * 1000; + return; + } + } + if (terminated) { + return; + } + boolean registered = false; + try { + storeInfo.setState(Metapb.StoreState.Up); + registered = pdProvider.registerStore(storeInfo) != 0; + } finally { + if (!registered) { + storeInfo.setState(Metapb.StoreState.Offline); + } + } + if (!registered) { + return; + } this.clusterStats = pdProvider.getClusterStats(); if (clusterStats.getState() == Metapb.ClusterState.Cluster_OK) { timerNextDelay = options.getStoreHBInterval() * 1000; @@ -407,6 +433,7 @@ public void shutdown() { // Never hold this monitor while joining the producer threads. synchronized (stateChangeLock) { // Acquiring the monitor waits for any in-flight state callback to finish. + reportOffline(); } } finally { if (interrupted) { @@ -415,6 +442,29 @@ public void shutdown() { } } + private void reportOffline() { + if (fatalExitRequested.get() || storeInfo == null) { + return; + } + storeInfo.setState(Metapb.StoreState.Offline); + if (storeInfo.getId() == 0 || pdProvider == null || + pdProvider.getPDClient() == null) { + return; + } + try { + // Preserve PD's current metadata, and never use Exiting for a temporary restart. + var client = pdProvider.getPDClient(); + client.getCache().removeStore(storeInfo.getId()); + Metapb.Store current = client.getStore(storeInfo.getId()); + if (current.getState() == Metapb.StoreState.Up) { + client.updateStore(current.toBuilder().setState(Metapb.StoreState.Offline).build()); + } + } catch (PDException | RuntimeException e) { + // PD may be unavailable during cluster shutdown; heartbeat expiry remains the fallback. + log.warn("Failed to report Store {} offline during shutdown", storeInfo.getId(), e); + } + } + private void requestFatalExit() { requestExit(-1); } diff --git a/hugegraph-store/hg-store-core/src/main/java/org/apache/hugegraph/store/HgStoreEngine.java b/hugegraph-store/hg-store-core/src/main/java/org/apache/hugegraph/store/HgStoreEngine.java index 4299e79cee..358764d4ea 100644 --- a/hugegraph-store/hg-store-core/src/main/java/org/apache/hugegraph/store/HgStoreEngine.java +++ b/hugegraph-store/hg-store-core/src/main/java/org/apache/hugegraph/store/HgStoreEngine.java @@ -170,7 +170,6 @@ public synchronized boolean init(final HgStoreEngineOptions opts) { heartbeatService = new HeartbeatService(this); heartbeatService.setStoreMetadata(partitionManager.getStoreMetadata()); - heartbeatService.addStateListener(this).init(options); metricService = HgMetricService.getInstance(); metricService.setHgStoreEngine(this).init(null); @@ -181,6 +180,8 @@ public synchronized boolean init(final HgStoreEngineOptions opts) { dataManager.setMetaManager(partitionManager); dataManager.setCmdClient(hgCmdClient); } + // Registration can restore partition engines immediately, so initialize their owners first. + heartbeatService.init(options); return true; } @@ -277,12 +278,6 @@ public void stateChanged(Store store, Metapb.StoreState oldState, Metapb.StoreSt * 1. Need to check the partition saved this time, delete the invalid partitions. */ public void restoreLocalPartitionEngine() { - // TODO: surface the outcome of this restore (a per-group ready signal, or a failed state - // reported to PD) instead of logging only; a Store is marked Up before this runs and a - // failed restore leaves it Up with missing shard groups. Paired with the TODO in - // StoreNodeService, which also says why this signal alone does not retire the manual - // Store rollout barrier in the Helm chart (helm/hugegraph) cluster preset. - // https://github.com/apache/hugegraph/issues/3229 try { if (!options.isFakePD()) { // FakePD mode does not require synchronization partitionManager.syncPartitionsFromPD(partition -> { @@ -296,13 +291,13 @@ public void restoreLocalPartitionEngine() { try { createPartitionEngine(p); } catch (Exception e) { - log.error("Partition {}-{} restore exception {}", p.getGraphName(), - p.getId(), e); + throw new IllegalStateException("Failed to restore partition " + + p.getGraphName() + "-" + p.getId(), e); } }); }); } catch (PDException e) { - log.error("HgStoreEngine restoreLocalPartitionEngine error {}", e); + throw new IllegalStateException("Failed to synchronize partitions before restore", e); } } diff --git a/hugegraph-store/hg-store-core/src/main/java/org/apache/hugegraph/store/pd/DefaultPdProvider.java b/hugegraph-store/hg-store-core/src/main/java/org/apache/hugegraph/store/pd/DefaultPdProvider.java index 1a99f27feb..b2fd24128d 100644 --- a/hugegraph-store/hg-store-core/src/main/java/org/apache/hugegraph/store/pd/DefaultPdProvider.java +++ b/hugegraph-store/hg-store-core/src/main/java/org/apache/hugegraph/store/pd/DefaultPdProvider.java @@ -133,7 +133,10 @@ public long registerStore(Store store) throws PDException { store); long storeId = 0; - Metapb.Store protoObj = store.getProtoObj(); + Metapb.Store protoObj = store.getProtoObj().toBuilder() + .setState(store.getState() == Metapb.StoreState.Up ? + Metapb.StoreState.Up : Metapb.StoreState.Offline) + .build(); try { storeId = pdClient.registerStore(protoObj); store.setId(storeId); diff --git a/hugegraph-store/hg-store-test/pom.xml b/hugegraph-store/hg-store-test/pom.xml index 28ca7102bf..c23309db84 100644 --- a/hugegraph-store/hg-store-test/pom.xml +++ b/hugegraph-store/hg-store-test/pom.xml @@ -248,6 +248,7 @@ **/PartitionEngineErrorTest.java **/BatchGraphIsolationTest.java **/GraphIDManagerTest.java + **/StoreRegistrationTest.java **/SchemaDriverTest.java diff --git a/hugegraph-store/hg-store-test/src/main/java/org/apache/hugegraph/store/HeartbeatShutdownTest.java b/hugegraph-store/hg-store-test/src/main/java/org/apache/hugegraph/store/HeartbeatShutdownTest.java index 7a89ddf241..bc40f20840 100644 --- a/hugegraph-store/hg-store-test/src/main/java/org/apache/hugegraph/store/HeartbeatShutdownTest.java +++ b/hugegraph-store/hg-store-test/src/main/java/org/apache/hugegraph/store/HeartbeatShutdownTest.java @@ -374,9 +374,10 @@ public void testShutdownWaitsForDerivedStateCallbackAndRejectsLateUpdates() thro assertStillClosing(shutdown); callbackRelease.countDown(); shutdown.get(5, TimeUnit.SECONDS); - service.onStateChanged(Metapb.StoreState.Offline); + Assert.assertEquals(Metapb.StoreState.Offline, service.getStoreInfo().getState()); + service.onStateChanged(Metapb.StoreState.Up); Assert.assertEquals(1, calls.get()); - Assert.assertEquals(Metapb.StoreState.Up, service.getStoreInfo().getState()); + Assert.assertEquals(Metapb.StoreState.Offline, service.getStoreInfo().getState()); } finally { callbackRelease.countDown(); service.shutdown(); diff --git a/hugegraph-store/hg-store-test/src/main/java/org/apache/hugegraph/store/StoreRegistrationTest.java b/hugegraph-store/hg-store-test/src/main/java/org/apache/hugegraph/store/StoreRegistrationTest.java new file mode 100644 index 0000000000..4d504c0803 --- /dev/null +++ b/hugegraph-store/hg-store-test/src/main/java/org/apache/hugegraph/store/StoreRegistrationTest.java @@ -0,0 +1,135 @@ +/* + * 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.store; + +import java.lang.reflect.Field; +import java.util.ArrayList; +import java.util.List; + +import org.apache.hugegraph.pd.client.ClientCache; +import org.apache.hugegraph.pd.client.PDClient; +import org.apache.hugegraph.pd.common.PDException; +import org.apache.hugegraph.pd.grpc.Metapb; +import org.apache.hugegraph.store.meta.Store; +import org.apache.hugegraph.store.meta.StoreMetadata; +import org.apache.hugegraph.store.options.HgStoreEngineOptions; +import org.apache.hugegraph.store.pd.PdProvider; +import org.junit.Assert; +import org.junit.Test; +import org.mockito.ArgumentCaptor; +import org.mockito.Mockito; + +public class StoreRegistrationTest { + + @Test + public void testRestoreCompletesBeforeUpRegistration() throws Exception { + List events = new ArrayList<>(); + HgStoreEngine engine = Mockito.mock(HgStoreEngine.class); + PdProvider provider = Mockito.mock(PdProvider.class); + Mockito.when(provider.registerStore(Mockito.any())).thenAnswer(call -> { + Store store = call.getArgument(0); + events.add(store.getState().name()); + return 7L; + }); + Mockito.doAnswer(call -> { + events.add("restore"); + return null; + }).when(engine).stateChanged(Mockito.any(), Mockito.any(), Mockito.eq(Metapb.StoreState.Up)); + HeartbeatService service = service(engine, provider); + try { + service.registerStore(); + Assert.assertEquals(List.of("Offline", "restore", "Up"), events); + } finally { + service.shutdown(); + } + } + + @Test + public void testFailedRestoreNeverRegistersUp() throws Exception { + HgStoreEngine engine = Mockito.mock(HgStoreEngine.class); + PdProvider provider = Mockito.mock(PdProvider.class); + Mockito.when(provider.registerStore(Mockito.any())).thenReturn(7L); + Mockito.doThrow(new IllegalStateException("restore failed")).when(engine) + .stateChanged(Mockito.any(), Mockito.any(), Mockito.eq(Metapb.StoreState.Up)); + HeartbeatService service = service(engine, provider); + try { + service.registerStore(); + Mockito.verify(provider, Mockito.times(1)).registerStore(Mockito.any()); + Assert.assertEquals(Metapb.StoreState.Offline, service.getStoreInfo().getState()); + } finally { + service.shutdown(); + } + } + + @Test + public void testShutdownReportsOfflineAndPreservesMetadata() throws Exception { + PdProvider provider = Mockito.mock(PdProvider.class); + PDClient client = Mockito.mock(PDClient.class); + Mockito.when(provider.getPDClient()).thenReturn(client); + Mockito.when(client.getCache()).thenReturn(Mockito.mock(ClientCache.class)); + Metapb.Store remote = Metapb.Store.newBuilder().setId(7).setState(Metapb.StoreState.Up) + .setRaftAddress("127.0.0.1:8510") + .addLabels(Metapb.StoreLabel.newBuilder() + .setKey("zone").setValue("a")).build(); + Mockito.when(client.getStore(7)).thenReturn(remote); + HeartbeatService service = service(null, provider); + service.getStoreInfo().setId(7); + service.shutdown(); + ArgumentCaptor request = ArgumentCaptor.forClass(Metapb.Store.class); + Mockito.verify(client).updateStore(request.capture()); + Assert.assertEquals(remote.toBuilder().setState(Metapb.StoreState.Offline).build(), + request.getValue()); + Assert.assertEquals(Metapb.StoreState.Offline, service.getStoreInfo().getState()); + } + + @Test + public void testShutdownContinuesWhenPdIsUnavailable() throws Exception { + PdProvider provider = Mockito.mock(PdProvider.class); + PDClient client = Mockito.mock(PDClient.class); + Mockito.when(provider.getPDClient()).thenReturn(client); + Mockito.when(client.getCache()).thenReturn(Mockito.mock(ClientCache.class)); + Mockito.when(client.getStore(7)).thenThrow(new PDException(1, "unavailable")); + HeartbeatService service = service(null, provider); + service.getStoreInfo().setId(7); + service.shutdown(); + Mockito.verify(client, Mockito.never()).updateStore(Mockito.any()); + Assert.assertEquals(Metapb.StoreState.Offline, service.getStoreInfo().getState()); + } + + private static HeartbeatService service(HgStoreEngine engine, PdProvider provider) throws Exception { + HeartbeatService service = new HeartbeatService(engine); + HgStoreEngineOptions options = new HgStoreEngineOptions(); + options.setGrpcAddress("127.0.0.1:8500"); + options.setRaftAddress("127.0.0.1:8510"); + Store store = new Store(); + store.setState(Metapb.StoreState.Unknown); + set(service, "options", options); + set(service, "storeInfo", store); + set(service, "pdProvider", provider); + service.setStoreMetadata(Mockito.mock(StoreMetadata.class)); + Mockito.when(provider.getClusterStats()).thenReturn(Metapb.ClusterStats.newBuilder() + .setState(Metapb.ClusterState.Cluster_OK).build()); + return service; + } + + private static void set(Object target, String name, Object value) throws Exception { + Field field = target.getClass().getDeclaredField(name); + field.setAccessible(true); + field.set(target, value); + } +} From 9c08bf6fb6e6c02e4a762172959d1df84a975564 Mon Sep 17 00:00:00 2001 From: pallavithegod Date: Fri, 16 Oct 2026 18:59:46 +0530 Subject: [PATCH 2/4] fix(store): preserve rolling restart readiness Retain existing keepalive leases during temporary Offline reporting. Restore once per process across reconnects and Up-registration retries. Attempt every local partition and cover restart and partial-failure paths. --- docker/README.md | 7 +- helm/hugegraph/docs/operations.md | 3 + hugegraph-pd/docs/api-reference.md | 7 +- .../apache/hugegraph/pd/StoreNodeService.java | 7 +- .../pd/core/StoreRestoreRegistrationTest.java | 63 ++++++++++++++- .../hugegraph/store/HeartbeatService.java | 19 +++-- .../apache/hugegraph/store/HgStoreEngine.java | 11 ++- .../store/StoreRegistrationTest.java | 80 +++++++++++++++++++ 8 files changed, 181 insertions(+), 16 deletions(-) diff --git a/docker/README.md b/docker/README.md index 559ef8d649..a068c8f26e 100644 --- a/docker/README.md +++ b/docker/README.md @@ -240,13 +240,18 @@ partition engines, and then register `Up`. A restore failure keeps the Store offline and the heartbeat loop retries it. Graceful shutdown stops heartbeat producers before reporting `Offline` to PD; if PD is unavailable, normal heartbeat expiry remains the fallback. A temporary stop does not request permanent retirement. +Offline reports retain the existing keepalive window without extending it, so a +quick restart does not fail the cluster's initial-store-count gate. A Store that +does not return is still removed by normal expiry. Stream reconnects after a +successful restore register `Up` directly without restoring the engines again. +Restoration attempts every local partition before reporting any failures. Upgrade PD before Store to enable this registration handshake. Older Stores keep their existing registration behavior; older PD versions ignore the restore gate. `Up` confirms local engine restore, not Raft catch-up or full shard-group health. PD's authenticated `/v1/stores` and `/v1/store/{id}` responses also expose `heartbeatAgeMillis`, computed with PD's clock (`-1` when no heartbeat is known). -Check it against the Store heartbeat interval to distinguish a stale `Up` record +Compare it with the fixed 30-second Store heartbeat period to distinguish a stale `Up` record from a recent heartbeat without waiting for keep-alive expiry. A future timestamp caused by a PD clock adjustment reports zero age. Keep the cluster preset's manual rollout barrier and verify fresh heartbeats, diff --git a/helm/hugegraph/docs/operations.md b/helm/hugegraph/docs/operations.md index b74ed7c589..b44db501fb 100644 --- a/helm/hugegraph/docs/operations.md +++ b/helm/hugegraph/docs/operations.md @@ -201,6 +201,9 @@ shard recovery before continuing; do not blindly retry writes. With PD and Store images containing #3229's registration changes, `Up` follows local engine restore and graceful stops report `Offline`. Use PD's `heartbeatAgeMillis` to reject stale `Up` records after an abrupt stop. +Temporary Offline reports preserve the existing keepalive window without renewing +it; quick restarts do not trip the initial-store-count gate, while expiry still +applies if the Store does not return. Stream reconnects do not repeat local restore. These signals do not prove Raft catch-up. diff --git a/hugegraph-pd/docs/api-reference.md b/hugegraph-pd/docs/api-reference.md index 5dffe754ab..d397734275 100644 --- a/hugegraph-pd/docs/api-reference.md +++ b/hugegraph-pd/docs/api-reference.md @@ -205,6 +205,9 @@ partition engines, then register `Up`. Upgrade PD before Store: older PD version ignore this restore gate, and older Stores retain their previous registration behavior. Pending Stores still require activation. `Up` does not guarantee Raft catch-up; graceful shutdown reports `Offline` without requesting permanent retirement. +Temporary Offline reports retain the existing keepalive entry without renewing it; +expiry still removes a Store that does not return. Once local restore succeeds, +stream reconnects register `Up` directly without repeating recovery. **Request**: ```protobuf @@ -925,8 +928,8 @@ curl http://localhost:8620/v1/partitions?graph_name=social_network The authenticated `/v1/stores` and `/v1/store/{id}` responses expose `heartbeatAgeMillis` alongside the existing `lastHeartBeat`. Age is computed with PD's clock, is `-1` if no heartbeat is recorded, and is clamped to zero for future -timestamps after a clock adjustment. Compare it with the configured Store -heartbeat interval to detect stale `Up` records before keep-alive expiry. +timestamps after a clock adjustment. Compare it with the fixed 30-second Store +heartbeat period (30000 ms) to detect stale `Up` records before keep-alive expiry. Read this state from the PD leader and continue to verify shard-group membership and leaders during rolling updates; fresh heartbeats do not prove Raft catch-up. 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 b0cb7c6e83..e8922fcc19 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 @@ -214,7 +214,7 @@ public Metapb.Store register(Metapb.Store store) throws PDException { checkStoreStatus(); } else if (storeState == Metapb.StoreState.Offline && lastStore.getState() == Metapb.StoreState.Up) { - storeInfoMeta.removeActiveStore(store); + // Keep the existing keepalive window during a temporary restart; do not renew it. onStoreStatusChanged(store, Metapb.StoreState.Up, Metapb.StoreState.Offline); checkStoreStatus(); } @@ -292,6 +292,8 @@ public synchronized Metapb.Store updateStore(Metapb.Store store) throws PDExcept } storeInfoMeta.updateStore(store); + // Temporary Offline reports retain the existing keepalive entry without renewing it. + // Removing it here makes a rolling restart fail the initial-store-count gate. if (store.getState() != Metapb.StoreState.Unknown && store.getState() != lastStore.getState()) { // If you want to take the store offline @@ -314,9 +316,6 @@ 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 diff --git a/hugegraph-pd/hg-pd-test/src/main/java/org/apache/hugegraph/pd/core/StoreRestoreRegistrationTest.java b/hugegraph-pd/hg-pd-test/src/main/java/org/apache/hugegraph/pd/core/StoreRestoreRegistrationTest.java index be187ecf22..159651170f 100644 --- a/hugegraph-pd/hg-pd-test/src/main/java/org/apache/hugegraph/pd/core/StoreRestoreRegistrationTest.java +++ b/hugegraph-pd/hg-pd-test/src/main/java/org/apache/hugegraph/pd/core/StoreRestoreRegistrationTest.java @@ -17,6 +17,8 @@ package org.apache.hugegraph.pd.core; +import java.util.ArrayList; +import java.util.HashMap; import java.util.List; import java.util.Map; import java.util.concurrent.atomic.AtomicReference; @@ -36,13 +38,14 @@ public class StoreRestoreRegistrationTest { private StoreNodeService service; private StoreInfoMeta metadata; private AtomicReference current; + private PDConfig config; @Before public void setUp() throws Exception { service = Mockito.mock(StoreNodeService.class, Mockito.CALLS_REAL_METHODS); metadata = Mockito.mock(StoreInfoMeta.class); current = new AtomicReference<>(store(Metapb.StoreState.Up)); - PDConfig config = Mockito.mock(PDConfig.class); + config = Mockito.mock(PDConfig.class); Mockito.when(config.getInitialStoreMap()).thenReturn(Map.of()); Whitebox.setInternalState(service, "pdConfig", config); Whitebox.setInternalState(service, "storeInfoMeta", metadata); @@ -60,7 +63,7 @@ public void setUp() throws Exception { public void testRestoringStoreRemainsOfflineUntilSecondRegistration() throws Exception { Metapb.Store restoring = service.register(store(Metapb.StoreState.Offline)); Assert.assertEquals(Metapb.StoreState.Offline, restoring.getState()); - Mockito.verify(metadata).removeActiveStore(restoring); + Mockito.verify(metadata, Mockito.never()).removeActiveStore(Mockito.any()); Mockito.verify(metadata, Mockito.never()).keepStoreAlive(Mockito.any()); Metapb.Store restored = service.register(store(Metapb.StoreState.Up)); @@ -85,6 +88,62 @@ public void testLegacyUpRegistrationStillActivatesOfflineStore() throws Exceptio service.register(store(Metapb.StoreState.Up)).getState()); } + @Test + public void testShutdownOfflinePreservesThreeStoreReadinessUntilExpiry() throws Exception { + assertTemporaryOfflinePreservesReadiness(false); + } + + @Test + public void testRestoreRegistrationPreservesThreeStoreReadinessUntilExpiry() throws Exception { + assertTemporaryOfflinePreservesReadiness(true); + } + + private void assertTemporaryOfflinePreservesReadiness(boolean registering) throws Exception { + Map stores = new HashMap<>(); + for (long id = 7; id < 10; id++) { + stores.put(id, store(Metapb.StoreState.Up).toBuilder().setId(id).build()); + } + Map active = new HashMap<>(stores); + Mockito.when(config.getMinStoreCount()).thenReturn(3); + Mockito.doCallRealMethod().when(service).checkStoreStatus(); + Mockito.when(metadata.getStore(Mockito.anyLong())).thenAnswer(call -> stores.get(call.getArgument(0))); + Mockito.when(metadata.getActiveStores()).thenAnswer(call -> new ArrayList<>(active.values())); + Mockito.when(metadata.getShardGroups()).thenReturn(List.of(Metapb.ShardGroup.newBuilder() + .setId(1).addShards(Metapb.Shard.newBuilder().setStoreId(7)) + .addShards(Metapb.Shard.newBuilder().setStoreId(8)) + .addShards(Metapb.Shard.newBuilder().setStoreId(9)).build())); + Mockito.doAnswer(call -> { + Metapb.Store store = call.getArgument(0); + stores.put(store.getId(), store); + return null; + }).when(metadata).updateStore(Mockito.any()); + Mockito.doAnswer(call -> { + active.remove(((Metapb.Store) call.getArgument(0)).getId()); + return null; + }).when(metadata).removeActiveStore(Mockito.any()); + Mockito.doAnswer(call -> { + Metapb.Store store = call.getArgument(0); + active.put(store.getId(), store); + return null; + }).when(metadata).keepStoreAlive(Mockito.any()); + + Metapb.Store offline = store(Metapb.StoreState.Offline); + if (registering) { + service.register(offline); + } else { + service.updateStore(offline); + } + Assert.assertEquals(Metapb.StoreState.Offline, service.getStore(7).getState()); + Metapb.StoreStats heartbeat = Metapb.StoreStats.newBuilder().setStoreId(8).build(); + Assert.assertEquals(Metapb.ClusterState.Cluster_OK, service.heartBeat(heartbeat).getState()); + Mockito.verify(metadata, Mockito.never()).removeActiveStore(Mockito.any()); + Mockito.verify(metadata, Mockito.never()).keepStoreAlive(Mockito.argThat(store -> store.getId() == 7)); + + // Expiry still removes a Store that never comes back; Offline must not renew the lease. + active.remove(7L); + Assert.assertEquals(Metapb.ClusterState.Cluster_Not_Ready, service.heartBeat(heartbeat).getState()); + } + private static Metapb.Store store(Metapb.StoreState state) { return Metapb.Store.newBuilder().setId(7).setAddress("127.0.0.1:8500") .setRaftAddress("127.0.0.1:8510").setState(state).build(); diff --git a/hugegraph-store/hg-store-core/src/main/java/org/apache/hugegraph/store/HeartbeatService.java b/hugegraph-store/hg-store-core/src/main/java/org/apache/hugegraph/store/HeartbeatService.java index 06105f7d7c..921b5f8cf2 100644 --- a/hugegraph-store/hg-store-core/src/main/java/org/apache/hugegraph/store/HeartbeatService.java +++ b/hugegraph-store/hg-store-core/src/main/java/org/apache/hugegraph/store/HeartbeatService.java @@ -70,6 +70,7 @@ public class HeartbeatService implements Lifecycle, Partit // Thread sleep time private volatile int timerNextDelay = 1000; private volatile boolean terminated = false; + private boolean restored = false; private final AtomicBoolean fatalExitRequested = new AtomicBoolean(); private Thread storeHeartbeatThread; private Thread partitionHeartbeatThread; @@ -200,13 +201,20 @@ protected void registerStore() { this.storeInfo.setStoreAddress(IpUtil.getNearestAddress(options.getGrpcAddress())); this.storeInfo.setRaftAddress(IpUtil.getNearestAddress(options.getRaftAddress())); - // Obtain the identity needed by partition recovery without advertising this Store Up. - this.storeInfo.setState(Metapb.StoreState.Offline); - long storeId = pdProvider.registerStore(this.storeInfo); + // Recovery is needed once per process, not when an established stream reconnects. + boolean restoring = !restored; + this.storeInfo.setState(restoring ? Metapb.StoreState.Offline : Metapb.StoreState.Up); + long storeId; + try { + storeId = pdProvider.registerStore(this.storeInfo); + } finally { + // Retry registration after a failure instead of heartbeating an unregistered Store. + this.storeInfo.setState(Metapb.StoreState.Offline); + } if (storeId != 0) { storeInfo.setId(storeId); storeMetadata.save(storeInfo); - if (storeEngine != null) { + if (restoring && storeEngine != null) { try { storeEngine.stateChanged(storeInfo, Metapb.StoreState.Offline, Metapb.StoreState.Up); } catch (RuntimeException e) { @@ -215,13 +223,14 @@ protected void registerStore() { return; } } + restored = true; if (terminated) { return; } boolean registered = false; try { storeInfo.setState(Metapb.StoreState.Up); - registered = pdProvider.registerStore(storeInfo) != 0; + registered = !restoring || pdProvider.registerStore(storeInfo) != 0; } finally { if (!registered) { storeInfo.setState(Metapb.StoreState.Offline); diff --git a/hugegraph-store/hg-store-core/src/main/java/org/apache/hugegraph/store/HgStoreEngine.java b/hugegraph-store/hg-store-core/src/main/java/org/apache/hugegraph/store/HgStoreEngine.java index 358764d4ea..f6f25a5042 100644 --- a/hugegraph-store/hg-store-core/src/main/java/org/apache/hugegraph/store/HgStoreEngine.java +++ b/hugegraph-store/hg-store-core/src/main/java/org/apache/hugegraph/store/HgStoreEngine.java @@ -286,16 +286,23 @@ public void restoreLocalPartitionEngine() { "Please delete the redundant data manually, {}", partition); }); } + List failures = new ArrayList<>(); partitionManager.getPartitions().forEach((k, g) -> { g.forEach((id, p) -> { try { createPartitionEngine(p); } catch (Exception e) { - throw new IllegalStateException("Failed to restore partition " + - p.getGraphName() + "-" + p.getId(), e); + failures.add(new IllegalStateException("Failed to restore partition " + + p.getGraphName() + "-" + p.getId(), e)); } }); }); + if (!failures.isEmpty()) { + IllegalStateException failure = new IllegalStateException("Failed to restore " + + failures.size() + " partitions"); + failures.forEach(failure::addSuppressed); + throw failure; + } } catch (PDException e) { throw new IllegalStateException("Failed to synchronize partitions before restore", e); } diff --git a/hugegraph-store/hg-store-test/src/main/java/org/apache/hugegraph/store/StoreRegistrationTest.java b/hugegraph-store/hg-store-test/src/main/java/org/apache/hugegraph/store/StoreRegistrationTest.java index 4d504c0803..799c3100a9 100644 --- a/hugegraph-store/hg-store-test/src/main/java/org/apache/hugegraph/store/StoreRegistrationTest.java +++ b/hugegraph-store/hg-store-test/src/main/java/org/apache/hugegraph/store/StoreRegistrationTest.java @@ -19,12 +19,16 @@ import java.lang.reflect.Field; import java.util.ArrayList; +import java.util.LinkedHashMap; import java.util.List; +import java.util.Map; import org.apache.hugegraph.pd.client.ClientCache; import org.apache.hugegraph.pd.client.PDClient; import org.apache.hugegraph.pd.common.PDException; import org.apache.hugegraph.pd.grpc.Metapb; +import org.apache.hugegraph.store.meta.Partition; +import org.apache.hugegraph.store.meta.PartitionManager; import org.apache.hugegraph.store.meta.Store; import org.apache.hugegraph.store.meta.StoreMetadata; import org.apache.hugegraph.store.options.HgStoreEngineOptions; @@ -76,6 +80,82 @@ public void testFailedRestoreNeverRegistersUp() throws Exception { } } + @Test + public void testReconnectRegistersUpWithoutRestoringAgain() throws Exception { + List events = new ArrayList<>(); + HgStoreEngine engine = Mockito.mock(HgStoreEngine.class); + PdProvider provider = Mockito.mock(PdProvider.class); + Mockito.when(provider.registerStore(Mockito.any())).thenAnswer(call -> { + events.add(((Store) call.getArgument(0)).getState().name()); + return 7L; + }); + HeartbeatService service = service(engine, provider); + try { + service.registerStore(); + Assert.assertEquals(List.of("Offline", "Up"), events); + events.clear(); + service.onStateChanged(Metapb.StoreState.Offline); + service.registerStore(); + Assert.assertEquals(List.of("Up"), events); + Assert.assertEquals(Metapb.StoreState.Up, service.getStoreInfo().getState()); + Mockito.verify(engine).stateChanged(Mockito.any(), Mockito.any(), Mockito.eq(Metapb.StoreState.Up)); + } finally { + service.shutdown(); + } + } + + @Test + public void testFailedUpRegistrationRetriesWithoutRestoringAgain() throws Exception { + HgStoreEngine engine = Mockito.mock(HgStoreEngine.class); + PdProvider provider = Mockito.mock(PdProvider.class); + Mockito.when(provider.registerStore(Mockito.any())).thenReturn(7L) + .thenThrow(new PDException(1, "unavailable")).thenReturn(7L); + HeartbeatService service = service(engine, provider); + try { + service.registerStore(); + Assert.assertEquals(Metapb.StoreState.Offline, service.getStoreInfo().getState()); + service.registerStore(); + Assert.assertEquals(Metapb.StoreState.Up, service.getStoreInfo().getState()); + Mockito.verify(engine).stateChanged(Mockito.any(), Mockito.any(), Mockito.eq(Metapb.StoreState.Up)); + } finally { + service.shutdown(); + } + } + + @Test + public void testRestoreAttemptsHealthyPartitionAfterFailure() throws Exception { + HgStoreEngine engine = Mockito.mock(HgStoreEngine.class, Mockito.CALLS_REAL_METHODS); + HgStoreEngineOptions options = new HgStoreEngineOptions(); + options.setFakePD(true); + Field optionsField = HgStoreEngine.class.getDeclaredField("options"); + optionsField.setAccessible(true); + optionsField.set(engine, options); + PartitionManager manager = Mockito.mock(PartitionManager.class); + engine.setPartitionManager(manager); + Partition first = new Partition(); + first.setId(1); + first.setGraphName("graph"); + Partition second = new Partition(); + second.setId(2); + second.setGraphName("graph"); + Map partitions = new LinkedHashMap<>(); + partitions.put(1, first); + partitions.put(2, second); + Mockito.when(manager.getPartitions()).thenReturn(Map.of("graph", partitions)); + Mockito.doThrow(new IllegalStateException("broken raft group")).when(engine).createPartitionEngine(first); + PartitionEngine healthy = Mockito.mock(PartitionEngine.class); + Mockito.doReturn(healthy).when(engine).createPartitionEngine(second); + try { + engine.restoreLocalPartitionEngine(); + Assert.fail("partial restore must not publish the Store Up"); + } catch (IllegalStateException e) { + Assert.assertEquals(1, e.getSuppressed().length); + Assert.assertTrue(e.getSuppressed()[0].getMessage().contains("graph-1")); + Assert.assertEquals("broken raft group", e.getSuppressed()[0].getCause().getMessage()); + } + Mockito.verify(engine).createPartitionEngine(second); + } + @Test public void testShutdownReportsOfflineAndPreservesMetadata() throws Exception { PdProvider provider = Mockito.mock(PdProvider.class); From 5511593ea0508bf6599e907a8d710341c7b1a89f Mon Sep 17 00:00:00 2001 From: pallavithegod Date: Sun, 18 Oct 2026 16:34:18 +0530 Subject: [PATCH 3/4] fix(store): separate shutdown from recovery leases Use a dedicated shutdown RPC while preserving administrative Offline checks. Renew leases and cluster readiness while an owned worker restores partitions. Return authoritative registration state and cover lifecycle and compatibility paths. --- docker/README.md | 10 +- helm/hugegraph/docs/operations.md | 5 +- hugegraph-pd/docs/api-reference.md | 14 +- .../apache/hugegraph/pd/client/PDClient.java | 31 +++- .../apache/hugegraph/pd/StoreNodeService.java | 29 ++- .../hg-pd-grpc/src/main/proto/pdpb.proto | 3 + .../hugegraph/pd/service/PDService.java | 19 ++ .../hugegraph/pd/core/PDCoreSuiteTest.java | 2 + .../pd/core/StoreRestoreRegistrationTest.java | 33 +++- .../pd/service/StoreShutdownReportTest.java | 115 ++++++++++++ .../hugegraph/store/HeartbeatService.java | 57 ++++-- .../hugegraph/store/pd/DefaultPdProvider.java | 7 +- .../store/HeartbeatShutdownTest.java | 14 +- .../store/StoreRegistrationTest.java | 168 ++++++++++++++++-- 14 files changed, 451 insertions(+), 56 deletions(-) create mode 100644 hugegraph-pd/hg-pd-test/src/main/java/org/apache/hugegraph/pd/service/StoreShutdownReportTest.java diff --git a/docker/README.md b/docker/README.md index a068c8f26e..2b6554abd4 100644 --- a/docker/README.md +++ b/docker/README.md @@ -240,14 +240,22 @@ partition engines, and then register `Up`. A restore failure keeps the Store offline and the heartbeat loop retries it. Graceful shutdown stops heartbeat producers before reporting `Offline` to PD; if PD is unavailable, normal heartbeat expiry remains the fallback. A temporary stop does not request permanent retirement. -Offline reports retain the existing keepalive window without extending it, so a +Shutdown reports retain the existing keepalive window without extending it, so a quick restart does not fail the cluster's initial-store-count gate. A Store that does not return is still removed by normal expiry. Stream reconnects after a successful restore register `Up` directly without restoring the engines again. Restoration attempts every local partition before reporting any failures. +Recovery runs on an owned worker while the heartbeat thread renews the restoring +Store's lease through Offline registrations. A slow or failed partition therefore +does not let a returned Store expire from the initial-store-count gate. +Partition heartbeat reports wait until the pulse stream starts after restoration. Upgrade PD before Store to enable this registration handshake. Older Stores keep their existing registration behavior; older PD versions ignore the restore gate. +Temporary shutdown uses a dedicated `ReportStoreShutdown` RPC, independent of +cluster readiness. Older PDs do not implement it, so Store logs the failed report +and relies on expiry; it never falls back to administrative `setStore(Offline)`. +Administrative Offline requests retain their readiness gate and active-entry removal. `Up` confirms local engine restore, not Raft catch-up or full shard-group health. PD's authenticated `/v1/stores` and `/v1/store/{id}` responses also expose `heartbeatAgeMillis`, computed with PD's clock (`-1` when no heartbeat is known). diff --git a/helm/hugegraph/docs/operations.md b/helm/hugegraph/docs/operations.md index b44db501fb..e940a2d102 100644 --- a/helm/hugegraph/docs/operations.md +++ b/helm/hugegraph/docs/operations.md @@ -201,9 +201,12 @@ shard recovery before continuing; do not blindly retry writes. With PD and Store images containing #3229's registration changes, `Up` follows local engine restore and graceful stops report `Offline`. Use PD's `heartbeatAgeMillis` to reject stale `Up` records after an abrupt stop. -Temporary Offline reports preserve the existing keepalive window without renewing +Temporary shutdown reports preserve the existing keepalive window without renewing it; quick restarts do not trip the initial-store-count gate, while expiry still applies if the Store does not return. Stream reconnects do not repeat local restore. +While recovery runs, Offline registrations renew the returned Store's lease. +The dedicated shutdown RPC falls back to normal expiry on older PDs without +using the administrative Offline path, so a PD rollback does not remove that lease. These signals do not prove Raft catch-up. diff --git a/hugegraph-pd/docs/api-reference.md b/hugegraph-pd/docs/api-reference.md index d397734275..4004766304 100644 --- a/hugegraph-pd/docs/api-reference.md +++ b/hugegraph-pd/docs/api-reference.md @@ -205,9 +205,19 @@ partition engines, then register `Up`. Upgrade PD before Store: older PD version ignore this restore gate, and older Stores retain their previous registration behavior. Pending Stores still require activation. `Up` does not guarantee Raft catch-up; graceful shutdown reports `Offline` without requesting permanent retirement. -Temporary Offline reports retain the existing keepalive entry without renewing it; -expiry still removes a Store that does not return. Once local restore succeeds, +Temporary shutdown reports retain the existing keepalive entry without renewing it; +expiry still removes a Store that does not return. During recovery, repeated Offline +registrations renew the returned Store's lease while an owned worker restores engines. +Once local restore succeeds, stream reconnects register `Up` directly without repeating recovery. +`RegisterStoreResponse.store` supplies the authoritative registration state; older +responses without this additive field fall back to a fresh Store lookup. + +`ReportStoreShutdown(SetStoreRequest)` returns `SetStoreResponse` and records +temporary Offline by Store ID, preserving PD metadata and the existing lease even +when the cluster is not ready. Older PDs return `UNIMPLEMENTED`; updated Stores +fall back to expiry, never to `SetStore`. Administrative `SetStore(Offline)` keeps +its cluster-readiness guard and removes the active entry as before. **Request**: ```protobuf diff --git a/hugegraph-pd/hg-pd-client/src/main/java/org/apache/hugegraph/pd/client/PDClient.java b/hugegraph-pd/hg-pd-client/src/main/java/org/apache/hugegraph/pd/client/PDClient.java index e616e27c41..5c5df55b4e 100644 --- a/hugegraph-pd/hg-pd-client/src/main/java/org/apache/hugegraph/pd/client/PDClient.java +++ b/hugegraph-pd/hg-pd-client/src/main/java/org/apache/hugegraph/pd/client/PDClient.java @@ -382,6 +382,26 @@ private String getLeaderIp(PDBlockingStub stub) { * @return */ public long registerStore(Metapb.Store store) throws PDException { + return registerStoreResponse(store).getStoreId(); + } + + public Metapb.Store registerStoreWithState(Metapb.Store store) throws PDException { + Pdpb.RegisterStoreResponse response = registerStoreResponse(store); + if (response.hasStore()) { + if (config.isEnableCache() && response.getStore().getId() != 0) { + cache.addStore(response.getStore().getId(), response.getStore()); + } + return response.getStore(); + } + if (response.getStoreId() == 0) { + return store.toBuilder().setId(0).build(); + } + // Older PDs return only the identity. Avoid a stale Offline entry from an earlier handshake. + cache.removeStore(response.getStoreId()); + return getStore(response.getStoreId()); + } + + private Pdpb.RegisterStoreResponse registerStoreResponse(Metapb.Store store) throws PDException { Pdpb.RegisterStoreRequest request = Pdpb.RegisterStoreRequest.newBuilder() .setHeader(header) .setStore(store).build(); @@ -389,7 +409,7 @@ public long registerStore(Metapb.Store store) throws PDException { Pdpb.RegisterStoreResponse response = blockingUnaryCall(PDGrpc.getRegisterStoreMethod(), request); handleResponseError(response.getHeader()); - return response.getStoreId(); + return response; } /** @@ -435,6 +455,15 @@ public Metapb.Store updateStore(Metapb.Store store) throws PDException { return store; } + public Metapb.Store reportStoreShutdown(long storeId) throws PDException { + Pdpb.SetStoreRequest request = Pdpb.SetStoreRequest.newBuilder().setHeader(header) + .setStore(Metapb.Store.newBuilder().setId(storeId)).build(); + Pdpb.SetStoreResponse response = getStub().reportStoreShutdown(request); + handleResponseError(response.getHeader()); + cache.removeStore(storeId); + return response.getStore(); + } + /** * Return to the active store * 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 e8922fcc19..a2ecae6ede 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 @@ -212,10 +212,12 @@ 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.Offline && - lastStore.getState() == Metapb.StoreState.Up) { - // Keep the existing keepalive window during a temporary restart; do not renew it. - onStoreStatusChanged(store, Metapb.StoreState.Up, Metapb.StoreState.Offline); + } else if (restoring && storeState == Metapb.StoreState.Offline) { + // A returned Store is alive even while its engines recover. Renew on each retry. + storeInfoMeta.keepStoreAlive(store); + if (lastStore.getState() == Metapb.StoreState.Up) { + onStoreStatusChanged(store, Metapb.StoreState.Up, Metapb.StoreState.Offline); + } checkStoreStatus(); } @@ -272,6 +274,7 @@ 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 { + boolean administrativeOffline = store.getState() == Metapb.StoreState.Offline; log.info("updateStore storeId: {}, address: {}, state: {}", store.getId(), store.getAddress(), store.getState()); Metapb.Store lastStore = storeInfoMeta.getStore(store.getId()); @@ -292,8 +295,10 @@ public synchronized Metapb.Store updateStore(Metapb.Store store) throws PDExcept } storeInfoMeta.updateStore(store); - // Temporary Offline reports retain the existing keepalive entry without renewing it. - // Removing it here makes a rolling restart fail the initial-store-count gate. + if (administrativeOffline) { + // A crashed restoring Store may already be Offline while its lease is still present. + storeInfoMeta.removeActiveStore(store); + } if (store.getState() != Metapb.StoreState.Unknown && store.getState() != lastStore.getState()) { // If you want to take the store offline @@ -332,6 +337,18 @@ public synchronized Metapb.Store updateStore(Metapb.Store store) throws PDExcept return store; } + public synchronized Metapb.Store reportStoreShutdown(long storeId) throws PDException { + Metapb.Store current = getStore(storeId); + if (current.getState() != Metapb.StoreState.Up) { + return current; + } + Metapb.Store offline = current.toBuilder().setState(Metapb.StoreState.Offline).build(); + storeInfoMeta.updateStore(offline); + // Preserve the existing lease, including its expiry; administrative Offline still removes it. + onStoreStatusChanged(offline, Metapb.StoreState.Up, Metapb.StoreState.Offline); + return offline; + } + /** * The shard of the shardGroup is reassigned * diff --git a/hugegraph-pd/hg-pd-grpc/src/main/proto/pdpb.proto b/hugegraph-pd/hg-pd-grpc/src/main/proto/pdpb.proto index 4e6c855322..5f977de4de 100644 --- a/hugegraph-pd/hg-pd-grpc/src/main/proto/pdpb.proto +++ b/hugegraph-pd/hg-pd-grpc/src/main/proto/pdpb.proto @@ -29,6 +29,8 @@ service PD { rpc GetStore(GetStoreRequest) returns (GetStoreResponse) {} // Modify Store status and other information. rpc SetStore(SetStoreRequest) returns (SetStoreResponse) {} + // A Store reports temporary shutdown without administrative offlining or lease renewal. + rpc ReportStoreShutdown(SetStoreRequest) returns (SetStoreResponse) {} // Look up the partition based on the key rpc DelStore(DetStoreRequest) returns (DetStoreResponse) {} rpc GetAllStores(GetAllStoresRequest) returns (GetAllStoresResponse) {} @@ -202,6 +204,7 @@ message RegisterStoreResponse { ResponseHeader header = 1; // Upon initial registration, returns the new store_id uint64 store_id = 2; + metapb.Store store = 3; } message SetStoreRequest { diff --git a/hugegraph-pd/hg-pd-service/src/main/java/org/apache/hugegraph/pd/service/PDService.java b/hugegraph-pd/hg-pd-service/src/main/java/org/apache/hugegraph/pd/service/PDService.java index f40b2f5801..55d0c0ea3c 100644 --- a/hugegraph-pd/hg-pd-service/src/main/java/org/apache/hugegraph/pd/service/PDService.java +++ b/hugegraph-pd/hg-pd-service/src/main/java/org/apache/hugegraph/pd/service/PDService.java @@ -403,6 +403,7 @@ public void registerStore(Pdpb.RegisterStoreRequest request, Metapb.Store store = storeNodeService.register(request.getStore()); response = Pdpb.RegisterStoreResponse.newBuilder().setHeader(okHeader) .setStoreId(store.getId()) + .setStore(store) .build(); } catch (PDException e) { response = Pdpb.RegisterStoreResponse.newBuilder().setHeader(newErrorHeader(e)).build(); @@ -438,6 +439,24 @@ public void getStore(Pdpb.GetStoreRequest request, observer.onCompleted(); } + @Override + public void reportStoreShutdown(Pdpb.SetStoreRequest request, + StreamObserver observer) { + if (!isLeader()) { + redirectToLeader(PDGrpc.getReportStoreShutdownMethod(), request, observer); + return; + } + Pdpb.SetStoreResponse response; + try { + Metapb.Store store = storeNodeService.reportStoreShutdown(request.getStore().getId()); + response = Pdpb.SetStoreResponse.newBuilder().setHeader(okHeader).setStore(store).build(); + } catch (PDException e) { + response = Pdpb.SetStoreResponse.newBuilder().setHeader(newErrorHeader(e)).build(); + } + observer.onNext(response); + observer.onCompleted(); + } + /** * 
      * Modify information such as the status of the store.
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 5ec2deede5..356a7b9b7a 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
@@ -28,6 +28,7 @@
 import org.apache.hugegraph.pd.raft.RaftEngineReadinessTest;
 import org.apache.hugegraph.pd.raft.RaftReflectionUtilTest;
 import org.apache.hugegraph.pd.raft.ZipUtilsTest;
+import org.apache.hugegraph.pd.service.StoreShutdownReportTest;
 import org.junit.runner.RunWith;
 import org.junit.runners.Suite;
 
@@ -48,6 +49,7 @@
         StoreMonitorDataServiceTest.class,
         StoreServiceTest.class,
         StoreRestoreRegistrationTest.class,
+        StoreShutdownReportTest.class,
         TaskScheduleServiceTest.class,
         IpAuthHandlerTest.class,
         RaftEngineIpAuthIntegrationTest.class,
diff --git a/hugegraph-pd/hg-pd-test/src/main/java/org/apache/hugegraph/pd/core/StoreRestoreRegistrationTest.java b/hugegraph-pd/hg-pd-test/src/main/java/org/apache/hugegraph/pd/core/StoreRestoreRegistrationTest.java
index 159651170f..85f38f9b89 100644
--- a/hugegraph-pd/hg-pd-test/src/main/java/org/apache/hugegraph/pd/core/StoreRestoreRegistrationTest.java
+++ b/hugegraph-pd/hg-pd-test/src/main/java/org/apache/hugegraph/pd/core/StoreRestoreRegistrationTest.java
@@ -64,7 +64,7 @@ public void testRestoringStoreRemainsOfflineUntilSecondRegistration() throws Exc
         Metapb.Store restoring = service.register(store(Metapb.StoreState.Offline));
         Assert.assertEquals(Metapb.StoreState.Offline, restoring.getState());
         Mockito.verify(metadata, Mockito.never()).removeActiveStore(Mockito.any());
-        Mockito.verify(metadata, Mockito.never()).keepStoreAlive(Mockito.any());
+        Mockito.verify(metadata).keepStoreAlive(restoring);
 
         Metapb.Store restored = service.register(store(Metapb.StoreState.Up));
         Assert.assertEquals(Metapb.StoreState.Up, restored.getState());
@@ -131,17 +131,44 @@ private void assertTemporaryOfflinePreservesReadiness(boolean registering) throw
         if (registering) {
             service.register(offline);
         } else {
-            service.updateStore(offline);
+            service.reportStoreShutdown(7);
         }
         Assert.assertEquals(Metapb.StoreState.Offline, service.getStore(7).getState());
         Metapb.StoreStats heartbeat = Metapb.StoreStats.newBuilder().setStoreId(8).build();
         Assert.assertEquals(Metapb.ClusterState.Cluster_OK, service.heartBeat(heartbeat).getState());
         Mockito.verify(metadata, Mockito.never()).removeActiveStore(Mockito.any());
-        Mockito.verify(metadata, Mockito.never()).keepStoreAlive(Mockito.argThat(store -> store.getId() == 7));
+        if (!registering) {
+            Mockito.verify(metadata, Mockito.never()).keepStoreAlive(Mockito.argThat(store -> store.getId() == 7));
+        }
 
         // Expiry still removes a Store that never comes back; Offline must not renew the lease.
         active.remove(7L);
         Assert.assertEquals(Metapb.ClusterState.Cluster_Not_Ready, service.heartBeat(heartbeat).getState());
+        if (registering) {
+            service.register(offline);
+            Assert.assertEquals(Metapb.ClusterState.Cluster_OK, service.heartBeat(heartbeat).getState());
+        }
+    }
+
+    @Test
+    public void testAdministrativeOfflineStillRemovesActiveEntry() throws Exception {
+        service.updateStore(store(Metapb.StoreState.Offline));
+        Mockito.verify(metadata).removeActiveStore(Mockito.any());
+    }
+
+    @Test
+    public void testAdministrativeOfflineRemovesAlreadyOfflineLease() throws Exception {
+        current.set(store(Metapb.StoreState.Offline));
+        service.updateStore(store(Metapb.StoreState.Offline));
+        Mockito.verify(metadata).removeActiveStore(Mockito.any());
+    }
+
+    @Test
+    public void testMetadataOnlyUpdateDoesNotRemoveRestoringLease() throws Exception {
+        current.set(store(Metapb.StoreState.Offline));
+        service.updateStore(Metapb.Store.newBuilder().setId(7)
+                .addLabels(Metapb.StoreLabel.newBuilder().setKey("zone").setValue("a")).build());
+        Mockito.verify(metadata, Mockito.never()).removeActiveStore(Mockito.any());
     }
 
     private static Metapb.Store store(Metapb.StoreState state) {
diff --git a/hugegraph-pd/hg-pd-test/src/main/java/org/apache/hugegraph/pd/service/StoreShutdownReportTest.java b/hugegraph-pd/hg-pd-test/src/main/java/org/apache/hugegraph/pd/service/StoreShutdownReportTest.java
new file mode 100644
index 0000000000..4bb3d0abfb
--- /dev/null
+++ b/hugegraph-pd/hg-pd-test/src/main/java/org/apache/hugegraph/pd/service/StoreShutdownReportTest.java
@@ -0,0 +1,115 @@
+/*
+ * 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.util.List;
+import java.util.concurrent.atomic.AtomicReference;
+
+import org.apache.hugegraph.pd.StoreNodeService;
+import org.apache.hugegraph.pd.grpc.Metapb;
+import org.apache.hugegraph.pd.grpc.Pdpb;
+import org.apache.hugegraph.pd.meta.StoreInfoMeta;
+import org.apache.hugegraph.testutil.Whitebox;
+import org.junit.Assert;
+import org.junit.Before;
+import org.junit.Test;
+import org.mockito.Mockito;
+
+import io.grpc.stub.StreamObserver;
+
+public class StoreShutdownReportTest {
+
+    private StoreNodeService stores;
+    private StoreInfoMeta metadata;
+    private PDService service;
+    private Metapb.Store original;
+
+    @Before
+    public void setUp() throws Exception {
+        stores = Mockito.mock(StoreNodeService.class, Mockito.CALLS_REAL_METHODS);
+        metadata = Mockito.mock(StoreInfoMeta.class);
+        original = Metapb.Store.newBuilder().setId(7).setState(Metapb.StoreState.Up)
+                .setRaftAddress("127.0.0.1:8510")
+                .addLabels(Metapb.StoreLabel.newBuilder().setKey("zone").setValue("a")).build();
+        Mockito.when(metadata.getStore(7L)).thenReturn(original);
+        Whitebox.setInternalState(stores, "storeInfoMeta", metadata);
+        Whitebox.setInternalState(stores, "statusListeners", List.of());
+        Mockito.doReturn(Metapb.ClusterStats.newBuilder()
+                .setState(Metapb.ClusterState.Cluster_Not_Ready).build()).when(stores).getClusterStats();
+        service = new PDService() {
+            @Override
+            public boolean isLeader() {
+                return true;
+            }
+        };
+        Whitebox.setInternalState(service, "storeNodeService", stores);
+    }
+
+    @Test
+    public void testShutdownReportBypassesReadinessAndPreservesMetadataAndLease() throws Exception {
+        AtomicReference response = new AtomicReference<>();
+        service.reportStoreShutdown(request(), observer(response));
+        Assert.assertEquals(Pdpb.ErrorType.OK, response.get().getHeader().getError().getType());
+        Metapb.Store expected = original.toBuilder().setState(Metapb.StoreState.Offline).build();
+        Assert.assertEquals(expected, response.get().getStore());
+        Mockito.verify(metadata).updateStore(expected);
+        Mockito.verify(metadata, Mockito.never()).removeActiveStore(Mockito.any());
+        Mockito.verify(metadata, Mockito.never()).keepStoreAlive(Mockito.any());
+    }
+
+    @Test
+    public void testAdministrativeOfflineStillRequiresHealthyCluster() throws Exception {
+        AtomicReference response = new AtomicReference<>();
+        service.setStore(request(), observer(response));
+        Assert.assertNotEquals(Pdpb.ErrorType.OK, response.get().getHeader().getError().getType());
+        Mockito.verify(metadata, Mockito.never()).updateStore(Mockito.any());
+    }
+
+    @Test
+    public void testRegistrationReturnsAuthoritativeState() throws Exception {
+        Metapb.Store offline = original.toBuilder().setState(Metapb.StoreState.Offline).build();
+        Mockito.doReturn(offline).when(stores).register(Mockito.any());
+        AtomicReference response = new AtomicReference<>();
+        service.registerStore(Pdpb.RegisterStoreRequest.newBuilder().setStore(offline).build(), observer(response));
+        Assert.assertEquals(7L, response.get().getStoreId());
+        Assert.assertEquals(offline, response.get().getStore());
+    }
+
+    private Pdpb.SetStoreRequest request() {
+        return Pdpb.SetStoreRequest.newBuilder().setStore(Metapb.Store.newBuilder()
+                .setId(7).setState(Metapb.StoreState.Offline).setRaftAddress("ignored")).build();
+    }
+
+    private static  StreamObserver observer(AtomicReference response) {
+        return new StreamObserver<>() {
+            @Override
+            public void onNext(T value) {
+                response.set(value);
+            }
+
+            @Override
+            public void onError(Throwable failure) {
+                throw new AssertionError(failure);
+            }
+
+            @Override
+            public void onCompleted() {
+            }
+        };
+    }
+}
diff --git a/hugegraph-store/hg-store-core/src/main/java/org/apache/hugegraph/store/HeartbeatService.java b/hugegraph-store/hg-store-core/src/main/java/org/apache/hugegraph/store/HeartbeatService.java
index 921b5f8cf2..8e38109673 100644
--- a/hugegraph-store/hg-store-core/src/main/java/org/apache/hugegraph/store/HeartbeatService.java
+++ b/hugegraph-store/hg-store-core/src/main/java/org/apache/hugegraph/store/HeartbeatService.java
@@ -23,6 +23,8 @@
 import java.util.Collections;
 import java.util.List;
 import java.util.Map;
+import java.util.concurrent.ExecutionException;
+import java.util.concurrent.FutureTask;
 import java.util.concurrent.atomic.AtomicBoolean;
 
 import org.apache.hugegraph.pd.common.PDException;
@@ -59,7 +61,7 @@ public class HeartbeatService implements Lifecycle, Partit
     private HgStoreEngineOptions options;
     private PdProvider pdProvider;
     private Store storeInfo;
-    private Metapb.ClusterStats clusterStats;
+    private volatile Metapb.ClusterStats clusterStats;
     private StoreMetadata storeMetadata;
     private final List stateListeners;
     private final Object partitionThreadLock = new Object();
@@ -71,6 +73,9 @@ public class HeartbeatService implements Lifecycle, Partit
     private volatile int timerNextDelay = 1000;
     private volatile boolean terminated = false;
     private boolean restored = false;
+    private FutureTask restoreTask;
+    private Thread restoreThread;
+    private volatile boolean pulseConnected;
     private final AtomicBoolean fatalExitRequested = new AtomicBoolean();
     private Thread storeHeartbeatThread;
     private Thread partitionHeartbeatThread;
@@ -118,10 +123,12 @@ public void setStoreMetadata(StoreMetadata storeMetadata) {
 
     // Whether the cluster is ready
     public boolean isClusterReady() {
-        if (clusterStats == null) {
-            clusterStats = pdProvider.getClusterStats();
+        Metapb.ClusterStats stats = clusterStats;
+        if (stats == null) {
+            stats = pdProvider.getClusterStats();
+            clusterStats = stats;
         }
-        return clusterStats != null && clusterStats.getState() == Metapb.ClusterState.Cluster_OK;
+        return stats != null && stats.getState() == Metapb.ClusterState.Cluster_OK;
     }
 
     /**
@@ -215,11 +222,28 @@ protected void registerStore() {
                 storeInfo.setId(storeId);
                 storeMetadata.save(storeInfo);
                 if (restoring && storeEngine != null) {
+                    this.clusterStats = pdProvider.getClusterStats();
+                    timerNextDelay = REGISTER_RETRY_INTERVAL * 1000;
+                    if (restoreTask == null) {
+                        restoreTask = new FutureTask<>(() -> {
+                            storeEngine.stateChanged(storeInfo, Metapb.StoreState.Offline, Metapb.StoreState.Up);
+                            return null;
+                        });
+                        restoreThread = new Thread(restoreTask, "store-partition-restore");
+                        restoreThread.start();
+                        return;
+                    }
+                    if (!restoreTask.isDone()) {
+                        return;
+                    }
                     try {
-                        storeEngine.stateChanged(storeInfo, Metapb.StoreState.Offline, Metapb.StoreState.Up);
-                    } catch (RuntimeException e) {
+                        restoreTask.get();
+                    } catch (InterruptedException e) {
+                        Thread.currentThread().interrupt();
+                        return;
+                    } catch (ExecutionException e) {
                         log.error("Store {} remains Offline after partition restore failed", storeId, e);
-                        timerNextDelay = REGISTER_RETRY_INTERVAL * 1000;
+                        restoreTask = null;
                         return;
                     }
                 }
@@ -248,7 +272,8 @@ protected void registerStore() {
                 log.info("Register Store id= {} successfully. store = {}, clusterStats {}",
                          storeInfo.getId(), storeInfo, this.clusterStats);
                 // Listen to partition messages
-                pdProvider.startHeartbeatStream(error -> {
+                pulseConnected = pdProvider.startHeartbeatStream(error -> {
+                    pulseConnected = false;
                     onStateChanged(Metapb.StoreState.Offline);
                     timerNextDelay = REGISTER_RETRY_INTERVAL * 1000;
                     wakeupHeartbeatThread();
@@ -350,7 +375,7 @@ protected synchronized void onStateChanged(Metapb.StoreState newState) {
     }
 
     protected void partitionHeartbeat() {
-        if (storeEngine == null) {
+        if (storeEngine == null || !pulseConnected) {
             return;
         }
 
@@ -438,8 +463,9 @@ public void shutdown() {
             stopProducers();
             interrupted |= joinHeartbeat(storeHeartbeatThread);
             interrupted |= joinHeartbeat(partitionHeartbeatThread);
-            // Producers may have submitted callbacks that create partition owners.
-            // Never hold this monitor while joining the producer threads.
+            interrupted |= joinHeartbeat(restoreThread);
+            // Finish restoration and queued state notifications before partition owners close.
+            // Never hold this monitor while joining owned threads.
             synchronized (stateChangeLock) {
                 // Acquiring the monitor waits for any in-flight state callback to finish.
                 reportOffline();
@@ -461,13 +487,8 @@ private void reportOffline() {
             return;
         }
         try {
-            // Preserve PD's current metadata, and never use Exiting for a temporary restart.
-            var client = pdProvider.getPDClient();
-            client.getCache().removeStore(storeInfo.getId());
-            Metapb.Store current = client.getStore(storeInfo.getId());
-            if (current.getState() == Metapb.StoreState.Up) {
-                client.updateStore(current.toBuilder().setState(Metapb.StoreState.Offline).build());
-            }
+            // Older PDs reject this dedicated RPC; never fall back to administrative Offline.
+            pdProvider.getPDClient().reportStoreShutdown(storeInfo.getId());
         } catch (PDException | RuntimeException e) {
             // PD may be unavailable during cluster shutdown; heartbeat expiry remains the fallback.
             log.warn("Failed to report Store {} offline during shutdown", storeInfo.getId(), e);
diff --git a/hugegraph-store/hg-store-core/src/main/java/org/apache/hugegraph/store/pd/DefaultPdProvider.java b/hugegraph-store/hg-store-core/src/main/java/org/apache/hugegraph/store/pd/DefaultPdProvider.java
index b2fd24128d..a3a739ceca 100644
--- a/hugegraph-store/hg-store-core/src/main/java/org/apache/hugegraph/store/pd/DefaultPdProvider.java
+++ b/hugegraph-store/hg-store-core/src/main/java/org/apache/hugegraph/store/pd/DefaultPdProvider.java
@@ -138,11 +138,12 @@ public long registerStore(Store store) throws PDException {
                                                Metapb.StoreState.Up : Metapb.StoreState.Offline)
                                      .build();
         try {
-            storeId = pdClient.registerStore(protoObj);
+            Metapb.Store registered = pdClient.registerStoreWithState(protoObj);
+            storeId = registered.getId();
             store.setId(storeId);
-            if (pdClient.getStore(storeId).getState() != Metapb.StoreState.Up) {
+            if (registered.getState() == Metapb.StoreState.Pending) {
                 LOG.warn("Store {} is not activated, state is {}", storeId,
-                         pdClient.getStore(storeId).getState());
+                         registered.getState());
             }
         } catch (PDException e) {
             LOG.error(
diff --git a/hugegraph-store/hg-store-test/src/main/java/org/apache/hugegraph/store/HeartbeatShutdownTest.java b/hugegraph-store/hg-store-test/src/main/java/org/apache/hugegraph/store/HeartbeatShutdownTest.java
index bc40f20840..e5441d1ad5 100644
--- a/hugegraph-store/hg-store-test/src/main/java/org/apache/hugegraph/store/HeartbeatShutdownTest.java
+++ b/hugegraph-store/hg-store-test/src/main/java/org/apache/hugegraph/store/HeartbeatShutdownTest.java
@@ -216,7 +216,7 @@ private static void runInvalidShardRecovery(Path marker) throws Exception {
         Constructor constructor = HgStoreEngine.class.getDeclaredConstructor();
         constructor.setAccessible(true);
         HgStoreEngine engine = constructor.newInstance();
-        Fixture service = new Fixture(false);
+        HeartbeatService service = new HeartbeatService(engine);
         RpcServer rpc = Mockito.mock(RpcServer.class);
         HgMetricService metrics = Mockito.mock(HgMetricService.class);
         setField(HgStoreEngine.class, engine, "partitionManager", manager);
@@ -227,7 +227,6 @@ private static void runInvalidShardRecovery(Path marker) throws Exception {
         Class configClass = Class.forName("org.apache.hugegraph.store.node.AppConfig");
         Object node = nodeClass.getConstructor(configClass).newInstance((Object) null);
         setField(nodeClass, node, "storeEngine", engine);
-        service.addStateListener(engine);
         Runtime.getRuntime().addShutdownHook(new Thread(() -> {
             System.out.println("NODE_DESTROY_ENTERED");
             try {
@@ -240,9 +239,14 @@ private static void runInvalidShardRecovery(Path marker) throws Exception {
                 throw new AssertionError(e);
             }
         }, "recovery-fixture-shutdown"));
-        service.start();
-        Assert.assertTrue(service.entered.await(5, TimeUnit.SECONDS));
-        service.onStateChanged(Metapb.StoreState.Up);
+        Mockito.when(provider.registerStore(Mockito.any())).thenReturn(1L);
+        Mockito.when(provider.getClusterStats()).thenReturn(Metapb.ClusterStats.newBuilder()
+                .setState(Metapb.ClusterState.Cluster_OK).build());
+        options.setGrpcAddress("127.0.0.1:8500");
+        options.setRaftAddress("127.0.0.1:8510");
+        options.setPdProvider(provider);
+        service.setStoreMetadata(metadata);
+        service.init(options);
     }
 
     private static void runInterruptedEngineShutdown(Path marker) throws Exception {
diff --git a/hugegraph-store/hg-store-test/src/main/java/org/apache/hugegraph/store/StoreRegistrationTest.java b/hugegraph-store/hg-store-test/src/main/java/org/apache/hugegraph/store/StoreRegistrationTest.java
index 799c3100a9..314ed7826d 100644
--- a/hugegraph-store/hg-store-test/src/main/java/org/apache/hugegraph/store/StoreRegistrationTest.java
+++ b/hugegraph-store/hg-store-test/src/main/java/org/apache/hugegraph/store/StoreRegistrationTest.java
@@ -22,8 +22,15 @@
 import java.util.LinkedHashMap;
 import java.util.List;
 import java.util.Map;
+import java.util.concurrent.CountDownLatch;
+import java.util.concurrent.ExecutionException;
+import java.util.concurrent.ExecutorService;
+import java.util.concurrent.Executors;
+import java.util.concurrent.Future;
+import java.util.concurrent.FutureTask;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.TimeoutException;
 
-import org.apache.hugegraph.pd.client.ClientCache;
 import org.apache.hugegraph.pd.client.PDClient;
 import org.apache.hugegraph.pd.common.PDException;
 import org.apache.hugegraph.pd.grpc.Metapb;
@@ -32,10 +39,10 @@
 import org.apache.hugegraph.store.meta.Store;
 import org.apache.hugegraph.store.meta.StoreMetadata;
 import org.apache.hugegraph.store.options.HgStoreEngineOptions;
+import org.apache.hugegraph.store.pd.DefaultPdProvider;
 import org.apache.hugegraph.store.pd.PdProvider;
 import org.junit.Assert;
 import org.junit.Test;
-import org.mockito.ArgumentCaptor;
 import org.mockito.Mockito;
 
 public class StoreRegistrationTest {
@@ -56,8 +63,8 @@ public void testRestoreCompletesBeforeUpRegistration() throws Exception {
         }).when(engine).stateChanged(Mockito.any(), Mockito.any(), Mockito.eq(Metapb.StoreState.Up));
         HeartbeatService service = service(engine, provider);
         try {
-            service.registerStore();
-            Assert.assertEquals(List.of("Offline", "restore", "Up"), events);
+            registerRestored(service);
+            Assert.assertEquals(List.of("Offline", "restore", "Offline", "Up"), events);
         } finally {
             service.shutdown();
         }
@@ -73,7 +80,9 @@ public void testFailedRestoreNeverRegistersUp() throws Exception {
         HeartbeatService service = service(engine, provider);
         try {
             service.registerStore();
-            Mockito.verify(provider, Mockito.times(1)).registerStore(Mockito.any());
+            Assert.assertThrows(ExecutionException.class, () -> restoreTask(service).get(5, TimeUnit.SECONDS));
+            service.registerStore();
+            Mockito.verify(provider, Mockito.times(2)).registerStore(Mockito.any());
             Assert.assertEquals(Metapb.StoreState.Offline, service.getStoreInfo().getState());
         } finally {
             service.shutdown();
@@ -91,11 +100,22 @@ public void testReconnectRegistersUpWithoutRestoringAgain() throws Exception {
         });
         HeartbeatService service = service(engine, provider);
         try {
-            service.registerStore();
-            Assert.assertEquals(List.of("Offline", "Up"), events);
+            registerRestored(service);
+            Assert.assertEquals(List.of("Offline", "Offline", "Up"), events);
             events.clear();
+            CountDownLatch offline = new CountDownLatch(1);
+            CountDownLatch up = new CountDownLatch(1);
+            service.addStateListener((store, previous, state) -> {
+                if (state == Metapb.StoreState.Offline) {
+                    offline.countDown();
+                } else if (state == Metapb.StoreState.Up) {
+                    up.countDown();
+                }
+            });
             service.onStateChanged(Metapb.StoreState.Offline);
+            Assert.assertTrue(offline.await(5, TimeUnit.SECONDS));
             service.registerStore();
+            Assert.assertTrue(up.await(5, TimeUnit.SECONDS));
             Assert.assertEquals(List.of("Up"), events);
             Assert.assertEquals(Metapb.StoreState.Up, service.getStoreInfo().getState());
             Mockito.verify(engine).stateChanged(Mockito.any(), Mockito.any(), Mockito.eq(Metapb.StoreState.Up));
@@ -108,10 +128,12 @@ public void testReconnectRegistersUpWithoutRestoringAgain() throws Exception {
     public void testFailedUpRegistrationRetriesWithoutRestoringAgain() throws Exception {
         HgStoreEngine engine = Mockito.mock(HgStoreEngine.class);
         PdProvider provider = Mockito.mock(PdProvider.class);
-        Mockito.when(provider.registerStore(Mockito.any())).thenReturn(7L)
+        Mockito.when(provider.registerStore(Mockito.any())).thenReturn(7L, 7L)
                .thenThrow(new PDException(1, "unavailable")).thenReturn(7L);
         HeartbeatService service = service(engine, provider);
         try {
+            service.registerStore();
+            restoreTask(service).get(5, TimeUnit.SECONDS);
             service.registerStore();
             Assert.assertEquals(Metapb.StoreState.Offline, service.getStoreInfo().getState());
             service.registerStore();
@@ -161,19 +183,17 @@ public void testShutdownReportsOfflineAndPreservesMetadata() throws Exception {
         PdProvider provider = Mockito.mock(PdProvider.class);
         PDClient client = Mockito.mock(PDClient.class);
         Mockito.when(provider.getPDClient()).thenReturn(client);
-        Mockito.when(client.getCache()).thenReturn(Mockito.mock(ClientCache.class));
         Metapb.Store remote = Metapb.Store.newBuilder().setId(7).setState(Metapb.StoreState.Up)
                                        .setRaftAddress("127.0.0.1:8510")
                                        .addLabels(Metapb.StoreLabel.newBuilder()
                                                                  .setKey("zone").setValue("a")).build();
-        Mockito.when(client.getStore(7)).thenReturn(remote);
+        Mockito.when(client.reportStoreShutdown(7)).thenReturn(remote.toBuilder()
+                .setState(Metapb.StoreState.Offline).build());
         HeartbeatService service = service(null, provider);
         service.getStoreInfo().setId(7);
         service.shutdown();
-        ArgumentCaptor request = ArgumentCaptor.forClass(Metapb.Store.class);
-        Mockito.verify(client).updateStore(request.capture());
-        Assert.assertEquals(remote.toBuilder().setState(Metapb.StoreState.Offline).build(),
-                            request.getValue());
+        Mockito.verify(client).reportStoreShutdown(7);
+        Mockito.verify(client, Mockito.never()).updateStore(Mockito.any());
         Assert.assertEquals(Metapb.StoreState.Offline, service.getStoreInfo().getState());
     }
 
@@ -182,8 +202,7 @@ public void testShutdownContinuesWhenPdIsUnavailable() throws Exception {
         PdProvider provider = Mockito.mock(PdProvider.class);
         PDClient client = Mockito.mock(PDClient.class);
         Mockito.when(provider.getPDClient()).thenReturn(client);
-        Mockito.when(client.getCache()).thenReturn(Mockito.mock(ClientCache.class));
-        Mockito.when(client.getStore(7)).thenThrow(new PDException(1, "unavailable"));
+        Mockito.when(client.reportStoreShutdown(7)).thenThrow(new PDException(1, "unavailable"));
         HeartbeatService service = service(null, provider);
         service.getStoreInfo().setId(7);
         service.shutdown();
@@ -207,9 +226,126 @@ private static HeartbeatService service(HgStoreEngine engine, PdProvider provide
         return service;
     }
 
+    private static FutureTask restoreTask(HeartbeatService service) throws Exception {
+        Field field = HeartbeatService.class.getDeclaredField("restoreTask");
+        field.setAccessible(true);
+        return (FutureTask) field.get(service);
+    }
+
+    private static void registerRestored(HeartbeatService service) throws Exception {
+        CountDownLatch up = new CountDownLatch(1);
+        service.addStateListener((store, previous, state) -> {
+            if (state == Metapb.StoreState.Up) {
+                up.countDown();
+            }
+        });
+        service.registerStore();
+        restoreTask(service).get(5, TimeUnit.SECONDS);
+        service.registerStore();
+        Assert.assertTrue(up.await(5, TimeUnit.SECONDS));
+    }
+
+    @Test
+    public void testSlowRestoreKeepsRegisteringOfflineAndSkipsPartitionPulse() throws Exception {
+        HgStoreEngine engine = Mockito.mock(HgStoreEngine.class);
+        PdProvider provider = Mockito.mock(PdProvider.class);
+        Mockito.when(provider.registerStore(Mockito.any())).thenReturn(7L);
+        CountDownLatch entered = new CountDownLatch(1);
+        CountDownLatch release = new CountDownLatch(1);
+        Mockito.doAnswer(call -> {
+            entered.countDown();
+            Assert.assertTrue(release.await(5, TimeUnit.SECONDS));
+            return null;
+        }).when(engine).stateChanged(Mockito.any(), Mockito.any(), Mockito.any());
+        HeartbeatService service = service(engine, provider);
+        Mockito.when(provider.getClusterStats()).thenReturn(Metapb.ClusterStats.newBuilder()
+                .setState(Metapb.ClusterState.Cluster_OK).build(), Metapb.ClusterStats.newBuilder()
+                .setState(Metapb.ClusterState.Cluster_Not_Ready).build());
+        try {
+            service.registerStore();
+            Assert.assertTrue(entered.await(5, TimeUnit.SECONDS));
+            Assert.assertTrue(service.isClusterReady());
+            service.registerStore();
+            Assert.assertFalse(service.isClusterReady());
+            Assert.assertEquals(Metapb.StoreState.Offline, service.getStoreInfo().getState());
+            Mockito.verify(provider, Mockito.times(2)).registerStore(Mockito.argThat(
+                    store -> store.getState() == Metapb.StoreState.Offline));
+            service.partitionHeartbeat();
+            Mockito.verify(engine, Mockito.never()).getLeaderPartition();
+            Mockito.verify(provider, Mockito.never()).partitionHeartbeat(Mockito.anyList());
+        } finally {
+            release.countDown();
+            service.shutdown();
+        }
+    }
+
     private static void set(Object target, String name, Object value) throws Exception {
         Field field = target.getClass().getDeclaredField(name);
         field.setAccessible(true);
         field.set(target, value);
     }
+
+    @Test
+    public void testOlderPdShutdownDoesNotFallBackToAdministrativeOffline() throws Exception {
+        PdProvider provider = Mockito.mock(PdProvider.class);
+        PDClient client = Mockito.mock(PDClient.class);
+        Mockito.when(provider.getPDClient()).thenReturn(client);
+        Mockito.when(client.reportStoreShutdown(7)).thenThrow(io.grpc.Status.UNIMPLEMENTED.asRuntimeException());
+        HeartbeatService service = service(null, provider);
+        service.getStoreInfo().setId(7);
+        service.shutdown();
+        Mockito.verify(client, Mockito.never()).updateStore(Mockito.any());
+        Assert.assertEquals(Metapb.StoreState.Offline, service.getStoreInfo().getState());
+    }
+
+    @Test
+    public void testShutdownWaitsForRestoreWorker() throws Exception {
+        HgStoreEngine engine = Mockito.mock(HgStoreEngine.class);
+        PdProvider provider = Mockito.mock(PdProvider.class);
+        Mockito.when(provider.registerStore(Mockito.any())).thenReturn(7L);
+        CountDownLatch entered = new CountDownLatch(1);
+        CountDownLatch release = new CountDownLatch(1);
+        Mockito.doAnswer(call -> {
+            entered.countDown();
+            Assert.assertTrue(release.await(5, TimeUnit.SECONDS));
+            return null;
+        }).when(engine).stateChanged(Mockito.any(), Mockito.any(), Mockito.any());
+        HeartbeatService service = service(engine, provider);
+        ExecutorService closer = Executors.newSingleThreadExecutor();
+        try {
+            service.registerStore();
+            Assert.assertTrue(entered.await(5, TimeUnit.SECONDS));
+            Future closed = closer.submit(service::shutdown);
+            Assert.assertThrows(TimeoutException.class, () -> closed.get(100, TimeUnit.MILLISECONDS));
+            release.countDown();
+            closed.get(5, TimeUnit.SECONDS);
+            Assert.assertEquals(Metapb.StoreState.Offline, service.getStoreInfo().getState());
+        } finally {
+            release.countDown();
+            service.shutdown();
+            closer.shutdownNow();
+        }
+    }
+
+    @Test
+    public void testProviderUsesRegistrationStateWithoutCachedLookup() throws Exception {
+        DefaultPdProvider provider = Mockito.mock(DefaultPdProvider.class, Mockito.CALLS_REAL_METHODS);
+        PDClient client = Mockito.mock(PDClient.class);
+        Field field = DefaultPdProvider.class.getDeclaredField("pdClient");
+        field.setAccessible(true);
+        field.set(provider, client);
+        Metapb.Store offline = Metapb.Store.newBuilder().setId(7).setState(Metapb.StoreState.Offline).build();
+        Mockito.when(client.registerStoreWithState(Mockito.any())).thenReturn(offline,
+                offline.toBuilder().setState(Metapb.StoreState.Up).build());
+        Store store = new Store();
+        store.setStoreAddress("127.0.0.1:8500");
+        store.setRaftAddress("127.0.0.1:8510");
+        store.setDeployPath("");
+        store.setDataPath("");
+        store.setState(Metapb.StoreState.Offline);
+        Assert.assertEquals(7L, provider.registerStore(store));
+        store.setState(Metapb.StoreState.Up);
+        Assert.assertEquals(7L, provider.registerStore(store));
+        Mockito.verify(client, Mockito.never()).getStore(Mockito.anyLong());
+    }
 }

From ebdc83863e384e2f0b899d0c64c766ec29ae66c4 Mon Sep 17 00:00:00 2001
From: pallavithegod 
Date: Mon, 19 Oct 2026 23:01:06 +0530
Subject: [PATCH 4/4] fix(store): separate leases from placement

Exclude restoring Stores from new placement while retaining their leases.

Require a privileged internal credential for shutdown reports.

Back off restore retries and throttle renewal and repeated traces.

Cover placement, gRPC credentials, lifecycle and deployment references.
---
 docker/README.md                              |   6 +
 docker/docker-compose-3pd-3store-3server.yml  |   3 +
 docker/docker-compose-hstore.yml              |   1 +
 helm/hugegraph/docs/operations.md             |   6 +
 .../templates/store-statefulset.yaml          |   5 +
 .../tests/credential_contract_test.yaml       |  15 ++
 hugegraph-pd/docs/api-reference.md            |  16 ++
 .../apache/hugegraph/pd/PartitionService.java |   2 +-
 .../apache/hugegraph/pd/StoreNodeService.java |  31 +++-
 .../hugegraph/pd/TaskScheduleService.java     |   8 +-
 .../hugegraph/pd/service/PDService.java       |   4 +-
 .../StoreShutdownAuthentication.java          |  56 +++++++
 .../pd/util/grpc/GRpcServerConfig.java        |   5 +
 .../hugegraph/pd/core/PDCoreSuiteTest.java    |   2 +
 .../pd/core/StoreRestoreRegistrationTest.java |  90 ++++++++++++
 .../StoreShutdownAuthenticationTest.java      | 138 ++++++++++++++++++
 .../hugegraph/store/HeartbeatService.java     |  36 ++++-
 .../hugegraph/store/pd/DefaultPdProvider.java |   2 +-
 .../store/StoreRegistrationTest.java          |  65 ++++++++-
 19 files changed, 470 insertions(+), 21 deletions(-)
 create mode 100644 hugegraph-pd/hg-pd-service/src/main/java/org/apache/hugegraph/pd/service/interceptor/StoreShutdownAuthentication.java
 create mode 100644 hugegraph-pd/hg-pd-test/src/main/java/org/apache/hugegraph/pd/service/StoreShutdownAuthenticationTest.java

diff --git a/docker/README.md b/docker/README.md
index 2b6554abd4..440b43109e 100644
--- a/docker/README.md
+++ b/docker/README.md
@@ -264,6 +264,12 @@ from a recent heartbeat without waiting for keep-alive expiry. A future timestam
 caused by a PD clock adjustment reports zero age.
 Keep the cluster preset's manual rollout barrier and verify fresh heartbeats,
 membership and leaders before replacing another Store.
+The shutdown report succeeds only if PD accepts a privileged internal credential:
+set `HG_PD_AUTH_SECRET_KEY` in Store to match PD's `auth.secret-key`. Without it,
+or if the RPC fails, shutdown continues and the existing lease expires normally.
+Restoring Offline Stores retain their liveness lease but cannot receive new
+placements until `Up`; recovery retries back off to the heartbeat interval.
+
 
 ### Select image versions
 
diff --git a/docker/docker-compose-3pd-3store-3server.yml b/docker/docker-compose-3pd-3store-3server.yml
index 999f2126fb..7b3e25ad8e 100644
--- a/docker/docker-compose-3pd-3store-3server.yml
+++ b/docker/docker-compose-3pd-3store-3server.yml
@@ -161,6 +161,7 @@ services:
     container_name: hg-store0
     hostname: store0
     environment:
+      HG_PD_AUTH_SECRET_KEY: ${HG_PD_AUTH_SECRET_KEY:?set HG_PD_AUTH_SECRET_KEY in .env; see docker/README.md}
       HG_STORE_PD_ADDRESS: pd0:8686,pd1:8686,pd2:8686
       HG_STORE_GRPC_HOST: store0
       HG_STORE_GRPC_PORT: "8500"
@@ -176,6 +177,7 @@ services:
     container_name: hg-store1
     hostname: store1
     environment:
+      HG_PD_AUTH_SECRET_KEY: ${HG_PD_AUTH_SECRET_KEY:?set HG_PD_AUTH_SECRET_KEY in .env; see docker/README.md}
       HG_STORE_PD_ADDRESS: pd0:8686,pd1:8686,pd2:8686
       HG_STORE_GRPC_HOST: store1
       HG_STORE_GRPC_PORT: "8500"
@@ -191,6 +193,7 @@ services:
     container_name: hg-store2
     hostname: store2
     environment:
+      HG_PD_AUTH_SECRET_KEY: ${HG_PD_AUTH_SECRET_KEY:?set HG_PD_AUTH_SECRET_KEY in .env; see docker/README.md}
       HG_STORE_PD_ADDRESS: pd0:8686,pd1:8686,pd2:8686
       HG_STORE_GRPC_HOST: store2
       HG_STORE_GRPC_PORT: "8500"
diff --git a/docker/docker-compose-hstore.yml b/docker/docker-compose-hstore.yml
index 0fcc3697f6..bab5ef832d 100644
--- a/docker/docker-compose-hstore.yml
+++ b/docker/docker-compose-hstore.yml
@@ -60,6 +60,7 @@ services:
       pd:
         condition: service_healthy
     environment:
+      HG_PD_AUTH_SECRET_KEY: ${HG_PD_AUTH_SECRET_KEY:?set HG_PD_AUTH_SECRET_KEY in .env; see docker/README.md}
       HG_STORE_PD_ADDRESS: pd:8686
       HG_STORE_GRPC_HOST: store
       HG_STORE_GRPC_PORT: "8500"
diff --git a/helm/hugegraph/docs/operations.md b/helm/hugegraph/docs/operations.md
index e940a2d102..cd10736185 100644
--- a/helm/hugegraph/docs/operations.md
+++ b/helm/hugegraph/docs/operations.md
@@ -210,6 +210,12 @@ using the administrative Offline path, so a PD rollback does not remove that lea
 These signals do not prove Raft catch-up.
 
+The shutdown report succeeds only if PD accepts a privileged internal credential:
+set `HG_PD_AUTH_SECRET_KEY` in Store to match PD's `auth.secret-key`. Without it,
+or if the RPC fails, shutdown continues and the existing lease expires normally.
+Restoring Offline Stores retain their liveness lease but cannot receive new
+placements until `Up`; recovery retries back off to the heartbeat interval.
+
 
 ### PD patrol and leader balance
 
diff --git a/helm/hugegraph/templates/store-statefulset.yaml b/helm/hugegraph/templates/store-statefulset.yaml
index 81c143ee8a..1fb1e5ebd6 100644
--- a/helm/hugegraph/templates/store-statefulset.yaml
+++ b/helm/hugegraph/templates/store-statefulset.yaml
@@ -153,6 +153,11 @@ spec:
               valueFrom:
                 fieldRef:
                   fieldPath: metadata.namespace
+            - name: HG_PD_AUTH_SECRET_KEY
+              valueFrom:
+                secretKeyRef:
+                  name: {{ include "hugegraph.pd.authSecretName" . | quote }}
+                  key: {{ include "hugegraph.pd.authSecretKey" . | quote }}
             - name: HG_STORE_PD_ADDRESS
               value: {{ include "hugegraph.pd.grpcPeersList" . | quote }}
             - name: HG_STORE_GRPC_HOST
diff --git a/helm/hugegraph/tests/credential_contract_test.yaml b/helm/hugegraph/tests/credential_contract_test.yaml
index 5fbddb830b..16d5820f89 100644
--- a/helm/hugegraph/tests/credential_contract_test.yaml
+++ b/helm/hugegraph/tests/credential_contract_test.yaml
@@ -131,3 +131,18 @@ tests:
               secretKeyRef:
                 name: provided-admin
                 key: password
+
+  - it: gives Store the configured PD credential for privileged shutdown reports
+    template: store-statefulset.yaml
+    set:
+      pd.auth.existingSecret: provided-pd
+      pd.auth.key: shutdown-key
+    asserts:
+      - contains:
+          path: spec.template.spec.containers[0].env
+          content:
+            name: HG_PD_AUTH_SECRET_KEY
+            valueFrom:
+              secretKeyRef:
+                name: provided-pd
+                key: shutdown-key
diff --git a/hugegraph-pd/docs/api-reference.md b/hugegraph-pd/docs/api-reference.md
index 4004766304..e9dbc953ae 100644
--- a/hugegraph-pd/docs/api-reference.md
+++ b/hugegraph-pd/docs/api-reference.md
@@ -219,6 +219,22 @@ when the cluster is not ready. Older PDs return `UNIMPLEMENTED`; updated Stores
 fall back to expiry, never to `SetStore`. Administrative `SetStore(Offline)` keeps
 its cluster-readiness guard and removes the active entry as before.
 
+Shutdown reporting requires a privileged internal Basic credential matching PD's
+`auth.secret-key`. Set the same `HG_PD_AUTH_SECRET_KEY` in each Store process;
+this shared credential authorizes internal callers across Store IDs, rather than
+identifying an individual Store. Keep it restricted to trusted internal services.
+Missing, invalid or unconfigured credentials refuse the RPC. Store still shuts
+down and PD retains its previous record until the existing lease expires if the
+report is refused, unreachable, or unsupported by an older PD.
+
+A live restoration lease is not placement readiness: new allocation, replica
+addition, balancing and capacity checks select only Stores whose canonical state
+is `Up`. Existing shard membership and liveness/expiry tracking remain separate.
+Offline renewal occurs at most every five seconds while restoration is polled;
+failed restore attempts back off from one second to the Store heartbeat interval,
+with a full failure trace only on the first or a changed partition failure set.
+
+
 **Request**:
 ```protobuf
 message RegisterStoreRequest {
diff --git a/hugegraph-pd/hg-pd-core/src/main/java/org/apache/hugegraph/pd/PartitionService.java b/hugegraph-pd/hg-pd-core/src/main/java/org/apache/hugegraph/pd/PartitionService.java
index 8d39006d45..511a1af8c7 100644
--- a/hugegraph-pd/hg-pd-core/src/main/java/org/apache/hugegraph/pd/PartitionService.java
+++ b/hugegraph-pd/hg-pd-core/src/main/java/org/apache/hugegraph/pd/PartitionService.java
@@ -648,7 +648,7 @@ public synchronized void splitPartition(Metapb.Graph graph, int toCount) throws
         var maxShardsPerStore = pdConfig.getPartition().getMaxShardsPerStore();
         var shardCount = pdConfig.getPartition().getShardCount();
 
-        if (shardCount * toCount > storeService.getActiveStores().size() * maxShardsPerStore) {
+        if (shardCount * toCount > storeService.getPlacementStores().size() * maxShardsPerStore) {
             throw new PDException(Pdpb.ErrorType.Too_Many_Partitions_Per_Store_VALUE,
                                   "can't satisfy target shard group count, reached the upper " +
                                   "limit of the cluster");
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 a2ecae6ede..b6d86904d5 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
@@ -430,6 +430,21 @@ public List getActiveStores() throws PDException {
         return storeInfoMeta.getActiveStores();
     }
 
+    /**
+     * Placement requires both a live lease and the current canonical Up state.
+     * Shutdown reports deliberately retain the old TTL entry until expiry.
+     */
+    public List getPlacementStores() throws PDException {
+        List stores = new ArrayList<>();
+        for (Metapb.Store active : getActiveStores()) {
+            Metapb.Store current = storeInfoMeta.getStore(active.getId());
+            if (current != null && current.getState() == Metapb.StoreState.Up) {
+                stores.add(current);
+            }
+        }
+        return stores;
+    }
+
     public List getTombStores() throws PDException {
         List stores = new ArrayList<>();
         for (Metapb.Store store : this.getStores()) {
@@ -458,7 +473,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,
@@ -518,7 +533,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,
@@ -532,7 +547,8 @@ public synchronized List reallocShards(Metapb.ShardGroup shardGrou
         }
 
         int shardCount = pdConfig.getPartition().getShardCount();
-        shardCount = Math.min(shardCount, stores.size());
+        // A leased but restoring member must not make an existing group shrink.
+        shardCount = Math.min(shardCount, getActiveStores().size());
         if (shardCount == 2 || shardCount < 1) {
             // Two shards could not elect a leader
             // It cannot be 0
@@ -543,6 +559,11 @@ public synchronized List reallocShards(Metapb.ShardGroup shardGrou
         shards.addAll(shardGroup.getShardsList());
 
         if (shardCount > shards.size()) {
+            long candidates = stores.stream().filter(store -> !isStoreInShards(shards, store.getId())).count();
+            if (candidates < shardCount - shards.size()) {
+                throw new PDException(Pdpb.ErrorType.LESS_ACTIVE_STORE_VALUE,
+                                      "Not enough restored stores to add replicas");
+            }
             // Need to add shards
             log.info("reallocShards ShardGroup {}, add shards from {} to {}",
                      shardGroup.getId(), shards.size(), shardCount);
@@ -604,7 +625,7 @@ public synchronized int splitShardGroups(List> groups)
                                                                                     PDException {
         int sum = groups.stream().map(pair -> pair.getValue()).reduce(0, Integer::sum);
         // shard group is too big
-        if (sum > getActiveStores().size() * pdConfig.getPartition().getMaxShardsPerStore()) {
+        if (sum > getPlacementStores().size() * pdConfig.getPartition().getMaxShardsPerStore()) {
             throw new PDException(Pdpb.ErrorType.Too_Many_Partitions_Per_Store_VALUE,
                                   "can't satisfy target shard group count");
         }
@@ -942,7 +963,7 @@ protected void onShardGroupOp(Metapb.ShardGroup shardGroup) {
     public boolean checkStoreCanOffline(Metapb.Store currentStore) {
         try {
             long currentStoreId = currentStore.getId();
-            List activeStores = this.getActiveStores();
+            List activeStores = this.getPlacementStores();
             Map storeMap = new HashMap<>();
             activeStores.forEach(store -> {
                 if (store.getId() != currentStoreId) {
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..97c51684e0 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,7 @@ public synchronized Map> balancePartitionShard() thr
             return null;
         }
 
-        int activeStores = storeService.getActiveStores().size();
+        int activeStores = storeService.getPlacementStores().size();
         if (activeStores == 0) {
             log.warn("balancePartitionShard non active stores, skip to balancePartitionShard");
             return null;
@@ -318,7 +318,7 @@ public synchronized Map> balancePartitionShard() thr
 
         // Count the partitions on each store, StoreId -> PartitionID, ShardRole
         Map> partitionMap = new HashMap<>();
-        storeService.getActiveStores().forEach(store -> {
+        storeService.getPlacementStores().forEach(store -> {
             partitionMap.put(store.getId(), new HashMap<>());
         });
 
@@ -651,7 +651,7 @@ public List autoSplitPartition() throws PDException {
 
         // The maximum split count that a compute cluster can support
         int splitCount = pdConfig.getPartition().getMaxShardsPerStore() *
-                         storeService.getActiveStores().size() /
+                         storeService.getPlacementStores().size() /
                          (storeService.getShardGroups().size() *
                           pdConfig.getPartition().getShardCount());
 
@@ -750,7 +750,7 @@ public Map canAllPartitionsMovedOut(Metapb.Store sourceStore) th
         // Record the amount of data in the partition to be migrated
         Map partitionDataSize = new HashMap<>();
 
-        storeService.getActiveStores().forEach(store -> {
+        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
diff --git a/hugegraph-pd/hg-pd-service/src/main/java/org/apache/hugegraph/pd/service/PDService.java b/hugegraph-pd/hg-pd-service/src/main/java/org/apache/hugegraph/pd/service/PDService.java
index 55d0c0ea3c..1e8237a705 100644
--- a/hugegraph-pd/hg-pd-service/src/main/java/org/apache/hugegraph/pd/service/PDService.java
+++ b/hugegraph-pd/hg-pd-service/src/main/java/org/apache/hugegraph/pd/service/PDService.java
@@ -1483,7 +1483,7 @@ public void delStore(Pdpb.DetStoreRequest request,
     private boolean checkShardCount(int newShardCount) {
         try {
             var maxCount = pdConfig.getPartition().getMaxShardsPerStore() *
-                           storeNodeService.getActiveStores().size() /
+                           storeNodeService.getPlacementStores().size() /
                            pdConfig.getConfigService().getPartitionCount();
 
             if (newShardCount > maxCount) {
@@ -1521,7 +1521,7 @@ public boolean isResourceEnough(int oldShardCount, int newShardCount) {
             // Estimate the storage space consumed after data bloat
             newDataSize = (long) Math.ceil(currentDataSize * expansionRatio);
             // Count the available space in all active stores
-            List activeStores = storeNodeService.getActiveStores();
+            List activeStores = storeNodeService.getPlacementStores();
             for (Metapb.Store store : activeStores) {
                 Metapb.StoreStats storeStats = store.getStats();
                 totalAvaible += storeStats.getAvailable();
diff --git a/hugegraph-pd/hg-pd-service/src/main/java/org/apache/hugegraph/pd/service/interceptor/StoreShutdownAuthentication.java b/hugegraph-pd/hg-pd-service/src/main/java/org/apache/hugegraph/pd/service/interceptor/StoreShutdownAuthentication.java
new file mode 100644
index 0000000000..314a0d6dff
--- /dev/null
+++ b/hugegraph-pd/hg-pd-service/src/main/java/org/apache/hugegraph/pd/service/interceptor/StoreShutdownAuthentication.java
@@ -0,0 +1,56 @@
+/*
+ * 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.interceptor;
+
+import org.apache.hugegraph.pd.common.Consts;
+import org.apache.hugegraph.pd.grpc.PDGrpc;
+import org.springframework.stereotype.Component;
+
+import io.grpc.Metadata;
+import io.grpc.ServerCall;
+import io.grpc.ServerCallHandler;
+import io.grpc.ServerInterceptor;
+import io.grpc.Status;
+
+/**
+ * Require a privileged internal credential for shutdown reports. Existing RPCs
+ * retain their transport policy until the general gRPC authentication migration.
+ */
+@Component
+public class StoreShutdownAuthentication extends Authentication implements ServerInterceptor {
+
+    @Override
+    public  ServerCall.Listener interceptCall(
+            ServerCall call, Metadata headers,
+            ServerCallHandler next) {
+        if (!PDGrpc.getReportStoreShutdownMethod().getFullMethodName()
+                   .equals(call.getMethodDescriptor().getFullMethodName())) {
+            return next.startCall(call, headers);
+        }
+        try {
+            // Holders of auth.secret-key are privileged internal callers, not
+            // individual Store identities. Validate on followers before redirect.
+            authenticate(headers.get(Consts.CREDENTIAL_KEY), null, token -> true, () -> true);
+        } catch (RuntimeException denied) {
+            call.close(Status.UNAUTHENTICATED.withDescription("Shutdown report requires an internal credential"),
+                       new Metadata());
+            return new ServerCall.Listener<>() { };
+        }
+        return next.startCall(call, headers);
+    }
+}
diff --git a/hugegraph-pd/hg-pd-service/src/main/java/org/apache/hugegraph/pd/util/grpc/GRpcServerConfig.java b/hugegraph-pd/hg-pd-service/src/main/java/org/apache/hugegraph/pd/util/grpc/GRpcServerConfig.java
index 224d417cc4..173de981a8 100644
--- a/hugegraph-pd/hg-pd-service/src/main/java/org/apache/hugegraph/pd/util/grpc/GRpcServerConfig.java
+++ b/hugegraph-pd/hg-pd-service/src/main/java/org/apache/hugegraph/pd/util/grpc/GRpcServerConfig.java
@@ -18,6 +18,7 @@
 package org.apache.hugegraph.pd.util.grpc;
 
 import org.apache.hugegraph.pd.config.PDConfig;
+import org.apache.hugegraph.pd.service.interceptor.StoreShutdownAuthentication;
 import org.apache.hugegraph.pd.util.HgExecutorUtil;
 import org.lognet.springboot.grpc.GRpcServerBuilderConfigurer;
 import org.springframework.beans.factory.annotation.Autowired;
@@ -33,6 +34,9 @@ public class GRpcServerConfig extends GRpcServerBuilderConfigurer {
     @Autowired
     private PDConfig pdConfig;
 
+    @Autowired
+    private StoreShutdownAuthentication storeShutdownAuthentication;
+
     @Override
     public void configure(ServerBuilder serverBuilder) {
         PDConfig.ThreadPoolGrpc poolGrpc = pdConfig.getThreadPoolGrpc();
@@ -40,6 +44,7 @@ public void configure(ServerBuilder serverBuilder) {
                 HgExecutorUtil.createExecutor(EXECUTOR_NAME, poolGrpc.getCore(), poolGrpc.getMax(),
                                               poolGrpc.getQueue()));
         serverBuilder.maxInboundMessageSize(MAX_INBOUND_MESSAGE_SIZE);
+        serverBuilder.intercept(storeShutdownAuthentication);
         // TODO: GrpcAuthentication is instantiated as a Spring bean but never registered
         // here - add serverBuilder.intercept(grpcAuthentication) once auth is refactored.
         // It extends Authentication, which now also checks the Basic password against
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 356a7b9b7a..4835c2cedb 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
@@ -29,6 +29,7 @@
 import org.apache.hugegraph.pd.raft.RaftReflectionUtilTest;
 import org.apache.hugegraph.pd.raft.ZipUtilsTest;
 import org.apache.hugegraph.pd.service.StoreShutdownReportTest;
+import org.apache.hugegraph.pd.service.StoreShutdownAuthenticationTest;
 import org.junit.runner.RunWith;
 import org.junit.runners.Suite;
 
@@ -50,6 +51,7 @@
         StoreServiceTest.class,
         StoreRestoreRegistrationTest.class,
         StoreShutdownReportTest.class,
+        StoreShutdownAuthenticationTest.class,
         TaskScheduleServiceTest.class,
         IpAuthHandlerTest.class,
         RaftEngineIpAuthIntegrationTest.class,
diff --git a/hugegraph-pd/hg-pd-test/src/main/java/org/apache/hugegraph/pd/core/StoreRestoreRegistrationTest.java b/hugegraph-pd/hg-pd-test/src/main/java/org/apache/hugegraph/pd/core/StoreRestoreRegistrationTest.java
index 85f38f9b89..dc23c54ddb 100644
--- a/hugegraph-pd/hg-pd-test/src/main/java/org/apache/hugegraph/pd/core/StoreRestoreRegistrationTest.java
+++ b/hugegraph-pd/hg-pd-test/src/main/java/org/apache/hugegraph/pd/core/StoreRestoreRegistrationTest.java
@@ -22,8 +22,13 @@
 import java.util.List;
 import java.util.Map;
 import java.util.concurrent.atomic.AtomicReference;
+import java.util.concurrent.atomic.AtomicLong;
 
 import org.apache.hugegraph.pd.StoreNodeService;
+import org.apache.hugegraph.pd.ConfigService;
+import org.apache.hugegraph.pd.PartitionService;
+import org.apache.hugegraph.pd.common.PDException;
+import org.apache.hugegraph.pd.grpc.Pdpb;
 import org.apache.hugegraph.pd.config.PDConfig;
 import org.apache.hugegraph.pd.grpc.Metapb;
 import org.apache.hugegraph.pd.meta.StoreInfoMeta;
@@ -50,6 +55,7 @@ public void setUp() throws Exception {
         Whitebox.setInternalState(service, "pdConfig", config);
         Whitebox.setInternalState(service, "storeInfoMeta", metadata);
         Whitebox.setInternalState(service, "statusListeners", List.of());
+        Whitebox.setInternalState(service, "shardGroupStatusListeners", List.of());
         Mockito.doNothing().when(service).checkStoreStatus();
         Mockito.when(metadata.storeExists(7L)).thenReturn(true);
         Mockito.when(metadata.getStore(7L)).thenAnswer(call -> current.get());
@@ -171,6 +177,90 @@ public void testMetadataOnlyUpdateDoesNotRemoveRestoringLease() throws Exception
         Mockito.verify(metadata, Mockito.never()).removeActiveStore(Mockito.any());
     }
 
+    @Test
+    public void testLiveRestoreLeaseDoesNotPermitPlacementUntilUp() throws Exception {
+        Mockito.when(config.getMinStoreCount()).thenReturn(3);
+        PDConfig.Partition partition = config.new Partition();
+        partition.setShardCount(3);
+        Mockito.when(config.getPartition()).thenReturn(partition);
+        ConfigService configuration = Mockito.mock(ConfigService.class);
+        Mockito.when(config.getConfigService()).thenReturn(configuration);
+        Mockito.when(configuration.getPartitionCount()).thenReturn(1);
+        PartitionService partitions = Mockito.mock(PartitionService.class);
+        Whitebox.setInternalState(service, "partitionService", partitions);
+        Metapb.Store second = store(Metapb.StoreState.Up).toBuilder().setId(8).build();
+        Metapb.Store third = second.toBuilder().setId(9).build();
+        // The retained shutdown lease can still contain the previous Up snapshot.
+        AtomicLong clock = new AtomicLong();
+        AtomicLong leaseExpiry = new AtomicLong(300_000L);
+        Mockito.when(metadata.getActiveStores()).thenAnswer(call -> clock.get() < leaseExpiry.get() ?
+                List.of(store(Metapb.StoreState.Up), second, third) : List.of(second, third));
+        Mockito.doAnswer(call -> {
+            leaseExpiry.set(clock.get() + 300_000L);
+            return null;
+        }).when(metadata).keepStoreAlive(Mockito.any());
+        Mockito.when(metadata.getStore(8L)).thenReturn(second);
+        Mockito.when(metadata.getStore(9L)).thenReturn(third);
+        AtomicReference group = new AtomicReference<>();
+        Mockito.when(metadata.getShardGroup(0)).thenAnswer(call -> group.get());
+        Mockito.doAnswer(call -> {
+            group.set(call.getArgument(0));
+            return null;
+        }).when(metadata).updateShardGroup(Mockito.any());
+        Metapb.ShardGroup underReplicated = Metapb.ShardGroup.newBuilder().setId(0)
+                .addShards(Metapb.Shard.newBuilder().setStoreId(8).setRole(Metapb.ShardRole.Leader))
+                .addShards(Metapb.Shard.newBuilder().setStoreId(9)).build();
+
+        // Controlled metadata clock: recovery lasts longer than the original five-minute lease.
+        for (int renewal = 0; renewal < 4; renewal++) {
+            clock.addAndGet(200_000L);
+            service.register(store(Metapb.StoreState.Offline));
+            Assert.assertEquals(3, service.getActiveStores().size());
+            Assert.assertEquals(2, service.getPlacementStores().size());
+            PDException allocation = Assert.assertThrows(PDException.class, () -> service.allocShards(null, 0));
+            Assert.assertEquals(Pdpb.ErrorType.LESS_ACTIVE_STORE_VALUE, allocation.getErrorCode());
+            Assert.assertThrows(PDException.class, () -> service.reallocShards(underReplicated));
+        }
+        Mockito.verify(metadata, Mockito.never()).updateShardGroup(Mockito.any());
+        Mockito.verify(metadata, Mockito.times(4)).keepStoreAlive(Mockito.any());
+
+        service.register(store(Metapb.StoreState.Up));
+        Assert.assertEquals(3, service.getPlacementStores().size());
+        Assert.assertTrue(service.reallocShards(underReplicated).stream().anyMatch(shard -> shard.getStoreId() == 7));
+        group.set(null);
+        Assert.assertTrue(service.allocShards(null, 0).stream().anyMatch(shard -> shard.getStoreId() == 7));
+    }
+
+    @Test
+    public void testTemporaryOfflineDoesNotShrinkExistingGroup() throws Exception {
+        Mockito.when(config.getMinStoreCount()).thenReturn(1);
+        PDConfig.Partition partition = config.new Partition();
+        partition.setShardCount(3);
+        Mockito.when(config.getPartition()).thenReturn(partition);
+        current.set(store(Metapb.StoreState.Offline));
+        Metapb.Store second = store(Metapb.StoreState.Up).toBuilder().setId(8).build();
+        Mockito.when(metadata.getStore(8L)).thenReturn(second);
+        Metapb.Store third = second.toBuilder().setId(9).build();
+        Mockito.when(metadata.getStore(9L)).thenReturn(third);
+        Mockito.when(metadata.getActiveStores()).thenReturn(List.of(store(Metapb.StoreState.Up), second, third));
+        Metapb.ShardGroup group = Metapb.ShardGroup.newBuilder().setId(0)
+                .addShards(Metapb.Shard.newBuilder().setStoreId(7))
+                .addShards(Metapb.Shard.newBuilder().setStoreId(8))
+                .addShards(Metapb.Shard.newBuilder().setStoreId(9).setRole(Metapb.ShardRole.Leader)).build();
+        Assert.assertEquals(group.getShardsList(), service.reallocShards(group));
+        Mockito.verify(metadata, Mockito.never()).updateShardGroup(Mockito.any());
+    }
+
+    @Test
+    public void testOnlyRestoringStoreCannotReceiveInitialLeader() throws Exception {
+        Mockito.when(config.getMinStoreCount()).thenReturn(1);
+        current.set(store(Metapb.StoreState.Offline));
+        Mockito.when(metadata.getActiveStores()).thenReturn(List.of(store(Metapb.StoreState.Up)));
+        PDException failure = Assert.assertThrows(PDException.class, () -> service.allocShards(null, 0));
+        Assert.assertEquals(Pdpb.ErrorType.NO_ACTIVE_STORE_VALUE, failure.getErrorCode());
+        Mockito.verify(metadata, Mockito.never()).updateShardGroup(Mockito.any());
+    }
+
     private static Metapb.Store store(Metapb.StoreState state) {
         return Metapb.Store.newBuilder().setId(7).setAddress("127.0.0.1:8500")
                            .setRaftAddress("127.0.0.1:8510").setState(state).build();
diff --git a/hugegraph-pd/hg-pd-test/src/main/java/org/apache/hugegraph/pd/service/StoreShutdownAuthenticationTest.java b/hugegraph-pd/hg-pd-test/src/main/java/org/apache/hugegraph/pd/service/StoreShutdownAuthenticationTest.java
new file mode 100644
index 0000000000..d41279a708
--- /dev/null
+++ b/hugegraph-pd/hg-pd-test/src/main/java/org/apache/hugegraph/pd/service/StoreShutdownAuthenticationTest.java
@@ -0,0 +1,138 @@
+/*
+ * 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.nio.charset.StandardCharsets;
+import java.util.Base64;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicInteger;
+
+import org.apache.hugegraph.pd.common.Consts;
+import org.apache.hugegraph.pd.config.PDConfig;
+import org.apache.hugegraph.pd.grpc.PDGrpc;
+import org.apache.hugegraph.pd.grpc.Pdpb;
+import org.apache.hugegraph.pd.service.interceptor.StoreShutdownAuthentication;
+import org.apache.hugegraph.testutil.Whitebox;
+import org.junit.After;
+import org.junit.Assert;
+import org.junit.Before;
+import org.junit.Test;
+import org.mockito.Mockito;
+
+import io.grpc.ManagedChannel;
+import io.grpc.ManagedChannelBuilder;
+import io.grpc.Metadata;
+import io.grpc.Server;
+import io.grpc.ServerBuilder;
+import io.grpc.Status;
+import io.grpc.StatusRuntimeException;
+import io.grpc.stub.MetadataUtils;
+import io.grpc.stub.StreamObserver;
+
+public class StoreShutdownAuthenticationTest {
+
+    private PDConfig config;
+    private Server server;
+    private ManagedChannel channel;
+    private final AtomicInteger shutdownCalls = new AtomicInteger();
+
+    @Before
+    public void setUp() throws Exception {
+        config = Mockito.mock(PDConfig.class);
+        Mockito.when(config.getSecretKey()).thenReturn("test-only-secret");
+        StoreShutdownAuthentication authentication = new StoreShutdownAuthentication();
+        Whitebox.setInternalState(authentication, "pdConfig", config);
+        server = ServerBuilder.forPort(0).intercept(authentication).addService(new PDGrpc.PDImplBase() {
+            @Override
+            public void reportStoreShutdown(Pdpb.SetStoreRequest request,
+                                            StreamObserver observer) {
+                shutdownCalls.incrementAndGet();
+                observer.onNext(Pdpb.SetStoreResponse.getDefaultInstance());
+                observer.onCompleted();
+            }
+
+            @Override
+            public void registerStore(Pdpb.RegisterStoreRequest request,
+                                      StreamObserver observer) {
+                observer.onNext(Pdpb.RegisterStoreResponse.getDefaultInstance());
+                observer.onCompleted();
+            }
+        }).build().start();
+        channel = ManagedChannelBuilder.forAddress("127.0.0.1", server.getPort()).usePlaintext().build();
+    }
+
+    @After
+    public void tearDown() throws Exception {
+        if (channel != null) {
+            channel.shutdownNow().awaitTermination(5, TimeUnit.SECONDS);
+        }
+        if (server != null) {
+            server.shutdownNow().awaitTermination(5, TimeUnit.SECONDS);
+        }
+    }
+
+    @Test
+    public void testMissingCredentialIsRefused() {
+        assertRefused(null);
+    }
+
+    @Test
+    public void testWrongCredentialIsRefused() {
+        assertRefused(credential("store:wrong"));
+        assertRefused("not-base64!");
+        assertRefused(credential("external:test-only-secret"));
+    }
+
+    @Test
+    public void testMissingConfiguredSecretFailsClosed() {
+        Mockito.when(config.getSecretKey()).thenReturn("");
+        assertRefused(credential("store:test-only-secret"));
+    }
+
+    @Test
+    public void testPrivilegedInternalCredentialPermitsReport() {
+        stub(credential("store:test-only-secret")).reportStoreShutdown(Pdpb.SetStoreRequest.getDefaultInstance());
+        Assert.assertEquals(1, shutdownCalls.get());
+    }
+
+    @Test
+    public void testExistingRegistrationRemainsCompatible() {
+        stub(null).registerStore(Pdpb.RegisterStoreRequest.getDefaultInstance());
+        Assert.assertEquals(0, shutdownCalls.get());
+    }
+
+    private void assertRefused(String credential) {
+        StatusRuntimeException failure = Assert.assertThrows(StatusRuntimeException.class,
+                () -> stub(credential).reportStoreShutdown(Pdpb.SetStoreRequest.getDefaultInstance()));
+        Assert.assertEquals(Status.Code.UNAUTHENTICATED, failure.getStatus().getCode());
+        Assert.assertEquals(0, shutdownCalls.get());
+    }
+
+    private PDGrpc.PDBlockingStub stub(String credential) {
+        Metadata headers = new Metadata();
+        if (credential != null) {
+            headers.put(Consts.CREDENTIAL_KEY, credential);
+        }
+        return PDGrpc.newBlockingStub(channel).withDeadlineAfter(5, TimeUnit.SECONDS)
+                     .withInterceptors(MetadataUtils.newAttachHeadersInterceptor(headers));
+    }
+
+    private static String credential(String value) {
+        return Base64.getEncoder().encodeToString(value.getBytes(StandardCharsets.UTF_8));
+    }
+}
diff --git a/hugegraph-store/hg-store-core/src/main/java/org/apache/hugegraph/store/HeartbeatService.java b/hugegraph-store/hg-store-core/src/main/java/org/apache/hugegraph/store/HeartbeatService.java
index 8e38109673..9ee70ceb95 100644
--- a/hugegraph-store/hg-store-core/src/main/java/org/apache/hugegraph/store/HeartbeatService.java
+++ b/hugegraph-store/hg-store-core/src/main/java/org/apache/hugegraph/store/HeartbeatService.java
@@ -20,12 +20,14 @@
 import static org.apache.hugegraph.pd.grpc.Pdpb.ErrorType.PD_UNREACHABLE_VALUE;
 
 import java.util.ArrayList;
+import java.util.Arrays;
 import java.util.Collections;
 import java.util.List;
 import java.util.Map;
 import java.util.concurrent.ExecutionException;
 import java.util.concurrent.FutureTask;
 import java.util.concurrent.atomic.AtomicBoolean;
+import java.util.stream.Collectors;
 
 import org.apache.hugegraph.pd.common.PDException;
 import org.apache.hugegraph.pd.common.PDRuntimeException;
@@ -75,6 +77,10 @@ public class HeartbeatService implements Lifecycle, Partit
     private boolean restored = false;
     private FutureTask restoreTask;
     private Thread restoreThread;
+    private long nextOfflineRegistration = Long.MIN_VALUE;
+    private long nextRestoreAttempt = Long.MIN_VALUE;
+    private long restoreRetryDelayMillis = 1000;
+    private String lastRestoreFailure;
     private volatile boolean pulseConnected;
     private final AtomicBoolean fatalExitRequested = new AtomicBoolean();
     private Thread storeHeartbeatThread;
@@ -211,9 +217,14 @@ protected void registerStore() {
             // Recovery is needed once per process, not when an established stream reconnects.
             boolean restoring = !restored;
             this.storeInfo.setState(restoring ? Metapb.StoreState.Offline : Metapb.StoreState.Up);
-            long storeId;
+            long now = monotonicTimeMillis();
+            long storeId = storeInfo.getId();
             try {
-                storeId = pdProvider.registerStore(this.storeInfo);
+                if (!restoring || storeId == 0 || now >= nextOfflineRegistration) {
+                    storeId = pdProvider.registerStore(this.storeInfo);
+                    nextOfflineRegistration = storeId == 0 ? Long.MIN_VALUE :
+                                              now + Math.min(5000L, options.getStoreHBInterval() * 1000L);
+                }
             } finally {
                 // Retry registration after a failure instead of heartbeating an unregistered Store.
                 this.storeInfo.setState(Metapb.StoreState.Offline);
@@ -225,6 +236,9 @@ protected void registerStore() {
                     this.clusterStats = pdProvider.getClusterStats();
                     timerNextDelay = REGISTER_RETRY_INTERVAL * 1000;
                     if (restoreTask == null) {
+                        if (now < nextRestoreAttempt) {
+                            return;
+                        }
                         restoreTask = new FutureTask<>(() -> {
                             storeEngine.stateChanged(storeInfo, Metapb.StoreState.Offline, Metapb.StoreState.Up);
                             return null;
@@ -242,7 +256,19 @@ protected void registerStore() {
                         Thread.currentThread().interrupt();
                         return;
                     } catch (ExecutionException e) {
-                        log.error("Store {} remains Offline after partition restore failed", storeId, e);
+                        Throwable failure = e.getCause();
+                        String signature = failure.getMessage() + Arrays.stream(failure.getSuppressed())
+                                .map(Throwable::toString).sorted().collect(Collectors.joining(";"));
+                        if (!signature.equals(lastRestoreFailure)) {
+                            log.error("Store {} remains Offline after partition restore failed", storeId, failure);
+                            lastRestoreFailure = signature;
+                        } else {
+                            log.warn("Store {} partition restore still failing; retrying after {} ms",
+                                     storeId, restoreRetryDelayMillis);
+                        }
+                        nextRestoreAttempt = now + restoreRetryDelayMillis;
+                        restoreRetryDelayMillis = Math.min(restoreRetryDelayMillis * 2,
+                                                           options.getStoreHBInterval() * 1000L);
                         restoreTask = null;
                         return;
                     }
@@ -305,6 +331,10 @@ protected void registerStore() {
         }
     }
 
+    protected long monotonicTimeMillis() {
+        return System.nanoTime() / 1_000_000L;
+    }
+
     protected void storeHeartbeat() {
         if (log.isDebugEnabled()) {
             log.debug("storeHeartbeat ... ");
diff --git a/hugegraph-store/hg-store-core/src/main/java/org/apache/hugegraph/store/pd/DefaultPdProvider.java b/hugegraph-store/hg-store-core/src/main/java/org/apache/hugegraph/store/pd/DefaultPdProvider.java
index a3a739ceca..352d179efd 100644
--- a/hugegraph-store/hg-store-core/src/main/java/org/apache/hugegraph/store/pd/DefaultPdProvider.java
+++ b/hugegraph-store/hg-store-core/src/main/java/org/apache/hugegraph/store/pd/DefaultPdProvider.java
@@ -67,7 +67,7 @@ public class DefaultPdProvider implements PdProvider {
     private GraphManager graphManager = null;
 
     public static String name = "store";
-    public static String authority = "default";
+    public static String authority = System.getenv().getOrDefault("HG_PD_AUTH_SECRET_KEY", "default");
 
     PDEventListener listener = new PDEventListener() {
         // Listening to pd change information listener
diff --git a/hugegraph-store/hg-store-test/src/main/java/org/apache/hugegraph/store/StoreRegistrationTest.java b/hugegraph-store/hg-store-test/src/main/java/org/apache/hugegraph/store/StoreRegistrationTest.java
index 314ed7826d..9ffca5c499 100644
--- a/hugegraph-store/hg-store-test/src/main/java/org/apache/hugegraph/store/StoreRegistrationTest.java
+++ b/hugegraph-store/hg-store-test/src/main/java/org/apache/hugegraph/store/StoreRegistrationTest.java
@@ -64,7 +64,7 @@ public void testRestoreCompletesBeforeUpRegistration() throws Exception {
         HeartbeatService service = service(engine, provider);
         try {
             registerRestored(service);
-            Assert.assertEquals(List.of("Offline", "restore", "Offline", "Up"), events);
+            Assert.assertEquals(List.of("Offline", "restore", "Up"), events);
         } finally {
             service.shutdown();
         }
@@ -82,7 +82,7 @@ public void testFailedRestoreNeverRegistersUp() throws Exception {
             service.registerStore();
             Assert.assertThrows(ExecutionException.class, () -> restoreTask(service).get(5, TimeUnit.SECONDS));
             service.registerStore();
-            Mockito.verify(provider, Mockito.times(2)).registerStore(Mockito.any());
+            Mockito.verify(provider).registerStore(Mockito.any());
             Assert.assertEquals(Metapb.StoreState.Offline, service.getStoreInfo().getState());
         } finally {
             service.shutdown();
@@ -101,7 +101,7 @@ public void testReconnectRegistersUpWithoutRestoringAgain() throws Exception {
         HeartbeatService service = service(engine, provider);
         try {
             registerRestored(service);
-            Assert.assertEquals(List.of("Offline", "Offline", "Up"), events);
+            Assert.assertEquals(List.of("Offline", "Up"), events);
             events.clear();
             CountDownLatch offline = new CountDownLatch(1);
             CountDownLatch up = new CountDownLatch(1);
@@ -128,7 +128,7 @@ public void testReconnectRegistersUpWithoutRestoringAgain() throws Exception {
     public void testFailedUpRegistrationRetriesWithoutRestoringAgain() throws Exception {
         HgStoreEngine engine = Mockito.mock(HgStoreEngine.class);
         PdProvider provider = Mockito.mock(PdProvider.class);
-        Mockito.when(provider.registerStore(Mockito.any())).thenReturn(7L, 7L)
+        Mockito.when(provider.registerStore(Mockito.any())).thenReturn(7L)
                .thenThrow(new PDException(1, "unavailable")).thenReturn(7L);
         HeartbeatService service = service(engine, provider);
         try {
@@ -268,7 +268,7 @@ public void testSlowRestoreKeepsRegisteringOfflineAndSkipsPartitionPulse() throw
             service.registerStore();
             Assert.assertFalse(service.isClusterReady());
             Assert.assertEquals(Metapb.StoreState.Offline, service.getStoreInfo().getState());
-            Mockito.verify(provider, Mockito.times(2)).registerStore(Mockito.argThat(
+            Mockito.verify(provider).registerStore(Mockito.argThat(
                     store -> store.getState() == Metapb.StoreState.Offline));
             service.partitionHeartbeat();
             Mockito.verify(engine, Mockito.never()).getLeaderPartition();
@@ -279,6 +279,61 @@ public void testSlowRestoreKeepsRegisteringOfflineAndSkipsPartitionPulse() throw
         }
     }
 
+    @Test
+    public void testRefusedRegistrationIsRetriedBeforeStartingRestore() throws Exception {
+        HgStoreEngine engine = Mockito.mock(HgStoreEngine.class);
+        PdProvider provider = Mockito.mock(PdProvider.class);
+        Mockito.when(provider.registerStore(Mockito.any())).thenReturn(0L);
+        HeartbeatService service = service(engine, provider);
+        service.getStoreInfo().setId(7);
+        try {
+            service.registerStore();
+            service.registerStore();
+            Assert.assertNull(restoreTask(service));
+            Mockito.verify(provider, Mockito.times(2)).registerStore(Mockito.any());
+            Mockito.verifyNoInteractions(engine);
+        } finally {
+            service.shutdown();
+        }
+    }
+
+    @Test
+    public void testRestoreRetryBackoffKeepsLeaseRenewalIndependent() throws Exception {
+        HgStoreEngine engine = Mockito.mock(HgStoreEngine.class);
+        PdProvider provider = Mockito.mock(PdProvider.class);
+        Mockito.when(provider.registerStore(Mockito.any())).thenReturn(7L);
+        Mockito.doThrow(new IllegalStateException("broken partition")).when(engine)
+               .stateChanged(Mockito.any(), Mockito.any(), Mockito.any());
+        HeartbeatService service = Mockito.spy(service(engine, provider));
+        Mockito.doReturn(1000L).when(service).monotonicTimeMillis();
+        try {
+            service.registerStore();
+            Assert.assertThrows(ExecutionException.class, () -> restoreTask(service).get(5, TimeUnit.SECONDS));
+            service.registerStore();
+            Assert.assertNull(restoreTask(service));
+            service.registerStore();
+            Assert.assertNull(restoreTask(service));
+            Mockito.verify(engine).stateChanged(Mockito.any(), Mockito.any(), Mockito.any());
+
+            Mockito.doReturn(2000L).when(service).monotonicTimeMillis();
+            service.registerStore();
+            Assert.assertThrows(ExecutionException.class, () -> restoreTask(service).get(5, TimeUnit.SECONDS));
+            service.registerStore();
+            Mockito.doReturn(3000L).when(service).monotonicTimeMillis();
+            service.registerStore();
+            Assert.assertNull(restoreTask(service));
+            Mockito.verify(engine, Mockito.times(2)).stateChanged(Mockito.any(), Mockito.any(), Mockito.any());
+
+            // Renewal remains due after five seconds, independently of the worker's retry deadline.
+            Mockito.doReturn(7000L).when(service).monotonicTimeMillis();
+            service.registerStore();
+            Mockito.verify(provider, Mockito.times(2)).registerStore(Mockito.any());
+            Assert.assertEquals(Metapb.StoreState.Offline, service.getStoreInfo().getState());
+        } finally {
+            service.shutdown();
+        }
+    }
+
     private static void set(Object target, String name, Object value) throws Exception {
         Field field = target.getClass().getDeclaredField(name);
         field.setAccessible(true);