[CELEBORN-2470] Add backoff with jitter to push retries and log hard split once per file

client/src/main/java/org/apache/celeborn/client/ShuffleClientImpl.java

@@ -22,7 +22,8 @@
 import java.util.*;
 import java.util.concurrent.Callable;
 import java.util.concurrent.ConcurrentHashMap;
-import java.util.concurrent.ExecutorService;
+import java.util.concurrent.ScheduledExecutorService;
+import java.util.concurrent.ThreadLocalRandom;
 import java.util.concurrent.TimeUnit;
 
 import scala.Tuple2;
@@ -131,7 +132,10 @@ public class ShuffleClientImpl extends ShuffleClient {
   private boolean pushReplicateEnabled;
   private boolean fetchExcludeWorkerOnFailureEnabled;
 
-  private final ExecutorService pushDataRetryPool;
+  private final ScheduledExecutorService pushDataRetryPool;
+  private final long pushRetryBackoffInitialMs;
+  private final long pushRetryBackoffMaxMs;
+  private final double pushRetryBackoffJitter;
 
   private final Map<Integer, Set<Integer>> splitting = JavaUtils.newConcurrentHashMap();
 
@@ -235,7 +239,11 @@ public ShuffleClientImpl(String appUniqueId, CelebornConf conf, UserIdentifier u
     initDataClientFactoryIfNeeded();
     int pushDataRetryThreads = conf.clientPushRetryThreads();
     pushDataRetryPool =
-        ThreadUtils.newDaemonCachedThreadPool("celeborn-retry-sender", pushDataRetryThreads, 60);
+        ThreadUtils.newDaemonThreadPoolScheduledExecutor(
+            "celeborn-retry-sender", pushDataRetryThreads);
+    pushRetryBackoffInitialMs = conf.clientPushRetryBackoffInitial();
+    pushRetryBackoffMaxMs = conf.clientPushRetryBackoffMax();
+    pushRetryBackoffJitter = conf.clientPushRetryBackoffJitter();
 
     reviveManager = new ReviveManager(this, conf);
 
@@ -304,6 +312,40 @@ private boolean isPushTargetWorkerExcluded(
     }
   }
 
+  /**
+   * Delay before the {@code retryTimes}-th push retry: an exponentially growing base capped at
+   * {@code celeborn.client.push.retry.backoff.max}, spread by jitter so that map tasks blocked on
+   * the same partition do not retry in lockstep. Returns 0 when backoff is disabled.
+   */
+  @VisibleForTesting
+  long pushRetryBackoffMs(int retryTimes) {
+    if (pushRetryBackoffInitialMs <= 0 || retryTimes <= 0) {
+      return 0;
+    }
+    // Shift-loop rather than a power, because the hard split retry counter is not bounded.
+    long delay = pushRetryBackoffInitialMs;
+    for (int i = 1; i < retryTimes && delay < pushRetryBackoffMaxMs; i++) {
+      delay = delay > pushRetryBackoffMaxMs / 2 ? pushRetryBackoffMaxMs : delay << 1;
+    }
+    delay = Math.min(delay, pushRetryBackoffMaxMs);
+    long jitter = (long) (delay * pushRetryBackoffJitter);
+    if (jitter <= 0) {
+      return delay;
+    }
+    long minDelay = delay - jitter;
+    long maxDelay = delay > pushRetryBackoffMaxMs - jitter ? pushRetryBackoffMaxMs : delay + jitter;
+    long jitterRange = maxDelay - minDelay;
+    if (jitterRange == Long.MAX_VALUE) {
+      return ThreadLocalRandom.current().nextLong(Long.MAX_VALUE);
+    }
+    return minDelay + ThreadLocalRandom.current().nextLong(jitterRange + 1);
+  }
+
+  @VisibleForTesting
+  void scheduleRetry(Runnable retry, long delayMs) {
+    pushDataRetryPool.schedule(retry, delayMs, TimeUnit.MILLISECONDS);
+  }
+
   private void submitRetryPushData(
       int shuffleId,
       byte[] body,
@@ -532,7 +574,7 @@ private void submitRetryPushMergedData(
     } else {
       ReviveRequest[] requests =
           addAndGetReviveRequests(shuffleId, mapId, attemptId, reviveFailedBatchesMap, cause);
-      pushDataRetryPool.submit(
+      scheduleRetry(
           () ->
               submitRetryPushMergedData(
                   pushState,
@@ -545,7 +587,8 @@ private void submitRetryPushMergedData(
                   requests,
                   remainReviveTimes - 1,
                   System.currentTimeMillis()
-                      + conf.clientRpcRequestPartitionLocationAskTimeout().duration().toMillis()));
+                      + conf.clientRpcRequestPartitionLocationAskTimeout().duration().toMillis()),
+          pushRetryBackoffMs(maxReviveTimes - (remainReviveTimes - 1)));
     }
   }
 
@@ -1133,6 +1176,9 @@ public void onFailure(Throwable e) {
       RpcResponseCallback wrappedCallback =
           new PushDataRpcResponseCallback() {
             int remainReviveTimes = maxReviveTimes;
+            // HARD_SPLIT retries deliberately do not consume remainReviveTimes, so they need their
+            // own counter to grow the backoff.
+            int hardSplitRetryTimes = 0;
             PartitionLocation latest = loc;
 
             @Override
@@ -1199,7 +1245,7 @@ public void onSuccess(ByteBuffer response) {
                           + conf.clientRpcRequestPartitionLocationAskTimeout()
                               .duration()
                               .toMillis();
-                  pushDataRetryPool.submit(
+                  scheduleRetry(
                       () ->
                           submitRetryPushData(
                               shuffleId,
@@ -1209,7 +1255,8 @@ public void onSuccess(ByteBuffer response) {
                               pushState,
                               reviveRequest,
                               remainReviveTimes,
-                              dueTime));
+                              dueTime),
+                      pushRetryBackoffMs(++hardSplitRetryTimes));
                 } else if (reason == StatusCode.PUSH_DATA_SUCCESS_PRIMARY_CONGESTED.getValue()) {
                   logger.debug(
                       "Push data to {} primary congestion required for shuffle {} map {} attempt {} partition {} batch {}.",
@@ -1291,7 +1338,7 @@ public void onFailure(Throwable e) {
                 long dueTime =
                     System.currentTimeMillis()
                         + conf.clientRpcRequestPartitionLocationAskTimeout().duration().toMillis();
-                pushDataRetryPool.submit(
+                scheduleRetry(
                     () ->
                         submitRetryPushData(
                             shuffleId,
@@ -1301,7 +1348,8 @@ public void onFailure(Throwable e) {
                             pushState,
                             reviveRequest,
                             remainReviveTimes,
-                            dueTime));
+                            dueTime),
+                    pushRetryBackoffMs(maxReviveTimes - remainReviveTimes));
               } else {
                 pushState.removeBatch(nextBatchId, latest.hostAndPushPort());
                 logger.info(
@@ -1643,7 +1691,11 @@ public void onSuccess(ByteBuffer response) {
                 ReviveRequest[] requests =
                     addAndGetReviveRequests(
                         shuffleId, mapId, attemptId, batchesNeedResubmit, StatusCode.HARD_SPLIT);
-                pushDataRetryPool.submit(
+                // remainReviveTimes is a parameter here and is not decremented on the HARD_SPLIT
+                // path, so there is no per-batch attempt counter to grow from. Apply the initial
+                // delay with jitter, which is what decorrelates the herd; growing this one too
+                // needs the counter threaded through doPushMergedData and is left as follow-up.
+                scheduleRetry(
                     () ->
                         submitRetryPushMergedData(
                             pushState,
@@ -1658,7 +1710,8 @@ public void onSuccess(ByteBuffer response) {
                             System.currentTimeMillis()
                                 + conf.clientRpcRequestPartitionLocationAskTimeout()
                                     .duration()
-                                    .toMillis()));
+                                    .toMillis()),
+                    pushRetryBackoffMs(1));
               }
             } else if (reason == StatusCode.PUSH_DATA_SUCCESS_PRIMARY_CONGESTED.getValue()) {
               logger.debug(
@@ -1731,7 +1784,7 @@ public void onFailure(Throwable e) {
             if (!mapperEnded(shuffleId, mapId)) {
               ReviveRequest[] requests =
                   addAndGetReviveRequests(shuffleId, mapId, attemptId, batches, cause);
-              pushDataRetryPool.submit(
+              scheduleRetry(
                   () ->
                       submitRetryPushMergedData(
                           pushState,
@@ -1746,7 +1799,8 @@ public void onFailure(Throwable e) {
                           System.currentTimeMillis()
                               + conf.clientRpcRequestPartitionLocationAskTimeout()
                                   .duration()
-                                  .toMillis()));
+                                  .toMillis()),
+                  pushRetryBackoffMs(maxReviveTimes - (remainReviveTimes - 1)));
             } else {
               pushState.removeBatch(groupedBatchId, hostPort);
               logger.info(