[Feature] PD: optional periodic partition leader balancing in TaskScheduleService
- Dominant language
- Java
- Stars
- 3.2k
- Forks
- 636
- Avg merge
- 3d 11h
- Merged PRs (30d)
- 14
Description
## 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):
```yaml
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 https://github.com/apache/hugegraph/pull/3132 (issue https://github.com/apache/hugegraph/issues/3131), whose documentation currently lists the manual balance endpoints as the operational workaround.
Contributor guide
Research direction
Start with hugegraph-pd/hg-pd-core/src/main/java/org/apache/hugegraph/pd/TaskScheduleService.java, especially init(), balancePartitionLeader(...), and the existing patrolStores() scheduling pattern. Then read PDConfig.java and the manual balance entry points in TaskAPI.java and PDService.java; done means an opt-in periodic task applies the configured interval and skew threshold while preserving the existing safety and rate-limit behavior.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- java
- Domain
- backend, databases, distributed-systems
- Issue type
- Feature
- Difficulty
- 4/5
- Estimated time
- 3-5 days
- Activity status
- Quiet
- Clarity
- Clearly specified
- Newbie friendliness
- 58/100