From b84bf81b8962507027d811206d8fa8a590e8b9d0 Mon Sep 17 00:00:00 2001 From: orion Date: Tue, 27 Oct 2026 07:52:38 +0000 Subject: [PATCH 1/8] fix(pd): preserve store drain state - Separate live Stores from Up placement destinations using current metadata. - Keep Exiting leases alive until local engines and PD references are empty. - Serialize heartbeat, registration and state changes and preserve drain on restart. - Cover drain, stale leases, shared groups and concurrent updates in regressions. --- hugegraph-pd/docs/api-reference.md | 14 + .../apache/hugegraph/pd/StoreNodeService.java | 54 ++- .../hugegraph/pd/TaskScheduleService.java | 19 +- .../hugegraph/pd/core/PDCoreSuiteTest.java | 1 + .../hugegraph/pd/core/StoreDrainTest.java | 345 ++++++++++++++++++ 5 files changed, 406 insertions(+), 27 deletions(-) create mode 100644 hugegraph-pd/hg-pd-test/src/main/java/org/apache/hugegraph/pd/core/StoreDrainTest.java diff --git a/hugegraph-pd/docs/api-reference.md b/hugegraph-pd/docs/api-reference.md index 122ae6c8b4..88cbe5e461 100644 --- a/hugegraph-pd/docs/api-reference.md +++ b/hugegraph-pd/docs/api-reference.md @@ -90,6 +90,20 @@ 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. +Heartbeats, registration and management state changes serialize their Store-state updates. +Re-registering an `Exiting` Store preserves that state even when its address is in the initial +Store list. This preserves the drain intent, but does not replace checking restored engines and +actual Raft configuration after a process restart. + ```protobuf message Store { uint64 id = 1; diff --git a/hugegraph-pd/hg-pd-core/src/main/java/org/apache/hugegraph/pd/StoreNodeService.java b/hugegraph-pd/hg-pd-core/src/main/java/org/apache/hugegraph/pd/StoreNodeService.java index 0f23195cf8..a76b442d3d 100644 --- a/hugegraph-pd/hg-pd-core/src/main/java/org/apache/hugegraph/pd/StoreNodeService.java +++ b/hugegraph-pd/hg-pd-core/src/main/java/org/apache/hugegraph/pd/StoreNodeService.java @@ -131,7 +131,7 @@ public boolean isOK() { * * @param store */ - public Metapb.Store register(Metapb.Store store) throws PDException { + public synchronized Metapb.Store register(Metapb.Store store) throws PDException { if (store.getId() == 0) { // Initial registration, generate a new ID, and ensure that the ID is not duplicated. store = newStoreNode(store); @@ -167,11 +167,15 @@ public Metapb.Store register(Metapb.Store store) throws PDException { // also exposes liveness that tells stale membership apart. // https://github.com/apache/hugegraph/issues/3229 Metapb.StoreState storeState = lastStore.getState(); - if (storeState == Metapb.StoreState.Offline || storeState == Metapb.StoreState.Up - || inInitialStoreList(store)) { - storeState = Metapb.StoreState.Up; - } else { - storeState = Metapb.StoreState.Pending; + // Restarting an exiting Store must not cancel its persisted drain intent, including + // when its address is still present in the bootstrap allowlist. + if (storeState != Metapb.StoreState.Exiting) { + if (storeState == Metapb.StoreState.Offline || storeState == Metapb.StoreState.Up + || inInitialStoreList(store)) { + storeState = Metapb.StoreState.Up; + } else { + storeState = Metapb.StoreState.Pending; + } } store = Metapb.Store.newBuilder(lastStore) @@ -212,6 +216,8 @@ public Metapb.Store register(Metapb.Store store) throws PDException { storeInfoMeta.keepStoreAlive(store); onStoreStatusChanged(store, Metapb.StoreState.Offline, Metapb.StoreState.Up); checkStoreStatus(); + } else if (storeState == Metapb.StoreState.Exiting) { + storeInfoMeta.keepStoreAlive(store); } // Wait for the store information to be saved before sending the changes @@ -409,6 +415,21 @@ public List getActiveStores() throws PDException { return storeInfoMeta.getActiveStores(); } + /** + * Return live Stores that may receive new replicas. The TTL entry records liveness and can + * still contain an older Up state after a Store starts exiting, so consult Store metadata. + */ + public List getPlacementStores() throws PDException { + List stores = new ArrayList<>(); + for (Metapb.Store activeStore : storeInfoMeta.getActiveStores()) { + Metapb.Store store = storeInfoMeta.getStore(activeStore.getId()); + if (store != null && store.getState() == Metapb.StoreState.Up) { + stores.add(store); + } + } + return stores; + } + public List getTombStores() throws PDException { List stores = new ArrayList<>(); for (Metapb.Store store : this.getStores()) { @@ -437,7 +458,7 @@ public synchronized List allocShards(Metapb.Graph graph, int partI // number cannot exceed the number of raft groups if (storeInfoMeta.getShardGroup(partId) == null) { // Get active store key - List stores = storeInfoMeta.getActiveStores(); + List stores = getPlacementStores(); if (stores.size() == 0) { throw new PDException(Pdpb.ErrorType.NO_ACTIVE_STORE_VALUE, @@ -497,7 +518,7 @@ public synchronized List allocShards(Metapb.Graph graph, int partI */ public synchronized List reallocShards(Metapb.ShardGroup shardGroup) throws PDException { - List stores = storeInfoMeta.getActiveStores(); + List stores = getPlacementStores(); if (stores.size() == 0) { throw new PDException(Pdpb.ErrorType.NO_ACTIVE_STORE_VALUE, @@ -727,7 +748,7 @@ public synchronized void updateShardGroupState(int groupId, Metapb.PartitionStat * @param storeStats * @throws PDException */ - public Metapb.ClusterStats heartBeat(Metapb.StoreStats storeStats) throws PDException { + public synchronized Metapb.ClusterStats heartBeat(Metapb.StoreStats storeStats) throws PDException { this.storeInfoMeta.updateStoreStats(storeStats); Metapb.Store lastStore = this.getStore(storeStats.getStoreId()); if (lastStore == null) { @@ -745,21 +766,16 @@ public Metapb.ClusterStats heartBeat(Metapb.StoreStats storeStats) throws PDExce Metapb.Store nowStore; // If you are going to take the store offline if (lastStore.getState() == Metapb.StoreState.Exiting) { - List activeStores = this.getActiveStores(); - Map storeMap = new HashMap<>(); - activeStores.forEach(store -> { - storeMap.put(store.getId(), store); - }); - // If the partition of the offline store is 0, it means that the migration has been - // completed and can be taken offline, if it is not 0, the migration is still in - // progress and you need to wait - if (storeStats.getPartitionCount() > 0 && - storeMap.containsKey(storeStats.getStoreId())) { + // TTL expiry only means the last heartbeat is stale, not that migration finished. + // Require both the source and PD's group view to be empty before retiring it. + if (storeStats.getPartitionCount() > 0 || + !getShardGroupsByStore(storeStats.getStoreId()).isEmpty()) { nowStore = Metapb.Store.newBuilder(lastStore) .setStats(storeStats) .setLastHeartbeat(System.currentTimeMillis()) .setState(Metapb.StoreState.Exiting).build(); this.storeInfoMeta.updateStore(nowStore); + this.storeInfoMeta.keepStoreAlive(nowStore); return this.clusterStats; } else { nowStore = Metapb.Store.newBuilder(lastStore) 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..5f63d7dafb 100644 --- a/hugegraph-pd/hg-pd-core/src/main/java/org/apache/hugegraph/pd/TaskScheduleService.java +++ b/hugegraph-pd/hg-pd-core/src/main/java/org/apache/hugegraph/pd/TaskScheduleService.java @@ -298,7 +298,8 @@ public synchronized Map> balancePartitionShard() thr return null; } - int activeStores = storeService.getActiveStores().size(); + List placementStores = storeService.getPlacementStores(); + int activeStores = placementStores.size(); if (activeStores == 0) { log.warn("balancePartitionShard non active stores, skip to balancePartitionShard"); return null; @@ -318,7 +319,7 @@ public synchronized Map> balancePartitionShard() thr // Count the partitions on each store, StoreId -> PartitionID, ShardRole Map> partitionMap = new HashMap<>(); - storeService.getActiveStores().forEach(store -> { + placementStores.forEach(store -> { partitionMap.put(store.getId(), new HashMap<>()); }); @@ -751,12 +752,15 @@ public Map canAllPartitionsMovedOut(Metapb.Store sourceStore) th Map partitionDataSize = new HashMap<>(); storeService.getActiveStores().forEach(store -> { + if (store.getId() == sourceStore.getId()) { + resultMap.put("current_store_is_online", true); + } + }); + storeService.getPlacementStores().forEach(store -> { if (store.getId() != sourceStore.getId()) { otherPartitionMap.put(store.getId(), new HashMap<>()); // Records the remaining disk space of other stores, in bytes availableDiskSpace.put(store.getId(), store.getStats().getAvailable()); - } else { - resultMap.put("current_store_is_online", true); } }); // Count the size of the partition to be migrated (from storeStats in KB) @@ -791,10 +795,9 @@ public Map canAllPartitionsMovedOut(Metapb.Store sourceStore) th } // Count the number of partitions of other stores and save them with a small top heap, so // that stores with fewer partitions are always prioritized - PriorityQueue> minHeap = new PriorityQueue<>(otherPartitionMap.size(), - (o1, o2) -> o1.getValue() - .compareTo( - o2.getValue())); + PriorityQueue> minHeap = + new PriorityQueue<>(Math.max(1, otherPartitionMap.size()), + (o1, o2) -> o1.getValue().compareTo(o2.getValue())); otherPartitionMap.forEach((storeId, shards) -> { minHeap.add(new KVPair(storeId, shards.size())); }); diff --git a/hugegraph-pd/hg-pd-test/src/main/java/org/apache/hugegraph/pd/core/PDCoreSuiteTest.java b/hugegraph-pd/hg-pd-test/src/main/java/org/apache/hugegraph/pd/core/PDCoreSuiteTest.java index ad97e37d55..0cfa69156b 100644 --- a/hugegraph-pd/hg-pd-test/src/main/java/org/apache/hugegraph/pd/core/PDCoreSuiteTest.java +++ b/hugegraph-pd/hg-pd-test/src/main/java/org/apache/hugegraph/pd/core/PDCoreSuiteTest.java @@ -56,6 +56,7 @@ PartitionServiceTest.class, StoreMonitorDataServiceTest.class, StoreServiceTest.class, + StoreDrainTest.class, TaskScheduleServiceTest.class, IpAuthHandlerTest.class, RaftEngineCliProcessorsTest.class, diff --git a/hugegraph-pd/hg-pd-test/src/main/java/org/apache/hugegraph/pd/core/StoreDrainTest.java b/hugegraph-pd/hg-pd-test/src/main/java/org/apache/hugegraph/pd/core/StoreDrainTest.java new file mode 100644 index 0000000000..906a49a25b --- /dev/null +++ b/hugegraph-pd/hg-pd-test/src/main/java/org/apache/hugegraph/pd/core/StoreDrainTest.java @@ -0,0 +1,345 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.hugegraph.pd.core; + +import java.lang.reflect.Field; +import java.util.ArrayList; +import java.util.HashMap; +import java.util.List; +import java.util.Map; +import java.util.concurrent.CountDownLatch; +import java.util.concurrent.ExecutorService; +import java.util.concurrent.Executors; +import java.util.concurrent.Future; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicBoolean; + +import org.apache.hugegraph.pd.ConfigService; +import org.apache.hugegraph.pd.PartitionService; +import org.apache.hugegraph.pd.StoreNodeService; +import org.apache.hugegraph.pd.TaskScheduleService; +import org.apache.hugegraph.pd.common.PDException; +import org.apache.hugegraph.pd.config.PDConfig; +import org.apache.hugegraph.pd.grpc.Metapb; +import org.apache.hugegraph.pd.grpc.Pdpb; +import org.apache.hugegraph.pd.grpc.pulse.ConfChangeType; +import org.apache.hugegraph.pd.meta.StoreInfoMeta; +import org.apache.hugegraph.pd.raft.RaftEngine; +import org.junit.Assert; +import org.junit.Test; +import org.mockito.Mockito; +import org.mockito.invocation.Invocation; + +import com.alipay.sofa.jraft.Node; + +public class StoreDrainTest { + + @Test + public void testExpiredLeaseDoesNotRetireStoreWithLocalGroups() throws Exception { + Fixture fixture = new Fixture(); + Metapb.Store exiting = fixture.addStore(4, Metapb.StoreState.Exiting); + // No active entry remains, as when a migration outlasts the heartbeat TTL. + Mockito.when(fixture.meta.getActiveStores()).thenReturn(List.of()); + fixture.heartbeat(4, 3); + Assert.assertEquals(Metapb.StoreState.Exiting, fixture.stores.get(4L).getState()); + Mockito.verify(fixture.meta).keepStoreAlive(fixture.stores.get(4L)); + Mockito.verify(fixture.meta, Mockito.never()).removeActiveStore(Mockito.any()); + Assert.assertEquals(exiting.getId(), fixture.stores.get(4L).getStats().getStoreId()); + } + + @Test + public void testZeroLocalGroupsDoesNotRetireReferencedStore() throws Exception { + Fixture fixture = new Fixture(); + fixture.addStore(4, Metapb.StoreState.Exiting); + fixture.groups.put(0, group(0, 1, 2, 4)); + fixture.heartbeat(4, 0); + Assert.assertEquals(Metapb.StoreState.Exiting, fixture.stores.get(4L).getState()); + Mockito.verify(fixture.meta).keepStoreAlive(fixture.stores.get(4L)); + Mockito.verify(fixture.meta, Mockito.never()).removeActiveStore(Mockito.any()); + } + + @Test + public void testRetiresOnlyAfterBothViewsAreEmpty() throws Exception { + Fixture fixture = new Fixture(); + fixture.addStore(4, Metapb.StoreState.Exiting); + fixture.heartbeat(4, 0); + Assert.assertEquals(Metapb.StoreState.Tombstone, fixture.stores.get(4L).getState()); + Mockito.verify(fixture.meta).removeActiveStore(fixture.stores.get(4L)); + Mockito.verify(fixture.meta, Mockito.never()).keepStoreAlive(Mockito.any()); + } + + @Test + public void testPlacementUsesCurrentStateInsteadOfStaleLease() throws Exception { + Fixture fixture = new Fixture(); + Metapb.Store up = fixture.addStore(1, Metapb.StoreState.Up); + fixture.addStore(2, Metapb.StoreState.Exiting); + fixture.addStore(3, Metapb.StoreState.Pending); + fixture.addStore(4, Metapb.StoreState.Tombstone); + Metapb.Store staleUp = Metapb.Store.newBuilder().setId(2) + .setState(Metapb.StoreState.Up).build(); + Mockito.when(fixture.meta.getActiveStores()).thenReturn(List.of( + up, staleUp, fixture.stores.get(3L), fixture.stores.get(4L), + Metapb.Store.newBuilder().setId(5).setState(Metapb.StoreState.Up).build())); + Assert.assertEquals(List.of(up), fixture.service.getPlacementStores()); + // Liveness remains available for an exiting source that still serves its groups. + Assert.assertEquals(5, fixture.service.getActiveStores().size()); + } + + @Test + public void testDrainCannotPlaceReplicaBackOnSourceAndNotifiesSharedGroupOnce() throws Exception { + Fixture fixture = new Fixture(); + for (long id = 1; id <= 3; id++) { + fixture.addStore(id, Metapb.StoreState.Up); + } + Metapb.Store source = fixture.addStore(4, Metapb.StoreState.Exiting); + fixture.groups.put(0, group(0, 1, 2, 4)); + Mockito.when(fixture.meta.getActiveStores()).thenReturn(List.of( + fixture.stores.get(1L), fixture.stores.get(2L), fixture.stores.get(3L), + source.toBuilder().setState(Metapb.StoreState.Up).build())); + List sharedPartitions = List.of( + Metapb.Partition.newBuilder().setId(0).setGraphName("graph-a").build(), + Metapb.Partition.newBuilder().setId(0).setGraphName("graph-b").build()); + Mockito.when(fixture.partitions.getPartitionById(0)).thenReturn(sharedPartitions); + + fixture.service.storeTurnoff(source); + + List result = fixture.groups.get(0).getShardsList(); + Assert.assertEquals(3, result.size()); + Assert.assertTrue(result.stream().noneMatch(shard -> shard.getStoreId() == 4)); + Assert.assertTrue(result.stream().anyMatch(shard -> shard.getStoreId() == 3)); + List changes = Mockito.mockingDetails(fixture.partitions).getInvocations().stream() + .filter(invocation -> invocation.getMethod().getName() + .equals("fireChangeShard")) + .toList(); + Assert.assertEquals(1, changes.size()); + Assert.assertEquals(sharedPartitions.get(0), changes.get(0).getArgument(0)); + Assert.assertEquals(result, changes.get(0).getArgument(1)); + Assert.assertEquals(ConfChangeType.CONF_CHANGE_TYPE_ADJUST, changes.get(0).getArgument(2)); + } + + @Test + public void testInitialAllocationExcludesExitingStoreWithStaleUpLease() throws Exception { + Fixture fixture = new Fixture(); + for (long id = 1; id <= 3; id++) { + fixture.addStore(id, Metapb.StoreState.Up); + } + Metapb.Store source = fixture.addStore(4, Metapb.StoreState.Exiting); + Mockito.when(fixture.meta.getActiveStores()).thenReturn(List.of( + source.toBuilder().setState(Metapb.StoreState.Up).build(), + fixture.stores.get(1L), fixture.stores.get(2L), fixture.stores.get(3L))); + List shards = fixture.service.allocShards( + Metapb.Graph.newBuilder().setGraphName("new-graph").build(), 0); + Assert.assertEquals(3, shards.size()); + Assert.assertTrue(shards.stream().noneMatch(shard -> shard.getStoreId() == 4)); + } + + @Test + public void testDrainingStoreCannotSatisfyPlacementMinimum() throws Exception { + Fixture fixture = new Fixture(); + fixture.addStore(1, Metapb.StoreState.Up); + fixture.addStore(2, Metapb.StoreState.Up); + Metapb.Store source = fixture.addStore(3, Metapb.StoreState.Exiting); + Mockito.when(fixture.meta.getActiveStores()).thenReturn(new ArrayList<>(fixture.stores.values())); + Metapb.ShardGroup original = group(0, 1, 2, 3); + fixture.groups.put(0, original); + + PDException error = Assert.assertThrows(PDException.class, () -> fixture.service.storeTurnoff(source)); + + Assert.assertEquals(Pdpb.ErrorType.LESS_ACTIVE_STORE_VALUE, error.getErrorCode()); + Assert.assertEquals(original, fixture.groups.get(0)); + Assert.assertEquals(3, fixture.service.getActiveStores().size()); + Assert.assertEquals(2, fixture.service.getPlacementStores().size()); + } + + @Test + public void testDrainIsBlockedWhenOtherLiveStoresCannotReceiveReplicas() throws Exception { + Fixture fixture = new Fixture(); + Metapb.Store source = fixture.addStore(1, Metapb.StoreState.Up); + fixture.addStore(2, Metapb.StoreState.Exiting); + fixture.addStore(3, Metapb.StoreState.Pending); + Mockito.when(fixture.meta.getActiveStores()).thenReturn(new ArrayList<>(fixture.stores.values())); + fixture.groups.put(0, group(0, 1)); + Mockito.when(fixture.partitions.getPartitions()).thenReturn(List.of( + Metapb.Partition.newBuilder().setId(0).setGraphName("graph-a").build())); + TaskScheduleService tasks = Mockito.mock(TaskScheduleService.class, Mockito.CALLS_REAL_METHODS); + Field storesField = TaskScheduleService.class.getDeclaredField("storeService"); + storesField.setAccessible(true); + storesField.set(tasks, fixture.service); + Field partitionsField = TaskScheduleService.class.getDeclaredField("partitionService"); + partitionsField.setAccessible(true); + partitionsField.set(tasks, fixture.partitions); + RaftEngine raft = RaftEngine.getInstance(); + Field nodeField = RaftEngine.class.getDeclaredField("raftNode"); + nodeField.setAccessible(true); + Object previous = nodeField.get(raft); + Node node = Mockito.mock(Node.class); + Mockito.when(node.isLeader(true)).thenReturn(true); + nodeField.set(raft, node); + try { + Map plan = tasks.canAllPartitionsMovedOut(source); + Assert.assertEquals(true, plan.get("current_store_is_online")); + Assert.assertEquals(false, plan.get("flag")); + Assert.assertNull(plan.get("movedPartitions")); + } finally { + nodeField.set(raft, previous); + } + } + + @Test + public void testReregisteringExitingStorePreservesDrainIntent() throws Exception { + Fixture fixture = new Fixture(); + Metapb.Store source = fixture.addStore(4, Metapb.StoreState.Exiting); + fixture.config.setInitialStoreList("source:8500"); + source = source.toBuilder().setAddress("source:8500").build(); + fixture.stores.put(4L, source); + + Metapb.Store registered = fixture.service.register(source); + + Assert.assertEquals(Metapb.StoreState.Exiting, registered.getState()); + Mockito.verify(fixture.meta).keepStoreAlive(registered); + Assert.assertTrue(fixture.service.getPlacementStores().isEmpty()); + } + + @Test + public void testHeartbeatCannotOverwriteConcurrentDrainIntent() throws Exception { + assertStateUpdateSerializesWithDrain(false); + } + + @Test + public void testRegistrationCannotOverwriteConcurrentDrainIntent() throws Exception { + assertStateUpdateSerializesWithDrain(true); + } + + private void assertStateUpdateSerializesWithDrain(boolean registration) throws Exception { + Fixture fixture = new Fixture(); + Metapb.Store source = fixture.addStore(4, Metapb.StoreState.Up); + CountDownLatch readUp = new CountDownLatch(1); + CountDownLatch continueUpdate = new CountDownLatch(1); + CountDownLatch drainRequested = new CountDownLatch(1); + CountDownLatch drainDone = new CountDownLatch(1); + AtomicBoolean firstRead = new AtomicBoolean(true); + Mockito.when(fixture.meta.getStore(4L)).thenAnswer(call -> { + Metapb.Store observed = fixture.stores.get(4L); + if (firstRead.compareAndSet(true, false)) { + readUp.countDown(); + Assert.assertTrue(continueUpdate.await(5, TimeUnit.SECONDS)); + } + return observed; + }); + ExecutorService workers = Executors.newFixedThreadPool(2); + try { + Future update = workers.submit(() -> { + if (registration) { + fixture.service.register(source); + } else { + fixture.heartbeat(4, 3); + } + return null; + }); + Assert.assertTrue(readUp.await(5, TimeUnit.SECONDS)); + Future drain = workers.submit(() -> { + drainRequested.countDown(); + fixture.service.updateStore(source.toBuilder().setState(Metapb.StoreState.Exiting).build()); + drainDone.countDown(); + return null; + }); + Assert.assertTrue(drainRequested.await(5, TimeUnit.SECONDS)); + Assert.assertFalse(drainDone.await(200, TimeUnit.MILLISECONDS)); + continueUpdate.countDown(); + update.get(5, TimeUnit.SECONDS); + drain.get(5, TimeUnit.SECONDS); + Assert.assertEquals(Metapb.StoreState.Exiting, fixture.stores.get(4L).getState()); + } finally { + continueUpdate.countDown(); + workers.shutdownNow(); + } + } + + private static Metapb.ShardGroup group(int id, long... storeIds) { + Metapb.ShardGroup.Builder group = Metapb.ShardGroup.newBuilder().setId(id); + for (int i = 0; i < storeIds.length; i++) { + group.addShards(Metapb.Shard.newBuilder().setStoreId(storeIds[i]) + .setRole(i == 0 ? Metapb.ShardRole.Leader : + Metapb.ShardRole.Follower)); + } + return group.build(); + } + + private static void setField(StoreNodeService service, String name, Object value) + throws Exception { + Field field = StoreNodeService.class.getDeclaredField(name); + field.setAccessible(true); + field.set(service, value); + } + + private static class Fixture { + private final StoreNodeService service = + Mockito.mock(StoreNodeService.class, Mockito.CALLS_REAL_METHODS); + private final StoreInfoMeta meta = Mockito.mock(StoreInfoMeta.class); + private final PartitionService partitions = Mockito.mock(PartitionService.class); + private final Map stores = new HashMap<>(); + private final Map groups = new HashMap<>(); + private final PDConfig config = new PDConfig(); + + private Fixture() throws Exception { + config.setMinStoreCount(3); + PDConfig.Partition partition = config.new Partition(); + partition.setShardCount(3); + config.setPartition(partition); + ConfigService configService = Mockito.mock(ConfigService.class); + Mockito.when(configService.getPartitionCount()).thenReturn(1); + config.setConfigService(configService); + setField(service, "pdConfig", config); + setField(service, "storeInfoMeta", meta); + setField(service, "partitionService", partitions); + setField(service, "shardGroupStatusListeners", List.of()); + setField(service, "statusListeners", List.of()); + setField(service, "clusterStats", Metapb.ClusterStats.getDefaultInstance()); + Mockito.when(meta.getActiveStores()).thenReturn(List.of()); + Mockito.when(meta.getStore(Mockito.anyLong())).thenAnswer(call -> + stores.get((Long) call.getArgument(0))); + Mockito.when(meta.storeExists(Mockito.anyLong())).thenAnswer(call -> + stores.containsKey((Long) call.getArgument(0))); + Mockito.when(meta.getShardGroup(Mockito.anyInt())).thenAnswer(call -> + groups.get((Integer) call.getArgument(0))); + Mockito.when(meta.getShardGroups()).thenAnswer(call -> new ArrayList<>(groups.values())); + Mockito.doAnswer(call -> { + Metapb.Store store = call.getArgument(0); + stores.put(store.getId(), store); + return null; + }).when(meta).updateStore(Mockito.any()); + Mockito.doAnswer(call -> { + Metapb.ShardGroup group = call.getArgument(0); + groups.put(group.getId(), group); + return null; + }).when(meta).updateShardGroup(Mockito.any()); + } + + private Metapb.Store addStore(long id, Metapb.StoreState state) { + Metapb.Store store = Metapb.Store.newBuilder().setId(id).setState(state).build(); + stores.put(id, store); + return store; + } + + private void heartbeat(long id, int groups) throws Exception { + service.heartBeat(Metapb.StoreStats.newBuilder().setStoreId(id) + .setPartitionCount(groups).build()); + } + } +} + From 30db68f4646fe7476a0812a8cee4cbcabb973c61 Mon Sep 17 00:00:00 2001 From: imbajin Date: Wed, 28 Oct 2026 04:33:21 +0800 Subject: [PATCH 2/8] fix(pd): revalidate placement targets before migration --- .../hugegraph/pd/TaskScheduleService.java | 79 +++++++++++---- .../hugegraph/pd/core/PDCoreSuiteTest.java | 1 + .../pd/core/StorePlacementSubmissionTest.java | 97 +++++++++++++++++++ 3 files changed, 159 insertions(+), 18 deletions(-) create mode 100644 hugegraph-pd/hg-pd-test/src/main/java/org/apache/hugegraph/pd/core/StorePlacementSubmissionTest.java 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 5f63d7dafb..bc63149eea 100644 --- a/hugegraph-pd/hg-pd-core/src/main/java/org/apache/hugegraph/pd/TaskScheduleService.java +++ b/hugegraph-pd/hg-pd-core/src/main/java/org/apache/hugegraph/pd/TaskScheduleService.java @@ -450,22 +450,61 @@ public synchronized Map> balancePartitionShard() thr maxHeap.addAll(tmpList); } - kvService.put(BALANCE_SHARD_KEY, "DOING", 180 * 1000); - - // Start the migration - movedPartitions.forEach((partId, storePair) -> { - // Neither the source nor destination storeID is 0 - if (storePair.getKey() > 0 && storePair.getValue() > 0) { - partitionService.movePartitionsShard(partId, storePair.getKey(), - storePair.getValue()); - } else { + boolean submitted = false; + Iterator>> submission = + movedPartitions.entrySet().iterator(); + while (submission.hasNext()) { + Map.Entry> move = submission.next(); + KVPair stores = move.getValue(); + if (stores.getKey() <= 0 || stores.getValue() <= 0) { log.warn("balancePartitionShard key or value is zero, partId:{} storePair:{}", - partId, storePair); + move.getKey(), stores); + submission.remove(); + continue; } - }); + if (!movePartitionIfPlacementStore(move.getKey(), stores.getKey(), + stores.getValue())) { + submission.remove(); + continue; + } + submitted = true; + } + if (submitted) { + kvService.put(BALANCE_SHARD_KEY, "DOING", 180 * 1000); + } return movedPartitions; } + /** + * Revalidate destination liveness and state while serializing with Store state changes. + * Preserve the existing service lock order: PartitionService then StoreNodeService. + */ + private boolean movePartitionIfPlacementStore(int partitionId, long sourceStoreId, + long destStoreId) { + synchronized (partitionService) { + synchronized (storeService) { + try { + boolean eligible = storeService.getPlacementStores().stream() + .anyMatch(store -> + store.getId() == destStoreId); + if (!eligible) { + log.warn("Skip moving partition {} from store {} to store {}: " + + "destination is no longer an Up placement store", + partitionId, sourceStoreId, destStoreId); + return false; + } + partitionService.movePartitionsShard(partitionId, sourceStoreId, + destStoreId); + return true; + } catch (PDException e) { + log.error("Cannot revalidate destination store {} for partition {}", + destStoreId, partitionId, e); + return false; + } + } + } + } + /** * Balance the number of leaders of partitions between stores */ @@ -873,15 +912,19 @@ public Map> movePartitions( if (!isLeader()) { return null; } - // Start the migration + // Start the migration, but revalidate every destination because a plan may be stale. log.info("begin move partitions:"); - movedPartitions.forEach((partId, storePair) -> { - // Neither the source nor destination storeID is 0 - if (storePair.getKey() > 0 && storePair.getValue() > 0) { - partitionService.movePartitionsShard(partId, storePair.getKey(), - storePair.getValue()); + Iterator>> iterator = + movedPartitions.entrySet().iterator(); + while (iterator.hasNext()) { + Map.Entry> move = iterator.next(); + KVPair stores = move.getValue(); + if (stores.getKey() <= 0 || stores.getValue() <= 0 || + !movePartitionIfPlacementStore(move.getKey(), stores.getKey(), + stores.getValue())) { + iterator.remove(); } - }); + } return movedPartitions; } diff --git a/hugegraph-pd/hg-pd-test/src/main/java/org/apache/hugegraph/pd/core/PDCoreSuiteTest.java b/hugegraph-pd/hg-pd-test/src/main/java/org/apache/hugegraph/pd/core/PDCoreSuiteTest.java index 0cfa69156b..9377d62de8 100644 --- a/hugegraph-pd/hg-pd-test/src/main/java/org/apache/hugegraph/pd/core/PDCoreSuiteTest.java +++ b/hugegraph-pd/hg-pd-test/src/main/java/org/apache/hugegraph/pd/core/PDCoreSuiteTest.java @@ -57,6 +57,7 @@ StoreMonitorDataServiceTest.class, StoreServiceTest.class, StoreDrainTest.class, + StorePlacementSubmissionTest.class, TaskScheduleServiceTest.class, IpAuthHandlerTest.class, RaftEngineCliProcessorsTest.class, diff --git a/hugegraph-pd/hg-pd-test/src/main/java/org/apache/hugegraph/pd/core/StorePlacementSubmissionTest.java b/hugegraph-pd/hg-pd-test/src/main/java/org/apache/hugegraph/pd/core/StorePlacementSubmissionTest.java new file mode 100644 index 0000000000..5492551dc7 --- /dev/null +++ b/hugegraph-pd/hg-pd-test/src/main/java/org/apache/hugegraph/pd/core/StorePlacementSubmissionTest.java @@ -0,0 +1,97 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.hugegraph.pd.core; + +import java.lang.reflect.Field; +import java.util.HashMap; +import java.util.List; +import java.util.Map; + +import org.apache.hugegraph.pd.PartitionService; +import org.apache.hugegraph.pd.StoreNodeService; +import org.apache.hugegraph.pd.TaskScheduleService; +import org.apache.hugegraph.pd.common.KVPair; +import org.apache.hugegraph.pd.grpc.Metapb; +import org.apache.hugegraph.pd.meta.StoreInfoMeta; +import org.apache.hugegraph.pd.raft.RaftEngine; +import org.junit.Assert; +import org.junit.Test; +import org.mockito.Mockito; + +import com.alipay.sofa.jraft.Node; + +public class StorePlacementSubmissionTest { + + @Test + public void testMovePlanRevalidatesDestinationBeforeSubmission() throws Exception { + StoreNodeService stores = Mockito.mock(StoreNodeService.class, + Mockito.CALLS_REAL_METHODS); + StoreInfoMeta meta = Mockito.mock(StoreInfoMeta.class); + PartitionService partitions = Mockito.mock(PartitionService.class); + Map current = new HashMap<>(); + Metapb.Store source = store(1L, Metapb.StoreState.Up); + Metapb.Store destinationLease = store(2L, Metapb.StoreState.Up); + current.put(source.getId(), source); + current.put(destinationLease.getId(), + destinationLease.toBuilder().setState(Metapb.StoreState.Exiting).build()); + Mockito.when(meta.getActiveStores()).thenReturn(List.of(source, destinationLease)); + Mockito.when(meta.getStore(Mockito.anyLong())).thenAnswer(call -> + current.get((Long) call.getArgument(0))); + setField(StoreNodeService.class, stores, "storeInfoMeta", meta); + + TaskScheduleService tasks = Mockito.mock(TaskScheduleService.class, + Mockito.CALLS_REAL_METHODS); + setField(TaskScheduleService.class, tasks, "storeService", stores); + setField(TaskScheduleService.class, tasks, "partitionService", partitions); + + RaftEngine raft = RaftEngine.getInstance(); + Field nodeField = RaftEngine.class.getDeclaredField("raftNode"); + nodeField.setAccessible(true); + Object previous = nodeField.get(raft); + Node node = Mockito.mock(Node.class); + Mockito.when(node.isLeader(true)).thenReturn(true); + nodeField.set(raft, node); + try { + Map> stalePlan = new HashMap<>(); + stalePlan.put(0, new KVPair<>(source.getId(), destinationLease.getId())); + Assert.assertTrue(tasks.movePartitions(stalePlan).isEmpty()); + Mockito.verify(partitions, Mockito.never()) + .movePartitionsShard(Mockito.anyInt(), Mockito.anyLong(), Mockito.anyLong()); + + current.put(destinationLease.getId(), destinationLease); + Map> validPlan = new HashMap<>(); + validPlan.put(0, new KVPair<>(source.getId(), destinationLease.getId())); + Assert.assertEquals(1, tasks.movePartitions(validPlan).size()); + Mockito.verify(partitions).movePartitionsShard(0, source.getId(), + destinationLease.getId()); + } finally { + nodeField.set(raft, previous); + } + } + + private static Metapb.Store store(long id, Metapb.StoreState state) { + return Metapb.Store.newBuilder().setId(id).setState(state).build(); + } + + private static void setField(Class owner, Object target, String name, Object value) + throws Exception { + Field field = owner.getDeclaredField(name); + field.setAccessible(true); + field.set(target, value); + } +} From 3a6f2c0d7fb9b3440b24ea2e22c8a47dbe82f766 Mon Sep 17 00:00:00 2001 From: orion Date: Wed, 28 Oct 2026 05:10:38 +0000 Subject: [PATCH 3/8] fix(pd): resume drains without lock inversion Retry explicit Exiting requests after eligible capacity is restored. Dispatch management state notifications after releasing the Store lock. Preserve Tombstone as terminal across management state updates. Cover capacity recovery and concurrent placement with regressions. --- hugegraph-pd/docs/api-reference.md | 6 + .../apache/hugegraph/pd/StoreNodeService.java | 102 +++++++++------- .../hugegraph/pd/core/StoreDrainTest.java | 80 +++++++++++++ .../pd/core/StorePlacementSubmissionTest.java | 112 ++++++++++++++++++ 4 files changed, 259 insertions(+), 41 deletions(-) diff --git a/hugegraph-pd/docs/api-reference.md b/hugegraph-pd/docs/api-reference.md index 88cbe5e461..f1c2ad9be4 100644 --- a/hugegraph-pd/docs/api-reference.md +++ b/hugegraph-pd/docs/api-reference.md @@ -99,6 +99,12 @@ management state or a successful state-update response alone is not a complete m 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 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 a76b442d3d..10a0b1692a 100644 --- a/hugegraph-pd/hg-pd-core/src/main/java/org/apache/hugegraph/pd/StoreNodeService.java +++ b/hugegraph-pd/hg-pd-core/src/main/java/org/apache/hugegraph/pd/StoreNodeService.java @@ -272,36 +272,44 @@ public Metapb.Store getStore(long id) throws PDException { /** * Update the store information, detect the change of store status, and notify Hugestore */ - public synchronized Metapb.Store updateStore(Metapb.Store store) throws PDException { - log.info("updateStore storeId: {}, address: {}, state: {}", store.getId(), - store.getAddress(), store.getState()); - Metapb.Store lastStore = storeInfoMeta.getStore(store.getId()); - if (lastStore == null) { - return null; - } - Metapb.Store.Builder builder = - Metapb.Store.newBuilder(lastStore).clearLabels().clearStats(); - store = builder.mergeFrom(store).build(); - if (store.getState() == Metapb.StoreState.Tombstone) { - List activeStores = getStores(); - if (lastStore.getState() == Metapb.StoreState.Up - && activeStores.size() - 1 < pdConfig.getMinStoreCount()) { - throw new PDException(Pdpb.ErrorType.LESS_ACTIVE_STORE_VALUE, - "The number of active stores is less then " + - pdConfig.getMinStoreCount()); + public Metapb.Store updateStore(Metapb.Store store) throws PDException { + Metapb.Store lastStore; + boolean stateChanged; + synchronized (this) { + log.info("updateStore storeId: {}, address: {}, state: {}", store.getId(), + store.getAddress(), store.getState()); + lastStore = storeInfoMeta.getStore(store.getId()); + if (lastStore == null) { + return null; } - } - - storeInfoMeta.updateStore(store); - if (store.getState() != Metapb.StoreState.Unknown && - store.getState() != lastStore.getState()) { - // If you want to take the store offline - if (store.getState() == Metapb.StoreState.Exiting) { - if (lastStore.getState() == Metapb.StoreState.Exiting) { - // If it is already in the offline state, no further processing will be made - return lastStore; + boolean requestedDrain = store.getState() == Metapb.StoreState.Exiting; + Metapb.Store.Builder builder = + Metapb.Store.newBuilder(lastStore).clearLabels().clearStats(); + store = builder.mergeFrom(store).build(); + // Tombstone is terminal for this Store identity, as in registration and + // heartbeat handling. A deferred retirement notification must not race a + // management request that revives the same identity. + if (lastStore.getState() == Metapb.StoreState.Tombstone && + store.getState() != Metapb.StoreState.Tombstone) { + throw new PDException(Pdpb.ErrorType.STORE_HAS_BEEN_REMOVED_VALUE, + String.format("Store id %d has been removed.", store.getId())); + } + if (store.getState() == Metapb.StoreState.Tombstone) { + List activeStores = getStores(); + if (lastStore.getState() == Metapb.StoreState.Up + && activeStores.size() - 1 < pdConfig.getMinStoreCount()) { + throw new PDException(Pdpb.ErrorType.LESS_ACTIVE_STORE_VALUE, + "The number of active stores is less then " + + pdConfig.getMinStoreCount()); } + } + storeInfoMeta.updateStore(store); + stateChanged = store.getState() != Metapb.StoreState.Unknown && + store.getState() != lastStore.getState(); + // A previous attempt can persist Exiting before capacity checks fail. Repeating + // the same intent must resume the remaining groups after capacity is restored. + if (requestedDrain) { List activeStores = this.getActiveStores(); Map storeMap = new HashMap<>(); activeStores.forEach(s -> { @@ -315,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; @@ -908,7 +922,13 @@ public void addShardGroupStatusListener(ShardGroupStatusListener listener) { protected void onStoreStatusChanged(Metapb.Store store, Metapb.StoreState old, Metapb.StoreState stats) { log.info("onStoreStatusChanged storeId = {} from {} to {}", store.getId(), old, stats); - statusListeners.forEach(e -> { + List listeners; + synchronized (statusListeners) { + listeners = new ArrayList<>(statusListeners); + } + // Do not hold the listener-list monitor across callbacks either: register() + // dispatches Up while holding this Store service's monitor. + listeners.forEach(e -> { e.onStoreStatusChanged(store, old, stats); }); } diff --git a/hugegraph-pd/hg-pd-test/src/main/java/org/apache/hugegraph/pd/core/StoreDrainTest.java b/hugegraph-pd/hg-pd-test/src/main/java/org/apache/hugegraph/pd/core/StoreDrainTest.java index 906a49a25b..5b4183b289 100644 --- a/hugegraph-pd/hg-pd-test/src/main/java/org/apache/hugegraph/pd/core/StoreDrainTest.java +++ b/hugegraph-pd/hg-pd-test/src/main/java/org/apache/hugegraph/pd/core/StoreDrainTest.java @@ -32,6 +32,7 @@ import org.apache.hugegraph.pd.ConfigService; import org.apache.hugegraph.pd.PartitionService; import org.apache.hugegraph.pd.StoreNodeService; +import org.apache.hugegraph.pd.StoreStatusListener; import org.apache.hugegraph.pd.TaskScheduleService; import org.apache.hugegraph.pd.common.PDException; import org.apache.hugegraph.pd.config.PDConfig; @@ -166,6 +167,85 @@ public void testDrainingStoreCannotSatisfyPlacementMinimum() throws Exception { Assert.assertEquals(2, fixture.service.getPlacementStores().size()); } + @Test + public void testRepeatedDrainResumesAfterCapacityIsRestored() throws Exception { + Fixture fixture = new Fixture(); + StoreStatusListener listener = Mockito.mock(StoreStatusListener.class); + setField(fixture.service, "statusListeners", List.of(listener)); + fixture.addStore(1, Metapb.StoreState.Up); + fixture.addStore(2, Metapb.StoreState.Up); + Metapb.Store source = fixture.addStore(3, Metapb.StoreState.Up); + Mockito.when(fixture.meta.getActiveStores()) + .thenAnswer(call -> new ArrayList<>(fixture.stores.values())); + Metapb.ShardGroup original = group(0, 1, 2, 3); + fixture.groups.put(0, original); + Metapb.Partition partition = Metapb.Partition.newBuilder() + .setId(0).setGraphName("graph").build(); + Mockito.when(fixture.partitions.getPartitionById(0)).thenReturn(List.of(partition)); + Metapb.Store drain = source.toBuilder().setState(Metapb.StoreState.Exiting).build(); + + PDException error = Assert.assertThrows(PDException.class, + () -> fixture.service.updateStore(drain)); + Assert.assertEquals(Pdpb.ErrorType.LESS_ACTIVE_STORE_VALUE, error.getErrorCode()); + Assert.assertEquals(Metapb.StoreState.Exiting, fixture.stores.get(3L).getState()); + Assert.assertEquals(original, fixture.groups.get(0)); + Assert.assertTrue(changeShardInvocations(fixture.partitions).isEmpty()); + + // A partial metadata update is not an explicit retry of the drain intent. + fixture.service.updateStore(Metapb.Store.newBuilder().setId(3) + .setVersion("metadata-only").build()); + Assert.assertEquals(original, fixture.groups.get(0)); + + fixture.addStore(4, Metapb.StoreState.Up); + fixture.service.updateStore(drain); + + List expected = fixture.groups.get(0).getShardsList(); + Assert.assertEquals(3, expected.size()); + Assert.assertTrue(expected.stream().noneMatch(shard -> shard.getStoreId() == 3)); + Assert.assertTrue(expected.stream().anyMatch(shard -> shard.getStoreId() == 4)); + Assert.assertEquals(Metapb.StoreState.Exiting, fixture.stores.get(3L).getState()); + List changes = changeShardInvocations(fixture.partitions); + Assert.assertEquals(1, changes.size()); + Assert.assertEquals(partition, changes.get(0).getArgument(0)); + Assert.assertEquals(expected, changes.get(0).getArgument(1)); + Assert.assertEquals(ConfChangeType.CONF_CHANGE_TYPE_ADJUST, changes.get(0).getArgument(2)); + + // Planning removed the source from the PD group; another identical request must + // not resubmit that group or repeat a Store-state notification. This is not a + // barrier proving that the actual Raft membership has committed the target. + fixture.service.updateStore(drain); + Assert.assertEquals(1, changeShardInvocations(fixture.partitions).size()); + Mockito.verifyNoInteractions(listener); + } + + @Test + public void testManagementCannotReviveTombstoneButAllowsPartialMetadataUpdate() throws Exception { + Fixture fixture = new Fixture(); + Metapb.Store retired = fixture.addStore(3, Metapb.StoreState.Tombstone); + for (Metapb.StoreState state : List.of(Metapb.StoreState.Up, + Metapb.StoreState.Exiting, + Metapb.StoreState.Offline, + Metapb.StoreState.Pending)) { + PDException error = Assert.assertThrows(PDException.class, () -> + fixture.service.updateStore(retired.toBuilder().setState(state).build())); + Assert.assertEquals(Pdpb.ErrorType.STORE_HAS_BEEN_REMOVED_VALUE, error.getErrorCode()); + Assert.assertEquals(retired, fixture.stores.get(3L)); + } + Metapb.Store partial = Metapb.Store.newBuilder().setId(3) + .setState(Metapb.StoreState.Unknown) + .setVersion("updated").build(); + Metapb.Store updated = fixture.service.updateStore(partial); + Assert.assertEquals(Metapb.StoreState.Tombstone, updated.getState()); + Assert.assertEquals("updated", updated.getVersion()); + Assert.assertEquals(updated, fixture.service.updateStore(updated)); + } + + private static List changeShardInvocations(PartitionService partitions) { + return Mockito.mockingDetails(partitions).getInvocations().stream() + .filter(invocation -> invocation.getMethod().getName().equals("fireChangeShard")) + .toList(); + } + @Test public void testDrainIsBlockedWhenOtherLiveStoresCannotReceiveReplicas() throws Exception { Fixture fixture = new Fixture(); diff --git a/hugegraph-pd/hg-pd-test/src/main/java/org/apache/hugegraph/pd/core/StorePlacementSubmissionTest.java b/hugegraph-pd/hg-pd-test/src/main/java/org/apache/hugegraph/pd/core/StorePlacementSubmissionTest.java index 5492551dc7..66b08fca79 100644 --- a/hugegraph-pd/hg-pd-test/src/main/java/org/apache/hugegraph/pd/core/StorePlacementSubmissionTest.java +++ b/hugegraph-pd/hg-pd-test/src/main/java/org/apache/hugegraph/pd/core/StorePlacementSubmissionTest.java @@ -18,15 +18,27 @@ package org.apache.hugegraph.pd.core; import java.lang.reflect.Field; +import java.util.ArrayList; +import java.util.Collections; import java.util.HashMap; import java.util.List; import java.util.Map; +import java.util.concurrent.CountDownLatch; +import java.util.concurrent.ExecutorService; +import java.util.concurrent.Executors; +import java.util.concurrent.Future; +import java.util.concurrent.TimeUnit; import org.apache.hugegraph.pd.PartitionService; import org.apache.hugegraph.pd.StoreNodeService; +import org.apache.hugegraph.pd.StoreStatusListener; import org.apache.hugegraph.pd.TaskScheduleService; import org.apache.hugegraph.pd.common.KVPair; +import org.apache.hugegraph.pd.common.PDException; +import org.apache.hugegraph.pd.config.PDConfig; import org.apache.hugegraph.pd.grpc.Metapb; +import org.apache.hugegraph.pd.grpc.Pdpb; +import org.apache.hugegraph.pd.meta.PartitionMeta; import org.apache.hugegraph.pd.meta.StoreInfoMeta; import org.apache.hugegraph.pd.raft.RaftEngine; import org.junit.Assert; @@ -37,6 +49,106 @@ public class StorePlacementSubmissionTest { + @Test + public void testTombstoneNotificationDoesNotInvertPlacementLocks() throws Exception { + StoreNodeService stores = Mockito.mock(StoreNodeService.class, + Mockito.CALLS_REAL_METHODS); + StoreInfoMeta meta = Mockito.mock(StoreInfoMeta.class); + PartitionService partitions = Mockito.mock(PartitionService.class, + Mockito.CALLS_REAL_METHODS); + PartitionMeta partitionMeta = Mockito.mock(PartitionMeta.class); + setField(PartitionService.class, partitions, "partitionMeta", partitionMeta); + setField(PartitionService.class, partitions, "storeService", stores); + Metapb.Store source = store(1L, Metapb.StoreState.Up); + Metapb.Store destination = store(2L, Metapb.StoreState.Up); + Metapb.Store retiring = store(3L, Metapb.StoreState.Exiting); + Map current = new HashMap<>(); + current.put(source.getId(), source); + current.put(destination.getId(), destination); + current.put(retiring.getId(), retiring); + Mockito.when(meta.getActiveStores()).thenReturn(List.of(source, destination)); + Mockito.when(meta.getStores(null)).thenReturn(List.of(source, destination, retiring)); + Mockito.when(meta.getShardGroups()).thenReturn(List.of()); + Mockito.when(meta.getStore(Mockito.anyLong())) + .thenAnswer(call -> current.get((Long) call.getArgument(0))); + Mockito.doAnswer(call -> { + Metapb.Store store = call.getArgument(0); + current.put(store.getId(), store); + return null; + }).when(meta).updateStore(Mockito.any()); + setField(StoreNodeService.class, stores, "storeInfoMeta", meta); + setField(StoreNodeService.class, stores, "pdConfig", new PDConfig()); + + CountDownLatch partitionHeld = new CountDownLatch(1); + CountDownLatch callbackEntered = new CountDownLatch(1); + List listeners = Collections.synchronizedList(new ArrayList<>()); + setField(StoreNodeService.class, stores, "statusListeners", listeners); + partitions.init(); + Metapb.Partition partition = Metapb.Partition.newBuilder() + .setId(0).setGraphName("graph").build(); + Metapb.Graph graph = Metapb.Graph.newBuilder().setGraphName("graph").build(); + Mockito.when(partitionMeta.getGraph("graph")).thenReturn(graph); + Mockito.doNothing().when(partitions).reallocPartitionShards(graph, partition); + // Keep the installed production Tombstone listener, storeOffline() and the + // synchronized shardOffline() real. Control only its metadata-query boundary. + Mockito.doAnswer(call -> { + callbackEntered.countDown(); + // Fail before blocking if the old lock order reappears, so the test always + // releases both threads instead of leaving deadlocked monitor owners behind. + Assert.assertFalse(Thread.holdsLock(stores)); + Assert.assertFalse(Thread.holdsLock(listeners)); + return List.of(partition); + }).when(partitions).getPartitionByStore(Mockito.any()); + + TaskScheduleService tasks = Mockito.mock(TaskScheduleService.class, + Mockito.CALLS_REAL_METHODS); + setField(TaskScheduleService.class, tasks, "storeService", stores); + setField(TaskScheduleService.class, tasks, "partitionService", partitions); + RaftEngine raft = RaftEngine.getInstance(); + Field nodeField = RaftEngine.class.getDeclaredField("raftNode"); + nodeField.setAccessible(true); + Object previous = nodeField.get(raft); + Node node = Mockito.mock(Node.class); + Mockito.when(node.isLeader(true)).thenReturn(true); + nodeField.set(raft, node); + ExecutorService workers = Executors.newFixedThreadPool(2); + try { + Future placement = workers.submit(() -> { + synchronized (partitions) { + partitionHeld.countDown(); + Assert.assertTrue(callbackEntered.await(5, TimeUnit.SECONDS)); + // The old retirement callback is waiting for this PartitionService + // monitor. A concurrent management update must not revive its ID. + PDException error = Assert.assertThrows(PDException.class, () -> + stores.updateStore(retiring.toBuilder() + .setState(Metapb.StoreState.Up).build())); + Assert.assertEquals(Pdpb.ErrorType.STORE_HAS_BEEN_REMOVED_VALUE, + error.getErrorCode()); + Map> plan = new HashMap<>(); + plan.put(0, new KVPair<>(source.getId(), destination.getId())); + Assert.assertEquals(1, tasks.movePartitions(plan).size()); + } + return null; + }); + Assert.assertTrue(partitionHeld.await(5, TimeUnit.SECONDS)); + Future retirement = workers.submit(() -> stores.updateStore( + retiring.toBuilder().setState(Metapb.StoreState.Tombstone).build())); + retirement.get(5, TimeUnit.SECONDS); + placement.get(5, TimeUnit.SECONDS); + Mockito.verify(partitions).shardOffline(partition, retiring.getId()); + Mockito.verify(partitions).movePartitionsShard(0, source.getId(), destination.getId()); + Assert.assertEquals(Metapb.StoreState.Tombstone, current.get(retiring.getId()).getState()); + } finally { + callbackEntered.countDown(); + workers.shutdownNow(); + try { + Assert.assertTrue(workers.awaitTermination(5, TimeUnit.SECONDS)); + } finally { + nodeField.set(raft, previous); + } + } + } + @Test public void testMovePlanRevalidatesDestinationBeforeSubmission() throws Exception { StoreNodeService stores = Mockito.mock(StoreNodeService.class, From 50515cb821c0ce1afd5ec3fc373104431cb7a881 Mon Sep 17 00:00:00 2001 From: orion Date: Wed, 28 Oct 2026 05:20:50 +0000 Subject: [PATCH 4/8] chore(pd): remove trailing test blank line --- .../main/java/org/apache/hugegraph/pd/core/StoreDrainTest.java | 1 - 1 file changed, 1 deletion(-) diff --git a/hugegraph-pd/hg-pd-test/src/main/java/org/apache/hugegraph/pd/core/StoreDrainTest.java b/hugegraph-pd/hg-pd-test/src/main/java/org/apache/hugegraph/pd/core/StoreDrainTest.java index 5b4183b289..e607b05926 100644 --- a/hugegraph-pd/hg-pd-test/src/main/java/org/apache/hugegraph/pd/core/StoreDrainTest.java +++ b/hugegraph-pd/hg-pd-test/src/main/java/org/apache/hugegraph/pd/core/StoreDrainTest.java @@ -422,4 +422,3 @@ private void heartbeat(long id, int groups) throws Exception { } } } - From 0376e08abecfc566ac5c904d042e39eb294bb2b2 Mon Sep 17 00:00:00 2001 From: imbajin Date: Wed, 28 Oct 2026 14:43:54 +0800 Subject: [PATCH 5/8] chore(ci): apply PR 343 drain fix --- .../workflows/fix-pr343-duplicate-drain.yml | 211 ++++++++++++++++++ 1 file changed, 211 insertions(+) create mode 100644 .github/workflows/fix-pr343-duplicate-drain.yml diff --git a/.github/workflows/fix-pr343-duplicate-drain.yml b/.github/workflows/fix-pr343-duplicate-drain.yml new file mode 100644 index 0000000000..a0964b4350 --- /dev/null +++ b/.github/workflows/fix-pr343-duplicate-drain.yml @@ -0,0 +1,211 @@ +name: Apply PR 343 drain fix + +on: + push: + branches: + - codex/pd-store-drain-state-20261027 + paths: + - .github/workflows/fix-pr343-duplicate-drain.yml + +permissions: + contents: write + +jobs: + patch-and-test: + if: github.event.head_commit.message == 'chore(ci): apply PR 343 drain fix' + runs-on: ubuntu-latest + timeout-minutes: 30 + steps: + - name: Checkout PR branch + uses: actions/checkout@v4 + with: + ref: codex/pd-store-drain-state-20261027 + fetch-depth: 0 + + - name: Set up Java 17 + uses: actions/setup-java@v4 + with: + distribution: temurin + java-version: '17' + cache: maven + + - name: Apply focused fix and regression + shell: bash + run: | + python3 - <<'PY' + from pathlib import Path + + def replace_once(path: str, old: str, new: str) -> None: + file = Path(path) + text = file.read_text() + count = text.count(old) + if count != 1: + raise SystemExit(f"Expected one match in {path}, found {count}: {old[:80]!r}") + file.write_text(text.replace(old, new, 1)) + + store = "hugegraph-pd/hg-pd-core/src/main/java/org/apache/hugegraph/pd/StoreNodeService.java" + replace_once( + store, + " public Metapb.Store updateStore(Metapb.Store store) throws PDException {\n", + " public Metapb.Store updateStore(Metapb.Store store) throws PDException {\n" + " return updateStore(store, Collections.emptyList());\n" + " }\n\n" + " public Metapb.Store updateStore(Metapb.Store store,\n" + " List submittedGroups) throws PDException {\n", + ) + replace_once( + store, + " storeTurnoff(store);\n" + " } else if (stateChanged) {\n", + " storeTurnoff(store, submittedGroups);\n" + " } else if (stateChanged) {\n", + ) + replace_once( + store, + " public synchronized void storeTurnoff(Metapb.Store store) throws PDException {\n" + " // Traverse ShardGroup,redistribution\n" + " for (Metapb.ShardGroup group : getShardGroupsByStore(store.getId())) {\n" + " Metapb.ShardGroup.Builder builder = Metapb.ShardGroup.newBuilder(group);\n" + " builder.clearShards();\n" + " group.getShardsList().forEach(shard -> {\n" + " if (shard.getStoreId() != store.getId()) {\n" + " builder.addShards(shard);\n" + " }\n" + " });\n" + " reallocShards(builder.build());\n" + " }\n" + " }\n", + " public synchronized void storeTurnoff(Metapb.Store store) throws PDException {\n" + " storeTurnoff(store, Collections.emptyList());\n" + " }\n\n" + " public synchronized void storeTurnoff(Metapb.Store store,\n" + " List submittedGroups)\n" + " throws PDException {\n" + " // Traverse ShardGroup,redistribution\n" + " for (Metapb.ShardGroup group : getShardGroupsByStore(store.getId())) {\n" + " // PDService can submit a capacity-aware migration before persisting\n" + " // Exiting. Keep that accepted destination authoritative instead of\n" + " // immediately issuing a second, independently planned change.\n" + " if (submittedGroups.contains(group.getId())) {\n" + " log.info(\"Skip duplicate drain submission for shard group {}\",\n" + " group.getId());\n" + " continue;\n" + " }\n" + " Metapb.ShardGroup.Builder builder = Metapb.ShardGroup.newBuilder(group);\n" + " builder.clearShards();\n" + " group.getShardsList().forEach(shard -> {\n" + " if (shard.getStoreId() != store.getId()) {\n" + " builder.addShards(shard);\n" + " }\n" + " });\n" + " reallocShards(builder.build());\n" + " }\n" + " }\n", + ) + + service = "hugegraph-pd/hg-pd-service/src/main/java/org/apache/hugegraph/pd/service/PDService.java" + replace_once( + service, + " Metapb.Store store = request.getStore();\n" + " // Before going offline, check whether the number of active machines is greater than\n", + " Metapb.Store store = request.getStore();\n" + " List submittedGroups = new ArrayList<>();\n" + " // Before going offline, check whether the number of active machines is greater than\n", + ) + replace_once( + service, + " // Perform partition migration operations\n" + " taskService.movePartitions((Map>) resultMap.get(\n" + " \"movedPartitions\"));\n", + " // Perform partition migration operations. Preserve the groups whose\n" + " // final placement was accepted so updateStore does not immediately\n" + " // submit a second independently planned membership change.\n" + " Map> submitted =\n" + " taskService.movePartitions(\n" + " (Map>) resultMap.get(\n" + " \"movedPartitions\"));\n" + " if (submitted != null) {\n" + " submittedGroups.addAll(submitted.keySet());\n" + " }\n", + ) + replace_once( + service, + " // Replace license using grpc\n" + " store = storeNodeService.updateStore(store);\n", + " // Replace license using grpc\n" + " store = storeNodeService.updateStore(store, submittedGroups);\n", + ) + + test = "hugegraph-pd/hg-pd-test/src/main/java/org/apache/hugegraph/pd/core/StoreDrainTest.java" + marker = ( + " @Test\n" + " public void testInitialAllocationExcludesExitingStoreWithStaleUpLease() throws Exception {\n" + ) + regression = ( + " @Test\n" + " public void testPreSubmittedDrainGroupIsNotDispatchedAgain() throws Exception {\n" + " Fixture fixture = new Fixture();\n" + " for (long id = 1; id <= 4; id++) {\n" + " fixture.addStore(id, Metapb.StoreState.Up);\n" + " }\n" + " Metapb.Store source = fixture.stores.get(4L);\n" + " Mockito.when(fixture.meta.getActiveStores())\n" + " .thenAnswer(call -> new ArrayList<>(fixture.stores.values()));\n" + " Metapb.ShardGroup submitted = group(0, 1, 2, 4);\n" + " fixture.groups.put(0, submitted);\n" + " fixture.groups.put(1, group(1, 1, 2, 4));\n" + " Metapb.Partition fallback = Metapb.Partition.newBuilder()\n" + " .setId(1)\n" + " .setGraphName(\"graph\").build();\n" + " Mockito.when(fixture.partitions.getPartitionById(1))\n" + " .thenReturn(List.of(fallback));\n" + " Metapb.Store drain = source.toBuilder()\n" + " .setState(Metapb.StoreState.Exiting).build();\n\n" + " fixture.service.updateStore(drain, List.of(0));\n\n" + " Assert.assertEquals(submitted, fixture.groups.get(0));\n" + " Assert.assertTrue(fixture.groups.get(1).getShardsList().stream()\n" + " .noneMatch(shard -> shard.getStoreId() == 4));\n" + " List changes = changeShardInvocations(fixture.partitions);\n" + " Assert.assertEquals(1, changes.size());\n" + " Assert.assertEquals(fallback, changes.get(0).getArgument(0));\n" + " }\n\n" + ) + replace_once(test, marker, regression + marker) + PY + + - name: Format and validate the focused change + shell: bash + run: | + mvn editorconfig:format -ntp + git diff --check + changed="$(git diff --name-only | sort)" + expected="$(printf '%s\n' \ + hugegraph-pd/hg-pd-core/src/main/java/org/apache/hugegraph/pd/StoreNodeService.java \ + hugegraph-pd/hg-pd-service/src/main/java/org/apache/hugegraph/pd/service/PDService.java \ + hugegraph-pd/hg-pd-test/src/main/java/org/apache/hugegraph/pd/core/StoreDrainTest.java | sort)" + test "$changed" = "$expected" || { + printf 'Unexpected changed files:\n%s\n' "$changed" >&2 + exit 1 + } + + - name: Run focused PD drain tests + shell: bash + run: | + mvn test -ntp \ + -pl hugegraph-pd/hg-pd-test -am \ + -Dtest=StoreDrainTest,StorePlacementSubmissionTest \ + -Dsurefire.failIfNoSpecifiedTests=false + + - name: Commit fix and remove one-off workflow + shell: bash + run: | + git config user.name 'chatgpt-codex-connector[bot]' + git config user.email 'chatgpt-codex-connector[bot]@users.noreply.github.com' + git add \ + hugegraph-pd/hg-pd-core/src/main/java/org/apache/hugegraph/pd/StoreNodeService.java \ + hugegraph-pd/hg-pd-service/src/main/java/org/apache/hugegraph/pd/service/PDService.java \ + hugegraph-pd/hg-pd-test/src/main/java/org/apache/hugegraph/pd/core/StoreDrainTest.java + git rm .github/workflows/fix-pr343-duplicate-drain.yml + git diff --cached --check + git commit -m 'fix(pd): avoid duplicate drain migration submission' + git push origin HEAD:codex/pd-store-drain-state-20261027 From 7a792bc850cf943a83b996635bcda0a929c5ec91 Mon Sep 17 00:00:00 2001 From: imbajin Date: Wed, 28 Oct 2026 14:45:31 +0800 Subject: [PATCH 6/8] chore(ci): fix PR 343 patch workflow --- .github/workflows/fix-pr343-duplicate-drain.yml | 1 - 1 file changed, 1 deletion(-) diff --git a/.github/workflows/fix-pr343-duplicate-drain.yml b/.github/workflows/fix-pr343-duplicate-drain.yml index a0964b4350..a0af67a731 100644 --- a/.github/workflows/fix-pr343-duplicate-drain.yml +++ b/.github/workflows/fix-pr343-duplicate-drain.yml @@ -12,7 +12,6 @@ permissions: jobs: patch-and-test: - if: github.event.head_commit.message == 'chore(ci): apply PR 343 drain fix' runs-on: ubuntu-latest timeout-minutes: 30 steps: From a85d49f265e950e505c5d11d44d4caece5190934 Mon Sep 17 00:00:00 2001 From: imbajin Date: Wed, 28 Oct 2026 14:49:10 +0800 Subject: [PATCH 7/8] chore(ci): stage tested PR 343 fix --- .github/workflows/stage-pr343-tested-fix.yml | 95 ++++++++++++++++++++ 1 file changed, 95 insertions(+) create mode 100644 .github/workflows/stage-pr343-tested-fix.yml diff --git a/.github/workflows/stage-pr343-tested-fix.yml b/.github/workflows/stage-pr343-tested-fix.yml new file mode 100644 index 0000000000..ae1849e1aa --- /dev/null +++ b/.github/workflows/stage-pr343-tested-fix.yml @@ -0,0 +1,95 @@ +name: Stage PR 343 tested fix + +on: + push: + branches: + - codex/pd-store-drain-state-20261027 + paths: + - .github/workflows/stage-pr343-tested-fix.yml + +permissions: + contents: write + +jobs: + stage-fix: + runs-on: ubuntu-latest + timeout-minutes: 30 + steps: + - name: Checkout PR branch + uses: actions/checkout@v4 + with: + ref: codex/pd-store-drain-state-20261027 + fetch-depth: 0 + + - name: Set up Java 17 + uses: actions/setup-java@v4 + with: + distribution: temurin + java-version: '17' + cache: maven + + - name: Reuse the reviewed patch and validation steps + shell: bash + run: | + python3 - <<'PY' + from pathlib import Path + import subprocess + + source = Path('.github/workflows/fix-pr343-duplicate-drain.yml') + lines = source.read_text().splitlines() + wanted = [ + 'Apply focused fix and regression', + 'Format and validate the focused change', + 'Run focused PD drain tests', + ] + scripts = {} + current = None + index = 0 + while index < len(lines): + line = lines[index] + if line.startswith(' - name: '): + current = line[len(' - name: '):] + if current in wanted and line == ' run: |': + index += 1 + block = [] + while index < len(lines): + item = lines[index] + if item.startswith(' '): + block.append(item[10:]) + index += 1 + continue + if not item.strip(): + block.append('') + index += 1 + continue + break + scripts[current] = '\n'.join(block) + '\n' + continue + index += 1 + + for name in wanted: + script = scripts.get(name) + if not script: + raise SystemExit(f'Missing workflow step: {name}') + print(f'Running reused step: {name}', flush=True) + subprocess.run( + ['bash', '-e', '-u', '-o', 'pipefail', '-c', script], + check=True, + ) + PY + + - name: Push the tested fix to a review branch + shell: bash + run: | + git config user.name 'chatgpt-codex-connector[bot]' + git config user.email 'chatgpt-codex-connector[bot]@users.noreply.github.com' + git add \ + hugegraph-pd/hg-pd-core/src/main/java/org/apache/hugegraph/pd/StoreNodeService.java \ + hugegraph-pd/hg-pd-service/src/main/java/org/apache/hugegraph/pd/service/PDService.java \ + hugegraph-pd/hg-pd-test/src/main/java/org/apache/hugegraph/pd/core/StoreDrainTest.java + git rm \ + .github/workflows/fix-pr343-duplicate-drain.yml \ + .github/workflows/stage-pr343-tested-fix.yml + git diff --cached --check + git commit -m 'fix(pd): avoid duplicate drain migration submission' + git push origin HEAD:refs/heads/codex/pr343-duplicate-drain-fix-20261028 From 476e80d603e79bcc7611714243654e491ad30574 Mon Sep 17 00:00:00 2001 From: orion Date: Wed, 28 Oct 2026 07:34:29 +0000 Subject: [PATCH 8/8] chore(ci): remove one-off drain workflows Remove the two branch-specific patch and staging workflows.\nKeep regression tests and production changes in their normal modules.\nRetain validation evidence outside the source tree. --- .../workflows/fix-pr343-duplicate-drain.yml | 210 ------------------ .github/workflows/stage-pr343-tested-fix.yml | 95 -------- 2 files changed, 305 deletions(-) delete mode 100644 .github/workflows/fix-pr343-duplicate-drain.yml delete mode 100644 .github/workflows/stage-pr343-tested-fix.yml diff --git a/.github/workflows/fix-pr343-duplicate-drain.yml b/.github/workflows/fix-pr343-duplicate-drain.yml deleted file mode 100644 index a0af67a731..0000000000 --- a/.github/workflows/fix-pr343-duplicate-drain.yml +++ /dev/null @@ -1,210 +0,0 @@ -name: Apply PR 343 drain fix - -on: - push: - branches: - - codex/pd-store-drain-state-20261027 - paths: - - .github/workflows/fix-pr343-duplicate-drain.yml - -permissions: - contents: write - -jobs: - patch-and-test: - runs-on: ubuntu-latest - timeout-minutes: 30 - steps: - - name: Checkout PR branch - uses: actions/checkout@v4 - with: - ref: codex/pd-store-drain-state-20261027 - fetch-depth: 0 - - - name: Set up Java 17 - uses: actions/setup-java@v4 - with: - distribution: temurin - java-version: '17' - cache: maven - - - name: Apply focused fix and regression - shell: bash - run: | - python3 - <<'PY' - from pathlib import Path - - def replace_once(path: str, old: str, new: str) -> None: - file = Path(path) - text = file.read_text() - count = text.count(old) - if count != 1: - raise SystemExit(f"Expected one match in {path}, found {count}: {old[:80]!r}") - file.write_text(text.replace(old, new, 1)) - - store = "hugegraph-pd/hg-pd-core/src/main/java/org/apache/hugegraph/pd/StoreNodeService.java" - replace_once( - store, - " public Metapb.Store updateStore(Metapb.Store store) throws PDException {\n", - " public Metapb.Store updateStore(Metapb.Store store) throws PDException {\n" - " return updateStore(store, Collections.emptyList());\n" - " }\n\n" - " public Metapb.Store updateStore(Metapb.Store store,\n" - " List submittedGroups) throws PDException {\n", - ) - replace_once( - store, - " storeTurnoff(store);\n" - " } else if (stateChanged) {\n", - " storeTurnoff(store, submittedGroups);\n" - " } else if (stateChanged) {\n", - ) - replace_once( - store, - " public synchronized void storeTurnoff(Metapb.Store store) throws PDException {\n" - " // Traverse ShardGroup,redistribution\n" - " for (Metapb.ShardGroup group : getShardGroupsByStore(store.getId())) {\n" - " Metapb.ShardGroup.Builder builder = Metapb.ShardGroup.newBuilder(group);\n" - " builder.clearShards();\n" - " group.getShardsList().forEach(shard -> {\n" - " if (shard.getStoreId() != store.getId()) {\n" - " builder.addShards(shard);\n" - " }\n" - " });\n" - " reallocShards(builder.build());\n" - " }\n" - " }\n", - " public synchronized void storeTurnoff(Metapb.Store store) throws PDException {\n" - " storeTurnoff(store, Collections.emptyList());\n" - " }\n\n" - " public synchronized void storeTurnoff(Metapb.Store store,\n" - " List submittedGroups)\n" - " throws PDException {\n" - " // Traverse ShardGroup,redistribution\n" - " for (Metapb.ShardGroup group : getShardGroupsByStore(store.getId())) {\n" - " // PDService can submit a capacity-aware migration before persisting\n" - " // Exiting. Keep that accepted destination authoritative instead of\n" - " // immediately issuing a second, independently planned change.\n" - " if (submittedGroups.contains(group.getId())) {\n" - " log.info(\"Skip duplicate drain submission for shard group {}\",\n" - " group.getId());\n" - " continue;\n" - " }\n" - " Metapb.ShardGroup.Builder builder = Metapb.ShardGroup.newBuilder(group);\n" - " builder.clearShards();\n" - " group.getShardsList().forEach(shard -> {\n" - " if (shard.getStoreId() != store.getId()) {\n" - " builder.addShards(shard);\n" - " }\n" - " });\n" - " reallocShards(builder.build());\n" - " }\n" - " }\n", - ) - - service = "hugegraph-pd/hg-pd-service/src/main/java/org/apache/hugegraph/pd/service/PDService.java" - replace_once( - service, - " Metapb.Store store = request.getStore();\n" - " // Before going offline, check whether the number of active machines is greater than\n", - " Metapb.Store store = request.getStore();\n" - " List submittedGroups = new ArrayList<>();\n" - " // Before going offline, check whether the number of active machines is greater than\n", - ) - replace_once( - service, - " // Perform partition migration operations\n" - " taskService.movePartitions((Map>) resultMap.get(\n" - " \"movedPartitions\"));\n", - " // Perform partition migration operations. Preserve the groups whose\n" - " // final placement was accepted so updateStore does not immediately\n" - " // submit a second independently planned membership change.\n" - " Map> submitted =\n" - " taskService.movePartitions(\n" - " (Map>) resultMap.get(\n" - " \"movedPartitions\"));\n" - " if (submitted != null) {\n" - " submittedGroups.addAll(submitted.keySet());\n" - " }\n", - ) - replace_once( - service, - " // Replace license using grpc\n" - " store = storeNodeService.updateStore(store);\n", - " // Replace license using grpc\n" - " store = storeNodeService.updateStore(store, submittedGroups);\n", - ) - - test = "hugegraph-pd/hg-pd-test/src/main/java/org/apache/hugegraph/pd/core/StoreDrainTest.java" - marker = ( - " @Test\n" - " public void testInitialAllocationExcludesExitingStoreWithStaleUpLease() throws Exception {\n" - ) - regression = ( - " @Test\n" - " public void testPreSubmittedDrainGroupIsNotDispatchedAgain() throws Exception {\n" - " Fixture fixture = new Fixture();\n" - " for (long id = 1; id <= 4; id++) {\n" - " fixture.addStore(id, Metapb.StoreState.Up);\n" - " }\n" - " Metapb.Store source = fixture.stores.get(4L);\n" - " Mockito.when(fixture.meta.getActiveStores())\n" - " .thenAnswer(call -> new ArrayList<>(fixture.stores.values()));\n" - " Metapb.ShardGroup submitted = group(0, 1, 2, 4);\n" - " fixture.groups.put(0, submitted);\n" - " fixture.groups.put(1, group(1, 1, 2, 4));\n" - " Metapb.Partition fallback = Metapb.Partition.newBuilder()\n" - " .setId(1)\n" - " .setGraphName(\"graph\").build();\n" - " Mockito.when(fixture.partitions.getPartitionById(1))\n" - " .thenReturn(List.of(fallback));\n" - " Metapb.Store drain = source.toBuilder()\n" - " .setState(Metapb.StoreState.Exiting).build();\n\n" - " fixture.service.updateStore(drain, List.of(0));\n\n" - " Assert.assertEquals(submitted, fixture.groups.get(0));\n" - " Assert.assertTrue(fixture.groups.get(1).getShardsList().stream()\n" - " .noneMatch(shard -> shard.getStoreId() == 4));\n" - " List changes = changeShardInvocations(fixture.partitions);\n" - " Assert.assertEquals(1, changes.size());\n" - " Assert.assertEquals(fallback, changes.get(0).getArgument(0));\n" - " }\n\n" - ) - replace_once(test, marker, regression + marker) - PY - - - name: Format and validate the focused change - shell: bash - run: | - mvn editorconfig:format -ntp - git diff --check - changed="$(git diff --name-only | sort)" - expected="$(printf '%s\n' \ - hugegraph-pd/hg-pd-core/src/main/java/org/apache/hugegraph/pd/StoreNodeService.java \ - hugegraph-pd/hg-pd-service/src/main/java/org/apache/hugegraph/pd/service/PDService.java \ - hugegraph-pd/hg-pd-test/src/main/java/org/apache/hugegraph/pd/core/StoreDrainTest.java | sort)" - test "$changed" = "$expected" || { - printf 'Unexpected changed files:\n%s\n' "$changed" >&2 - exit 1 - } - - - name: Run focused PD drain tests - shell: bash - run: | - mvn test -ntp \ - -pl hugegraph-pd/hg-pd-test -am \ - -Dtest=StoreDrainTest,StorePlacementSubmissionTest \ - -Dsurefire.failIfNoSpecifiedTests=false - - - name: Commit fix and remove one-off workflow - shell: bash - run: | - git config user.name 'chatgpt-codex-connector[bot]' - git config user.email 'chatgpt-codex-connector[bot]@users.noreply.github.com' - git add \ - hugegraph-pd/hg-pd-core/src/main/java/org/apache/hugegraph/pd/StoreNodeService.java \ - hugegraph-pd/hg-pd-service/src/main/java/org/apache/hugegraph/pd/service/PDService.java \ - hugegraph-pd/hg-pd-test/src/main/java/org/apache/hugegraph/pd/core/StoreDrainTest.java - git rm .github/workflows/fix-pr343-duplicate-drain.yml - git diff --cached --check - git commit -m 'fix(pd): avoid duplicate drain migration submission' - git push origin HEAD:codex/pd-store-drain-state-20261027 diff --git a/.github/workflows/stage-pr343-tested-fix.yml b/.github/workflows/stage-pr343-tested-fix.yml deleted file mode 100644 index ae1849e1aa..0000000000 --- a/.github/workflows/stage-pr343-tested-fix.yml +++ /dev/null @@ -1,95 +0,0 @@ -name: Stage PR 343 tested fix - -on: - push: - branches: - - codex/pd-store-drain-state-20261027 - paths: - - .github/workflows/stage-pr343-tested-fix.yml - -permissions: - contents: write - -jobs: - stage-fix: - runs-on: ubuntu-latest - timeout-minutes: 30 - steps: - - name: Checkout PR branch - uses: actions/checkout@v4 - with: - ref: codex/pd-store-drain-state-20261027 - fetch-depth: 0 - - - name: Set up Java 17 - uses: actions/setup-java@v4 - with: - distribution: temurin - java-version: '17' - cache: maven - - - name: Reuse the reviewed patch and validation steps - shell: bash - run: | - python3 - <<'PY' - from pathlib import Path - import subprocess - - source = Path('.github/workflows/fix-pr343-duplicate-drain.yml') - lines = source.read_text().splitlines() - wanted = [ - 'Apply focused fix and regression', - 'Format and validate the focused change', - 'Run focused PD drain tests', - ] - scripts = {} - current = None - index = 0 - while index < len(lines): - line = lines[index] - if line.startswith(' - name: '): - current = line[len(' - name: '):] - if current in wanted and line == ' run: |': - index += 1 - block = [] - while index < len(lines): - item = lines[index] - if item.startswith(' '): - block.append(item[10:]) - index += 1 - continue - if not item.strip(): - block.append('') - index += 1 - continue - break - scripts[current] = '\n'.join(block) + '\n' - continue - index += 1 - - for name in wanted: - script = scripts.get(name) - if not script: - raise SystemExit(f'Missing workflow step: {name}') - print(f'Running reused step: {name}', flush=True) - subprocess.run( - ['bash', '-e', '-u', '-o', 'pipefail', '-c', script], - check=True, - ) - PY - - - name: Push the tested fix to a review branch - shell: bash - run: | - git config user.name 'chatgpt-codex-connector[bot]' - git config user.email 'chatgpt-codex-connector[bot]@users.noreply.github.com' - git add \ - hugegraph-pd/hg-pd-core/src/main/java/org/apache/hugegraph/pd/StoreNodeService.java \ - hugegraph-pd/hg-pd-service/src/main/java/org/apache/hugegraph/pd/service/PDService.java \ - hugegraph-pd/hg-pd-test/src/main/java/org/apache/hugegraph/pd/core/StoreDrainTest.java - git rm \ - .github/workflows/fix-pr343-duplicate-drain.yml \ - .github/workflows/stage-pr343-tested-fix.yml - git diff --cached --check - git commit -m 'fix(pd): avoid duplicate drain migration submission' - git push origin HEAD:refs/heads/codex/pr343-duplicate-drain-fix-20261028