[Feature] PD: optional periodic partition leader balancing in TaskScheduleService

[bitflicker64]

Feature Description (功能描述)

Problem

Under the current implementation, leader balancing is only performed when explicitly requested; there is no periodic or event-driven invocation of the leader balance logic. After a store restart or rebuild, the raft groups that lost their leader elect new leaders on the surviving stores, and those leaders stay where they landed. As a result, leader distribution can remain skewed until a manual balance operation is triggered (for example after noticing the skew in Hubble or via /v1/shards).

I would like PD to optionally run its existing leader balance logic on a schedule, gated by configuration, so that clusters converge back to an even leader distribution without operator intervention.

All file and line references below are against the current master branch of apache/hugegraph and were verified there at the time of writing.

Current behavior (verified on master)

1. TaskScheduleService.init() schedules four periodic jobs, none of which balances leaders (verified on master, hugegraph-pd/hg-pd-core/src/main/java/org/apache/hugegraph/pd/TaskScheduleService.java, lines 121-189):

   • patrolStores() every 60s (lines 122-129), which only flags stores Offline when they miss heartbeats,
   • kvService.clearTTLData() every 1s (lines 130-134),
   • storeService.getQuota() every 30s (lines 135-145),
   • expired monitor data cleanup every 600s, gated by store.monitor_data_enabled (lines 147-166).

   Neither balancePartitionLeader(...) (lines 455-567) nor balancePartitionShard() (lines 273-450) is referenced from init(). There is even an unused constant KEY_ENABLE_AUTO_BALANCE = "key/ENABLE_AUTO_BALANCE" at line 58 that no code reads, which suggests automatic balancing was considered at some point but never wired up.

2. The only triggers for leader balancing today are manual, request-driven entry points (verified on master):

   • REST: GET /v1/task/balanceLeaders in hugegraph-pd/hg-pd-service/src/main/java/org/apache/hugegraph/pd/rest/TaskAPI.java (mapping at line 38, endpoint at lines 92-95), and a second mapping GET /v1/balanceLeaders in rest/StoreAPI.java (lines 55, 149-151). Both call PDRestService.balancePartitionLeader() (service/PDRestService.java, lines 264-266), which invokes TaskScheduleService.balancePartitionLeader(true).
   • REST: GET /v1/task/balancePartitions (rest/TaskAPI.java, lines 68-78) for shard balancing.
   • gRPC: PDService.balanceLeaders(...) (hugegraph-pd/hg-pd-service/src/main/java/org/apache/hugegraph/pd/service/PDService.java, lines 1353-1372) calls balancePartitionLeader(true); PDService.movePartition(...) (lines 1182-1200) calls balancePartitionShard(). The client wrapper is PDClient.balanceLeaders() (hugegraph-pd/hg-pd-client/src/main/java/org/apache/hugegraph/pd/client/PDClient.java, lines 1200-1206).

   None of these are invoked by a timer or a store status event inside PD.

3. How leaders end up skewed (verified on master):

   • Initial allocation is balanced by construction: StoreNodeService.allocShards(...) assigns shard groups round-robin and marks the first shard of each group as Leader (hugegraph-pd/hg-pd-core/src/main/java/org/apache/hugegraph/pd/StoreNodeService.java, lines 456-468).
   • After that, leadership is decided by raft elections on the store side. PD only learns the outcome through partition heartbeats: PartitionService.partitionHeartbeat(...) updates the shard group when the leader term or shard set changes (hugegraph-pd/hg-pd-core/src/main/java/org/apache/hugegraph/pd/PartitionService.java, lines 945-970).
   • The only mechanism PD has to move a leader is PartitionService.transferLeader(...) (lines 774-791), which fires a TransferLeader instruction to the store, and the only caller that computes a balanced target distribution is TaskScheduleService.balancePartitionLeader(...).

   So when a store goes down, its leaders fail over to the surviving replicas; when it comes back, nothing moves any leader back to it until balancePartitionLeader is invoked by hand.

4. Operational consequences. These are expected behavior based on the design above, not measured claims: with 3-way replication and one store restarted, the surviving two stores end up carrying all leaders, so leader-side work (reads served from leaders, raft log append and apply coordination) concentrates on them, while the recovered store runs mostly as a follower. In the worst case a rolling restart can walk all leaders onto the last-restarted-first-recovered nodes. Under the current design, nothing in PD converges the distribution back on its own; a manual balance call is the intended mechanism.

Proposal

Add a config-gated periodic leader balance task to TaskScheduleService.init(), reusing the existing balancePartitionLeader(...) logic rather than introducing a new algorithm:

• Schedule with the same executor.scheduleWithFixedDelay(...) pattern used by patrolStores(), guarded by isLeader() like the other jobs.
• Before acting, compute the current leader distribution from storeService.getShardGroups() shard roles and skip the run when the skew is within a threshold (for example, when max leaders per store minus min leaders per store is at or below the threshold). This avoids constant churn when the cluster is already essentially balanced, which matters because balancePartitionLeader today transfers toward an exact target regardless of how small the deviation is.
• When the threshold is exceeded, call balancePartitionLeader(false). The false path keeps the built-in 30s rate limit (BalanceLeaderInterval, TaskScheduleService.java lines 62 and 463-467) as an extra guard. The existing safety checks already refuse to run during split or move tasks and while a shard balance is in progress (lines 471-479); the periodic wrapper should catch that PDException and log at info level rather than error.

Suggested configuration, following the existing flat pd.* key style of pd.patrol-interval (hugegraph-pd/hg-pd-core/src/main/java/org/apache/hugegraph/pd/config/PDConfig.java, lines 47-49, verified on master; also note the file's store.monitor_data_enabled precedent for an enable flag, lines 178-179):

pd:
  # existing
  patrol-interval: 300
  # new, disabled by default
  balance-leader-enable: false
  balance-leader-interval: 300
  balance-leader-max-skew: 2

Defaults are deliberately conservative: disabled by default so existing deployments see no behavior change, a 300s interval matching the patrol-interval default, and a skew threshold that tolerates small imbalances. I am happy to adjust naming and defaults to whatever the maintainers prefer, and I can work on a PR for this if the direction is acceptable.

Alternatives considered

• Operator-side cron (Kubernetes CronJob or similar) calling GET /v1/task/balanceLeaders. This works and is what I do today, but it pushes cluster-internal scheduling policy out to every operator, has no view of PD's internal state beyond what the endpoint enforces, and has to be rediscovered and reimplemented by each deployment (bare metal, Docker Compose, Helm). PD already owns a scheduler, the balance logic, and the safety guards, so it is the natural place for this.
• Event-driven rebalance on store recovery (trigger when a store transitions back to Up, similar to how init() already registers a StoreStatusListener for Tombstone transitions at TaskScheduleService.java lines 168-188). This is attractive but needs care around raft catch-up time on the recovered store; a periodic task with a skew threshold covers the same cases with simpler semantics and also handles skew from causes other than restarts. The two approaches are compatible and the event trigger could be a follow-up.

Context

While testing distributed deployments on Kubernetes, I observed that after a Store pod restarted, leader distribution remained skewed until /v1/task/balanceLeaders was called manually. On Kubernetes, pod restarts are routine rather than exceptional, which is what motivated this request. For reference, the deployment tooling used for these tests is the Helm chart contributed in #3132 (issue #3131), whose documentation currently lists the manual balance endpoints as the operational workaround.