Skip to content

fix(pd): preserve store drain state #3418

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
imbajin merged 8 commits into from
Oct 28, 2026
Merged

fix(pd): preserve store drain state #3418

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
8 commits
Select commit Hold shift + click to select a range
b84bf81
fix(pd): preserve store drain state
Oct 27, 2026
30db68f
fix(pd): revalidate placement targets before migration
imbajin Oct 27, 2026
3a6f2c0
fix(pd): resume drains without lock inversion
Oct 28, 2026
50515cb
chore(pd): remove trailing test blank line
Oct 28, 2026
0376e08
chore(ci): apply PR 343 drain fix
imbajin Oct 28, 2026
7a792bc
chore(ci): fix PR 343 patch workflow
imbajin Oct 28, 2026
a85d49f
chore(ci): stage tested PR 343 fix
imbajin Oct 28, 2026
476e80d
chore(ci): remove one-off drain workflows
Oct 28, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
20 changes: 20 additions & 0 deletions hugegraph-pd/docs/api-reference.md
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -90,6 +90,26 @@ enum ShardRole {

Represents a Store node in the cluster.

An `Exiting` Store continues renewing its heartbeat lease while draining, but is excluded from
new replica allocation and migration destinations. Expiry of an old heartbeat does not complete
retirement. PD keeps the Store `Exiting` until its latest heartbeat reports no local partition
engines and PD has no shard group referencing its ID. Before deleting its Pod or volume, also
verify the remaining groups' actual Raft membership, full replicas and fresh heartbeats; a
management state or a successful state-update response alone is not a complete maintenance check.
Replica allocation still enforces the configured minimum Store count using eligible destinations.
If the remaining eligible Stores do not satisfy that minimum, draining remains blocked; expiry of
the source lease must not be used to force retirement or reduce the configured replica policy.
If a drain request fails this capacity check, add enough eligible capacity and repeat the same
`Exiting` management request to resume planning for groups that still reference the source in PD.
The persisted intent stays `Exiting`; there is no need to toggle it to `Up`. This retry does not
establish that an already dispatched migration has completed or recover a lost Store-side task.
`Tombstone` is permanent for that Store identity: management updates cannot return it to an active
state, just as registration and heartbeat already reject a retired identity.
Heartbeats, registration and management state changes serialize their Store-state updates.
Re-registering an `Exiting` Store preserves that state even when its address is in the initial
Store list. This preserves the drain intent, but does not replace checking restored engines and
actual Raft configuration after a process restart.

```protobuf
message Store {
uint64 id = 1;
Expand Down
156 changes: 96 additions & 60 deletions hugegraph-pd/hg-pd-core/src/main/java/org/apache/hugegraph/pd/StoreNodeService.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -131,7 +131,7 @@ public boolean isOK() {
*
* @param store
*/
public Metapb.Store register(Metapb.Store store) throws PDException {
public synchronized Metapb.Store register(Metapb.Store store) throws PDException {
if (store.getId() == 0) {
// Initial registration, generate a new ID, and ensure that the ID is not duplicated.
store = newStoreNode(store);
Expand Down Expand Up @@ -167,11 +167,15 @@ public Metapb.Store register(Metapb.Store store) throws PDException {
// also exposes liveness that tells stale membership apart.
// https://github.com/apache/hugegraph/issues/3229
Metapb.StoreState storeState = lastStore.getState();
if (storeState == Metapb.StoreState.Offline || storeState == Metapb.StoreState.Up
|| inInitialStoreList(store)) {
storeState = Metapb.StoreState.Up;
} else {
storeState = Metapb.StoreState.Pending;
// Restarting an exiting Store must not cancel its persisted drain intent, including
// when its address is still present in the bootstrap allowlist.
if (storeState != Metapb.StoreState.Exiting) {
if (storeState == Metapb.StoreState.Offline || storeState == Metapb.StoreState.Up
|| inInitialStoreList(store)) {
storeState = Metapb.StoreState.Up;
} else {
storeState = Metapb.StoreState.Pending;
}
}

store = Metapb.Store.newBuilder(lastStore)
Expand Down Expand Up @@ -212,6 +216,8 @@ public Metapb.Store register(Metapb.Store store) throws PDException {
storeInfoMeta.keepStoreAlive(store);
onStoreStatusChanged(store, Metapb.StoreState.Offline, Metapb.StoreState.Up);
checkStoreStatus();
} else if (storeState == Metapb.StoreState.Exiting) {
storeInfoMeta.keepStoreAlive(store);
}

// Wait for the store information to be saved before sending the changes
Expand Down Expand Up @@ -266,36 +272,44 @@ public Metapb.Store getStore(long id) throws PDException {
/**
* Update the store information, detect the change of store status, and notify Hugestore
*/
public synchronized Metapb.Store updateStore(Metapb.Store store) throws PDException {
log.info("updateStore storeId: {}, address: {}, state: {}", store.getId(),
store.getAddress(), store.getState());
Metapb.Store lastStore = storeInfoMeta.getStore(store.getId());
if (lastStore == null) {
return null;
}
Metapb.Store.Builder builder =
Metapb.Store.newBuilder(lastStore).clearLabels().clearStats();
store = builder.mergeFrom(store).build();
if (store.getState() == Metapb.StoreState.Tombstone) {
List<Metapb.Store> activeStores = getStores();
if (lastStore.getState() == Metapb.StoreState.Up
&& activeStores.size() - 1 < pdConfig.getMinStoreCount()) {
throw new PDException(Pdpb.ErrorType.LESS_ACTIVE_STORE_VALUE,
"The number of active stores is less then " +
pdConfig.getMinStoreCount());
public Metapb.Store updateStore(Metapb.Store store) throws PDException {
Metapb.Store lastStore;
boolean stateChanged;
synchronized (this) {
log.info("updateStore storeId: {}, address: {}, state: {}", store.getId(),
store.getAddress(), store.getState());
lastStore = storeInfoMeta.getStore(store.getId());
if (lastStore == null) {
return null;
}
}

storeInfoMeta.updateStore(store);
if (store.getState() != Metapb.StoreState.Unknown &&
store.getState() != lastStore.getState()) {
// If you want to take the store offline
if (store.getState() == Metapb.StoreState.Exiting) {
if (lastStore.getState() == Metapb.StoreState.Exiting) {
// If it is already in the offline state, no further processing will be made
return lastStore;
boolean requestedDrain = store.getState() == Metapb.StoreState.Exiting;
Metapb.Store.Builder builder =
Metapb.Store.newBuilder(lastStore).clearLabels().clearStats();
store = builder.mergeFrom(store).build();
// Tombstone is terminal for this Store identity, as in registration and
// heartbeat handling. A deferred retirement notification must not race a
// management request that revives the same identity.
if (lastStore.getState() == Metapb.StoreState.Tombstone &&
store.getState() != Metapb.StoreState.Tombstone) {
throw new PDException(Pdpb.ErrorType.STORE_HAS_BEEN_REMOVED_VALUE,
String.format("Store id %d has been removed.", store.getId()));
}
if (store.getState() == Metapb.StoreState.Tombstone) {
List<Metapb.Store> activeStores = getStores();
if (lastStore.getState() == Metapb.StoreState.Up
&& activeStores.size() - 1 < pdConfig.getMinStoreCount()) {
throw new PDException(Pdpb.ErrorType.LESS_ACTIVE_STORE_VALUE,
"The number of active stores is less then " +
pdConfig.getMinStoreCount());
}
}

storeInfoMeta.updateStore(store);
stateChanged = store.getState() != Metapb.StoreState.Unknown &&
store.getState() != lastStore.getState();
// A previous attempt can persist Exiting before capacity checks fail. Repeating
// the same intent must resume the remaining groups after capacity is restored.
if (requestedDrain) {
List<Metapb.Store> activeStores = this.getActiveStores();
Map<Long, Metapb.Store> storeMap = new HashMap<>();
activeStores.forEach(s -> {
Expand All @@ -309,20 +323,26 @@ public synchronized Metapb.Store updateStore(Metapb.Store store) throws PDExcept
storeInfoMeta.removeActiveStore(store);
}
storeTurnoff(store);
} else if (store.getState() == Metapb.StoreState.Offline) {
// Monitor that the store has gone offline and is removed from the active
storeInfoMeta.removeActiveStore(store);
} else if (store.getState() == Metapb.StoreState.Tombstone) {
// When the status changes, the store is shut down, the shardGroup is modified,
// and the replica is migrated
log.info("updateStore removeActiveStores store {}", store.getId());
storeInfoMeta.removeActiveStore(store);
// Storage goes offline
storeTurnoff(store);
} else if (store.getState() == Metapb.StoreState.Up) {
storeInfoMeta.keepStoreAlive(store);
checkStoreStatus();
} else if (stateChanged) {
if (store.getState() == Metapb.StoreState.Offline) {
// Monitor that the store has gone offline and is removed from the active
storeInfoMeta.removeActiveStore(store);
} else if (store.getState() == Metapb.StoreState.Tombstone) {
// When the status changes, the store is shut down, the shardGroup is modified,
// and the replica is migrated
log.info("updateStore removeActiveStores store {}", store.getId());
storeInfoMeta.removeActiveStore(store);
// Storage goes offline
storeTurnoff(store);
} else if (store.getState() == Metapb.StoreState.Up) {
storeInfoMeta.keepStoreAlive(store);
checkStoreStatus();
}
}
}
// Tombstone listeners enter synchronized PartitionService methods. Never call
// them with this monitor held: placement takes PartitionService before this.
if (stateChanged) {
onStoreStatusChanged(lastStore, lastStore.getState(), store.getState());
}
return store;
Expand Down Expand Up @@ -409,6 +429,21 @@ public List<Metapb.Store> getActiveStores() throws PDException {
return storeInfoMeta.getActiveStores();
}

/**
* Return live Stores that may receive new replicas. The TTL entry records liveness and can
* still contain an older Up state after a Store starts exiting, so consult Store metadata.
*/
public List<Metapb.Store> getPlacementStores() throws PDException {
List<Metapb.Store> stores = new ArrayList<>();
for (Metapb.Store activeStore : storeInfoMeta.getActiveStores()) {
Metapb.Store store = storeInfoMeta.getStore(activeStore.getId());
if (store != null && store.getState() == Metapb.StoreState.Up) {
stores.add(store);
}
}
return stores;
}

public List<Metapb.Store> getTombStores() throws PDException {
List<Metapb.Store> stores = new ArrayList<>();
for (Metapb.Store store : this.getStores()) {
Expand Down Expand Up @@ -437,7 +472,7 @@ public synchronized List<Metapb.Shard> allocShards(Metapb.Graph graph, int partI
// number cannot exceed the number of raft groups
if (storeInfoMeta.getShardGroup(partId) == null) {
// Get active store key
List<Metapb.Store> stores = storeInfoMeta.getActiveStores();
List<Metapb.Store> stores = getPlacementStores();

if (stores.size() == 0) {
throw new PDException(Pdpb.ErrorType.NO_ACTIVE_STORE_VALUE,
Expand Down Expand Up @@ -497,7 +532,7 @@ public synchronized List<Metapb.Shard> allocShards(Metapb.Graph graph, int partI
*/
public synchronized List<Metapb.Shard> reallocShards(Metapb.ShardGroup shardGroup) throws
PDException {
List<Metapb.Store> stores = storeInfoMeta.getActiveStores();
List<Metapb.Store> stores = getPlacementStores();
Comment thread
imbajin marked this conversation as resolved.

if (stores.size() == 0) {
throw new PDException(Pdpb.ErrorType.NO_ACTIVE_STORE_VALUE,
Expand Down Expand Up @@ -727,7 +762,7 @@ public synchronized void updateShardGroupState(int groupId, Metapb.PartitionStat
* @param storeStats
* @throws PDException
*/
public Metapb.ClusterStats heartBeat(Metapb.StoreStats storeStats) throws PDException {
public synchronized Metapb.ClusterStats heartBeat(Metapb.StoreStats storeStats) throws PDException {
this.storeInfoMeta.updateStoreStats(storeStats);
Metapb.Store lastStore = this.getStore(storeStats.getStoreId());
if (lastStore == null) {
Expand All @@ -745,21 +780,16 @@ public Metapb.ClusterStats heartBeat(Metapb.StoreStats storeStats) throws PDExce
Metapb.Store nowStore;
// If you are going to take the store offline
if (lastStore.getState() == Metapb.StoreState.Exiting) {
List<Metapb.Store> activeStores = this.getActiveStores();
Map<Long, Metapb.Store> storeMap = new HashMap<>();
activeStores.forEach(store -> {
storeMap.put(store.getId(), store);
});
// If the partition of the offline store is 0, it means that the migration has been
// completed and can be taken offline, if it is not 0, the migration is still in
// progress and you need to wait
if (storeStats.getPartitionCount() > 0 &&
storeMap.containsKey(storeStats.getStoreId())) {
// TTL expiry only means the last heartbeat is stale, not that migration finished.
// Require both the source and PD's group view to be empty before retiring it.
if (storeStats.getPartitionCount() > 0 ||
!getShardGroupsByStore(storeStats.getStoreId()).isEmpty()) {
nowStore = Metapb.Store.newBuilder(lastStore)
.setStats(storeStats)
.setLastHeartbeat(System.currentTimeMillis())
.setState(Metapb.StoreState.Exiting).build();
this.storeInfoMeta.updateStore(nowStore);
this.storeInfoMeta.keepStoreAlive(nowStore);
return this.clusterStats;
} else {
nowStore = Metapb.Store.newBuilder(lastStore)
Expand Down Expand Up @@ -892,7 +922,13 @@ public void addShardGroupStatusListener(ShardGroupStatusListener listener) {
protected void onStoreStatusChanged(Metapb.Store store, Metapb.StoreState old,
Metapb.StoreState stats) {
log.info("onStoreStatusChanged storeId = {} from {} to {}", store.getId(), old, stats);
statusListeners.forEach(e -> {
List<StoreStatusListener> listeners;
synchronized (statusListeners) {
listeners = new ArrayList<>(statusListeners);
}
// Do not hold the listener-list monitor across callbacks either: register()
// dispatches Up while holding this Store service's monitor.
listeners.forEach(e -> {
e.onStoreStatusChanged(store, old, stats);
});
}
Expand Down
Loading
Loading