[yew1eb]

What changes were proposed in this pull request?

In HA mode, aggregate worker/application heartbeats on the master leader over a short time window (default 1s) and replicate them as ONE BatchHeartbeat raft log entry, instead of one entry per heartbeat.

• New HeartbeatAggregator on the leader: pending heartbeats are kept in per-worker/per-app maps. Worker snapshots keep the latest heartbeat; application heartbeat deltas (written bytes, file/shuffle/application counts, and fallback counts) are accumulated, while timestamps, gauges, and metric labels keep the latest snapshot. A scheduled thread flushes one batch per window; heartbeat RPCs are enqueued asynchronously and replied immediately from leader-local memory.
• AppLost, WorkerLost, and RegisterWorker serialize submission with batch flushing and flush the affected pending heartbeat first, preserving ordering with both buffered and already-drained batches. Master shutdown performs a best-effort final flush before stopping Raft.
• New internal protobuf messages BatchHeartbeatRequest / PbBatchHeartbeatRequest (Type.BatchHeartbeat = 31); MetaHandler expands the batch and applies each child heartbeat through the same handleWorkerHeartbeat / handleAppHeartbeat helpers as the single-entry path.
• Fully additive and off by default:
  celeborn.master.ha.heartbeat.batch.enabled (false),
  celeborn.master.ha.heartbeat.batch.interval (1s).

Design doc (CIP-23 Batched Heartbeat Raft Writes): https://docs.google.com/document/d/1YXDCjk_kR_5jyRGmUz1DqImjnAtsr7oFEwpYKK1NxAE/edit?usp=sharing

Why are the changes needed?

Heartbeats are the most frequent metadata writes: A/10 + W/30 raft entries/s for A apps and W workers. On one of our production clusters at peak (900+ running apps + 200 workers, ~97 heartbeat entries/s), follower apply lag (RatisApplyCompletedIndexDiff) grows in lockstep with RunningApplicationCount (~50 -> 1000). A follower master flame graph shows the cost is the per-entry raft pipeline (appendEntries ~35%, log flush/fsync ~20%), not apply compute (~7.7%).

Heartbeat liveness and worker snapshots are refreshed periodically, and replies are built from leader-local memory. Batching avoids a synchronous Raft commit for each heartbeat and can cut Raft write volume by ~100x at the reported peak workload.

Known trade-offs when batching is enabled:

• Heartbeat replies do not wait for replication. If a batch is lost during failover or submission fails, liveness/snapshot state can converge on later heartbeats, but delta counters in the lost batch are not replayed.
• A lost or timed-out response can cause the client to retry the same heartbeat with the same request ID. Batched submissions use new Raft request IDs, so this path does not preserve the original heartbeat retry deduplication. Repeated deltas can overcount cumulative shuffle/application/fallback metrics and temporarily bias partition-size estimation. Cumulative metric overcounts do not automatically self-correct; partition-size inputs are reset each estimation cycle. This is an accepted trade-off for the opt-in feature, which is disabled by default.

JIRA: https://issues.apache.org/jira/browse/CELEBORN-2429

Does this PR resolve a correctness bug?

• Yes

Does this PR introduce any user-facing change?

• Yes

Adds the two configuration options above. Existing synchronous heartbeat submission remains the default.

How was this patch tested?

• HeartbeatAggregatorSuiteJ (single-node Raft, full offer → submit → replicate → apply chain): verifies aggregation, empty windows, application delta accumulation with the latest timestamp, and ordering before AppLost / WorkerLost / RegisterWorker. All 5 tests passed in JDK 17 CI on 39542082.
• MasterStateMachineSuiteJ#testBatchHeartbeat: batched entry applies correctly and survives the ResourceRequest → PbMetaRequest wire round-trip.
• MasterStateMachineSuiteJ: all 6 tests passed in the same JDK 17 CI run. All 109 PR checks passed on 39542082.

[yew1eb]

Ping @pan3793 @cxzl25 @RexXiong @SteNicholas @zaynt4606

[AmandeepSingh285]

There is no final flush before the master stops. Those heartbeats are always lost. Should we try to flush once before shut down.

[yew1eb]

Good catch. Master.stop() stops the aggregator before the raft server and before leadereship transfer, so a final flush can still commit. Added a best-effort flushSafely() in stop(): on the leader it commits the pending batch synchronously; on a non-leader it's a no-op.

[AmandeepSingh285]

If the active leader dies before the batched worker heartbeat is committed to Ratis, the new leader keeps the last applied worker snapshot. Load-aware slot allocation then ranks disks from stale flush/fetch times, usable space, and active slots, so traffic can keep landing on workers that have since become hot or unhealthy.
highWorkload is applied only when that heartbeat is applied. If the lost heartbeat was the one that marked the worker overloaded, the new leader will still treat it as available and offer new slots until the next heartbeat (~30s).
The same gap exists for decommission: the worker may already be draining from an event it received, but the status it reported in the lost heartbeat never reaches the new leader. Until the next heartbeat applies, the master can still give applications slots on a worker that is leaving the cluster.

The impact of losing a worker/app heartbeat after this change could be significantly higher, as losing the subsequent heartbeat would further extend the period during which the leader operates on stale worker state.

[yew1eb]

This is a valid concern: if the leader dies before the batch commits, the new leader loses up to one flush window of heartbeats for all workers/apps. That said, the incremental risk is bounded by the flush interval T (default 1s), not the heartbeat period: batched or not, a lost heartbeat is only refreshed by that worker's next heartbeat (30s), so the same stale gap already exists today whenever a heartbeat is lost in a failover. With the final flush added above, graceful shutdowns (rolling restart, the common case in production) lose nothing regardless of T. After a leader change, the extended worker/app timeout deadlines prevent false WorkerLost/AppLost during the convergence window, and workers re-report their highWorkload/decommission status in the next heartbeat, so the state self-heals the same way the non-batch path relies on.

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 58.55%. Comparing base (90954f7) to head (6a329a5).
⚠️ Report is 33 commits behind head on main.

Additional details and impacted files

@@             Coverage Diff              @@
##               main    #3810      +/-   ##
============================================
- Coverage     58.74%   58.55%   -0.18%     
  Complexity      231      231              
============================================
  Files           399      399              
  Lines         28058    28080      +22     
  Branches       2740     2745       +5     
============================================
- Hits          16479    16440      -39     
- Misses        10380    10437      +57     
- Partials       1199     1203       +4     

☔ View full report in Codecov by Harness.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.
• 📦 JS Bundle Analysis: Save yourself from yourself by tracking and limiting bundle sizes in JS merges.

[AmandeepSingh285]

From the underlying implementation, submitRequest() appears to be synchronous since it waits on

raftClientReply = server.submitClientRequestAsync(raftClientRequest).get();

If the Ratis submission takes a long time, the single flush thread will be blocked and subsequent flushes will be delayed. During this time, new worker/app heartbeats will continue accumulating in the pending maps.
Should we consider putting an upper bound on the number of pending heartbeat entries to prevent unbounded memory growth in case the Ratis call is blocked for an extended period? If the limit is reached, we could drop subsequent heartbeats and emit a warning log so that this condition is visible.

[yew1eb]

You're right that submitRequest() is synchronous, so a blocked Ratis call would stall the single flush thread and delay subsequent flushes. But an upper bound on pending entries shouldn't be necessary:

1. Entries don't accumulate while the thread is blocked. The maps are keyed by worker/app id with newest-wins semantics, and flush() drains and clears them under the lock before submitRequest — so heartbeats arriving during a blocked submit just overwrite into fresh maps, bounded by #workers + #apps, the same order as the master's own metadata.
2. This blocking is pre-existing, with a larger blast radius on the non-batch path. There, every heartbeat synchronously calls the same submitRequest on an RPC handler thread — a slow Ratis can stall many RPC threads at once, one per in-flight heartbeat. Batching reduces this to a single dedicated flush thread, and collapses per-heartbeat raft submions into one per flush window.
3. The impact stays small: a delayed flush means at most a few extra seconds of stale heartbeat state — far below the 120s/300s timeouts, self-heals next window — and the existingn log already fires when a submit exceeds the flush interval, so the condition is visible.

[github-actions]

This PR is stale because it has been open 20 days with no activity. Remove stale label or comment or this will be closed in 10 days.

[zaynt4606]

I found two correctness issues with heartbeat batching enabled. Reproduction details and suggested fixes are in the inline comments.

[zaynt4606]

[P2] Preserve application heartbeat deltas when coalescing

The application heartbeat fields are not all snapshots. CommitHandler.commitMetrics() uses sumThenReset() for written bytes/file counts, and LifecycleManager also resets the shuffle/application counters and fallback counts after collecting them. updateAppHeartbeatMeta() adds these deltas.

Consequently, replacing the pending heartbeat for the same app loses the earlier heartbeat's counters permanently, even though it has already been acknowledged successfully. This can happen when a slow Raft submission delays subsequent flushes, or the configured batch interval exceeds the application's heartbeat interval. The next heartbeat cannot recover the discarded deltas.

I reproduced this with the actual metadata manager/handler and a mocked synchronous Raft submission: two heartbeats reporting 100 and 200 written bytes produce 200 with batching enabled, versus 300 with batching disabled. File, shuffle, application and fallback counts are also undercounted. Besides metrics, missing written bytes/file counts can bias partition-size estimation and disk-slot capacity.

Could we accumulate the delta fields while retaining the latest timestamp, and update testDuplicateAppHeartbeatKeepsNewest to verify that distinct heartbeat deltas are conserved?

[yew1eb]

Good catch, and thanks for the detailed reproduction. Fixed in the latest push: offerAppHeartbeat now coalesces by accumulating the delta fields (totalWritten, fileCount, shuffleCount, applicationCount, and shuffle/application fallback counts) onto the newest heartbeat, while time and clientGauges remain newest-wins since gauges are snapshots. testDuplicateAppHeartbeatKeepsNewest is replaced by testDuplicateAppHeartbeatAccumulatesDeltas, which verifies that two heartbeats reporting 100/200 written bytes apply 300, with fallback counts summed and the newest timestamp kept.

[zaynt4606]

[P2] Preserve ordering between buffered heartbeats and lifecycle operations

Returning after enqueueing allows a successfully acknowledged heartbeat to be applied after a subsequent lifecycle operation, because AppLost, WorkerLost and RegisterWorker still submit directly to Raft.

Two reproducible sequences with batching enabled are:

• AppHeartbeat -> AppLost -> flush: AppLost removes the application, but the buffered heartbeat then recreates its appHeartbeatTime entry. A normal client shutdown can hit this sequence after stopping heartbeats and unregistering. The running-application count becomes incorrect and the app expires again later; its removed shuffle registrations are not restored.
• WorkerHeartbeat -> WorkerLost -> RegisterWorker using the same host/ports, followed by flush: the old heartbeat overwrites the newly registered worker's workload state and lastHeartbeat. Other snapshot fields, including disk/status information, follow the same update path, potentially affecting slot allocation.

I reproduced both using the actual metadata manager/handler with a mocked synchronous Raft submission. The corresponding non-batched controls preserve application removal and the new worker state. These sequences do not require a leader failure: the old synchronous path completes the heartbeat submission before acknowledging it, whereas the batched path leaves it pending.

Could we preserve ordering across heartbeat batches and these lifecycle operations, including already-drained batches, and add regression coverage for both sequences?

[yew1eb]

Both sequences are addressed in the latest push. HeartbeatAggregator now serializes all raft submissions through a dedicated submissionLock: the periodic flush() holds it across drain + submit, and AppLost/WorkerLost/RegisterWorker in HAMasterMetaManager go through a new submitWithOrdering(...), which first drains and flushes pending heartbeats of the affected app/worker, then submits the lifecycle request — all under the same lock. This covers both pending entries and already-drained in-flight batches, and costs no extra raft write when no matching heartbeat is buffered. Regression coverage added: testAppLostFlushesPendingAppHeartbeatFirst (buffered heartbeat applied before AppLost, app not resurrected) and testWorkerLostAndRegisterWorkerFlushPendingWorkerHeartbeatFirst (stale heartbeat cannot overwrite freshly registered worker state).

[zaynt4606]

Non-blocking: the new delta accumulation can count a retried application heartbeat twice.

MasterClient.sendMessageInner() retains the same requestId across retries. If the master handles a heartbeat but its response is lost or times out, the same heartbeat can be resent. The batched path in HAMasterMetaManager.handleAppHeartbeat() discards the original ID, accumulates the deltas again, and submits batches with new Raft request IDs, bypassing the deduplication preserved by the non-batched path.

I reproduced this on 39542082 with a real single-node Raft server: submitting the same heartbeat twice with the same ID produces 100 bytes / 10 files / 1 shuffle with batching disabled, versus 200 bytes / 20 files / 2 shuffles with batching enabled. This happens both within one batch window and when the retry arrives after the first flush.

Could we retain the original heartbeat identity and deduplicate retries, including retries after flushing, while continuing to accumulate distinct heartbeat deltas? Regression tests should cover both same-window and cross-window retries.

[yew1eb]

@zaynt4606 Thanks for the careful reproduction — the mechanism is real: MasterClient keeps the same requestId across retries, and the batched path bypasses the Ratis (clientId, callId) dedup that the non-batched path gets for free.

That said, after walking through what a double-applied app heartbeat actually does in updateAppHeartbeatMeta, I lean toward accepting this as a documented trade-off rather than adding dedup state:

1. The non-idempotent fields are all transient, self-healing metrics. partitionTotalWritten/partitionTotalFileCount are consumed via sumThenReset() on every updatePartitionSize() cycle, so a duplicated delta biases at most one estimation window of estimatedPartitionSize (further bounded by the min/max clamp). shuffleTotalCount/applicationTotalCount/fallback counts are metrics only. appHeartbeatTime and clientGauges are idempotent for a retried (identical) heartbeat. No metadata corruption, no liveness impact.
2. The trigger requires a lost/timed-out response plus a retry — a rare RPC-layer event, bounded by maxRetries, for a feature that is off by default and explicitly opted into.
3. A leader-side dedup cache would still be best-effort. It cannot cover the case where the batch is committed, the leader crashes, and the retry reaches the new leader with an empty cache — so we would add permanent state and a capacity knob without actually restoring the exactly-once guarantee Ratis gives the non-batched path. Doing it correctly (carrying child requestIds in the batch proto and deduplicating in the state machine) is significantly more complexity for a metrics-only, transient effect.

If you are OK with it, I would prefer to document this as a known trade-off in the PR description (and the CIP), or track it in a follow-up JIRA, rather than grow the aggregator for an incomplete fix. WDYT?

[zaynt4606]

Thanks for the clarification. Since batching is disabled by default and this requires a lost/timed-out response followed by a retry, I don't consider it a blocker. I'm fine with documenting the trade-off, distinguishing cumulative metric overcounts from temporary partition-size estimation bias. LGTM.

[zaynt4606]

Merge Summary:

• merged into main d8a4c71c

Posted by merge_pr.py

[zaynt4606]

Thanks @yew1eb for the contribution!
Merged to main(v1.0.0).

Commit: d8a4c71