[CELEBORN-2429][CIP-22] Batch worker/application heartbeats into aggregated raft log entries

common/src/main/proto/TransportMessages.proto

@@ -463,6 +463,11 @@ message PbMetaBatchUnregisterShuffles {
   repeated string shuffleKeys = 1;
 }
 
+message PbBatchHeartbeatRequest {
+  repeated PbMetaWorkerHeartbeatRequest workerHeartbeats = 1;
+  repeated PbMetaAppHeartbeatRequest appHeartbeats = 2;
+}
+
 message PbUnregisterShuffleResponse {
   int32 status = 1;
 }
@@ -998,6 +1003,7 @@ enum PbMetaRequestType {
   BatchUnRegisterShuffle = 28;
   ReviseLostShuffles = 29;
   RegisterApplicationInfo = 30;
+  BatchHeartbeat = 31;
 }
 
 message PbMetaRequest {
@@ -1019,6 +1025,7 @@ message PbMetaRequest {
   PbReportWorkerDecommission reportWorkerDecommissionRequest = 24;
   PbMetaBatchUnregisterShuffles batchUnregisterShuffleRequest = 25;
   PbRegisterApplicationInfo registerApplicationInfoRequest = 26;
+  PbBatchHeartbeatRequest batchHeartbeatRequest = 27;
   PbReviseLostShuffles reviseLostShufflesRequest = 102;
 }
 

common/src/main/scala/org/apache/celeborn/common/CelebornConf.scala

@@ -815,6 +815,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] = {
@@ -3086,6 +3088,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")

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.&lt;id&gt;.host | &lt;required&gt; | false | Host to bind of master node <id> in HA mode. | 0.3.0 | celeborn.ha.master.node.&lt;id&gt;.host | 
 | celeborn.master.ha.node.&lt;id&gt;.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.&lt;id&gt;.port | 9097 | false | Port to bind of master node <id> in HA mode. | 0.3.0 | celeborn.ha.master.node.&lt;id&gt;.port | 

master/src/main/java/org/apache/celeborn/service/deploy/master/clustermeta/ha/HAMasterMetaManager.java

@@ -46,6 +46,8 @@ public class HAMasterMetaManager extends AbstractMetaManager {
 
   protected HARaftServer ratisServer;
 
+  private HeartbeatAggregator heartbeatAggregator;
+
   public HAMasterMetaManager(RpcEnv rpcEnv, CelebornConf conf) {
     this(rpcEnv, conf, new CelebornRackResolver(conf));
   }
@@ -68,6 +70,13 @@ public HARaftServer getRatisServer() {
 
   public void setRatisServer(HARaftServer ratisServer) {
     this.ratisServer = ratisServer;
+    if (conf.masterHaHeartbeatBatchEnabled()) {
+      this.heartbeatAggregator = new HeartbeatAggregator(ratisServer, conf);
+    }
+  }
+
+  public HeartbeatAggregator getHeartbeatAggregator() {
+    return heartbeatAggregator;
   }
 
   @Override
@@ -168,22 +177,27 @@ public void handleAppHeartbeat(
       Map<String, Long> applicationFallbackCounts,
       long time,
       String requestId) {
+    ResourceProtos.AppHeartbeatRequest appHeartbeatRequest =
+        ResourceProtos.AppHeartbeatRequest.newBuilder()
+            .setAppId(appId)
+            .setTime(time)
+            .setTotalWritten(totalWritten)
+            .setFileCount(fileCount)
+            .setShuffleCount(shuffleCount)
+            .setApplicationCount(applicationCount)
+            .putAllShuffleFallbackCounts(shuffleFallbackCounts)
+            .putAllApplicationFallbackCounts(applicationFallbackCounts)
+            .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)
-                      .build())
+              .setAppHeartbeatRequest(appHeartbeatRequest)
               .build());
     } catch (CelebornRuntimeException e) {
       LOG.error("Handle heartbeat for {} failed!", appId, e);
@@ -312,23 +326,30 @@ public void handleWorkerHeartbeat(
       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);

master/src/main/java/org/apache/celeborn/service/deploy/master/clustermeta/ha/HeartbeatAggregator.java

@@ -0,0 +1,152 @@
+/*
+ * 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.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 Map<String, ResourceProtos.WorkerHeartbeatRequest> workerHeartbeats = new HashMap<>();
+  private Map<String, ResourceProtos.AppHeartbeatRequest> appHeartbeats = new HashMap<>();
+
+  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), heartbeat);
+    } finally {
+      pendingLock.unlock();
+    }
+  }
+
+  public void offerAppHeartbeat(ResourceProtos.AppHeartbeatRequest heartbeat) {
+    pendingLock.lock();
+    try {
+      appHeartbeats.put(heartbeat.getAppId(), heartbeat);
+    } finally {
+      pendingLock.unlock();
+    }
+  }
+
+  private static String workerKey(ResourceProtos.WorkerHeartbeatRequest heartbeat) {
+    return heartbeat.getHost()
+        + ":"
+        + heartbeat.getRpcPort()
+        + ":"
+        + heartbeat.getPushPort()
+        + ":"
+        + heartbeat.getFetchPort()
+        + ":"
+        + heartbeat.getReplicatePort();
+  }
+
+  public void stop() {
+    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() {
+    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();
+    }
+
+    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());
+    }
+  }
+}