This is an automated email from the ASF dual-hosted git repository.
zaynt4606 pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/celeborn.git
The following commit(s) were added to refs/heads/main by this push:
new d8a4c71c50 [CELEBORN-2429][CIP-23] Batch worker/application heartbeats
into aggregated raft log entries
d8a4c71c50 is described below
commit d8a4c71c50a46c62b3ab6b667f689d8b4f27e284
Author: yew1eb <[email protected]>
AuthorDate: Sun Oct 25 17:16:20 2026 +0800
[CELEBORN-2429][CIP-23] Batch worker/application heartbeats into aggregated
raft log entries
### 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%).
<img width="1942" height="701" alt="image"
src="https://github.com/user-attachments/assets/7712bfea-9008-4e67-8423-528e1351298b";
/>
<img width="1732" height="902" alt="image"
src="https://github.com/user-attachments/assets/feb1834d-3668-49f9-ae90-639bd909e04d";
/>
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 accept [...]
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?
- [x] 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`.
Closes #3810 from yew1eb/CELEBORN-2429.
Authored-by: yew1eb <[email protected]>
Signed-off-by: zhengtao <[email protected]>
Co-Authored-By: Codex <[email protected]>
AI-Model: gpt-6
AI-Contributed/Feature: 0/553
AI-Contributed/UT: 0/434
---
common/src/main/proto/TransportMessages.proto | 7 +
.../org/apache/celeborn/common/CelebornConf.scala | 26 ++
docs/configuration/ha.md | 2 +
.../master/clustermeta/ha/HAMasterMetaManager.java | 115 +++++---
.../master/clustermeta/ha/HeartbeatAggregator.java | 248 ++++++++++++++++
.../deploy/master/clustermeta/ha/MetaHandler.java | 138 +++++----
master/src/main/proto/Resource.proto | 8 +
.../celeborn/service/deploy/master/Master.scala | 9 +
.../clustermeta/ha/HeartbeatAggregatorSuiteJ.java | 321 +++++++++++++++++++++
.../clustermeta/ha/MasterStateMachineSuiteJ.java | 110 +++++++
.../master/clustermeta/ha/RatisBaseSuiteJ.java | 3 +-
11 files changed, 891 insertions(+), 96 deletions(-)
diff --git a/common/src/main/proto/TransportMessages.proto
b/common/src/main/proto/TransportMessages.proto
index dd671914a0..b7b04d6cf2 100644
--- a/common/src/main/proto/TransportMessages.proto
+++ b/common/src/main/proto/TransportMessages.proto
@@ -490,6 +490,11 @@ message PbMetaBatchUnregisterShuffles {
repeated string shuffleKeys = 1;
}
+message PbBatchHeartbeatRequest {
+ repeated PbMetaWorkerHeartbeatRequest workerHeartbeats = 1;
+ repeated PbMetaAppHeartbeatRequest appHeartbeats = 2;
+}
+
message PbUnregisterShuffleResponse {
int32 status = 1;
}
@@ -1036,6 +1041,7 @@ enum PbMetaRequestType {
BatchUnRegisterShuffle = 28;
ReviseLostShuffles = 29;
RegisterApplicationInfo = 30;
+ BatchHeartbeat = 31;
}
message PbMetaRequest {
@@ -1057,6 +1063,7 @@ message PbMetaRequest {
PbReportWorkerDecommission reportWorkerDecommissionRequest = 24;
PbMetaBatchUnregisterShuffles batchUnregisterShuffleRequest = 25;
PbRegisterApplicationInfo registerApplicationInfoRequest = 26;
+ PbBatchHeartbeatRequest batchHeartbeatRequest = 27;
PbReviseLostShuffles reviseLostShufflesRequest = 102;
}
diff --git
a/common/src/main/scala/org/apache/celeborn/common/CelebornConf.scala
b/common/src/main/scala/org/apache/celeborn/common/CelebornConf.scala
index 5f535fb344..c1dc8aed87 100644
--- a/common/src/main/scala/org/apache/celeborn/common/CelebornConf.scala
+++ b/common/src/main/scala/org/apache/celeborn/common/CelebornConf.scala
@@ -822,6 +822,8 @@ class CelebornConf(loadDefaults: Boolean) extends Cloneable
with Logging with Se
def haMasterRatisSnapshotAutoTriggerThreshold: Long =
get(HA_MASTER_RATIS_SNAPSHOT_AUTO_TRIGGER_THRESHOLD)
def haMasterRatisSnapshotRetentionFileNum: Int =
get(HA_MASTER_RATIS_SNAPSHOT_RETENTION_FILE_NUM)
+ def masterHaHeartbeatBatchEnabled: Boolean =
get(MASTER_HA_HEARTBEAT_BATCH_ENABLED)
+ def masterHaHeartbeatBatchIntervalMs: Long =
get(MASTER_HA_HEARTBEAT_BATCH_INTERVAL)
def masterPersistWorkerNetworkLocation: Boolean =
get(MASTER_PERSIST_WORKER_NETWORK_LOCATION)
def haRatisCustomConfigs: JMap[String, String] = {
@@ -3138,6 +3140,30 @@ object CelebornConf extends Logging {
.intConf
.createWithDefault(3)
+ val MASTER_HA_HEARTBEAT_BATCH_ENABLED: ConfigEntry[Boolean] =
+ buildConf("celeborn.master.ha.heartbeat.batch.enabled")
+ .categories("ha")
+ .version("1.0.0")
+ .doc(
+ "Whether to aggregate worker/app heartbeats on the raft leader and
flush them " +
+ "periodically as a single BatchHeartbeat raft log entry, reducing
raft log entries, " +
+ "fsyncs and state machine applies by roughly the number of
heartbeats per flush window.")
+ .booleanConf
+ .createWithDefault(false)
+
+ val MASTER_HA_HEARTBEAT_BATCH_INTERVAL: ConfigEntry[Long] =
+ buildConf("celeborn.master.ha.heartbeat.batch.interval")
+ .categories("ha")
+ .version("1.0.0")
+ .doc(
+ "The interval at which aggregated heartbeats are flushed as a single
raft log entry. " +
+ "Heartbeat replies do not wait for raft replication when batching is
enabled, so a " +
+ "larger interval only means the leader may lose at most one interval
of heartbeat " +
+ "metadata on failover, which is negligible compared to the
worker/app heartbeat " +
+ "timeouts (120s/300s by default).")
+ .timeConf(TimeUnit.MILLISECONDS)
+ .createWithDefaultString("1s")
+
val MASTER_PERSIST_WORKER_NETWORK_LOCATION: ConfigEntry[Boolean] =
buildConf("celeborn.master.persist.workerNetworkLocation")
.categories("master")
diff --git a/docs/configuration/ha.md b/docs/configuration/ha.md
index 51007ed055..bac7348775 100644
--- a/docs/configuration/ha.md
+++ b/docs/configuration/ha.md
@@ -22,6 +22,8 @@ license: |
| celeborn.master.ha.enabled | false | false | When true, master nodes run as
Raft cluster mode. | 0.3.0 | celeborn.ha.enabled |
| celeborn.master.ha.graceful.shutdown.enabled | false | false | When true,
the master will transfer Raft leadership before shutting down gracefully. This
reduces chances of client side failures by avoiding the Raft election window
where no leader is available. | 0.7.0 | |
| celeborn.master.ha.graceful.shutdown.timeout | 30s | false | Timeout for the
master graceful shutdown process including Raft leadership transfer. Used as
the shutdown hook timeout and the transfer-leadership request timeout. | 0.7.0
| |
+| celeborn.master.ha.heartbeat.batch.enabled | false | false | Whether to
aggregate worker/app heartbeats on the raft leader and flush them periodically
as a single BatchHeartbeat raft log entry, reducing raft log entries, fsyncs
and state machine applies by roughly the number of heartbeats per flush window.
| 1.0.0 | |
+| celeborn.master.ha.heartbeat.batch.interval | 1s | false | The interval at
which aggregated heartbeats are flushed as a single raft log entry. Heartbeat
replies do not wait for raft replication when batching is enabled, so a larger
interval only means the leader may lose at most one interval of heartbeat
metadata on failover, which is negligible compared to the worker/app heartbeat
timeouts (120s/300s by default). | 1.0.0 | |
| celeborn.master.ha.node.<id>.host | <required> | false | Host to
bind of master node <id> in HA mode. | 0.3.0 |
celeborn.ha.master.node.<id>.host |
| celeborn.master.ha.node.<id>.internal.port | 8097 | false | Internal
port for the workers and other masters to bind to a master node <id> in HA
mode. | 0.5.0 | |
| celeborn.master.ha.node.<id>.port | 9097 | false | Port to bind of
master node <id> in HA mode. | 0.3.0 | celeborn.ha.master.node.<id>.port
|
diff --git
a/master/src/main/java/org/apache/celeborn/service/deploy/master/clustermeta/ha/HAMasterMetaManager.java
b/master/src/main/java/org/apache/celeborn/service/deploy/master/clustermeta/ha/HAMasterMetaManager.java
index 55f717d16f..5631467257 100644
---
a/master/src/main/java/org/apache/celeborn/service/deploy/master/clustermeta/ha/HAMasterMetaManager.java
+++
b/master/src/main/java/org/apache/celeborn/service/deploy/master/clustermeta/ha/HAMasterMetaManager.java
@@ -47,6 +47,8 @@ public class HAMasterMetaManager extends AbstractMetaManager {
protected HARaftServer ratisServer;
+ private HeartbeatAggregator heartbeatAggregator;
+
public HAMasterMetaManager(RpcEnv rpcEnv, CelebornConf conf) {
this(rpcEnv, conf, new CelebornRackResolver(conf));
}
@@ -78,6 +80,13 @@ public class HAMasterMetaManager extends AbstractMetaManager
{
public void setRatisServer(HARaftServer ratisServer) {
this.ratisServer = ratisServer;
+ if (conf.masterHaHeartbeatBatchEnabled()) {
+ this.heartbeatAggregator = new HeartbeatAggregator(ratisServer, conf);
+ }
+ }
+
+ public HeartbeatAggregator getHeartbeatAggregator() {
+ return heartbeatAggregator;
}
@Override
@@ -180,24 +189,29 @@ public class HAMasterMetaManager extends
AbstractMetaManager {
Map<String, Long> clientGauges,
Map<String, String> clientMetricLabels,
String requestId) {
+ ResourceProtos.AppHeartbeatRequest appHeartbeatRequest =
+ ResourceProtos.AppHeartbeatRequest.newBuilder()
+ .setAppId(appId)
+ .setTime(time)
+ .setTotalWritten(totalWritten)
+ .setFileCount(fileCount)
+ .setShuffleCount(shuffleCount)
+ .setApplicationCount(applicationCount)
+ .putAllShuffleFallbackCounts(shuffleFallbackCounts)
+ .putAllApplicationFallbackCounts(applicationFallbackCounts)
+ .putAllClientGauges(clientGauges)
+ .putAllClientMetricLabels(clientMetricLabels)
+ .build();
+ if (heartbeatAggregator != null) {
+ heartbeatAggregator.offerAppHeartbeat(appHeartbeatRequest);
+ return;
+ }
try {
ratisServer.submitRequest(
ResourceRequest.newBuilder()
.setCmdType(Type.AppHeartbeat)
.setRequestId(requestId)
- .setAppHeartbeatRequest(
- ResourceProtos.AppHeartbeatRequest.newBuilder()
- .setAppId(appId)
- .setTime(time)
- .setTotalWritten(totalWritten)
- .setFileCount(fileCount)
- .setShuffleCount(shuffleCount)
- .setApplicationCount(applicationCount)
- .putAllShuffleFallbackCounts(shuffleFallbackCounts)
-
.putAllApplicationFallbackCounts(applicationFallbackCounts)
- .putAllClientGauges(clientGauges)
- .putAllClientMetricLabels(clientMetricLabels)
- .build())
+ .setAppHeartbeatRequest(appHeartbeatRequest)
.build());
} catch (CelebornRuntimeException e) {
LOG.error("Handle heartbeat for {} failed!", appId, e);
@@ -207,13 +221,19 @@ public class HAMasterMetaManager extends
AbstractMetaManager {
@Override
public void handleAppLost(String appId, String requestId) {
+ ResourceRequest request =
+ ResourceRequest.newBuilder()
+ .setCmdType(Type.AppLost)
+ .setRequestId(requestId)
+
.setAppLostRequest(ResourceProtos.AppLostRequest.newBuilder().setAppId(appId).build())
+ .build();
try {
- ratisServer.submitRequest(
- ResourceRequest.newBuilder()
- .setCmdType(Type.AppLost)
- .setRequestId(requestId)
-
.setAppLostRequest(ResourceProtos.AppLostRequest.newBuilder().setAppId(appId).build())
- .build());
+ if (heartbeatAggregator != null) {
+ // Flush any pending heartbeat of this app first, so it is not applied
after AppLost.
+ heartbeatAggregator.submitAppLifecycle(appId, request);
+ } else {
+ ratisServer.submitRequest(request);
+ }
} catch (CelebornRuntimeException e) {
LOG.error("Handle app lost for {} failed!", appId, e);
throw e;
@@ -273,7 +293,7 @@ public class HAMasterMetaManager extends
AbstractMetaManager {
public void handleWorkerLost(
String host, int rpcPort, int pushPort, int fetchPort, int
replicatePort, String requestId) {
try {
- ratisServer.submitRequest(
+ ResourceRequest request =
ResourceRequest.newBuilder()
.setCmdType(Type.WorkerLost)
.setRequestId(requestId)
@@ -285,7 +305,16 @@ public class HAMasterMetaManager extends
AbstractMetaManager {
.setFetchPort(fetchPort)
.setReplicatePort(replicatePort)
.build())
- .build());
+ .build();
+ if (heartbeatAggregator != null) {
+ // Flush any pending heartbeat of this worker first, so it is not
applied after
+ // WorkerLost.
+ heartbeatAggregator.submitWorkerLifecycle(
+ HeartbeatAggregator.workerKey(host, rpcPort, pushPort, fetchPort,
replicatePort),
+ request);
+ } else {
+ ratisServer.submitRequest(request);
+ }
} catch (CelebornRuntimeException e) {
LOG.error("Handle worker lost for {} failed!", host, e);
throw e;
@@ -326,23 +355,30 @@ public class HAMasterMetaManager extends
AbstractMetaManager {
boolean highWorkload,
WorkerStatus workerStatus,
String requestId) {
+ ResourceProtos.WorkerHeartbeatRequest workerHeartbeatRequest =
+ ResourceProtos.WorkerHeartbeatRequest.newBuilder()
+ .setHost(host)
+ .setRpcPort(rpcPort)
+ .setPushPort(pushPort)
+ .setFetchPort(fetchPort)
+ .setReplicatePort(replicatePort)
+ .putAllDisks(MetaUtil.toPbDiskInfos(disks))
+ .setWorkerStatus(MetaUtil.toPbWorkerStatus(workerStatus))
+ .setTime(time)
+ .setHighWorkload(highWorkload)
+ .build();
+ if (heartbeatAggregator != null) {
+ heartbeatAggregator.offerWorkerHeartbeat(workerHeartbeatRequest);
+ updateWorkerResourceConsumptions(
+ host, rpcPort, pushPort, fetchPort, replicatePort,
userResourceConsumption);
+ return;
+ }
try {
ratisServer.submitRequest(
ResourceRequest.newBuilder()
.setCmdType(Type.WorkerHeartbeat)
.setRequestId(requestId)
- .setWorkerHeartbeatRequest(
- ResourceProtos.WorkerHeartbeatRequest.newBuilder()
- .setHost(host)
- .setRpcPort(rpcPort)
- .setPushPort(pushPort)
- .setFetchPort(fetchPort)
- .setReplicatePort(replicatePort)
- .putAllDisks(MetaUtil.toPbDiskInfos(disks))
- .setWorkerStatus(MetaUtil.toPbWorkerStatus(workerStatus))
- .setTime(time)
- .setHighWorkload(highWorkload)
- .build())
+ .setWorkerHeartbeatRequest(workerHeartbeatRequest)
.build());
updateWorkerResourceConsumptions(
host, rpcPort, pushPort, fetchPort, replicatePort,
userResourceConsumption);
@@ -365,7 +401,7 @@ public class HAMasterMetaManager extends
AbstractMetaManager {
Map<UserIdentifier, ResourceConsumption> userResourceConsumption,
String requestId) {
try {
- ratisServer.submitRequest(
+ ResourceRequest request =
ResourceRequest.newBuilder()
.setCmdType(Type.RegisterWorker)
.setRequestId(requestId)
@@ -380,7 +416,16 @@ public class HAMasterMetaManager extends
AbstractMetaManager {
.setNetworkLocation(networkLocation)
.putAllDisks(MetaUtil.toPbDiskInfos(disks))
.build())
- .build());
+ .build();
+ if (heartbeatAggregator != null) {
+ // Flush any pending heartbeat of this worker first, so a stale
heartbeat is not
+ // applied after the fresh registration.
+ heartbeatAggregator.submitWorkerLifecycle(
+ HeartbeatAggregator.workerKey(host, rpcPort, pushPort, fetchPort,
replicatePort),
+ request);
+ } else {
+ ratisServer.submitRequest(request);
+ }
updateWorkerResourceConsumptions(
host, rpcPort, pushPort, fetchPort, replicatePort,
userResourceConsumption);
} catch (CelebornRuntimeException e) {
diff --git
a/master/src/main/java/org/apache/celeborn/service/deploy/master/clustermeta/ha/HeartbeatAggregator.java
b/master/src/main/java/org/apache/celeborn/service/deploy/master/clustermeta/ha/HeartbeatAggregator.java
new file mode 100644
index 0000000000..1f45db8bf4
--- /dev/null
+++
b/master/src/main/java/org/apache/celeborn/service/deploy/master/clustermeta/ha/HeartbeatAggregator.java
@@ -0,0 +1,248 @@
+/*
+ * 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.celeborn.service.deploy.master.clustermeta.ha;
+
+import java.util.ArrayList;
+import java.util.Collection;
+import java.util.Collections;
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+import java.util.concurrent.ScheduledExecutorService;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.locks.ReentrantLock;
+
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+import org.apache.celeborn.common.CelebornConf;
+import org.apache.celeborn.common.client.MasterClient;
+import org.apache.celeborn.common.util.ThreadUtils;
+import org.apache.celeborn.service.deploy.master.clustermeta.ResourceProtos;
+import
org.apache.celeborn.service.deploy.master.clustermeta.ResourceProtos.ResourceRequest;
+import
org.apache.celeborn.service.deploy.master.clustermeta.ResourceProtos.Type;
+
+public class HeartbeatAggregator {
+ private static final Logger LOG =
LoggerFactory.getLogger(HeartbeatAggregator.class);
+
+ private final HARaftServer ratisServer;
+ private final long batchIntervalMs;
+
+ private final ReentrantLock pendingLock = new ReentrantLock();
+ private final Map<String, ResourceProtos.WorkerHeartbeatRequest>
workerHeartbeats =
+ new HashMap<>();
+ private final Map<String, ResourceProtos.AppHeartbeatRequest> appHeartbeats
= new HashMap<>();
+
+ // Serializes raft submissions so a lifecycle request is never applied
before heartbeats
+ // accepted earlier.
+ private final ReentrantLock submissionLock = new ReentrantLock();
+
+ private final ScheduledExecutorService flushExecutor;
+
+ public HeartbeatAggregator(HARaftServer ratisServer, CelebornConf conf) {
+ this.ratisServer = ratisServer;
+ this.batchIntervalMs = conf.masterHaHeartbeatBatchIntervalMs();
+ this.flushExecutor =
+
ThreadUtils.newDaemonSingleThreadScheduledExecutor("master-heartbeat-aggregator");
+ this.flushExecutor.scheduleWithFixedDelay(
+ this::flushSafely, batchIntervalMs, batchIntervalMs,
TimeUnit.MILLISECONDS);
+ LOG.info("HeartbeatAggregator started, flush interval {} ms.",
batchIntervalMs);
+ }
+
+ public void offerWorkerHeartbeat(ResourceProtos.WorkerHeartbeatRequest
heartbeat) {
+ pendingLock.lock();
+ try {
+ workerHeartbeats.put(
+ workerKey(
+ heartbeat.getHost(),
+ heartbeat.getRpcPort(),
+ heartbeat.getPushPort(),
+ heartbeat.getFetchPort(),
+ heartbeat.getReplicatePort()),
+ heartbeat);
+ } finally {
+ pendingLock.unlock();
+ }
+ }
+
+ public void offerAppHeartbeat(ResourceProtos.AppHeartbeatRequest heartbeat) {
+ pendingLock.lock();
+ try {
+ appHeartbeats.merge(heartbeat.getAppId(), heartbeat,
HeartbeatAggregator::mergeAppHeartbeat);
+ } finally {
+ pendingLock.unlock();
+ }
+ }
+
+ // Counter fields are client-side deltas and must be accumulated when
coalescing; time and
+ // client gauges are snapshots where the newest value wins.
+ private static ResourceProtos.AppHeartbeatRequest mergeAppHeartbeat(
+ ResourceProtos.AppHeartbeatRequest pending,
ResourceProtos.AppHeartbeatRequest latest) {
+ return latest
+ .toBuilder()
+ .setTotalWritten(pending.getTotalWritten() + latest.getTotalWritten())
+ .setFileCount(pending.getFileCount() + latest.getFileCount())
+ .setShuffleCount(pending.getShuffleCount() + latest.getShuffleCount())
+ .setApplicationCount(pending.getApplicationCount() +
latest.getApplicationCount())
+ .putAllShuffleFallbackCounts(
+ mergeCounts(
+ pending.getShuffleFallbackCountsMap(),
latest.getShuffleFallbackCountsMap()))
+ .putAllApplicationFallbackCounts(
+ mergeCounts(
+ pending.getApplicationFallbackCountsMap(),
+ latest.getApplicationFallbackCountsMap()))
+ .build();
+ }
+
+ private static Map<String, Long> mergeCounts(Map<String, Long> a,
Map<String, Long> b) {
+ if (a.isEmpty()) {
+ return b;
+ }
+ if (b.isEmpty()) {
+ return a;
+ }
+ Map<String, Long> merged = new HashMap<>(a);
+ b.forEach((key, value) -> merged.merge(key, value, Long::sum));
+ return merged;
+ }
+
+ static String workerKey(
+ String host, int rpcPort, int pushPort, int fetchPort, int
replicatePort) {
+ return host + ":" + rpcPort + ":" + pushPort + ":" + fetchPort + ":" +
replicatePort;
+ }
+
+ public void stop() {
+ flushSafely();
+ flushExecutor.shutdownNow();
+ }
+
+ private void flushSafely() {
+ try {
+ flush();
+ } catch (Throwable t) {
+ // Dropped batches self-heal next interval.
+ LOG.error("Failed to flush aggregated heartbeats, dropping this batch.",
t);
+ }
+ }
+
+ private void flush() {
+ submissionLock.lock();
+ try {
+ List<ResourceProtos.WorkerHeartbeatRequest> drainedWorkers;
+ List<ResourceProtos.AppHeartbeatRequest> drainedApps;
+ pendingLock.lock();
+ try {
+ if (workerHeartbeats.isEmpty() && appHeartbeats.isEmpty()) {
+ return;
+ }
+ drainedWorkers = new ArrayList<>(workerHeartbeats.values());
+ workerHeartbeats.clear();
+ drainedApps = new ArrayList<>(appHeartbeats.values());
+ appHeartbeats.clear();
+ } finally {
+ pendingLock.unlock();
+ }
+ submitBatch(drainedWorkers, drainedApps);
+ } finally {
+ submissionLock.unlock();
+ }
+ }
+
+ /** AppLost etc.: flush the app's pending heartbeat before submitting. */
+ public void submitAppLifecycle(String appId, ResourceRequest
lifecycleRequest) {
+ submitWithOrdering(Collections.singletonList(appId),
Collections.emptyList(), lifecycleRequest);
+ }
+
+ /** WorkerLost/RegisterWorker etc.: flush the worker's pending heartbeat
before submitting. */
+ public void submitWorkerLifecycle(String workerKey, ResourceRequest
lifecycleRequest) {
+ submitWithOrdering(
+ Collections.emptyList(), Collections.singletonList(workerKey),
lifecycleRequest);
+ }
+
+ // Flushes pending heartbeats of the affected apps/workers first, then
submits the lifecycle
+ // request, all under submissionLock so earlier accepted heartbeats apply
first.
+ private void submitWithOrdering(
+ Collection<String> appIds, Collection<String> workerKeys,
ResourceRequest lifecycleRequest) {
+ submissionLock.lock();
+ try {
+ List<ResourceProtos.WorkerHeartbeatRequest> drainedWorkers = new
ArrayList<>();
+ List<ResourceProtos.AppHeartbeatRequest> drainedApps = new ArrayList<>();
+ pendingLock.lock();
+ try {
+ for (String appId : appIds) {
+ ResourceProtos.AppHeartbeatRequest heartbeat =
appHeartbeats.remove(appId);
+ if (heartbeat != null) {
+ drainedApps.add(heartbeat);
+ }
+ }
+ for (String workerKey : workerKeys) {
+ ResourceProtos.WorkerHeartbeatRequest heartbeat =
workerHeartbeats.remove(workerKey);
+ if (heartbeat != null) {
+ drainedWorkers.add(heartbeat);
+ }
+ }
+ } finally {
+ pendingLock.unlock();
+ }
+ if (!drainedWorkers.isEmpty() || !drainedApps.isEmpty()) {
+ submitBatch(drainedWorkers, drainedApps);
+ }
+ ratisServer.submitRequest(lifecycleRequest);
+ } finally {
+ submissionLock.unlock();
+ }
+ }
+
+ // Caller must hold submissionLock.
+ private void submitBatch(
+ List<ResourceProtos.WorkerHeartbeatRequest> drainedWorkers,
+ List<ResourceProtos.AppHeartbeatRequest> drainedApps) {
+ if (!ratisServer.isLeader()) {
+ return;
+ }
+
+ ResourceRequest batchRequest =
+ ResourceRequest.newBuilder()
+ .setCmdType(Type.BatchHeartbeat)
+ .setRequestId(MasterClient.genRequestId())
+ .setBatchHeartbeatRequest(
+ ResourceProtos.BatchHeartbeatRequest.newBuilder()
+ .addAllWorkerHeartbeats(drainedWorkers)
+ .addAllAppHeartbeats(drainedApps)
+ .build())
+ .build();
+ long startNs = System.nanoTime();
+ ratisServer.submitRequest(batchRequest);
+ long elapsedMs = TimeUnit.NANOSECONDS.toMillis(System.nanoTime() -
startNs);
+ if (elapsedMs > batchIntervalMs) {
+ LOG.warn(
+ "Submitting aggregated heartbeats ({} worker, {} app) took {} ms;
raft commits are "
+ + "slower than the flush interval {} ms.",
+ drainedWorkers.size(),
+ drainedApps.size(),
+ elapsedMs,
+ batchIntervalMs);
+ }
+ if (LOG.isDebugEnabled()) {
+ LOG.debug(
+ "Flushed aggregated heartbeats, {} worker heartbeats, {} app
heartbeats.",
+ drainedWorkers.size(),
+ drainedApps.size());
+ }
+ }
+}
diff --git
a/master/src/main/java/org/apache/celeborn/service/deploy/master/clustermeta/ha/MetaHandler.java
b/master/src/main/java/org/apache/celeborn/service/deploy/master/clustermeta/ha/MetaHandler.java
index 834b201086..825047c0b2 100644
---
a/master/src/main/java/org/apache/celeborn/service/deploy/master/clustermeta/ha/MetaHandler.java
+++
b/master/src/main/java/org/apache/celeborn/service/deploy/master/clustermeta/ha/MetaHandler.java
@@ -156,34 +156,24 @@ public class MetaHandler {
break;
case AppHeartbeat:
- appId = request.getAppHeartbeatRequest().getAppId();
- long time = request.getAppHeartbeatRequest().getTime();
- long totalWritten =
request.getAppHeartbeatRequest().getTotalWritten();
- long fileCount = request.getAppHeartbeatRequest().getFileCount();
- long shuffleCount =
request.getAppHeartbeatRequest().getShuffleCount();
- long applicationCount =
request.getAppHeartbeatRequest().getApplicationCount();
- LOG.debug("Handle app heartbeat for {} with shuffle count {}",
appId, shuffleCount);
- Map<String, Long> shuffleFallbackCounts =
- request.getAppHeartbeatRequest().getShuffleFallbackCountsMap();
- if (CollectionUtils.isNotEmpty(shuffleFallbackCounts)) {
- LOG.warn(
- "{} shuffle fallbacks in app {}",
- shuffleFallbackCounts.values().stream().mapToLong(v ->
v).sum(),
- appId);
+ handleAppHeartbeat(request.getAppHeartbeatRequest());
+ break;
+
+ case BatchHeartbeat:
+ List<PbMetaWorkerHeartbeatRequest> workerHeartbeats =
+ request.getBatchHeartbeatRequest().getWorkerHeartbeatsList();
+ List<PbMetaAppHeartbeatRequest> appHeartbeats =
+ request.getBatchHeartbeatRequest().getAppHeartbeatsList();
+ LOG.debug(
+ "Handle batch heartbeat with {} worker heartbeats and {} app
heartbeats.",
+ workerHeartbeats.size(),
+ appHeartbeats.size());
+ for (PbMetaWorkerHeartbeatRequest workerHeartbeat :
workerHeartbeats) {
+ handleWorkerHeartbeat(workerHeartbeat);
+ }
+ for (PbMetaAppHeartbeatRequest appHeartbeat : appHeartbeats) {
+ handleAppHeartbeat(appHeartbeat);
}
- Map<String, Long> applicationFallbackCounts =
-
request.getAppHeartbeatRequest().getApplicationFallbackCountsMap();
- metaSystem.updateAppHeartbeatMeta(
- appId,
- time,
- totalWritten,
- fileCount,
- shuffleCount,
- applicationCount,
- shuffleFallbackCounts,
- applicationFallbackCounts,
- request.getAppHeartbeatRequest().getClientGaugesMap(),
- request.getAppHeartbeatRequest().getClientMetricLabelsMap());
break;
case AppLost:
@@ -215,39 +205,7 @@ public class MetaHandler {
break;
case WorkerHeartbeat:
- host = request.getWorkerHeartbeatRequest().getHost();
- rpcPort = request.getWorkerHeartbeatRequest().getRpcPort();
- pushPort = request.getWorkerHeartbeatRequest().getPushPort();
- fetchPort = request.getWorkerHeartbeatRequest().getFetchPort();
- Map<String, PbDiskInfo> pbDiskInfoMap =
request.getWorkerHeartbeatRequest().getDisksMap();
- diskInfos = MetaUtil.fromPbDiskInfoMap(pbDiskInfoMap);
- replicatePort =
request.getWorkerHeartbeatRequest().getReplicatePort();
- boolean highWorkload =
request.getWorkerHeartbeatRequest().getHighWorkload();
- if (request.getWorkerHeartbeatRequest().hasWorkerStatus()) {
- workerStatus =
-
MetaUtil.fromPbWorkerStatus(request.getWorkerHeartbeatRequest().getWorkerStatus());
- } else {
- workerStatus = WorkerStatus.normalWorkerStatus();
- }
-
- LOG.debug(
- "Handle worker heartbeat for {} {} {} {} {} {}",
- host,
- rpcPort,
- pushPort,
- fetchPort,
- replicatePort,
- diskInfos);
- metaSystem.updateWorkerHeartbeatMeta(
- host,
- rpcPort,
- pushPort,
- fetchPort,
- replicatePort,
- diskInfos,
- request.getWorkerHeartbeatRequest().getTime(),
- workerStatus,
- highWorkload);
+ handleWorkerHeartbeat(request.getWorkerHeartbeatRequest());
break;
case RegisterWorker:
@@ -337,6 +295,66 @@ public class MetaHandler {
return responseBuilder.build();
}
+ private void handleWorkerHeartbeat(PbMetaWorkerHeartbeatRequest
workerHeartbeat) {
+ String host = workerHeartbeat.getHost();
+ int rpcPort = workerHeartbeat.getRpcPort();
+ int pushPort = workerHeartbeat.getPushPort();
+ int fetchPort = workerHeartbeat.getFetchPort();
+ int replicatePort = workerHeartbeat.getReplicatePort();
+ Map<String, DiskInfo> diskInfos =
MetaUtil.fromPbDiskInfoMap(workerHeartbeat.getDisksMap());
+ boolean highWorkload = workerHeartbeat.getHighWorkload();
+ WorkerStatus workerStatus;
+ if (workerHeartbeat.hasWorkerStatus()) {
+ workerStatus =
MetaUtil.fromPbWorkerStatus(workerHeartbeat.getWorkerStatus());
+ } else {
+ workerStatus = WorkerStatus.normalWorkerStatus();
+ }
+
+ LOG.debug(
+ "Handle worker heartbeat for {} {} {} {} {} {}",
+ host,
+ rpcPort,
+ pushPort,
+ fetchPort,
+ replicatePort,
+ diskInfos);
+ metaSystem.updateWorkerHeartbeatMeta(
+ host,
+ rpcPort,
+ pushPort,
+ fetchPort,
+ replicatePort,
+ diskInfos,
+ workerHeartbeat.getTime(),
+ workerStatus,
+ highWorkload);
+ }
+
+ private void handleAppHeartbeat(PbMetaAppHeartbeatRequest appHeartbeat) {
+ String appId = appHeartbeat.getAppId();
+ long shuffleCount = appHeartbeat.getShuffleCount();
+ LOG.debug("Handle app heartbeat for {} with shuffle count {}", appId,
shuffleCount);
+ Map<String, Long> shuffleFallbackCounts =
appHeartbeat.getShuffleFallbackCountsMap();
+ if (CollectionUtils.isNotEmpty(shuffleFallbackCounts)) {
+ LOG.warn(
+ "{} shuffle fallbacks in app {}",
+ shuffleFallbackCounts.values().stream().mapToLong(v -> v).sum(),
+ appId);
+ }
+ Map<String, Long> applicationFallbackCounts =
appHeartbeat.getApplicationFallbackCountsMap();
+ metaSystem.updateAppHeartbeatMeta(
+ appId,
+ appHeartbeat.getTime(),
+ appHeartbeat.getTotalWritten(),
+ appHeartbeat.getFileCount(),
+ shuffleCount,
+ appHeartbeat.getApplicationCount(),
+ shuffleFallbackCounts,
+ applicationFallbackCounts,
+ appHeartbeat.getClientGaugesMap(),
+ appHeartbeat.getClientMetricLabelsMap());
+ }
+
public void writeToSnapShot(File file) throws IOException {
try {
metaSystem.writeMetaInfoToFile(file);
diff --git a/master/src/main/proto/Resource.proto
b/master/src/main/proto/Resource.proto
index 51eab97715..d3d71d9390 100644
--- a/master/src/main/proto/Resource.proto
+++ b/master/src/main/proto/Resource.proto
@@ -42,6 +42,8 @@ enum Type {
ReviseLostShuffles = 29;
RegisterApplicationInfo = 30;
+
+ BatchHeartbeat = 31;
}
enum WorkerEventType {
@@ -74,6 +76,7 @@ message ResourceRequest {
optional ReportWorkerDecommissionRequest reportWorkerDecommissionRequest =
24;
optional BatchUnregisterShuffleRequest batchUnregisterShuffleRequest = 25;
optional RegisterApplicationInfoRequest registerApplicationInfoRequest = 26;
+ optional BatchHeartbeatRequest batchHeartbeatRequest = 27;
optional ReviseLostShufflesRequest reviseLostShufflesRequest = 102;
}
@@ -106,6 +109,11 @@ message BatchUnregisterShuffleRequest {
repeated string shuffleKeys = 1;
}
+message BatchHeartbeatRequest {
+ repeated WorkerHeartbeatRequest workerHeartbeats = 1;
+ repeated AppHeartbeatRequest appHeartbeats = 2;
+}
+
message HeartbeatInfo {
string appId = 1;
int64 totalWritten = 2;
diff --git
a/master/src/main/scala/org/apache/celeborn/service/deploy/master/Master.scala
b/master/src/main/scala/org/apache/celeborn/service/deploy/master/Master.scala
index de2540e867..d74a3711d5 100644
---
a/master/src/main/scala/org/apache/celeborn/service/deploy/master/Master.scala
+++
b/master/src/main/scala/org/apache/celeborn/service/deploy/master/Master.scala
@@ -1709,6 +1709,15 @@ private[celeborn] class Master(
val transferLeadership = conf.haMasterGracefulShutdownEnabled
statusSystem match {
case ha: HAMasterMetaManager =>
+ val heartbeatAggregator = ha.getHeartbeatAggregator
+ if (heartbeatAggregator != null) {
+ try {
+ heartbeatAggregator.stop()
+ } catch {
+ case e: Exception =>
+ logError("Failed to stop heartbeat aggregator during Master
shutdown.", e)
+ }
+ }
val ratisServer = ha.getRatisServer
if (ratisServer != null) {
try {
diff --git
a/master/src/test/java/org/apache/celeborn/service/deploy/master/clustermeta/ha/HeartbeatAggregatorSuiteJ.java
b/master/src/test/java/org/apache/celeborn/service/deploy/master/clustermeta/ha/HeartbeatAggregatorSuiteJ.java
new file mode 100644
index 0000000000..276c5cf3b3
--- /dev/null
+++
b/master/src/test/java/org/apache/celeborn/service/deploy/master/clustermeta/ha/HeartbeatAggregatorSuiteJ.java
@@ -0,0 +1,321 @@
+/*
+ * 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.celeborn.service.deploy.master.clustermeta.ha;
+
+import java.io.File;
+import java.io.IOException;
+import java.util.Collections;
+import java.util.HashMap;
+import java.util.UUID;
+
+import org.junit.After;
+import org.junit.Assert;
+import org.junit.Before;
+import org.junit.Test;
+
+import org.apache.celeborn.common.CelebornConf;
+import org.apache.celeborn.common.meta.WorkerInfo;
+import org.apache.celeborn.common.meta.WorkerStatus;
+
+public class HeartbeatAggregatorSuiteJ {
+ private HARaftServer ratisServer;
+ private HAMasterMetaManager metaSystem;
+
+ @Before
+ public void init() throws Exception {
+ CelebornConf conf = new CelebornConf();
+ conf.set(CelebornConf.MASTER_HA_HEARTBEAT_BATCH_ENABLED().key(), "true");
+ // Use a 1s window (not shorter): on oversubscribed CI runners the handler
calls below can
+ // be scheduled far apart, and a short window would split the offers into
several batches.
+ conf.set(CelebornConf.MASTER_HA_HEARTBEAT_BATCH_INTERVAL().key(), "1s");
+ metaSystem = new HAMasterMetaManager(null, conf);
+ MetaHandler handler = new MetaHandler(metaSystem);
+ File tmpDir = File.createTempFile("celeborn-ratis-tmp", "for-test-only");
+ tmpDir.delete();
+ tmpDir.mkdirs();
+ conf.set(CelebornConf.HA_MASTER_RATIS_STORAGE_DIR().key(),
tmpDir.getAbsolutePath());
+ String id = UUID.randomUUID().toString();
+ MasterNode masterNode =
+ new
MasterNode.Builder().setNodeId(id).setHost("localhost").setRatisPort(9998).build();
+ ratisServer =
+ HARaftServer.newMasterRatisServer(handler, conf, masterNode,
Collections.emptyList());
+ metaSystem.setRatisServer(ratisServer);
+ ratisServer.start();
+ waitForLeader();
+ }
+
+ private void waitForLeader() throws InterruptedException {
+ // Wait for isLeaderReady(), not just isLeader(): a newly elected leader
rejects writes
+ // with LeaderNotReadyException until its no-op entry for the new term is
committed.
+ for (int i = 0; i < 100; i++) {
+ if (ratisServer.isLeader()) {
+ try {
+ if (ratisServer
+ .getServer()
+ .getDivision(ratisServer.getGroupId())
+ .getInfo()
+ .isLeaderReady()) {
+ return;
+ }
+ } catch (IOException e) {
+ // Division not available yet; keep polling.
+ }
+ }
+ Thread.sleep(200);
+ }
+ Assert.fail("Raft server did not become ready leader in time.");
+ }
+
+ @After
+ public void shutdown() {
+ if (metaSystem.getHeartbeatAggregator() != null) {
+ metaSystem.getHeartbeatAggregator().stop();
+ }
+ if (ratisServer != null) {
+ ratisServer.stop();
+ }
+ }
+
+ @Test
+ public void testBatchHeartbeatAggregation() throws Exception {
+ Assert.assertNotNull(metaSystem.getHeartbeatAggregator());
+
+ metaSystem.handleRegisterWorker(
+ "host1",
+ 1,
+ 2,
+ 3,
+ 4,
+ 5,
+ "networkLocation1",
+ new HashMap<>(),
+ new HashMap<>(),
+ UUID.randomUUID().toString() + "#1");
+ Thread.sleep(2000);
+ Assert.assertEquals(1, metaSystem.workersMap.size());
+ long appliedIndexAfterRegister = lastAppliedIndex();
+
+ // 4 heartbeat offers: 3 for the same worker (collapsed to the newest), 1
for an app.
+ long time1 = System.currentTimeMillis();
+ long time2 = time1 + 10;
+ metaSystem.handleWorkerHeartbeat(
+ "host1",
+ 1,
+ 2,
+ 3,
+ 4,
+ new HashMap<>(),
+ new HashMap<>(),
+ time1,
+ false,
+ WorkerStatus.normalWorkerStatus(),
+ UUID.randomUUID().toString() + "#2");
+ metaSystem.handleWorkerHeartbeat(
+ "host1",
+ 1,
+ 2,
+ 3,
+ 4,
+ new HashMap<>(),
+ new HashMap<>(),
+ time1,
+ false,
+ WorkerStatus.normalWorkerStatus(),
+ UUID.randomUUID().toString() + "#3");
+ metaSystem.handleAppHeartbeat(
+ "app-1",
+ 100,
+ 10,
+ 1,
+ 1,
+ new HashMap<>(),
+ new HashMap<>(),
+ time2,
+ new HashMap<>(),
+ new HashMap<>(),
+ UUID.randomUUID().toString() + "#4");
+ metaSystem.handleWorkerHeartbeat(
+ "host1",
+ 1,
+ 2,
+ 3,
+ 4,
+ new HashMap<>(),
+ new HashMap<>(),
+ time2,
+ false,
+ WorkerStatus.normalWorkerStatus(),
+ UUID.randomUUID().toString() + "#5");
+
+ // Wait comfortably past one 1s flush window plus the blocking
submit/apply.
+ Thread.sleep(4000);
+
+ // The 4 offers must have been flushed as (far) fewer raft log entries.
+ long newEntries = lastAppliedIndex() - appliedIndexAfterRegister;
+ Assert.assertTrue(
+ "Expected heartbeats to be merged into few raft log entries, but got "
+ newEntries,
+ newEntries >= 1 && newEntries <= 2);
+
+ // The newest heartbeat wins.
+ WorkerInfo workerInfo = metaSystem.workersMap.values().iterator().next();
+ Assert.assertEquals("host1", workerInfo.host());
+ Assert.assertEquals(time2, workerInfo.lastHeartbeat());
+ Assert.assertEquals(Long.valueOf(time2),
metaSystem.appHeartbeatTime.get("app-1"));
+ }
+
+ @Test
+ public void testEmptyWindowProducesNoRaftLog() throws Exception {
+ Thread.sleep(2000);
+ long appliedIndex = lastAppliedIndex();
+ Thread.sleep(500);
+ Assert.assertEquals(appliedIndex, lastAppliedIndex());
+ }
+
+ @Test
+ public void testDuplicateAppHeartbeatAccumulatesDeltas() throws Exception {
+ // Duplicate heartbeats of the same app within one window are coalesced:
delta counters
+ // are accumulated, while snapshots (time, client gauges) keep the newest
value.
+ java.util.Map<String, Long> fallback1 = new HashMap<>();
+ fallback1.put("shuffle-1", 1L);
+ java.util.Map<String, Long> fallback2 = new HashMap<>();
+ fallback2.put("shuffle-1", 2L);
+ java.util.Map<String, Long> gauges1 = new HashMap<>();
+ gauges1.put("gauge-1", 1L);
+ java.util.Map<String, Long> gauges2 = new HashMap<>();
+ gauges2.put("gauge-1", 2L);
+ long time1 = System.currentTimeMillis();
+ long time2 = time1 + 10;
+ metaSystem.handleAppHeartbeat(
+ "app-merge",
+ 100,
+ 10,
+ 1,
+ 1,
+ fallback1,
+ new HashMap<>(),
+ time1,
+ gauges1,
+ new HashMap<>(),
+ UUID.randomUUID().toString() + "#1");
+ metaSystem.handleAppHeartbeat(
+ "app-merge",
+ 200,
+ 20,
+ 2,
+ 3,
+ fallback2,
+ new HashMap<>(),
+ time2,
+ gauges2,
+ new HashMap<>(),
+ UUID.randomUUID().toString() + "#2");
+
+ Thread.sleep(4000);
+
+ Assert.assertEquals(300, metaSystem.partitionTotalWritten.sum());
+ Assert.assertEquals(30, metaSystem.partitionTotalFileCount.sum());
+ Assert.assertEquals(3, metaSystem.shuffleTotalCount.sum());
+ Assert.assertEquals(4, metaSystem.applicationTotalCount.sum());
+ Assert.assertEquals(Long.valueOf(3),
metaSystem.shuffleFallbackCounts.get("shuffle-1"));
+ Assert.assertEquals(Long.valueOf(time2),
metaSystem.appHeartbeatTime.get("app-merge"));
+ }
+
+ @Test
+ public void testAppLostFlushesPendingAppHeartbeatFirst() throws Exception {
+ // A buffered app heartbeat accepted before AppLost must be applied before
the AppLost,
+ // so the app is not resurrected by a heartbeat flushed afterwards.
+ long time1 = System.currentTimeMillis();
+ metaSystem.handleAppHeartbeat(
+ "app-lost",
+ 100,
+ 10,
+ 1,
+ 1,
+ new HashMap<>(),
+ new HashMap<>(),
+ time1,
+ new HashMap<>(),
+ new HashMap<>(),
+ UUID.randomUUID().toString() + "#1");
+ metaSystem.handleAppLost("app-lost", UUID.randomUUID().toString() + "#2");
+
+ Thread.sleep(2000);
+
+ Assert.assertFalse(metaSystem.appHeartbeatTime.containsKey("app-lost"));
+ // The pending heartbeat is flushed before AppLost, so its delta counters
are not lost.
+ Assert.assertEquals(100, metaSystem.partitionTotalWritten.sum());
+ }
+
+ @Test
+ public void
testWorkerLostAndRegisterWorkerFlushPendingWorkerHeartbeatFirst() throws
Exception {
+ metaSystem.handleRegisterWorker(
+ "host1",
+ 1,
+ 2,
+ 3,
+ 4,
+ 5,
+ "networkLocation1",
+ new HashMap<>(),
+ new HashMap<>(),
+ UUID.randomUUID().toString() + "#1");
+ Thread.sleep(2000);
+ Assert.assertEquals(1, metaSystem.workersMap.size());
+
+ // A stale buffered heartbeat followed by WorkerLost and RegisterWorker
for the same
+ // worker: the heartbeat must be applied before the lifecycle requests, so
it cannot
+ // overwrite the freshly registered worker state.
+ long staleTime = 1000L;
+ metaSystem.handleWorkerHeartbeat(
+ "host1",
+ 1,
+ 2,
+ 3,
+ 4,
+ new HashMap<>(),
+ new HashMap<>(),
+ staleTime,
+ false,
+ WorkerStatus.normalWorkerStatus(),
+ UUID.randomUUID().toString() + "#2");
+ metaSystem.handleWorkerLost("host1", 1, 2, 3, 4,
UUID.randomUUID().toString() + "#3");
+ metaSystem.handleRegisterWorker(
+ "host1",
+ 1,
+ 2,
+ 3,
+ 4,
+ 5,
+ "networkLocation1",
+ new HashMap<>(),
+ new HashMap<>(),
+ UUID.randomUUID().toString() + "#4");
+
+ Thread.sleep(2000);
+
+ Assert.assertEquals(1, metaSystem.workersMap.size());
+ WorkerInfo workerInfo = metaSystem.workersMap.values().iterator().next();
+ Assert.assertTrue(
+ "Stale heartbeat should not overwrite the registered worker state",
+ workerInfo.lastHeartbeat() > staleTime);
+ }
+
+ private long lastAppliedIndex() {
+ return
ratisServer.getMasterStateMachine().getLastAppliedTermIndex().getIndex();
+ }
+}
diff --git
a/master/src/test/java/org/apache/celeborn/service/deploy/master/clustermeta/ha/MasterStateMachineSuiteJ.java
b/master/src/test/java/org/apache/celeborn/service/deploy/master/clustermeta/ha/MasterStateMachineSuiteJ.java
index f272153664..bc9b46670b 100644
---
a/master/src/test/java/org/apache/celeborn/service/deploy/master/clustermeta/ha/MasterStateMachineSuiteJ.java
+++
b/master/src/test/java/org/apache/celeborn/service/deploy/master/clustermeta/ha/MasterStateMachineSuiteJ.java
@@ -581,4 +581,114 @@ public class MasterStateMachineSuiteJ extends
RatisBaseSuiteJ {
.getIndex());
stopRaftServers(raftServers);
}
+
+ @Test
+ public void testBatchHeartbeat() throws InvalidProtocolBufferException {
+ StateMachine stateMachine = ratisServer.getMasterStateMachine();
+
+ // Register two workers first, heartbeat apply is a no-op for unregistered
workers.
+ for (int i = 1; i <= 2; i++) {
+ PbMetaRequest registerRequest =
+ PbMetaRequest.newBuilder()
+ .setMetaRequestType(
+
org.apache.celeborn.common.protocol.PbMetaRequestType.RegisterWorker)
+ .setRequestId(UUID.randomUUID().toString())
+ .setRegisterWorkerRequest(
+
org.apache.celeborn.common.protocol.PbMetaRegisterWorkerRequest.newBuilder()
+ .setHost("host" + i)
+ .setRpcPort(1)
+ .setPushPort(2)
+ .setFetchPort(3)
+ .setReplicatePort(4)
+ .setInternalPort(5)
+ .build())
+ .build();
+ Assert.assertTrue(stateMachine.runCommand(registerRequest,
-1).getSuccess());
+ }
+ Assert.assertEquals(2, metaSystem.workersMap.size());
+
+ long heartbeatTime = System.currentTimeMillis();
+ org.apache.celeborn.common.protocol.PbMetaWorkerHeartbeatRequest
workerHeartbeat1 =
+
org.apache.celeborn.common.protocol.PbMetaWorkerHeartbeatRequest.newBuilder()
+ .setHost("host1")
+ .setRpcPort(1)
+ .setPushPort(2)
+ .setFetchPort(3)
+ .setReplicatePort(4)
+ .setTime(heartbeatTime)
+ .build();
+ org.apache.celeborn.common.protocol.PbMetaWorkerHeartbeatRequest
workerHeartbeat2 =
+
org.apache.celeborn.common.protocol.PbMetaWorkerHeartbeatRequest.newBuilder()
+ .setHost("host2")
+ .setRpcPort(1)
+ .setPushPort(2)
+ .setFetchPort(3)
+ .setReplicatePort(4)
+ .setTime(heartbeatTime)
+ .build();
+ org.apache.celeborn.common.protocol.PbMetaAppHeartbeatRequest appHeartbeat
=
+
org.apache.celeborn.common.protocol.PbMetaAppHeartbeatRequest.newBuilder()
+ .setAppId("app-1")
+ .setTime(heartbeatTime)
+ .setTotalWritten(100)
+ .setFileCount(10)
+ .build();
+
+ PbMetaRequest batchRequest =
+ PbMetaRequest.newBuilder()
+ .setMetaRequestType(
+
org.apache.celeborn.common.protocol.PbMetaRequestType.BatchHeartbeat)
+ .setRequestId(UUID.randomUUID().toString())
+ .setBatchHeartbeatRequest(
+
org.apache.celeborn.common.protocol.PbBatchHeartbeatRequest.newBuilder()
+ .addWorkerHeartbeats(workerHeartbeat1)
+ .addWorkerHeartbeats(workerHeartbeat2)
+ .addAppHeartbeats(appHeartbeat)
+ .build())
+ .build();
+
+ PbMetaRequestResponse response = stateMachine.runCommand(batchRequest, -1);
+ Assert.assertTrue(response.getSuccess());
+
+ // Each worker heartbeat must be applied with its own time field (see
CELEBORN-2399).
+ for (WorkerInfo workerInfo : metaSystem.workersMap.values()) {
+ Assert.assertEquals(heartbeatTime, workerInfo.lastHeartbeat());
+ }
+ Assert.assertEquals(Long.valueOf(heartbeatTime),
metaSystem.appHeartbeatTime.get("app-1"));
+
+ // The batched log entry written by the leader (ResourceRequest) must stay
wire-compatible
+ // with the apply side (PbMetaRequest).
+ ResourceProtos.ResourceRequest resourceRequest =
+ ResourceProtos.ResourceRequest.newBuilder()
+ .setCmdType(ResourceProtos.Type.BatchHeartbeat)
+ .setRequestId(UUID.randomUUID().toString())
+ .setBatchHeartbeatRequest(
+ ResourceProtos.BatchHeartbeatRequest.newBuilder()
+ .addWorkerHeartbeats(
+ ResourceProtos.WorkerHeartbeatRequest.newBuilder()
+ .setHost("host1")
+ .setRpcPort(1)
+ .setPushPort(2)
+ .setFetchPort(3)
+ .setReplicatePort(4)
+ .setTime(heartbeatTime)
+ .build())
+ .addAppHeartbeats(
+ ResourceProtos.AppHeartbeatRequest.newBuilder()
+ .setAppId("app-1")
+ .setTime(heartbeatTime)
+ .setTotalWritten(100)
+ .setFileCount(10)
+ .build())
+ .build())
+ .build();
+ PbMetaRequest parsed =
PbMetaRequest.parseFrom(resourceRequest.toByteString());
+ Assert.assertEquals(
+ org.apache.celeborn.common.protocol.PbMetaRequestType.BatchHeartbeat,
+ parsed.getMetaRequestType());
+ Assert.assertEquals(1,
parsed.getBatchHeartbeatRequest().getWorkerHeartbeatsCount());
+ Assert.assertEquals(1,
parsed.getBatchHeartbeatRequest().getAppHeartbeatsCount());
+ Assert.assertEquals(
+ heartbeatTime,
parsed.getBatchHeartbeatRequest().getWorkerHeartbeats(0).getTime());
+ }
}
diff --git
a/master/src/test/java/org/apache/celeborn/service/deploy/master/clustermeta/ha/RatisBaseSuiteJ.java
b/master/src/test/java/org/apache/celeborn/service/deploy/master/clustermeta/ha/RatisBaseSuiteJ.java
index b727758e4a..3a03c63618 100644
---
a/master/src/test/java/org/apache/celeborn/service/deploy/master/clustermeta/ha/RatisBaseSuiteJ.java
+++
b/master/src/test/java/org/apache/celeborn/service/deploy/master/clustermeta/ha/RatisBaseSuiteJ.java
@@ -28,11 +28,12 @@ import org.apache.celeborn.common.CelebornConf;
public class RatisBaseSuiteJ {
HARaftServer ratisServer;
+ HAMasterMetaManager metaSystem;
@Before
public void init() throws Exception {
CelebornConf conf = new CelebornConf();
- HAMasterMetaManager metaSystem = new HAMasterMetaManager(null, conf);
+ metaSystem = new HAMasterMetaManager(null, conf);
MetaHandler handler = new MetaHandler(metaSystem);
File tmpDir1 = File.createTempFile("celeborn-ratis-tmp", "for-test-only");
tmpDir1.delete();
