You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@spark.apache.org by we...@apache.org on 2026/07/10 17:45:08 UTC

(spark) branch branch-4.x updated: [SPARK-57994][CORE] Buffer shuffle-migration relocation that races map output registration

This is an automated email from the ASF dual-hosted git repository.

cloud-fan pushed a commit to branch branch-4.x
in repository https://gitbox.apache.org/repos/asf/spark.git


The following commit(s) were added to refs/heads/branch-4.x by this push:
     new 8eb789303c2f [SPARK-57994][CORE] Buffer shuffle-migration relocation that races map output registration
8eb789303c2f is described below

commit 8eb789303c2fc21d4d06cb0429b1cf1a26660300
Author: ChuckLin2025 <lz...@gmail.com>
AuthorDate: Sat Jul 11 01:44:02 2026 +0800

    [SPARK-57994][CORE] Buffer shuffle-migration relocation that races map output registration
    
    ### 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)
    
    Closes #57075 from ChuckLin2025/buffer-racing-shuffle-migration.
    
    Authored-by: ChuckLin2025 <lz...@gmail.com>
    Signed-off-by: Wenchen Fan <we...@databricks.com>
    (cherry picked from commit ec7173c58d25e66f6f5db87d494dc99ce8336e86)
    Signed-off-by: Wenchen Fan <we...@databricks.com>
---
 .../scala/org/apache/spark/MapOutputTracker.scala  | 69 +++++++++++++++++--
 .../org/apache/spark/internal/config/package.scala | 13 ++++
 .../org/apache/spark/MapOutputTrackerSuite.scala   | 78 ++++++++++++++++++++++
 3 files changed, 156 insertions(+), 4 deletions(-)

diff --git a/core/src/main/scala/org/apache/spark/MapOutputTracker.scala b/core/src/main/scala/org/apache/spark/MapOutputTracker.scala
index 51ba146f728f..742baf40ecbf 100644
--- a/core/src/main/scala/org/apache/spark/MapOutputTracker.scala
+++ b/core/src/main/scala/org/apache/spark/MapOutputTracker.scala
@@ -54,10 +54,16 @@ import org.apache.spark.util.io.{ChunkedByteBuffer, ChunkedByteBufferOutputStrea
  * serialized map statuses in order to speed up tasks' requests for map output statuses.
  *
  * All public methods of this class are thread-safe.
+ *
+ * @param bufferRacingMigrations if true, a migration relocation ([[updateMapOutput]]) that
+ *   arrives before this map id's output is registered is buffered and replayed on registration
+ *   instead of being dropped. See [[pendingMigratedOutputs]]. Read once at construction time so
+ *   the behavior is stable for the lifetime of this shuffle.
  */
 private class ShuffleStatus(
     numPartitions: Int,
-    numReducers: Int = -1) extends Logging {
+    numReducers: Int = -1,
+    bufferRacingMigrations: Boolean = false) extends Logging {
 
   private val (readLock, writeLock) = {
     val lock = new ReentrantReadWriteLock()
@@ -189,6 +195,25 @@ private class ShuffleStatus(
    */
   private[spark] val mapIdToMapIndex = new HashMap[Long, Int]()
 
+  /**
+   * Migration relocation reports (from executor decommission) that arrived before the map output
+   * was registered on the driver, buffered here to be applied once registration happens.
+   *
+   * The relocation report and the task-completion registration for the same map file race on
+   * separate threads: a migrated block is uploaded to a peer and reported back, routing to
+   * [[updateMapOutput]], while the task-completion event that registers the map output travels
+   * through the DAG scheduler independently. When the report wins the race the map output is not
+   * tracked yet, so there is no map status to relocate. Rather than drop the relocation, it is
+   * recorded here (keyed by mapId, the specific task attempt) and replayed by [[addMapOutput]]
+   * right after registration, so the migrated (peer) location wins over the just-registered
+   * origin location, which would otherwise be nulled when the decommissioned executor is removed.
+   *
+   * A buffered entry that never gets a matching registration (e.g. a superseded task attempt)
+   * simply never replays and is freed when the shuffle is unregistered. Guarded by the write
+   * lock, same as the canonical map-status state.
+   */
+  private[this] val pendingMigratedOutputs = new HashMap[Long, BlockManagerId]
+
   /**
    * Register a map output. If there is already a registered location for the map output then it
    * will be replaced by the new location. Returns true if the checksum in the new MapStatus is
@@ -215,6 +240,19 @@ private class ShuffleStatus(
     }
     mapStatuses(mapIndex) = status
     mapIdToMapIndex(status.mapId) = mapIndex
+
+    // If a migration relocation for this map output arrived before this registration (see
+    // updateMapOutput), apply it now so the migrated (peer) location wins over the origin
+    // location that was just registered and would otherwise be nulled when the decommissioned
+    // executor is removed.
+    if (pendingMigratedOutputs.nonEmpty) {
+      pendingMigratedOutputs.remove(status.mapId).foreach { bmAddress =>
+        status.updateLocation(bmAddress)
+        invalidateSerializedMapOutputStatusCache()
+        logInfo(log"Applied buffered migrated location ${MDC(BLOCK_MANAGER_ID, bmAddress)} " +
+          log"for ${MDC(MAP_ID, status.mapId)} on registration.")
+      }
+    }
     isChecksumMismatch
   }
 
@@ -229,7 +267,13 @@ private class ShuffleStatus(
   }
 
   /**
-   * Update the map output location (e.g. during migration).
+   * Update the map output location following a shuffle-data migration (e.g. during executor
+   * decommission).
+   *
+   * If the relocation report arrives before this map id's output has been registered by the
+   * task-completion event in the DAG scheduler, the relocation is buffered in
+   * [[pendingMigratedOutputs]] keyed by mapId and replayed by [[addMapOutput]] once registration
+   * happens, so the migrated location is not lost. See [[pendingMigratedOutputs]].
    */
   def updateMapOutput(mapId: Long, bmAddress: BlockManagerId): Unit = withWriteLock {
     try {
@@ -252,6 +296,15 @@ private class ShuffleStatus(
             mapStatusesDeleted(index) = null
             logInfo(log"Recover ${MDC(MAP_ID, mapStatus.mapId)}" +
               log" ${MDC(BLOCK_MANAGER_ID, mapStatus.location)}")
+          } else if (bufferRacingMigrations) {
+            // This map id's output is not committed/registered yet (e.g. the task-completion
+            // event has not landed on the driver). Buffer the relocation so it is applied once
+            // the map output is registered, instead of dropping it (which would leave the map
+            // pointing at the dead origin executor). A later relocation for the same untracked
+            // mapId overwrites the pending entry (latest migration wins).
+            logInfo(log"Map output ${MDC(MAP_ID, mapId)} is not registered yet; buffering the " +
+              log"migrated location ${MDC(BLOCK_MANAGER_ID, bmAddress)} to apply on registration.")
+            pendingMigratedOutputs(mapId) = bmAddress
           } else {
             logWarning(log"Asked to update map output ${MDC(MAP_ID, mapId)} " +
               log"for untracked map status.")
@@ -757,6 +810,12 @@ private[spark] class MapOutputTrackerMaster(
   private val shuffleMigrationEnabled = conf.get(DECOMMISSION_ENABLED) &&
     conf.get(STORAGE_DECOMMISSION_ENABLED) && conf.get(STORAGE_DECOMMISSION_SHUFFLE_BLOCKS_ENABLED)
 
+  // Whether to buffer a migration relocation that races map output registration and replay it on
+  // registration, instead of dropping it (see ShuffleStatus.pendingMigratedOutputs). Read once
+  // and passed to each ShuffleStatus at construction time.
+  private val bufferRacingMigrations =
+    conf.get(STORAGE_DECOMMISSION_SHUFFLE_BUFFER_RACING_MIGRATIONS)
+
   // Number of map and reduce tasks above which we do not assign preferred locations based on map
   // output sizes. We limit the size of jobs for which assign preferred locations as computing the
   // top locations by size becomes expensive.
@@ -872,11 +931,13 @@ private[spark] class MapOutputTrackerMaster(
 
   def registerShuffle(shuffleId: Int, numMaps: Int, numReduces: Int): Unit = {
     if (pushBasedShuffleEnabled) {
-      if (shuffleStatuses.put(shuffleId, new ShuffleStatus(numMaps, numReduces)).isDefined) {
+      if (shuffleStatuses.put(shuffleId,
+        new ShuffleStatus(numMaps, numReduces, bufferRacingMigrations)).isDefined) {
         throw new IllegalArgumentException("Shuffle ID " + shuffleId + " registered twice")
       }
     } else {
-      if (shuffleStatuses.put(shuffleId, new ShuffleStatus(numMaps)).isDefined) {
+      if (shuffleStatuses.put(shuffleId,
+        new ShuffleStatus(numMaps, bufferRacingMigrations = bufferRacingMigrations)).isDefined) {
         throw new IllegalArgumentException("Shuffle ID " + shuffleId + " registered twice")
       }
     }
diff --git a/core/src/main/scala/org/apache/spark/internal/config/package.scala b/core/src/main/scala/org/apache/spark/internal/config/package.scala
index c9f7654df70f..d8fc2a5c1fec 100644
--- a/core/src/main/scala/org/apache/spark/internal/config/package.scala
+++ b/core/src/main/scala/org/apache/spark/internal/config/package.scala
@@ -586,6 +586,19 @@ package object config {
       .checkValue(_ > 0, "The maximum number of threads should be positive")
       .createWithDefault(8)
 
+  private[spark] val STORAGE_DECOMMISSION_SHUFFLE_BUFFER_RACING_MIGRATIONS =
+    ConfigBuilder("spark.storage.decommission.shuffleBlocks.bufferRacingMigrations")
+      .internal()
+      .doc("Whether to buffer a shuffle-migration relocation report that arrives before the " +
+        "map output it relocates has been registered on the driver, and replay it once " +
+        "registration happens. When false, such a relocation is dropped (legacy behavior), " +
+        "which can leave the map output pointing at the decommissioned origin executor and " +
+        "surface downstream as a fetch failure once that executor is removed.")
+      .version("4.3.0")
+      .withBindingPolicy(ConfigBindingPolicy.NOT_APPLICABLE)
+      .booleanConf
+      .createWithDefault(true)
+
   private[spark] val STORAGE_DECOMMISSION_RDD_BLOCKS_ENABLED =
     ConfigBuilder("spark.storage.decommission.rddBlocks.enabled")
       .doc("Whether to transfer RDD blocks during block manager decommissioning.")
diff --git a/core/src/test/scala/org/apache/spark/MapOutputTrackerSuite.scala b/core/src/test/scala/org/apache/spark/MapOutputTrackerSuite.scala
index 2015ceb8b090..3e42c8a159c2 100644
--- a/core/src/test/scala/org/apache/spark/MapOutputTrackerSuite.scala
+++ b/core/src/test/scala/org/apache/spark/MapOutputTrackerSuite.scala
@@ -1167,6 +1167,84 @@ class MapOutputTrackerSuite extends SparkFunSuite with LocalSparkContext {
     }
   }
 
+  test("SPARK-57994: updateMapOutput buffers and replays a relocation that races registration") {
+    val newConf = new SparkConf
+    newConf.set(STORAGE_DECOMMISSION_SHUFFLE_BUFFER_RACING_MIGRATIONS, true)
+    val hostname = "localhost"
+    val rpcEnv = createRpcEnv("spark", hostname, 0, new SecurityManager(newConf))
+    val masterTracker = newTrackerMaster(newConf)
+    masterTracker.trackerEndpoint = rpcEnv.setupEndpoint(MapOutputTracker.ENDPOINT_NAME,
+      new MapOutputTrackerMasterEndpoint(rpcEnv, masterTracker, newConf))
+    val workerRpcEnv = createRpcEnv("spark-worker", hostname, 0, new SecurityManager(newConf))
+    val workerTracker = new MapOutputTrackerWorker(newConf)
+    workerTracker.trackerEndpoint =
+      workerRpcEnv.setupEndpointRef(rpcEnv.address, MapOutputTracker.ENDPOINT_NAME)
+
+    try {
+      masterTracker.registerShuffle(10, 1, MergeStatus.SHUFFLE_PUSH_DUMMY_NUM_REDUCES)
+      val origin = BlockManagerId("exec-1", "hostA", 1000)
+      val peer = BlockManagerId("exec-2", "hostB", 1000)
+
+      // The migrated block is reported by the peer before the task-completion event registers the
+      // map output on the driver. The map is untracked, so the relocation is buffered rather than
+      // dropped.
+      masterTracker.updateMapOutput(10, 100, peer)
+      assert(masterTracker.getNumAvailableOutputs(10) == 0,
+        "Relocation before registration must not create an output on its own")
+
+      // Task completion registers the map output at the origin (decommissioning) executor. The
+      // buffered relocation is replayed, so the location becomes the peer.
+      masterTracker.registerMapOutput(10, 0, MapStatus(origin, Array(2L), 100))
+      assert(masterTracker.getNumAvailableOutputs(10) == 1)
+
+      // The origin executor is decommissioned and removed. Without the buffer-and-replay fix the
+      // output would still point at the origin and be nulled here, producing a null MapStatus /
+      // fetch failure downstream. With the fix it points at the peer, so removal is a no-op.
+      masterTracker.removeOutputsOnExecutor("exec-1")
+      assert(masterTracker.getNumAvailableOutputs(10) == 1,
+        "The replayed peer location should survive removal of the origin executor")
+
+      masterTracker.incrementEpoch()
+      workerTracker.updateEpoch(masterTracker.getEpoch)
+      val result = workerTracker.getMapSizesByExecutorId(10, 0).toSeq
+      assert(result.nonEmpty, "Reducer should find the migrated map output at the peer")
+      assert(result.head._1 == peer)
+    } finally {
+      masterTracker.stop()
+      workerTracker.stop()
+      rpcEnv.shutdown()
+      workerRpcEnv.shutdown()
+    }
+  }
+
+  test("SPARK-57994: updateMapOutput drops a racing relocation when buffering is disabled") {
+    val newConf = new SparkConf
+    newConf.set(STORAGE_DECOMMISSION_SHUFFLE_BUFFER_RACING_MIGRATIONS, false)
+    val rpcEnv = createRpcEnv("spark", "localhost", 0, new SecurityManager(newConf))
+    val masterTracker = newTrackerMaster(newConf)
+    masterTracker.trackerEndpoint = rpcEnv.setupEndpoint(MapOutputTracker.ENDPOINT_NAME,
+      new MapOutputTrackerMasterEndpoint(rpcEnv, masterTracker, newConf))
+
+    try {
+      masterTracker.registerShuffle(10, 1, MergeStatus.SHUFFLE_PUSH_DUMMY_NUM_REDUCES)
+      val origin = BlockManagerId("exec-1", "hostA", 1000)
+      val peer = BlockManagerId("exec-2", "hostB", 1000)
+
+      // With buffering off, a relocation for an untracked map output is dropped (legacy behavior),
+      // so the later registration keeps the origin location unchanged.
+      masterTracker.updateMapOutput(10, 100, peer)
+      masterTracker.registerMapOutput(10, 0, MapStatus(origin, Array(2L), 100))
+
+      masterTracker.incrementEpoch()
+      val result = masterTracker.getMapSizesByExecutorId(10, 0).toSeq
+      assert(result.head._1 == origin,
+        "With buffering disabled the dropped relocation leaves the origin location in place")
+    } finally {
+      masterTracker.stop()
+      rpcEnv.shutdown()
+    }
+  }
+
   test("SPARK-57491: staleMapIndexes propagated from driver to worker via getStatuses") {
     val rpcEnv = createRpcEnv("test")
     val masterTracker = newTrackerMaster()


---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@spark.apache.org
For additional commands, e-mail: commits-help@spark.apache.org