Disable peer decommissioning when shuffle max disk size is 0

Author: EnricoMi

core/src/main/scala/org/apache/spark/internal/config/package.scala

@@ -627,7 +627,8 @@ package object config {
         "shuffle blocks. Rejecting remote shuffle blocks means that an executor will not receive " +
         "any shuffle migrations, and if there are no other executors available for migration " +
         "then shuffle blocks will be lost unless " +
-        s"${STORAGE_DECOMMISSION_FALLBACK_STORAGE_PATH.key} is configured.")
+        s"${STORAGE_DECOMMISSION_FALLBACK_STORAGE_PATH.key} is configured. " +
+        s"Set to 0 to migrate to fallback storage only.")
       .version("3.2.0")
       .bytesConf(ByteUnit.BYTE)
       .createOptional

core/src/main/scala/org/apache/spark/storage/BlockManagerDecommissioner.scala

@@ -28,6 +28,7 @@ import org.apache.spark._
 import org.apache.spark.errors.SparkCoreErrors
 import org.apache.spark.internal.{config, Logging}
 import org.apache.spark.internal.LogKeys._
+import org.apache.spark.internal.config.STORAGE_DECOMMISSION_SHUFFLE_MAX_DISK_SIZE
 import org.apache.spark.shuffle.ShuffleBlockInfo
 import org.apache.spark.storage.BlockManagerMessages.ReplicateBlock
 import org.apache.spark.util.{ThreadUtils, Utils}
@@ -319,8 +320,15 @@ private[storage] class BlockManagerDecommissioner(
       log"${MDC(TOTAL, localShuffles.size)} local shuffles are added. " +
       log"In total, ${MDC(NUM_REMAINED, remainedShuffles)} shuffles are remained.")
 
+    // migrate to fallback storage only if
+    // STORAGE_DECOMMISSION_FALLBACK_STORAGE_PATH is set and
+    // STORAGE_DECOMMISSION_SHUFFLE_MAX_DISK_SIZE is 0
+    val fallbackOnly = conf.get(config.STORAGE_DECOMMISSION_FALLBACK_STORAGE_PATH).isDefined &&
+      conf.get(STORAGE_DECOMMISSION_SHUFFLE_MAX_DISK_SIZE).contains(0)
+
     // Update the threads doing migrations
-    val livePeerSet = bm.getPeers(false).toSet
+    val livePeerSet = if (fallbackOnly) Set(FallbackStorage.FALLBACK_BLOCK_MANAGER_ID)
+      else bm.getPeers(false).toSet
     val currentPeerSet = migrationPeers.keys.toSet
     val deadPeers = currentPeerSet.diff(livePeerSet)
     // Randomize the orders of the peers to avoid hotspot nodes.

core/src/test/scala/org/apache/spark/storage/FallbackStorageSuite.scala

@@ -175,6 +175,7 @@ class FallbackStorageSuite extends SparkFunSuite with LocalSparkContext {
     val conf = new SparkConf(false)
       .set("spark.app.id", "testId")
       .set(STORAGE_DECOMMISSION_SHUFFLE_BLOCKS_ENABLED, true)
+      .set(STORAGE_DECOMMISSION_SHUFFLE_MAX_DISK_SIZE, 0L) // migrate to fallback storage only
       .set(STORAGE_DECOMMISSION_FALLBACK_STORAGE_PATH,
         Files.createTempDirectory("tmp").toFile.getAbsolutePath + "/")
 
@@ -200,8 +201,6 @@ class FallbackStorageSuite extends SparkFunSuite with LocalSparkContext {
       when(resolver.getDataFile(shuffleId, mapId)).thenReturn(dataFile)
     }
 
-    when(bm.getPeers(mc.any()))
-      .thenReturn(Seq(FallbackStorage.FALLBACK_BLOCK_MANAGER_ID))
     val bmm = new BlockManagerMaster(new NoopRpcEndpointRef(conf), null, conf, false)
     when(bm.master).thenReturn(bmm)
     val blockTransferService = mock(classOf[BlockTransferService])