imbajin commented on code in PR #3356:
URL: https://github.com/apache/hugegraph/pull/3356#discussion_r4301480191
##########
hugegraph-pd/hg-pd-core/src/main/java/org/apache/hugegraph/pd/TaskScheduleService.java:
##########
@@ -582,6 +626,113 @@ public synchronized Map<Integer, Long>
balancePartitionLeader(boolean immediatel
return results;
}
+ /**
+ * The periodic entry of leader balancing. Calls balancePartitionLeader
only when the
+ * leader count per store has drifted further apart than
pd.balance-leader-max-skew, so a
+ * balanced cluster sees no transfers. Skipped while a shard group is not
Normal, has no
+ * recent partition heartbeat or has a shard that is neither Leader nor
Follower, or while a
+ * store that holds a shard is not active, because the target distribution
would include a
+ * replica that cannot take a leader. The plan uses the groups checked
here; a group that
+ * changes before its transfer is sent, or whose target trails the leader,
keeps its leader.
+ *
+ * @return shard group id -> the store id its leader was moved to; empty
when skipped
+ */
+ public synchronized Map<Integer, Long> autoBalancePartitionLeader() throws
PDException {
+ // init() does not schedule this with a negative skew, which would
balance every run
+ int maxSkew = pdConfig.getBalanceLeaderMaxSkew();
+ if (!isLeader() || maxSkew < 0) {
Review Comment:
⚠️ Important: `isLeader()` is checked only at the start of this periodic
pass. If leadership changes while this method reads shard groups and index
stats, it still enters the planner; `canTakeLeader()` rechecks the group and
indexes but not PD leadership, and `PartitionService.transferLeader()` forwards
through this PD's local pulse subject. Please recheck leadership before
dispatch so a former PD cannot send stale transfer instructions.
##########
hugegraph-pd/hg-pd-test/src/main/java/org/apache/hugegraph/pd/core/TaskScheduleServiceTest.java:
##########
@@ -75,6 +83,217 @@ public void testBalancePartitionLeader() throws PDException
{
getStoreNodeService().getStoreInfoMeta().removeAll();
}
+ @Test
+ public void testAutoBalancePartitionLeaderWithinSkew() throws PDException {
+ // leaders spread 2/2/2 over the three stores
+ var rst = autoBalance(new long[]{1, 2, 3, 1, 2, 3}, 1, 2, 3);
+ assertTrue(rst.isEmpty());
+ }
+
+ @Test
+ public void testAutoBalancePartitionLeaderOverSkew() throws PDException {
+ // every leader on store 1, as after stores 2 and 3 restarted
+ var rst = autoBalance(new long[]{1, 1, 1, 1, 1, 1}, 1, 2, 3);
+ assertEquals(4, rst.size());
+ assertFalse(rst.containsValue(1L));
+ }
+
+ @Test
+ public void testAutoBalancePartitionLeaderWithInactiveStore() throws
PDException {
+ // store 3 holds shards but sends no heartbeat, so it cannot take a
leader
+ var rst = autoBalance(new long[]{1, 1, 1, 1, 1, 1}, 1, 2);
+ assertTrue(rst.isEmpty());
+ }
+
+ @Test
+ public void testAutoBalancePartitionLeaderWithWarnGroup() throws
PDException {
+ // store 3 just stopped: still in the active set, but its groups are
already Warn
+ var rst = autoBalance(new long[]{1, 1, 1, 1, 1, 1},
+ g ->
g.toBuilder().setState(Metapb.PartitionState.PState_Warn)
+ .build(), 0, 1, 2, 3);
+ assertTrue(rst.isEmpty());
+
+ // once the group is Normal again, the same imbalance is balanced
+ rst = autoBalance(new long[]{1, 1, 1, 1, 1, 1}, g -> g, 0, 1, 2, 3);
+ assertEquals(4, rst.size());
+ }
+
+ @Test
+ public void testAutoBalancePartitionLeaderWithoutQuorum() throws
PDException {
+ // group 0 lost its quorum: no partition heartbeat for a minute, while
its last state
+ // is still Normal and every store is still in the active set
+ var rst = autoBalance(new long[]{1, 1, 1, 1, 1, 1}, g -> g, 60 * 1000,
1, 2, 3);
+ assertTrue(rst.isEmpty());
+
+ // group 0 never reported stats at all
+ rst = autoBalance(new long[]{1, 1, 1, 1, 1, 1}, g -> g, -1, 1, 2, 3);
+ assertTrue(rst.isEmpty());
+
+ // a heartbeat within the window: the same imbalance is balanced
+ rst = autoBalance(new long[]{1, 1, 1, 1, 1, 1}, g -> g, 10 * 1000, 1,
2, 3);
+ assertEquals(4, rst.size());
+ }
+
+ @Test
+ public void testAutoBalancePartitionLeaderWithPendingMember() throws
PDException {
+ // reallocShards added the shard on store 3 and left its role None
+ for (var role : List.of(Metapb.ShardRole.None,
Metapb.ShardRole.Learner)) {
+ var rst = autoBalance(new long[]{1, 1, 1, 1, 1, 1},
+ g -> g.toBuilder().setShards(2,
g.getShards(2).toBuilder()
+
.setRole(role))
+ .build(), 0, 1, 2, 3);
+ assertTrue(role.name(), rst.isEmpty());
+ }
+ }
+
+ @Test
+ public void testAutoBalancePartitionLeaderWithChangeAfterCheck() throws
PDException {
+ // reallocShards turns store 3's shards into new None members right
after the check
+ // read the groups: the plan is made from the checked groups and no
transfer is sent
+ var rst = autoBalance(new long[]{1, 1, 1, 1, 1, 1}, g -> g, 0,
+ (storeId, groupId) -> 10_000,
+ g -> g.toBuilder().setShards(2,
g.getShards(2).toBuilder()
+
.setRole(Metapb.ShardRole.None))
+ .build(), 1, 2, 3);
+ assertTrue(rst.isEmpty());
+ }
+
+ @Test
+ public void testAutoBalancePartitionLeaderWithLaggingFollower() throws
PDException {
+ // store 3's replicas trail the leaders by 5000 entries or report no
index: its share
+ // of the leaders stays on store 1 and only store 2 takes leaders
+ for (long store3Index : new long[]{5_000, -1}) {
+ var rst = autoBalance(new long[]{1, 1, 1, 1, 1, 1}, g -> g, 0,
+ (storeId, groupId) -> storeId == 3 ?
store3Index : 10_000,
+ null, 1, 2, 3);
+ assertEquals(2, rst.size());
+ assertTrue(rst.values().stream().allMatch(storeId -> storeId ==
2L));
+ }
+
+ // 500 entries behind is close enough to take a leader
+ var rst = autoBalance(new long[]{1, 1, 1, 1, 1, 1}, g -> g, 0,
+ (storeId, groupId) -> storeId == 3 ? 9_500 :
10_000,
+ null, 1, 2, 3);
+ assertEquals(4, rst.size());
+ assertTrue(rst.containsValue(3L));
+ }
+
+ @Test
+ public void testAutoBalancePartitionLeaderWithNegativeSkew() throws
PDException {
+ // with the default skew this layout moves four leaders; -1 must not
run a plan
+ getPdConfig().setBalanceLeaderMaxSkew(-1);
+ try {
+ var rst = autoBalance(new long[]{1, 1, 1, 1, 1, 1}, 1, 2, 3);
+ assertTrue(rst.isEmpty());
+ } finally {
+ getPdConfig().setBalanceLeaderMaxSkew(2);
+ }
+ }
+
+ private Map<Integer, Long> autoBalance(long[] leaderStores, long...
activeStores) throws
+
PDException {
+ return autoBalance(leaderStores, g -> g, 0, activeStores);
+ }
+
+ private Map<Integer, Long> autoBalance(long[] leaderStores,
+ UnaryOperator<Metapb.ShardGroup>
group0,
+ long group0StatsAge,
+ long... activeStores) throws
PDException {
+ return autoBalance(leaderStores, group0, group0StatsAge, (storeId,
groupId) -> 10_000,
+ null, activeStores);
+ }
+
+ /**
+ * Stores shard groups 0..n-1 on stores 1, 2 and 3 with the leader of
group i on
+ * leaderStores[i], all Normal with fresh partition stats, then applies
group0 to group 0
+ * and ages its stats by group0StatsAge ms (-1: no stats). Marks
activeStores alive, each
+ * reporting appliedIndex(store, group) for every group (negative: not
reported), and runs
+ * one periodic balance. A non-null afterRead is applied to every stored
group right after
+ * the balance reads the group list, like a change landing between check
and transfer.
+ */
+ private Map<Integer, Long> autoBalance(long[] leaderStores,
+ UnaryOperator<Metapb.ShardGroup>
group0,
+ long group0StatsAge,
+ ToLongBiFunction<Long, Integer>
appliedIndex,
+ UnaryOperator<Metapb.ShardGroup>
afterRead,
+ long... activeStores) throws
PDException {
+ var storeMeta = getStoreNodeService().getStoreInfoMeta();
+ var partitionMeta = MetadataFactory.newPartitionMeta(getPdConfig());
+ // Set aside groups other tests in the suite left behind, so only ours
are planned
+ var otherGroups = storeMeta.getShardGroups();
+ for (var group : otherGroups) {
+ storeMeta.deleteShardGroup(group.getId());
+ }
+ var partitions = new ArrayList<Metapb.Partition>();
+ long now = System.currentTimeMillis();
+ for (int i = 0; i < leaderStores.length; i++) {
+ var group = genShardGroup(i, leaderStores[i],
Metapb.PartitionState.PState_Normal);
+ long statsAge = 0;
+ if (i == 0) {
+ group = group0.apply(group);
+ statsAge = group0StatsAge;
+ }
+ storeMeta.updateShardGroup(group);
+ if (statsAge >= 0) {
+
partitionMeta.updatePartitionStats(Metapb.PartitionStats.newBuilder()
+
.setId(i)
+
.setTimestamp(
+
now - statsAge)
+
.build());
+ }
+ partitions.add(genPartition(i));
+ }
+ for (long storeId = 1; storeId <= 3; storeId++) {
+ var stats = Metapb.StoreStats.newBuilder();
+ for (int i = 0; i < leaderStores.length; i++) {
+ long index = appliedIndex.applyAsLong(storeId, i);
+ if (index >= 0) {
+
stats.addRaftStats(Metapb.RaftStats.newBuilder().setPartitionId(i)
+
.setCommittedIndex(index));
+ }
+ }
+ var store =
Metapb.Store.newBuilder().setId(storeId).setStats(stats).build();
+ if (Arrays.stream(activeStores).anyMatch(id -> id ==
store.getId())) {
+ storeMeta.keepStoreAlive(store);
+ } else {
+ storeMeta.removeActiveStore(store);
+ }
+ }
+ getPdConfig().getPartition().setShardCount(3);
+ getPartitionService().updatePartition(partitions);
+ try {
+ StoreNodeService storeService = getStoreNodeService();
+ if (afterRead != null) {
+ storeService = new StoreNodeService(getPdConfig()) {
+ @Override
+ public List<Metapb.ShardGroup> getShardGroups() throws
PDException {
+ var groups = super.getShardGroups();
+ for (var group : groups) {
+ storeMeta.updateShardGroup(afterRead.apply(group));
+ }
+ return groups;
+ }
+ };
+ }
+ // A fresh service, so a balance run by another test does not hit
the rate limit
+ var service = new TaskScheduleService(getPdConfig(), storeService,
+ getPartitionService());
+ return service.autoBalancePartitionLeader();
Review Comment:
⚠️ Important: Every added case calls `autoBalancePartitionLeader()` directly
through this helper; none exercises `init()` or asserts that the timer is
registered only for valid enabled configuration. A regression in the new
schedule/config gate would leave the feature inactive (or schedule it under
invalid settings) while these tests pass. Please add deterministic coverage for
disabled, valid, and invalid configuration.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
