[SPARK-57994][CORE] Buffer shuffle-migration relocation that races map output registration

[ChuckLin2025]

What changes were proposed in this pull request?

During executor decommission, a migrated shuffle block's driver-side location can be silently dropped, leaving the map output pointing at the dead origin executor. That entry is then nulled when the origin executor is removed, surfacing downstream as a null MapStatus (NPE on MapStatus.location()) or a fetch failure.

The root cause is a race between two independently-threaded channels that both start when a shuffle map file is committed on the decommissioning executor:

• Registration: task completion -> DAG scheduler event loop -> registerMapOutput, recording the map output at the origin (decommissioning) executor.
• Relocation: the migrated block is uploaded to a peer; the peer reports it; the report is routed off the event loop to MapOutputTracker.updateMapOutput, relocating the map output to the peer.

The two run on different threads with no ordering guarantee. When the relocation beats the registration, the map output is not tracked yet, so ShuffleStatus.updateMapOutput hits the untracked branch, logs a warning, and drops the relocation with no retry. Registration then records the stale origin location; when the decommissioned executor is later removed (removeOutputsOnExecutor), that slot is nulled.

This PR buffers the racing relocation instead of dropping it:

• Adds pendingMigratedOutputs (a mapId -> BlockManagerId map) to ShuffleStatus, guarded by the same write lock as the canonical map-status state.
• When a relocation arrives for an untracked map output, updateMapOutput buffers it in pendingMigratedOutputs (keyed by mapId, the specific task attempt) instead of logging-and-dropping. A later relocation for the same untracked mapId overwrites the entry (latest migration wins).
• addMapOutput replays the buffered relocation immediately after the map output is registered, so the migrated (peer) location wins over the just-registered origin location. The replay runs only when the buffer is non-empty, so the normal registration path is unaffected.
• A buffered entry that never gets a matching registration (e.g. a superseded task attempt) simply never replays and is freed when the shuffle unregisters, so no explicit cleanup is needed.
• The buffering is gated on a new internal config spark.storage.decommission.shuffleBlocks.bufferRacingMigrations (default true); when disabled, the legacy log-and-drop behavior is preserved.

Why are the changes needed?

Without this fix, a shuffle-block migration that is reported before the corresponding map output is registered is lost. The map output keeps the origin (decommissioning) executor location, which is nulled once that executor is removed, so a downstream reducer hits a null MapStatus (NPE) or a fetch failure and the stage fails unnecessarily. Decommission is expected to preserve shuffle data by migrating it, so this defeats the purpose of shuffle migration during decommission.

Does this PR introduce any user-facing change?

No. This is an internal correctness fix. It adds an internal config (spark.storage.decommission.shuffleBlocks.bufferRacingMigrations, default true) that is not part of the public API. When the race does not occur the behavior is unchanged.

This should also be backported to branch-4.x (4.3.0), which carries the identical ShuffleStatus code; hence the 4.3.0 config version.

How was this patch tested?

Two new unit tests in MapOutputTrackerSuite:

• "SPARK-57994: updateMapOutput buffers and replays a relocation that races registration" (buffering enabled): drives the relocation-before-registration interleaving, registers the origin location, removes the origin executor, and asserts the reducer still finds the migrated output at the peer via getMapSizesByExecutorId. Without the fix the relocation is dropped and the assertion fails with a missing map output.
• "SPARK-57994: updateMapOutput drops a racing relocation when buffering is disabled" (buffering disabled): a negative control that verifies the legacy log-and-drop behavior is preserved when the gating config is off.

The full MapOutputTrackerSuite passes (35 tests), and core/scalastyle / core/Test/scalastyle report no issues.

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

Generated-by: Claude Code (Anthropic)

[ChuckLin2025]

@cloud-fan and @Ngone51 could you take a look at this PR ? Thanks.

[cloud-fan]

0 blocking, 1 non-blocking, 0 nits.
Clean, well-scoped concurrency fix; only a minor version-annotation nit.

Suggestions (1)

• config/package.scala:597: .version("5.0.0") should likely be .version("4.3.0") for a backportable fix — see inline

Verification

Traced both race orderings under ShuffleStatus's single write lock. Relocation-then-registration: updateMapOutput buffers into pendingMigratedOutputs (keyed by mapId), and addMapOutput replays it after registration so the peer location wins. Registration-then-relocation: the existing Some(mapStatus) branch updates in place. removeOutputsOnExecutor also takes the write lock, so it cannot interleave between register and replay. The buffer key (mapId = task attempt id) matches on both the report side (ShuffleIndexBlockId.mapId) and the replay side (MapStatus.mapId), consistent with the existing mapIdToMapIndex invariant. When the config is off the buffer is never populated, exactly preserving the legacy log-and-drop.

[cloud-fan]

This is a backportable correctness fix and the identical ShuffleStatus code already exists on branch-4.x (4.3.0-SNAPSHOT), so per Spark's versioning policy the config should carry the next-open-release version rather than master's:

Suggested change

	.version("5.0.0")
	.version("4.3.0")

Non-blocking — only matters if you intend this to backport (which the fix seems to warrant).

[ChuckLin2025]

Thanks! Applied 4.3.0. This is a correctness fix and the commit cherry-picks cleanly onto branch-4.x (4.3.0-SNAPSHOT), so it's good to backport there once this lands on master.

[cloud-fan]

LGTM if ci is green

[cloud-fan]

thanks, merging to master/4.x

[cloud-fan]

Merge Summary:

• merged into master ec7173c
• merged into branch-4.x 8eb7893

Posted by merge_spark_pr.py