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

[Dreamstick9]

What changes were proposed in this pull request?

SPARK-36206 added shuffle corruption diagnosis based on shuffle checksums, but only for shuffle
blocks that are not merged. When a push-merged shuffle chunk is detected as corrupt, the reducer
falls back to the original shuffle blocks and reports nothing about the cause, because no checksum
of the merged data exists. SPARK-37695 made that explicit by skipping the diagnosis of a
ShuffleBlockChunkId with a TODO SPARK-36284. This PR fills that gap.

Server side, RemoteBlockPushResolver now calculates a checksum per chunk of a merged shuffle
partition while it merges the pushed blocks, and stores them next to the merged data, index and
meta files as shuffleMerged_<appId>_<shuffleId>_<shuffleMergeId>_<reduceId>.checksum.<ALGORITHM>
(one long per chunk, so one entry less than the index file, which has a leading zero offset). The
checksum algorithm is part of the file name, as it is for the checksum file of a non-merged shuffle
block, so the checksums can never be compared against those of a different algorithm after the
shuffle service is reconfigured.

The running checksum of the current chunk is fed in writeBuf, right after the data of a block has
been written, so merging does not read the merged data back in the common case. Two situations make
the running checksum stop describing the current chunk, and both are detected by tracking the
position up to which the checksum has consumed the data file:

• a block that was partially written is abandoned, and its data is overwritten by the next block,
  since dataFilePos only advances when a block commits;
• a block is still being written when the shuffle merge is finalized, and its data is truncated
  away.

In both cases the checksum of the chunk is recalculated from the merged data file when the chunk
gets sealed. The checksum entry is written from updateChunkInfo, together with the chunk offset
and the chunk bitmap, and it follows the same tracked-position and rewind-on-retry protocol, so the
checksum file stays aligned with the index file even when the update of the index or meta file
fails and is retried.

Client side, a new DiagnoseShuffleChunkCorruption message asks the shuffle service which merged a
chunk for the cause of its corruption, and the response reuses CorruptionCause.
ShuffleBlockFetcherIterator now diagnoses a corrupt shuffle chunk before it falls back to the
original shuffle blocks. The diagnosis of a merged chunk is answered by comparing the checksum the
reducer calculated against the stored one and against a recalculation of the chunk on disk, exactly
like ExternalShuffleBlockResolver.diagnoseShuffleBlockCorruption does for a non-merged block, so
DISK_ISSUE, NETWORK_ISSUE and CHECKSUM_VERIFY_PASS keep their existing meanings. A chunk that
was merged with a different algorithm than the reducer uses is reported as
UNSUPPORTED_CHECKSUM_ALGORITHM.

Two shuffle service configurations are added:

Configuration	Default
spark.shuffle.push.server.mergedShuffleChecksum.enabled	true
spark.shuffle.push.server.mergedShuffleChecksum.algorithm	ADLER32

The default algorithm matches the default of spark.shuffle.checksum.algorithm, so the two line up
out of the box. An unsupported algorithm only disables the calculation with a warning instead of
failing the merge.

The checksums are strictly diagnostic, so nothing about them is allowed to affect the merge. A
failure to write or recompute them gives up on the checksums of that shuffle partition, logs a
warning and deletes the checksum file when the partition is finalized, instead of counting towards
spark.shuffle.push.server.ioExceptionsThresholdDuringMerge or failing the finalization of an
otherwise correctly merged partition. A chunk of such a partition then simply diagnoses as
UNKNOWN_ISSUE.

Notes for reviewers, mostly on the trade-offs I would like your opinion on:

1. The checksum is calculated over the data as the shuffle service received it, so it covers
   corruption of the merged data on the shuffle service disk and in transit to the reducer. It
   cannot distinguish corruption that already happened while a mapper pushed the block, which is
   noted in the code.
2. The diagnosis of a merged chunk is informational. A corrupt chunk still falls back to the
   original shuffle blocks exactly as before, since a chunk is never re-fetched. It does however
   delay that fallback by the drain of the chunk and one synchronous RPC. Like the existing
   diagnosis of a shuffle block, that RPC uses spark.network.timeout for both the connection and
   the request, so a shuffle service that stops responding between the fetch of a chunk and its
   diagnosis can delay by up to twice that timeout a fallback which is immediate today. The
   connection is normally still cached from the fetch of the chunk itself, so this only shows up
   when the shuffle service dies in that window, but it is the one behaviour change here that
   makes a recoverable path slower. Happy to gate it on a separate configuration, or to skip the
   diagnosis for chunks whose shuffle service just failed, if you prefer either.
3. The default of spark.shuffle.push.server.mergedShuffleChecksum.enabled is true, matching
   spark.shuffle.checksum.enabled. It costs one checksum pass over the merged data, one extra
   file, and one more file descriptor per actively merging partition, which is held for as long as
   the index and meta file descriptors are. The checksum file is written exactly as often as the
   index file, once per sealed chunk, so it follows the same MergeShuffleFile pattern, but it
   could be opened only when a chunk is sealed if the fourth descriptor is a concern. Please say
   so if you would rather have the whole thing off by default.
4. A chunk that was merged with a different algorithm than the reducer uses is reported as
   UNSUPPORTED_CHECKSUM_ALGORITHM even though both algorithms are supported on their own, since
   appending a new Cause would break a reducer running an older Spark version. The shuffle
   service logs which two algorithms disagree.
5. MergedShuffleFileManager gains a default method that returns UNKNOWN_ISSUE, so third party
   implementations and NoOpMergedShuffleFileManager keep working unchanged.

Why are the changes needed?

Push-based shuffle is the only place where Spark detects shuffle data corruption and cannot say
anything about its cause. When a merged chunk is corrupt, the reducer silently falls back to the
original blocks, so a shuffle service with a failing disk keeps corrupting merged data and the only
symptom is degraded performance from repeated fallbacks. With the merged chunks checksummed, the
same diagnosis that is available for regular shuffle blocks since Spark 3.2 tells the user whether
the corruption came from the shuffle service disk or from the network.

Does this PR introduce any user-facing change?

Yes. Two new shuffle service configurations, documented in docs/configuration.md. With the
default configuration a shuffle service with push-based shuffle enabled writes one additional
checksum file per merged shuffle partition, and the reducer logs the diagnosed cause of a corrupt
merged shuffle chunk, for example:

BlockChunk shuffleChunk_0_0_2_1 is corrupted due to DISK_ISSUE

Previously the corruption of a merged chunk was reported as:

BlockChunk shuffleChunk_0_0_2_1 is corrupted but corruption diagnosis is skipped due to lack of
shuffle checksum support for push-based shuffle.

How was this patch tested?

New unit tests, plus existing ones extended:

• RemoteBlockPushResolverSuite: the shared validateChunks helper now also verifies that every
  stored checksum matches the data its chunk ended up with in the merged shuffle data file, so all
  the existing merge tests cover the checksums, including the ones that inject IOExceptions into
  the index and meta file updates. New tests cover the checksums of the merged chunks and the entry
  count against the index file, a chunk whose data is overwritten after a block push failed, a
  chunk sealed while a block push is still in flight at finalization, the four diagnosis outcomes
  (CHECKSUM_VERIFY_PASS, NETWORK_ISSUE, DISK_ISSUE against a data file corrupted after the
  merge, and UNSUPPORTED_CHECKSUM_ALGORITHM), diagnosis of an unknown chunk, diagnosis when the
  checksum calculation is disabled, and a checksum file that cannot be written, which must leave
  the merge and its MergeStatuses untouched. Cleanup assertions for the checksum file were added
  to the outdated shuffle merge id and removeShuffleMerge tests.
• ExternalBlockHandlerSuite: the new RPC is routed to the MergedShuffleFileManager and answered
  with a CorruptionCause.
• BlockTransferMessagesSuite: round trip of DiagnoseShuffleChunkCorruption, and of
  DiagnoseCorruption and CorruptionCause, which were not covered before.
• ShuffleBlockFetcherIteratorSuite: the diagnosis of a ShuffleBlockChunkId reaches the shuffle
  service which merged the chunk, using the external shuffle service port for a push-merged-local
  chunk, and a corrupt merged chunk is diagnosed before it falls back to the original blocks.

Was this patch authored or co-authored using generative AI tooling?

Generated-by: Claude Code 2.1.220

[uros-b]

Just a drive-by note: branch-4.3 has already been cut, so double-check whether this is intended to be backported there as well

[Dreamstick9]

hmm nice catch.the version numbers in this PR were out of date. dev/next_version_candidates.py now shows 4.4.0 for branch-4.x, and branch-4.3 is cut. I changed the two configuration rows and the @since tag to 4.4.0.This change is a new feature. Therefore, it goes to master and branch-4.x only. I do not backport it to branch-4.3.