[CELEBORN-2364] Add support to cleanup stale shuffle dependencies across all Spark apps when cluster under load

client/src/main/scala/org/apache/celeborn/client/ApplicationHeartbeater.scala

@@ -23,6 +23,7 @@ import java.util.function.Consumer
 
 import scala.collection.JavaConverters._
 
+import com.google.common.annotations.VisibleForTesting
 import org.apache.commons.lang3.StringUtils
 
 import org.apache.celeborn.common.CelebornConf
@@ -47,6 +48,13 @@ class ApplicationHeartbeater(
   private var stopped = false
   private val reviseLostShuffles = conf.reviseLostShufflesEnabled
 
+  private val gcOnOverloadEnabled = conf.clientGcOnOverloadEnabled
+  private val gcOnOverloadMinIntervalMs = conf.clientGcOnOverloadMinIntervalMs
+  private val gcOnOverloadMinIntervalNs = TimeUnit.MILLISECONDS.toNanos(gcOnOverloadMinIntervalMs)
+
+  @VisibleForTesting
+  @volatile private[client] var lastGcTriggerTimeNs = 0L
+
   // Use independent app heartbeat threads to avoid being blocked by other operations.
   private val appHeartbeatIntervalMs = conf.appHeartbeatIntervalMs
   private val applicationUnregisterEnabled = conf.applicationUnregisterEnabled
@@ -91,6 +99,7 @@ class ApplicationHeartbeater(
               logDebug("Successfully send app heartbeat.")
               workerStatusTracker.handleHeartbeatResponse(response)
               checkQuotaExceeds(response.checkQuotaResponse)
+              handleGcSignal(response.shouldTriggerGc)
               // revise shuffle id if there are lost shuffles
               if (reviseLostShuffles) {
                 val masterRecordedShuffleIds = response.registeredShuffles
@@ -170,6 +179,21 @@ class ApplicationHeartbeater(
     }
   }
 
+  private[client] def handleGcSignal(shouldTriggerGc: Boolean): Unit = {
+    if (!gcOnOverloadEnabled || !shouldTriggerGc) return
+    val nowNs = System.nanoTime()
+    if (nowNs - lastGcTriggerTimeNs >= gcOnOverloadMinIntervalNs) {
+      logInfo(
+        "Cluster is overloaded; triggering System.gc() to release stale shuffle dependencies.")
+      lastGcTriggerTimeNs = nowNs
+      System.gc()
+    } else {
+      logDebug(
+        s"Cluster overload GC signal received but skipped: last GC was triggered " +
+          s"${(nowNs - lastGcTriggerTimeNs) / 1000000L}ms ago (min interval: ${gcOnOverloadMinIntervalMs}ms).")
+    }
+  }
+
   def stop(): Unit = {
     this.synchronized {
       if (!stopped) {

client/src/test/scala/org/apache/celeborn/client/ApplicationHeartbeaterSuite.scala

@@ -0,0 +1,112 @@
+/*
+ * 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.client
+
+import java.util.concurrent.ConcurrentHashMap
+
+import org.apache.celeborn.CelebornFunSuite
+import org.apache.celeborn.common.CelebornConf
+
+class ApplicationHeartbeaterSuite extends CelebornFunSuite {
+
+  private def makeHeartbeater(
+      gcEnabled: Boolean = true,
+      minIntervalMs: Long = 60000L): ApplicationHeartbeater = {
+    val conf = new CelebornConf()
+    conf.set(CelebornConf.CLIENT_GC_ON_CLUSTER_OVERLOAD_ENABLED, gcEnabled)
+    conf.set(CelebornConf.CLIENT_GC_ON_CLUSTER_OVERLOAD_MIN_INTERVAL, minIntervalMs)
+
+    val registeredShuffles = ConcurrentHashMap.newKeySet[Int]()
+      .asInstanceOf[ConcurrentHashMap.KeySetView[Int, java.lang.Boolean]]
+
+    new ApplicationHeartbeater(
+      "test-app",
+      conf,
+      null, // MasterClient not needed for handleGcSignal unit tests
+      () => (0L, 0L) -> (0L, 0L, Map.empty, Map.empty),
+      null, // WorkerStatusTracker not needed
+      registeredShuffles,
+      _ => ())
+  }
+
+  test("GC is not triggered when feature is disabled") {
+    val hb = makeHeartbeater(gcEnabled = false)
+    hb.handleGcSignal(shouldTriggerGc = true)
+    assert(hb.lastGcTriggerTimeNs == 0L)
+  }
+
+  test("GC is not triggered when signal is false") {
+    val hb = makeHeartbeater(gcEnabled = true)
+    hb.handleGcSignal(shouldTriggerGc = false)
+    assert(hb.lastGcTriggerTimeNs == 0L)
+  }
+
+  test("GC is triggered on first signal when enabled") {
+    val hb = makeHeartbeater(gcEnabled = true, minIntervalMs = 0L)
+    val before = System.nanoTime()
+    hb.handleGcSignal(shouldTriggerGc = true)
+    assert(hb.lastGcTriggerTimeNs >= before)
+  }
+
+  test("GC is skipped when called again within the cooldown interval") {
+    val hb = makeHeartbeater(gcEnabled = true, minIntervalMs = 60000L)
+
+    hb.handleGcSignal(shouldTriggerGc = true)
+    val firstTrigger = hb.lastGcTriggerTimeNs
+    assert(firstTrigger > 0L)
+
+    // Second call immediately — well within 60s cooldown
+    hb.handleGcSignal(shouldTriggerGc = true)
+    assert(
+      hb.lastGcTriggerTimeNs == firstTrigger,
+      "lastGcTriggerTimeNs should not change on second call")
+  }
+
+  test("GC fires again after cooldown interval has elapsed") {
+    // Use 0ms cooldown so any elapsed time satisfies it.
+    val hb = makeHeartbeater(gcEnabled = true, minIntervalMs = 0L)
+
+    hb.handleGcSignal(shouldTriggerGc = true)
+    val firstTrigger = hb.lastGcTriggerTimeNs
+
+    // With 0ms interval the next call should always be allowed.
+    hb.handleGcSignal(shouldTriggerGc = true)
+    assert(hb.lastGcTriggerTimeNs >= firstTrigger)
+  }
+
+  test("no GC when both signal is false and feature is disabled") {
+    val hb = makeHeartbeater(gcEnabled = false)
+    hb.handleGcSignal(shouldTriggerGc = false)
+    assert(hb.lastGcTriggerTimeNs == 0L)
+  }
+
+  test("cooldown is measured from the last successful GC trigger") {
+    val hb = makeHeartbeater(gcEnabled = true, minIntervalMs = 60000L)
+
+    // First trigger — sets lastGcTriggerTimeNs
+    hb.handleGcSignal(shouldTriggerGc = true)
+    val firstTrigger = hb.lastGcTriggerTimeNs
+    assert(firstTrigger > 0L)
+
+    // Three more calls within cooldown — none should update lastGcTriggerTimeNs
+    hb.handleGcSignal(shouldTriggerGc = true)
+    hb.handleGcSignal(shouldTriggerGc = true)
+    hb.handleGcSignal(shouldTriggerGc = true)
+    assert(hb.lastGcTriggerTimeNs == firstTrigger)
+  }
+}

common/src/main/proto/TransportMessages.proto

@@ -510,6 +510,7 @@ message PbHeartbeatFromApplicationResponse {
   repeated PbWorkerInfo shuttingWorkers = 4;
   repeated int32 registeredShuffles = 5;
   PbCheckQuotaResponse checkQuotaResponse = 6;
+  bool shouldTriggerGc = 7;
 }
 
 message PbCheckQuota {

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

@@ -936,6 +936,10 @@ class CelebornConf(loadDefaults: Boolean) extends Cloneable with Logging with Se
   def userQuotaEnabled: Boolean = get(USER_QUOTA_ENABLED)
   def quotaInterruptShuffleEnabled: Boolean = get(QUOTA_INTERRUPT_SHUFFLE_ENABLED)
 
+  def clusterOverloadGcEnabled: Boolean = get(MASTER_CLUSTER_OVERLOAD_GC_ENABLED)
+  def clientGcOnOverloadEnabled: Boolean = get(CLIENT_GC_ON_CLUSTER_OVERLOAD_ENABLED)
+  def clientGcOnOverloadMinIntervalMs: Long = get(CLIENT_GC_ON_CLUSTER_OVERLOAD_MIN_INTERVAL)
+
   // //////////////////////////////////////////////////////
   //                      Identity                       //
   // //////////////////////////////////////////////////////
@@ -2517,6 +2521,15 @@ object CelebornConf extends Logging {
       .timeConf(TimeUnit.MILLISECONDS)
       .createWithDefaultString("300s")
 
+  val MASTER_CLUSTER_OVERLOAD_GC_ENABLED: ConfigEntry[Boolean] =
+    buildConf("celeborn.master.clusterOverload.gc.enabled")
+      .categories("master")
+      .version("0.7.0")
+      .doc("Whether to enable the master signaling clients to trigger GC when the cluster " +
+        "storage is overloaded (any cluster quota dimension exceeds the overload threshold).")
+      .booleanConf
+      .createWithDefault(false)
+
   val DFS_EXPIRE_DIRS_TIMEOUT: ConfigEntry[Long] =
     buildConf("celeborn.master.dfs.expireDirs.timeout")
       .categories("master")
@@ -4776,6 +4789,26 @@ object CelebornConf extends Logging {
       .booleanConf
       .createWithDefault(true)
 
+  val CLIENT_GC_ON_CLUSTER_OVERLOAD_ENABLED: ConfigEntry[Boolean] =
+    buildConf("celeborn.client.clusterOverload.gc.enabled")
+      .categories("client")
+      .version("0.7.0")
+      .doc("When true, the client will trigger System.gc() upon receiving a GC signal from " +
+        "the master indicating cluster storage is overloaded. Disable to ignore the signal.")
+      .booleanConf
+      .createWithDefault(true)
+
+  val CLIENT_GC_ON_CLUSTER_OVERLOAD_MIN_INTERVAL: ConfigEntry[Long] =
+    buildConf("celeborn.client.clusterOverload.gc.minInterval")
+      .categories("client")
+      .version("0.7.0")
+      .doc("Minimum time that must elapse between consecutive GC triggers on the client side " +
+        "in response to master GC signals. This prevents excessive GC pressure when the " +
+        "cluster remains overloaded across multiple heartbeat intervals.")
+      .timeConf(TimeUnit.MILLISECONDS)
+      .checkValue(_ >= 0, "Should be >= 0.")
+      .createWithDefaultString("5m")
+
   val CLIENT_EXCLUDE_PEER_WORKER_ON_FAILURE_ENABLED: ConfigEntry[Boolean] =
     buildConf("celeborn.client.excludePeerWorkerOnFailure.enabled")
       .categories("client")
@@ -6895,6 +6928,18 @@ object CelebornConf extends Logging {
       .checkValue(v => v >= 0.0 && v < 1.0, "Should be in [0.0, 1).")
       .createWithDefault(0.3)
 
+  val QUOTA_CLUSTER_OVERLOAD_FACTOR: ConfigEntry[Double] =
+    buildConf("celeborn.quota.overload.factor")
+      .categories("quota")
+      .dynamic
+      .doc("This config decides the quota * factor at which to consider the cluster 'overloaded'." +
+        " When the cluster is overloaded and `celeborn.master.clusterOverload.gc.enabled` is true," +
+        " application heartbeat responses contain a signal to trigger a GC to clean up dangling shuffle dependencies")
+      .version("0.7.0")
+      .doubleConf
+      .checkValue(v => v > 0.0 && v <= 1.0, "Should be in (0.0, 1.0].")
+      .createWithDefault(0.8)
+
   val QUOTA_CLUSTER_DISK_BYTES_WRITTEN: ConfigEntry[Long] =
     buildConf("celeborn.quota.cluster.diskBytesWritten")
       .categories("quota")

common/src/main/scala/org/apache/celeborn/common/protocol/message/ControlMessages.scala

@@ -397,7 +397,8 @@ object ControlMessages extends Logging {
       unknownWorkers: util.List[WorkerInfo],
       shuttingWorkers: util.List[WorkerInfo],
       registeredShuffles: util.List[Integer],
-      checkQuotaResponse: CheckQuotaResponse) extends Message
+      checkQuotaResponse: CheckQuotaResponse,
+      shouldTriggerGc: Boolean = false) extends Message
 
   case class CheckQuota(userIdentifier: UserIdentifier) extends Message
 
@@ -890,7 +891,8 @@ object ControlMessages extends Logging {
           unknownWorkers,
           shuttingWorkers,
           registeredShuffles,
-          checkQuotaResponse) =>
+          checkQuotaResponse,
+          shouldTriggerGc) =>
       val pbCheckQuotaResponse = PbCheckQuotaResponse.newBuilder().setAvailable(
         checkQuotaResponse.isAvailable).setReason(checkQuotaResponse.reason)
       val payload = PbHeartbeatFromApplicationResponse.newBuilder()
@@ -903,6 +905,7 @@ object ControlMessages extends Logging {
           shuttingWorkers.asScala.map(PbSerDeUtils.toPbWorkerInfo(_, true, true)).toList.asJava)
         .addAllRegisteredShuffles(registeredShuffles)
         .setCheckQuotaResponse(pbCheckQuotaResponse)
+        .setShouldTriggerGc(shouldTriggerGc)
         .build().toByteArray
       new TransportMessage(MessageType.HEARTBEAT_FROM_APPLICATION_RESPONSE, payload)
 
@@ -1380,7 +1383,8 @@ object ControlMessages extends Logging {
           pbHeartbeatFromApplicationResponse.getShuttingWorkersList.asScala
             .map(PbSerDeUtils.fromPbWorkerInfo).toList.asJava,
           pbHeartbeatFromApplicationResponse.getRegisteredShufflesList,
-          checkQuotaResponse)
+          checkQuotaResponse,
+          pbHeartbeatFromApplicationResponse.getShouldTriggerGc)
 
       case CHECK_QUOTA_VALUE =>
         val pbCheckAvailable = PbCheckQuota.parseFrom(message.getPayload)

common/src/main/scala/org/apache/celeborn/common/quota/StorageQuota.scala

@@ -37,4 +37,5 @@ case class StorageQuota(
 
 object StorageQuota {
   val DEFAULT_QUOTA = StorageQuota(Long.MaxValue, Long.MaxValue, Long.MaxValue, Long.MaxValue)
+  val CLUSTER_OVERLOAD_LIMIT_DEFAULT_MULTIPLIER: Double = 0.8
 }

common/src/test/scala/org/apache/celeborn/common/util/PbSerDeUtilsTest.scala

@@ -814,14 +814,16 @@ class PbSerDeUtilsTest extends CelebornFunSuite {
       mockWorkers("host1").toList.asJava,
       mockWorkers("host2").toList.asJava,
       Array(Integer.valueOf(1)).toList.asJava,
-      CheckQuotaResponse(isAvailable = false, "test_reason"))
+      CheckQuotaResponse(isAvailable = false, "test_reason"),
+      shouldTriggerGc = true)
     val toTransportHeartbeatFromApplicationResponse =
       ControlMessages.toTransportMessage(heartbeatFromApplicationResponse)
     val fromTransportHeartbeatFromApplicationResponse =
       ControlMessages.fromTransportMessage(toTransportHeartbeatFromApplicationResponse)
         .asInstanceOf[HeartbeatFromApplicationResponse]
 
     assert(fromTransportHeartbeatFromApplicationResponse.equals(heartbeatFromApplicationResponse))
+    assert(fromTransportHeartbeatFromApplicationResponse.shouldTriggerGc)
   }
 
   test("HeartbeatFromApplicationResponse backward compatibility without checkQuotaResponse") {

docs/configuration/client.md

@@ -27,6 +27,8 @@ license: |
 | celeborn.client.application.uuidSuffix.enabled | false | false | Whether to add UUID suffix for application id for unique. When `true`, add UUID suffix for unique application id. Currently, this only applies to Spark and MR. | 0.6.0 |  | 
 | celeborn.client.chunk.prefetch.enabled | false | false | Whether to enable chunk prefetch when creating CelebornInputStream. | 0.5.1 |  | 
 | celeborn.client.closeIdleConnections | true | false | Whether client will close idle connections. | 0.3.0 |  | 
+| celeborn.client.clusterOverload.gc.enabled | true | false | When true, the client will trigger System.gc() upon receiving a GC signal from the master indicating cluster storage is overloaded. Disable to ignore the signal. | 0.7.0 |  | 
+| celeborn.client.clusterOverload.gc.minInterval | 5m | false | Minimum time that must elapse between consecutive GC triggers on the client side in response to master GC signals. This prevents excessive GC pressure when the cluster remains overloaded across multiple heartbeat intervals. | 0.7.0 |  | 
 | celeborn.client.commitFiles.ignoreExcludedWorker | false | false | When true, LifecycleManager will skip workers which are in the excluded list. | 0.3.0 |  | 
 | celeborn.client.eagerlyCreateInputStream.threads | 32 | false | Threads count for streamCreatorPool in CelebornShuffleReader. | 0.3.1 |  | 
 | celeborn.client.excludePeerWorkerOnFailure.enabled | true | false | When true, Celeborn will exclude partition's peer worker on failure when push data to replica failed. | 0.3.0 |  | 

docs/configuration/master.md

@@ -37,6 +37,7 @@ license: |
 | celeborn.internal.port.enabled | false | false | Whether to create a internal port on Masters/Workers for inter-Masters/Workers communication. This is beneficial when SASL authentication is enforced for all interactions between clients and Celeborn Services, but the services can exchange messages without being subject to SASL authentication. | 0.5.0 |  | 
 | celeborn.logConf.enabled | false | false | When `true`, log the CelebornConf for debugging purposes. | 0.5.0 |  | 
 | celeborn.master.allowWorkerHostPattern | <undefined> | false | Pattern of worker host that allowed to register with the master. If not set, all workers are allowed to register. | 0.6.0 |  | 
+| celeborn.master.clusterOverload.gc.enabled | false | false | Whether to enable the master signaling clients to trigger GC when the cluster storage is overloaded (any cluster quota dimension exceeds the overload threshold). | 0.7.0 |  | 
 | celeborn.master.denyWorkerHostPattern | <undefined> | false | Pattern of worker host that denied to register with the master. If not set, no workers are denied to register. | 0.6.0 |  | 
 | celeborn.master.dfs.expireDirs.timeout | 1h | false | The timeout for an expired dirs to be deleted on dfs like HDFS, S3, OSS. | 0.6.0 |  | 
 | celeborn.master.estimatedPartitionSize.initialSize | 64mb | false | Initial partition size for estimation, it will change according to runtime stats. | 0.3.0 | celeborn.shuffle.initialEstimatedPartitionSize | 

docs/configuration/quota.md

@@ -26,6 +26,7 @@ license: |
 | celeborn.quota.cluster.hdfsFileCount | 9223372036854775807 | true | Cluster level quota dynamic configuration for written hdfs file count. | 0.6.0 |  | 
 | celeborn.quota.enabled | true | false | When Master side sets to true, the master will enable to check the quota via QuotaManager. When Client side sets to true, LifecycleManager will request Master side to check whether the current user has enough quota before registration of shuffle. Fallback to the default shuffle service when Master side checks that there is no enough quota for current user. | 0.2.0 |  | 
 | celeborn.quota.interruptShuffle.enabled | false | false | Whether to enable interrupt shuffle when quota exceeds. | 0.6.0 |  | 
+| celeborn.quota.overload.factor | 0.8 | true | This config decides the quota * factor at which to consider the cluster 'overloaded'. When the cluster is overloaded and `celeborn.master.clusterOverload.gc.enabled` is true, application heartbeat responses contain a signal to trigger a GC to clean up dangling shuffle dependencies | 0.7.0 |  | 
 | celeborn.quota.tenant.diskBytesWritten | 9223372036854775807b | true | Tenant level quota dynamic configuration for written disk bytes. | 0.5.0 |  | 
 | celeborn.quota.tenant.diskFileCount | 9223372036854775807 | true | Tenant level quota dynamic configuration for written disk file count. | 0.5.0 |  | 
 | celeborn.quota.tenant.enabled | true | false | Whether to enable tenant-level quota. | 0.6.0 |  | 

master/src/main/scala/org/apache/celeborn/service/deploy/master/Master.scala

@@ -1255,7 +1255,8 @@ private[celeborn] class Master(
         new util.ArrayList[WorkerInfo](
           (statusSystem.shutdownWorkers.asScala ++ statusSystem.decommissionWorkers.asScala).asJava),
         new util.ArrayList(appRelatedShuffles),
-        quotaManager.checkApplicationQuotaStatus(appId)))
+        quotaManager.checkApplicationQuotaStatus(appId),
+        if (conf.clusterOverloadGcEnabled) quotaManager.isClusterOverloaded else false))
     } else {
       context.reply(OneWayMessageResponse)
     }