diff --git a/helm/hugegraph/docs/operations.md b/helm/hugegraph/docs/operations.md index 4c90e1e6ba..b2474854c6 100644 --- a/helm/hugegraph/docs/operations.md +++ b/helm/hugegraph/docs/operations.md @@ -245,12 +245,28 @@ curl --user "hg:${PD_SECRET}" "${PD_REST}/v1/task/balanceLeaders" ``` A busy task returns body `status: 1001`; inspect the result and wait for current -maintenance to finish rather than looping patrol/balance calls. A follower can -return an empty success, and leader transfers are asynchronous: HTTP 200 or a +maintenance to finish rather than looping patrol/balance calls. Images containing +the task response fix for #3231 return body `status: 100` on a follower, naming the +known leader's Raft address (`null` while no leader is known). Legacy images can +still return empty success from followers. + +On updated images, partition patrol retains `status` and `partitions` and adds +`reallocatedShardGroups` (groups whose shard allocation changed) and +`offlineShardRequests` (group ID to retired Store ID). The latter records requests, +not successful repairs: inspect group membership to verify their outcome. +An empty result means patrol ran without scheduling changes. Leader and partition +balances retain their transfer maps; a partition balance skipped because of its +maintenance gates returns `{"status":0,"outcome":"skipped"}` instead of an empty +body. Nonzero `status` identifies a refusal rather than a successful empty plan. + +Leader transfers are asynchronous: HTTP 200 or a returned transfer plan does not prove completion. Recheck `/v1/partitions` and `/v1/stores` for actual leaders, full membership and fresh heartbeats before continuing maintenance. + + Preserve volumes and recorded member identities during recovery. Deleting a RocksDB lock file or assigning an empty volume does not recover shard data. The chart does not provide backup or automatic disk-loss recovery; keep backups diff --git a/hugegraph-pd/docs/api-reference.md b/hugegraph-pd/docs/api-reference.md index 35d70fadd4..bf3e558b4b 100644 --- a/hugegraph-pd/docs/api-reference.md +++ b/hugegraph-pd/docs/api-reference.md @@ -842,6 +842,28 @@ Point Kubernetes readiness probes, `depends_on` healthchecks and any "wait for P Match on the body rather than on the status code alone. A PD that predates this endpoint does not reliably answer `404` for it: `RestAuthentication` refuses a request it does not exclude by writing an error envelope, and as of 1.7.0 it does so without setting a status, so an unknown path answers `200` with `{"status":-1,"error":"Unauthorized!"}`. A status-only probe therefore reads such a PD as ready. The body match holds whichever status a refusal carries: a shell gate should use `curl -fsS http://:8620/v1/ready | grep -q '"ready":true'`, and a Kubernetes `httpGet` probe should be paired with a PD image that carries the endpoint. +### Manual maintenance task results + +Task endpoints under `/v1/task` require the local PD to be leader. A follower +answers an error envelope with `status: 100` and the known leader's Raft address +in `error` (`null` if no leader is known). Check `status`, not HTTP 200 alone. +`/v1/balanceLeaders` applies the same check. A leadership change during execution +can also return this error; inspect cluster state before retrying. + +`/v1/task/patrolPartitions` retains `status` and `partitions`, and adds +`reallocatedShardGroups` (IDs whose allocation changed) and `offlineShardRequests` +(group ID to retired Store ID). Offline-shard requests do not prove completed +repairs. Empty collections distinguish a completed patrol with nothing to do. +The changed-group IDs include tombstone-triggered reallocations and are deduplicated +against the initial shard-count pass. + +Leader and partition balances retain their transfer maps. A skipped partition +balance instead returns `status: 0` with `outcome: skipped`; a completed balance +with no transfers returns `{}`. Transfers are asynchronous, so verify membership +and leaders before proceeding with maintenance. +`/v1/task/compact` preserves its `compact ok` success body as `text/plain`; +structured error bodies use `application/json`, including follower refusals. + ### Metrics ```bash diff --git a/hugegraph-pd/hg-pd-core/src/main/java/org/apache/hugegraph/pd/StoreNodeService.java b/hugegraph-pd/hg-pd-core/src/main/java/org/apache/hugegraph/pd/StoreNodeService.java index f8d8b81ed6..13d84183f6 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 @@ -336,6 +336,11 @@ public synchronized Metapb.Store updateStore(Metapb.Store store) throws PDExcept * @throws PDException */ public synchronized void storeTurnoff(Metapb.Store store) throws PDException { + storeTurnoffWithResult(store); + } + + public synchronized List storeTurnoffWithResult(Metapb.Store store) throws PDException { + List changed = new ArrayList<>(); // Traverse ShardGroup,redistribution for (Metapb.ShardGroup group : getShardGroupsByStore(store.getId())) { Metapb.ShardGroup.Builder builder = Metapb.ShardGroup.newBuilder(group); @@ -345,8 +350,12 @@ public synchronized void storeTurnoff(Metapb.Store store) throws PDException { builder.addShards(shard); } }); - reallocShards(builder.build()); + List shards = reallocShards(builder.build()); + if (!group.getShardsList().equals(shards)) { + changed.add(group.getId()); + } } + return changed; } /** 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..fb5352c813 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 @@ -41,6 +41,7 @@ import org.apache.hugegraph.pd.meta.TaskInfoMeta; import org.apache.hugegraph.pd.raft.RaftEngine; +import lombok.Getter; import lombok.extern.slf4j.Slf4j; /** @@ -249,14 +250,31 @@ public List patrolStores() throws PDException { * of replicas in the shard group */ public List patrolPartitions() throws PDException { + PartitionPatrolResult result = patrolPartitionsWithResult(); + return result == null ? null : result.getPartitions(); + } + + @Getter + public static class PartitionPatrolResult { + + private final List reallocatedShardGroups = new ArrayList<>(); + private final List partitions = new ArrayList<>(); + private final Map offlineShardRequests = new HashMap<>(); + } + + public PartitionPatrolResult patrolPartitionsWithResult() throws PDException { if (!isLeader()) { return null; } + PartitionPatrolResult result = new PartitionPatrolResult(); // If the number of replicas is inconsistent, reallocate replicas for (Metapb.ShardGroup group : storeService.getShardGroups()) { if (group.getShardsCount() != pdConfig.getPartition().getShardCount()) { - storeService.reallocShards(group); + List shards = storeService.reallocShards(group); + if (!group.getShardsList().equals(shards)) { + result.reallocatedShardGroups.add(group.getId()); + } kvService.put(BALANCE_SHARD_KEY, "DOING", 180 * 1000); } } @@ -273,13 +291,20 @@ public List patrolPartitions() throws PDException { } partIds.add(partition.getId()); - storeService.storeTurnoff(pair.getValue()); + for (Integer changed : storeService.storeTurnoffWithResult(pair.getValue())) { + if (!result.reallocatedShardGroups.contains(changed)) { + result.reallocatedShardGroups.add(changed); + } + } partitionService.shardOffline(partition, pair.getValue().getId()); + // shardOffline logs its own failures; report a request, not a completed repair. + result.offlineShardRequests.put(partition.getId(), pair.getValue().getId()); + result.partitions.add(partition); } } - return null; + return result; } /** diff --git a/hugegraph-pd/hg-pd-service/src/main/java/org/apache/hugegraph/pd/rest/TaskAPI.java b/hugegraph-pd/hg-pd-service/src/main/java/org/apache/hugegraph/pd/rest/TaskAPI.java index eaa53908d1..58ae7266b0 100644 --- a/hugegraph-pd/hg-pd-service/src/main/java/org/apache/hugegraph/pd/rest/TaskAPI.java +++ b/hugegraph-pd/hg-pd-service/src/main/java/org/apache/hugegraph/pd/rest/TaskAPI.java @@ -20,12 +20,14 @@ import java.util.List; import java.util.Map; +import org.apache.hugegraph.pd.TaskScheduleService.PartitionPatrolResult; import org.apache.hugegraph.pd.common.KVPair; import org.apache.hugegraph.pd.common.PDException; import org.apache.hugegraph.pd.grpc.Metapb; import org.apache.hugegraph.pd.service.PDRestService; import org.springframework.beans.factory.annotation.Autowired; import org.springframework.http.MediaType; +import org.springframework.http.ResponseEntity; import org.springframework.web.bind.annotation.GetMapping; import org.springframework.web.bind.annotation.RequestMapping; import org.springframework.web.bind.annotation.ResponseBody; @@ -57,8 +59,11 @@ public String patrolStores() { @ResponseBody public String patrolPartitions() { try { - List partitions = pdRestService.patrolPartitions(); - return toJSON(partitions, "partitions"); + PartitionPatrolResult result = pdRestService.patrolPartitionsWithResult(); + String partitions = toJSON(result.getPartitions(), "partitions"); + return partitions.substring(0, partitions.length() - 1) + + ",\"reallocatedShardGroups\":" + toJSON(result.getReallocatedShardGroups()) + + ",\"offlineShardRequests\":" + toJSON(result.getOfflineShardRequests()) + "}"; } catch (PDException e) { e.printStackTrace(); return toJSON(e); @@ -67,13 +72,16 @@ public String patrolPartitions() { @GetMapping(value = "/balancePartitions", produces = MediaType.APPLICATION_JSON_VALUE) @ResponseBody - public Map> balancePartitions() { + public String balancePartitions() { try { Map> partitions = pdRestService.balancePartitions(); - return partitions; + if (partitions == null) { + return toJSON(Map.of("status", 0, "outcome", "skipped")); + } + return toJSON(partitions); } catch (PDException e) { e.printStackTrace(); - return null; + return toJSON(e); } } @@ -89,12 +97,6 @@ public String splitPartitions() { } } - // TODO: make the task routes tell a real run from a no-op and a follower answer (a follower - // returns an empty success without doing anything), and return a body that names what was - // scheduled. Part 1 (#3233) replaced the bare 500 with a refusal body; this is part 2. The - // Helm chart (helm/hugegraph) documents the leader-first sequence and the 180 s spacing in its - // README until then. - // https://github.com/apache/hugegraph/issues/3231 @GetMapping(value = "/balanceLeaders", produces = MediaType.APPLICATION_JSON_VALUE) @ResponseBody public String balanceLeaders() { @@ -108,8 +110,12 @@ public String balanceLeaders() { } @GetMapping(value = "/compact") - public String dbCompaction() throws PDException { - pdRestService.dbCompaction(); - return "compact ok"; + public ResponseEntity dbCompaction() { + try { + pdRestService.dbCompaction(); + return ResponseEntity.ok().contentType(MediaType.TEXT_PLAIN).body("compact ok"); + } catch (PDException e) { + return ResponseEntity.ok().contentType(MediaType.APPLICATION_JSON).body(toJSON(e)); + } } } diff --git a/hugegraph-pd/hg-pd-service/src/main/java/org/apache/hugegraph/pd/service/PDRestService.java b/hugegraph-pd/hg-pd-service/src/main/java/org/apache/hugegraph/pd/service/PDRestService.java index c9d55370cf..014c34704e 100644 --- a/hugegraph-pd/hg-pd-service/src/main/java/org/apache/hugegraph/pd/service/PDRestService.java +++ b/hugegraph-pd/hg-pd-service/src/main/java/org/apache/hugegraph/pd/service/PDRestService.java @@ -40,6 +40,7 @@ import org.apache.hugegraph.pd.grpc.discovery.RegisterInfo; import org.apache.hugegraph.pd.model.RegistryRestRequest; import org.apache.hugegraph.pd.model.RegistryRestResponse; +import org.apache.hugegraph.pd.raft.RaftEngine; import org.springframework.beans.factory.InitializingBean; import org.springframework.beans.factory.annotation.Autowired; import org.springframework.stereotype.Service; @@ -137,11 +138,38 @@ public Map getShardGroupCache() { } public List patrolStores() throws PDException { - return monitorService.patrolStores(); + return runLeaderTask(() -> monitorService.patrolStores()); } public List patrolPartitions() throws PDException { - return monitorService.patrolPartitions(); + return runLeaderTask(() -> monitorService.patrolPartitions()); + } + + public TaskScheduleService.PartitionPatrolResult patrolPartitionsWithResult() throws PDException { + return runLeaderTask(() -> monitorService.patrolPartitionsWithResult()); + } + + private T runLeaderTask(LeaderTask task) throws PDException { + checkTaskLeader(); + T result = task.run(); + // Do not turn a scheduler's follower no-op into success if leadership changed mid-call. + checkTaskLeader(); + return result; + } + + private void checkTaskLeader() throws PDException { + RaftEngine engine = RaftEngine.getInstance(); + if (!engine.isLeader()) { + throw new PDException(Pdpb.ErrorType.NOT_LEADER_VALUE, + "Task requires the PD leader; leader Raft address: " + + engine.getLeader()); + } + } + + @FunctionalInterface + private interface LeaderTask { + + T run() throws PDException; } public Metapb.PartitionStats getPartitionStats(String graphName, int partitionId) throws @@ -154,11 +182,11 @@ public List getPartitionStatus(String graphName) throws P } public Map> balancePartitions() throws PDException { - return monitorService.balancePartitionShard(); + return runLeaderTask(() -> monitorService.balancePartitionShard()); } public List splitPartitions() throws PDException { - return monitorService.autoSplitPartition(); + return runLeaderTask(() -> monitorService.autoSplitPartition()); } public List getStoreStats(boolean isActive) throws PDException { @@ -262,11 +290,17 @@ public List getPartitionLog(Long start, Long end) throws PDExc } public Map balancePartitionLeader() throws PDException { - return monitorService.balancePartitionLeader(true); + return runLeaderTask(() -> monitorService.balancePartitionLeader(true)); } public void dbCompaction() throws PDException { - monitorService.dbCompaction(""); + runLeaderTask(() -> { + if (!monitorService.dbCompaction("")) { + throw new PDException(Pdpb.ErrorType.NOT_LEADER_VALUE, + "Compaction was refused by the scheduler; retry on the PD leader"); + } + return null; + }); } public List getShardList(int partitionId) throws PDException { diff --git a/hugegraph-pd/hg-pd-service/src/test/java/org/apache/hugegraph/pd/rest/TaskAPIResultTest.java b/hugegraph-pd/hg-pd-service/src/test/java/org/apache/hugegraph/pd/rest/TaskAPIResultTest.java new file mode 100644 index 0000000000..1ea0566ead --- /dev/null +++ b/hugegraph-pd/hg-pd-service/src/test/java/org/apache/hugegraph/pd/rest/TaskAPIResultTest.java @@ -0,0 +1,160 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.hugegraph.pd.rest; + +import java.lang.reflect.Proxy; +import java.util.Map; + +import org.apache.hugegraph.pd.TaskScheduleService.PartitionPatrolResult; +import org.apache.hugegraph.pd.common.KVPair; +import org.apache.hugegraph.pd.common.PDException; +import org.apache.hugegraph.pd.grpc.Metapb; +import org.apache.hugegraph.pd.grpc.Pdpb; +import org.apache.hugegraph.pd.raft.RaftEngine; +import org.apache.hugegraph.pd.service.PDRestService; +import org.apache.hugegraph.testutil.Whitebox; +import org.junit.Assert; +import org.junit.Test; +import org.springframework.http.MediaType; + +import com.alipay.sofa.jraft.Node; +import com.alipay.sofa.jraft.entity.PeerId; +import com.fasterxml.jackson.databind.JsonNode; +import com.fasterxml.jackson.databind.ObjectMapper; + +public class TaskAPIResultTest { + + @Test + public void testFollowerRefusesEveryTaskWithoutInvokingScheduler() throws Exception { + RaftEngine engine = RaftEngine.getInstance(); + Node original = engine.getRaftNode(); + Node follower = (Node) Proxy.newProxyInstance( + Node.class.getClassLoader(), new Class[]{Node.class}, + (proxy, method, args) -> { + if (method.getName().equals("isLeader")) { + return false; + } + if (method.getName().equals("getLeaderId")) { + return new PeerId("127.0.0.2", 8610); + } + throw new AssertionError("unexpected node call: " + method.getName()); + }); + Whitebox.setInternalState(engine, "raftNode", follower); + try { + TaskAPI api = new TaskAPI(); + // Its scheduler is deliberately unset: no follower request may reach it. + api.pdRestService = new PDRestService(); + for (String response : new String[]{api.patrolStores(), api.patrolPartitions(), + api.balancePartitions(), api.balanceLeaders(), + api.splitPartitions(), api.dbCompaction().getBody()}) { + JsonNode body = json(response); + Assert.assertEquals(Pdpb.ErrorType.NOT_LEADER_VALUE, body.get("status").asInt()); + Assert.assertTrue(body.get("error").asText().contains("127.0.0.2:8610")); + } + StoreAPI alias = new StoreAPI(); + alias.pdRestService = api.pdRestService; + Assert.assertEquals(Pdpb.ErrorType.NOT_LEADER_VALUE, + json(alias.balanceLeaders()).get("status").asInt()); + } finally { + Whitebox.setInternalState(engine, "raftNode", original); + } + } + + @Test + public void testPatrolReportsReallocationsAndOfflineRequests() throws Exception { + PartitionPatrolResult result = new PartitionPatrolResult(); + result.getReallocatedShardGroups().add(7); + result.getOfflineShardRequests().put(8, 42L); + result.getPartitions().add(Metapb.Partition.newBuilder().setId(8) + .setGraphName("graph").build()); + TaskAPI api = patrolAPI(result); + + JsonNode body = json(api.patrolPartitions()); + Assert.assertEquals(0, body.get("status").asInt()); + Assert.assertEquals(7, body.get("reallocatedShardGroups").get(0).asInt()); + Assert.assertEquals(42L, body.get("offlineShardRequests").get("8").asLong()); + Assert.assertEquals(8, body.get("partitions").get(0).get("id").asInt()); + } + + @Test + public void testPatrolWithNothingToDoHasEmptyResults() throws Exception { + JsonNode body = json(patrolAPI(new PartitionPatrolResult()).patrolPartitions()); + Assert.assertEquals(0, body.get("status").asInt()); + Assert.assertEquals(0, body.get("partitions").size()); + Assert.assertEquals(0, body.get("reallocatedShardGroups").size()); + Assert.assertEquals(0, body.get("offlineShardRequests").size()); + } + + @Test + public void testSkippedBalanceDiffersFromCompletedEmptyBalance() throws Exception { + TaskAPI skipped = new TaskAPI(); + skipped.pdRestService = balanceService(null); + Assert.assertEquals("skipped", json(skipped.balancePartitions()).get("outcome").asText()); + + TaskAPI completed = new TaskAPI(); + completed.pdRestService = balanceService(Map.of()); + Assert.assertEquals("{}", completed.balancePartitions()); + } + + @Test + public void testCompactionContentTypesPreservePlainTextSuccess() throws Exception { + TaskAPI api = new TaskAPI(); + api.pdRestService = new PDRestService() { + @Override + public void dbCompaction() { + } + }; + Assert.assertEquals("compact ok", api.dbCompaction().getBody()); + Assert.assertEquals(MediaType.TEXT_PLAIN, + api.dbCompaction().getHeaders().getContentType()); + api.pdRestService = new PDRestService() { + @Override + public void dbCompaction() throws PDException { + throw new PDException(Pdpb.ErrorType.NOT_LEADER_VALUE, "follower"); + } + }; + Assert.assertEquals(MediaType.APPLICATION_JSON, + api.dbCompaction().getHeaders().getContentType()); + Assert.assertEquals(Pdpb.ErrorType.NOT_LEADER_VALUE, + json(api.dbCompaction().getBody()).get("status").asInt()); + } + + private static TaskAPI patrolAPI(PartitionPatrolResult result) { + TaskAPI api = new TaskAPI(); + api.pdRestService = new PDRestService() { + @Override + public PartitionPatrolResult patrolPartitionsWithResult() { + return result; + } + }; + return api; + } + + private static PDRestService balanceService(Map> result) { + return new PDRestService() { + @Override + public Map> balancePartitions() { + return result; + } + }; + } + + private static JsonNode json(String response) throws Exception { + return new ObjectMapper().readTree(response); + } +} 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 c0c15975e8..5329a9e93a 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 @@ -53,6 +53,7 @@ StoreMonitorDataServiceTest.class, StoreServiceTest.class, TaskScheduleServiceTest.class, + PartitionPatrolResultTest.class, IpAuthHandlerTest.class, RaftEngineIpAuthIntegrationTest.class, RaftEngineLeaderAddressTest.class, diff --git a/hugegraph-pd/hg-pd-test/src/main/java/org/apache/hugegraph/pd/core/PartitionPatrolResultTest.java b/hugegraph-pd/hg-pd-test/src/main/java/org/apache/hugegraph/pd/core/PartitionPatrolResultTest.java new file mode 100644 index 0000000000..048bc304bd --- /dev/null +++ b/hugegraph-pd/hg-pd-test/src/main/java/org/apache/hugegraph/pd/core/PartitionPatrolResultTest.java @@ -0,0 +1,178 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.hugegraph.pd.core; + +import java.util.List; +import java.util.Map; + +import org.apache.hugegraph.pd.KvService; +import org.apache.hugegraph.pd.PartitionService; +import org.apache.hugegraph.pd.StoreNodeService; +import org.apache.hugegraph.pd.TaskScheduleService; +import org.apache.hugegraph.pd.TaskScheduleService.PartitionPatrolResult; +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.rest.TaskAPI; +import org.apache.hugegraph.pd.service.PDRestService; +import org.apache.hugegraph.pd.raft.RaftEngine; +import org.apache.hugegraph.testutil.Whitebox; +import org.junit.After; +import org.junit.Assert; +import org.junit.Before; +import org.junit.Test; +import org.mockito.Mockito; + +import com.alipay.sofa.jraft.Node; +import com.alipay.sofa.jraft.entity.PeerId; +import com.fasterxml.jackson.databind.JsonNode; +import com.fasterxml.jackson.databind.ObjectMapper; + +public class PartitionPatrolResultTest { + + private Node original; + private Node node; + private TaskScheduleService service; + private StoreNodeService stores; + private PartitionService partitions; + + @Before + public void setUp() { + original = RaftEngine.getInstance().getRaftNode(); + node = Mockito.mock(Node.class); + Mockito.when(node.isLeader(true)).thenReturn(true); + Whitebox.setInternalState(RaftEngine.getInstance(), "raftNode", node); + service = Mockito.mock(TaskScheduleService.class, Mockito.CALLS_REAL_METHODS); + stores = Mockito.mock(StoreNodeService.class); + partitions = Mockito.mock(PartitionService.class); + PDConfig config = new PDConfig(); + PDConfig.Partition partition = config.new Partition(); + partition.setShardCount(3); + config.setPartition(partition); + Whitebox.setInternalState(service, "pdConfig", config); + Whitebox.setInternalState(service, "storeService", stores); + Whitebox.setInternalState(service, "partitionService", partitions); + Whitebox.setInternalState(service, "kvService", Mockito.mock(KvService.class)); + } + + @After + public void tearDown() { + Whitebox.setInternalState(RaftEngine.getInstance(), "raftNode", original); + } + + @Test + public void testReportsChangedGroupsAndOfflineRequests() throws Exception { + Metapb.Shard oldShard = Metapb.Shard.newBuilder().setStoreId(1).build(); + Metapb.Shard added = Metapb.Shard.newBuilder().setStoreId(2).build(); + Metapb.ShardGroup changed = Metapb.ShardGroup.newBuilder().setId(7).addShards(oldShard).build(); + Metapb.ShardGroup unchanged = changed.toBuilder().setId(8).build(); + Mockito.when(stores.getShardGroups()).thenReturn(List.of(changed, unchanged)); + Mockito.when(stores.reallocShards(changed)).thenReturn(List.of(oldShard, added)); + Mockito.when(stores.reallocShards(unchanged)).thenReturn(List.of(oldShard)); + Metapb.Store retired = Metapb.Store.newBuilder().setId(9).build(); + Metapb.Partition affected = Metapb.Partition.newBuilder().setId(10).setGraphName("graph").build(); + Mockito.when(stores.getTombStores()).thenReturn(List.of(retired)); + Mockito.when(partitions.getPartitionByStore(retired)).thenReturn(List.of(affected, affected)); + + PartitionPatrolResult result = service.patrolPartitionsWithResult(); + + Assert.assertEquals(List.of(7), result.getReallocatedShardGroups()); + Assert.assertEquals(Map.of(10, 9L), result.getOfflineShardRequests()); + Assert.assertEquals(List.of(affected), result.getPartitions()); + Mockito.verify(partitions, Mockito.times(1)).shardOffline(affected, 9L); + } + + @Test + public void testFollowerDoesNotPatrol() throws Exception { + Mockito.when(node.isLeader(true)).thenReturn(false); + Assert.assertNull(service.patrolPartitionsWithResult()); + Mockito.verifyNoInteractions(stores, partitions); + } + + @Test + public void testLeadershipChangeAfterPatrolReturnsStructuredRefusal() throws Exception { + Mockito.when(node.isLeader(true)).thenReturn(true, false); + Mockito.when(node.getLeaderId()).thenReturn(new PeerId("127.0.0.2", 8610)); + TaskScheduleService scheduler = Mockito.mock(TaskScheduleService.class); + Mockito.when(scheduler.patrolPartitionsWithResult()).thenReturn(new PartitionPatrolResult()); + PDRestService rest = new PDRestService(); + Whitebox.setInternalState(rest, "monitorService", scheduler); + TaskAPI api = new TaskAPI(); + Whitebox.setInternalState(api, "pdRestService", rest); + JsonNode response = + new ObjectMapper().readTree(api.patrolPartitions()); + Assert.assertEquals(Pdpb.ErrorType.NOT_LEADER_VALUE, + response.get("status").asInt()); + Mockito.verify(scheduler).patrolPartitionsWithResult(); + } + + @Test + public void testCompactionRefusalSurvivesLeadershipRegain() throws Exception { + // The outer checks both see a leader, but the scheduler observed a follower in between. + Mockito.when(node.isLeader(true)).thenReturn(true); + TaskScheduleService scheduler = Mockito.mock(TaskScheduleService.class); + Mockito.when(scheduler.dbCompaction("")).thenReturn(false); + PDRestService rest = new PDRestService(); + Whitebox.setInternalState(rest, "monitorService", scheduler); + TaskAPI api = new TaskAPI(); + Whitebox.setInternalState(api, "pdRestService", rest); + JsonNode response = new ObjectMapper().readTree(api.dbCompaction().getBody()); + Assert.assertEquals(Pdpb.ErrorType.NOT_LEADER_VALUE, response.get("status").asInt()); + Mockito.verify(scheduler).dbCompaction(""); + Mockito.when(scheduler.dbCompaction("")).thenReturn(true); + Assert.assertEquals("compact ok", api.dbCompaction().getBody()); + } + + @Test + public void testReportsTombstoneReallocationWithCorrectInitialShardCount() throws Exception { + assertTombstoneReallocation(false); + } + + @Test + public void testDeduplicatesTombstoneAndInitialReallocations() throws Exception { + assertTombstoneReallocation(true); + } + + private void assertTombstoneReallocation(boolean initiallyUnderReplicated) throws Exception { + stores = Mockito.mock(StoreNodeService.class, Mockito.CALLS_REAL_METHODS); + Whitebox.setInternalState(service, "storeService", stores); + Metapb.Shard retiredShard = Metapb.Shard.newBuilder().setStoreId(9).build(); + Metapb.Shard second = Metapb.Shard.newBuilder().setStoreId(2).build(); + Metapb.Shard third = Metapb.Shard.newBuilder().setStoreId(3).build(); + Metapb.Shard replacement = Metapb.Shard.newBuilder().setStoreId(4).build(); + Metapb.ShardGroup before = Metapb.ShardGroup.newBuilder().setId(7) + .addAllShards(List.of(retiredShard, second, third)).build(); + Metapb.ShardGroup initial = initiallyUnderReplicated ? + before.toBuilder().clearShards().addShards(retiredShard).build() : before; + Mockito.doReturn(List.of(initial)).when(stores).getShardGroups(); + Mockito.doReturn(List.of(before)).when(stores).getShardGroupsByStore(9L); + if (initiallyUnderReplicated) { + Mockito.doReturn(before.getShardsList()).when(stores).reallocShards(initial); + } + Mockito.doReturn(List.of(second, third, replacement)).when(stores) + .reallocShards(before.toBuilder().clearShards().addAllShards(List.of(second, third)).build()); + Metapb.Store retired = Metapb.Store.newBuilder().setId(9).build(); + Mockito.doReturn(List.of(retired)).when(stores).getTombStores(); + Metapb.Partition affected = Metapb.Partition.newBuilder().setId(7).setGraphName("graph").build(); + Mockito.when(partitions.getPartitionByStore(retired)).thenReturn(List.of(affected)); + + PartitionPatrolResult result = service.patrolPartitionsWithResult(); + Assert.assertEquals(List.of(7), result.getReallocatedShardGroups()); + Assert.assertEquals(Map.of(7, 9L), result.getOfflineShardRequests()); + } +}