[SPARK-36284][CORE][SHUFFLE] Add shuffle checksum support for push-based shuffle

common/network-common/src/main/java/org/apache/spark/network/util/TransportConf.java

@@ -613,4 +613,24 @@ public long mergedShuffleCleanerShutdownTimeout() {
     return JavaUtils.timeStringAsSec(
       conf.get("spark.shuffle.push.server.mergedShuffleCleaner.shutdown.timeout", "60s"));
   }
+
+  /**
+   * Whether the shuffle server calculates a checksum for every chunk of a merged shuffle
+   * partition while merging pushed blocks. The checksums are stored alongside the merged shuffle
+   * data and are only used to diagnose the cause of a corrupted shuffle chunk.
+   */
+  public boolean mergedShuffleChecksumEnabled() {
+    return conf.getBoolean("spark.shuffle.push.server.mergedShuffleChecksum.enabled", true);
+  }
+
+  /**
+   * The algorithm used to calculate the checksums of the merged shuffle chunks. The reducer
+   * calculates the checksum of a corrupted chunk with spark.shuffle.checksum.algorithm, so
+   * corruption of a merged shuffle chunk can only be diagnosed when the two match.
+   */
+  public String mergedShuffleChecksumAlgorithm() {
+    // Upper cased like the spark.shuffle.checksum.algorithm of the application is
+    return conf.get("spark.shuffle.push.server.mergedShuffleChecksum.algorithm", "ADLER32")
+      .toUpperCase(Locale.ROOT);
+  }
 }

common/network-shuffle/src/main/java/org/apache/spark/network/shuffle/BlockStoreClient.java

@@ -90,6 +90,48 @@ public Cause diagnoseCorruption(
     }
   }
 
+  /**
+   * Send the diagnosis request for the corrupted chunk of a merged shuffle partition to the
+   * shuffle server which merged it.
+   *
+   * @param host the host of the shuffle server which merged the chunk.
+   * @param port the port of the shuffle server which merged the chunk.
+   * @param shuffleId the shuffleId of the corrupted shuffle chunk
+   * @param shuffleMergeId the shuffleMergeId of the corrupted shuffle chunk
+   * @param reduceId the reduceId of the corrupted shuffle chunk
+   * @param chunkId the chunkId of the corrupted shuffle chunk
+   * @param checksum the shuffle checksum which calculated at client side for the corrupted
+   *                 shuffle chunk
+   * @param algorithm the checksum algorithm which is used for calculating checksum
+   * @return The cause of the shuffle chunk corruption
+   */
+  public Cause diagnoseShuffleChunkCorruption(
+      String host,
+      int port,
+      int shuffleId,
+      int shuffleMergeId,
+      int reduceId,
+      int chunkId,
+      long checksum,
+      String algorithm) {
+    try {
+      TransportClient client = clientFactory.createClient(host, port);
+      ByteBuffer response = client.sendRpcSync(
+        new DiagnoseShuffleChunkCorruption(
+          appId, shuffleId, shuffleMergeId, reduceId, chunkId, checksum, algorithm).toByteBuffer(),
+        transportConf.connectionTimeoutMs()
+      );
+      CorruptionCause cause =
+        (CorruptionCause) BlockTransferMessage.Decoder.fromByteBuffer(response);
+      return cause.cause;
+    } catch (Exception e) {
+      // A shuffle service that does not support this request yet answers it with an error, so
+      // the exception is logged to tell that case apart from a genuinely unknown cause.
+      logger.warn("Failed to get the corruption cause of the shuffle chunk.", e);
+      return Cause.UNKNOWN_ISSUE;
+    }
+  }
+
   /**
    * Fetch a sequence of blocks from a remote node asynchronously,
    *

common/network-shuffle/src/main/java/org/apache/spark/network/shuffle/ExternalBlockHandler.java

@@ -280,6 +280,13 @@ protected void handleMessage(
       // In any cases of the error, diagnoseShuffleBlockCorruption should return UNKNOWN_ISSUE,
       // so it should always reply as success.
       callback.onSuccess(new CorruptionCause(cause).toByteBuffer());
+    } else if (msgObj instanceof DiagnoseShuffleChunkCorruption msg) {
+      checkAuth(client, msg.appId);
+      Cause cause = mergeManager.diagnoseShuffleChunkCorruption(msg.appId, msg.shuffleId,
+        msg.shuffleMergeId, msg.reduceId, msg.chunkId, msg.checksum, msg.algorithm);
+      // In any cases of the error, diagnoseShuffleChunkCorruption should return UNKNOWN_ISSUE,
+      // so it should always reply as success.
+      callback.onSuccess(new CorruptionCause(cause).toByteBuffer());
     } else {
       throw new UnsupportedOperationException("Unexpected message: " + msgObj);
     }

common/network-shuffle/src/main/java/org/apache/spark/network/shuffle/MergedShuffleFileManager.java

@@ -25,6 +25,7 @@
 import org.apache.spark.annotation.Evolving;
 import org.apache.spark.network.buffer.ManagedBuffer;
 import org.apache.spark.network.client.StreamCallbackWithID;
+import org.apache.spark.network.shuffle.checksum.Cause;
 import org.apache.spark.network.shuffle.protocol.ExecutorShuffleInfo;
 import org.apache.spark.network.shuffle.protocol.FinalizeShuffleMerge;
 import org.apache.spark.network.shuffle.protocol.MergeStatuses;
@@ -133,6 +134,34 @@ MergedBlockMeta getMergedBlockMeta(
    */
   void removeShuffleMerge(RemoveShuffleMerge removeShuffleMerge);
 
+  /**
+   * Diagnose the cause of the corruption of a merged shuffle chunk by comparing the checksum
+   * calculated by the reducer against the one calculated while the chunk was merged. This is
+   * best effort, so it returns {@link Cause#UNKNOWN_ISSUE} rather than failing when the
+   * checksum of the chunk is unavailable.
+   *
+   * @param appId application ID
+   * @param shuffleId shuffle ID
+   * @param shuffleMergeId shuffleMergeId is used to uniquely identify merging process
+   *                       of shuffle by an indeterminate stage attempt.
+   * @param reduceId reducer ID
+   * @param chunkId the ID of the corrupted chunk of the merged shuffle partition
+   * @param checksumByReader the checksum of the chunk calculated by the reducer
+   * @param algorithm the checksum algorithm the reducer used
+   * @return the cause of the corruption
+   * @since 4.4.0
+   */
+  default Cause diagnoseShuffleChunkCorruption(
+      String appId,
+      int shuffleId,
+      int shuffleMergeId,
+      int reduceId,
+      int chunkId,
+      long checksumByReader,
+      String algorithm) {
+    return Cause.UNKNOWN_ISSUE;
+  }
+
   /**
    * Optionally close any resources associated the MergedShuffleFileManager, such as the
    * leveldb for state persistence.