Code review feedback, remove regex we don't need, reduce some log levels, add a Since annotaiton.

Author: holdenk

core/src/main/scala/org/apache/spark/shuffle/IndexShuffleBlockResolver.scala

@@ -47,6 +47,7 @@ import org.apache.spark.util.Utils
 // org.apache.spark.network.shuffle.ExternalShuffleBlockResolver#getSortBasedShuffleBlockData().
 private[spark] class IndexShuffleBlockResolver(
     conf: SparkConf,
+    // var for testing
     var _blockManager: BlockManager = null)
   extends ShuffleBlockResolver
   with Logging with MigratableResolver {
@@ -61,19 +62,14 @@ private[spark] class IndexShuffleBlockResolver(
   /**
    * Get the shuffle files that are stored locally. Used for block migrations.
    */
-  override def getStoredShuffles(): Set[ShuffleBlockInfo] = {
-    // Matches ShuffleIndexBlockId name
-    val pattern = "shuffle_(\\d+)_(\\d+)_.+\\.index".r
-    val rootDirs = blockManager.diskBlockManager.localDirs
-    // ExecutorDiskUtil puts things inside one level hashed sub directories
-    val searchDirs = rootDirs.flatMap(_.listFiles()).filter(_.isDirectory()) ++ rootDirs
-    val filenames = searchDirs.flatMap(_.list())
-    logDebug(s"Got block files ${filenames.toList}")
-    filenames.flatMap { fname =>
-      pattern.findAllIn(fname).matchData.map {
-        matched => ShuffleBlockInfo(matched.group(1).toInt, matched.group(2).toLong)
-      }
-    }.toSet
+  override def getStoredShuffles(): Seq[ShuffleBlockInfo] = {
+    val allBlocks = blockManager.diskBlockManager.getAllBlocks()
+    allBlocks.flatMap {
+      case ShuffleIndexBlockId(shuffleId, mapId, _) =>
+        Some(ShuffleBlockInfo(shuffleId, mapId))
+      case _ =>
+        None
+    }
   }
 
   /**

core/src/main/scala/org/apache/spark/shuffle/MigratableResolver.scala

@@ -17,7 +17,7 @@
 
 package org.apache.spark.shuffle
 
-import org.apache.spark.annotation.Experimental
+import org.apache.spark.annotation.{Experimental, Since}
 import org.apache.spark.network.buffer.ManagedBuffer
 import org.apache.spark.network.client.StreamCallbackWithID
 import org.apache.spark.serializer.SerializerManager
@@ -28,11 +28,12 @@ import org.apache.spark.storage.BlockId
  * An experimental trait to allow Spark to migrate shuffle blocks.
  */
 @Experimental
+@Since("3.1.0")
 trait MigratableResolver {
   /**
    * Get the shuffle ids that are stored locally. Used for block migrations.
    */
-  def getStoredShuffles(): Set[ShuffleBlockInfo]
+  def getStoredShuffles(): Seq[ShuffleBlockInfo]
 
   /**
    * Write a provided shuffle block as a stream. Used for block migrations.

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

@@ -668,7 +668,7 @@ private[spark] class BlockManager(
     }
 
     if (blockId.isShuffle) {
-      logInfo(s"Putting shuffle block ${blockId}")
+      logDebug(s"Putting shuffle block ${blockId}")
       try {
         return migratableResolver.putShuffleBlockAsStream(blockId, serializerManager)
       } catch {
@@ -677,7 +677,7 @@ private[spark] class BlockManager(
           s"resolver ${shuffleManager.shuffleBlockResolver}")
       }
     }
-    logInfo(s"Putting regular block ${blockId}")
+    logDebug(s"Putting regular block ${blockId}")
     // All other blocks
     val (_, tmpFile) = diskBlockManager.createTempLocalBlock()
     val channel = new CountingWritableChannel(

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

@@ -35,8 +35,8 @@ import org.apache.spark.util.ThreadUtils
  * It creates a Thread to retry offloading all RDD cache and Shuffle blocks
  */
 private[storage] class BlockManagerDecommissioner(
-  conf: SparkConf,
-  bm: BlockManager) extends Logging {
+    conf: SparkConf,
+    bm: BlockManager) extends Logging {
 
   private val maxReplicationFailuresForDecommission =
     conf.get(config.STORAGE_DECOMMISSION_MAX_REPLICATION_FAILURE_PER_BLOCK)
@@ -195,7 +195,7 @@ private[storage] class BlockManagerDecommissioner(
   private[storage] def refreshOffloadingShuffleBlocks(): Unit = {
     // Update the queue of shuffles to be migrated
     logInfo("Offloading shuffle blocks")
-    val localShuffles = bm.migratableResolver.getStoredShuffles()
+    val localShuffles = bm.migratableResolver.getStoredShuffles().toSet
     val newShufflesToMigrate = localShuffles.diff(migratingShuffles).toSeq
     shufflesToMigrate.addAll(newShufflesToMigrate.map(x => (x, 0)).asJava)
     migratingShuffles ++= newShufflesToMigrate
@@ -223,7 +223,7 @@ private[storage] class BlockManagerDecommissioner(
   private[storage] def stopOffloadingShuffleBlocks(): Unit = {
     logInfo("Stopping offloading shuffle blocks.")
     // Stop as gracefully as possible.
-    migrationPeers.values.foreach{_.running = false}
+    migrationPeers.values.foreach{ _.running = false }
     shuffleMigrationPool.shutdown()
     shuffleMigrationPool.shutdownNow()
   }

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

@@ -44,7 +44,7 @@ class BlockManagerDecommissionUnitSuite extends SparkFunSuite with Matchers {
       ids: Set[(Int, Long, Int)]): Unit = {
 
     when(mockMigratableShuffleResolver.getStoredShuffles())
-      .thenReturn(ids.map(triple => ShuffleBlockInfo(triple._1, triple._2)).toSet)
+      .thenReturn(ids.map(triple => ShuffleBlockInfo(triple._1, triple._2)).toSeq)
 
     ids.foreach { case (shuffleId: Int, mapId: Long, reduceId: Int) =>
       when(mockMigratableShuffleResolver.getMigrationBlocks(mc.any()))