Make FallbackStorage cleanup multithreaded

Author: EnricoMi

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

@@ -623,6 +623,13 @@ package object config {
       .booleanConf
       .createWithDefault(false)
 
+  private[spark] val STORAGE_DECOMMISSION_FALLBACK_STORAGE_CLEANUP_THREADS =
+    ConfigBuilder("spark.storage.decommission.fallbackStorage.cleanUp.threads")
+      .doc("Number of threads that clean up fallback storage data.")
+      .version("4.2.0")
+      .intConf
+      .createWithDefault(5)
+
   private[spark] val STORAGE_DECOMMISSION_FALLBACK_STORAGE_CLEANUP_WAIT_ON_SHUTDOWN =
     ConfigBuilder("spark.storage.decommission.fallbackStorage.cleanUp.waitOnShutdown")
       .doc("If true, Spark waits for all fallback storage data to be cleaned up " +

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

@@ -19,10 +19,10 @@ package org.apache.spark.storage
 
 import java.io.DataInputStream
 import java.nio.ByteBuffer
-import java.util.concurrent.{BlockingQueue, LinkedBlockingQueue, TimeUnit}
+import java.util.concurrent.{ThreadPoolExecutor, TimeUnit}
 import java.util.concurrent.atomic.AtomicBoolean
 
-import scala.concurrent.Future
+import scala.concurrent.{ExecutionContext, ExecutionContextExecutor, Future}
 import scala.reflect.ClassTag
 
 import org.apache.hadoop.conf.Configuration
@@ -32,15 +32,15 @@ import org.apache.spark.{SparkConf, SparkException}
 import org.apache.spark.deploy.SparkHadoopUtil
 import org.apache.spark.internal.Logging
 import org.apache.spark.internal.LogKeys._
-import org.apache.spark.internal.config.{STORAGE_DECOMMISSION_FALLBACK_STORAGE_CLEANUP, STORAGE_DECOMMISSION_FALLBACK_STORAGE_CLEANUP_WAIT_ON_SHUTDOWN, STORAGE_DECOMMISSION_FALLBACK_STORAGE_PATH}
+import org.apache.spark.internal.config.{STORAGE_DECOMMISSION_FALLBACK_STORAGE_CLEANUP, STORAGE_DECOMMISSION_FALLBACK_STORAGE_CLEANUP_THREADS, STORAGE_DECOMMISSION_FALLBACK_STORAGE_CLEANUP_WAIT_ON_SHUTDOWN, STORAGE_DECOMMISSION_FALLBACK_STORAGE_PATH}
 import org.apache.spark.network.buffer.{ManagedBuffer, NioManagedBuffer}
 import org.apache.spark.network.util.JavaUtils
 import org.apache.spark.rpc.{RpcAddress, RpcEndpointRef, RpcTimeout}
 import org.apache.spark.shuffle.{IndexShuffleBlockResolver, ShuffleBlockInfo}
 import org.apache.spark.shuffle.IndexShuffleBlockResolver.NOOP_REDUCE_ID
 import org.apache.spark.storage.BlockManagerMessages.RemoveShuffle
-import org.apache.spark.storage.FallbackStorage.{cleanupShufflesThread, stopped}
-import org.apache.spark.util.{ShutdownHookManager, Utils}
+import org.apache.spark.storage.FallbackStorage.{cleanupShufflesThreadPool, stopped}
+import org.apache.spark.util.{ShutdownHookManager, ThreadUtils, Utils}
 
 /**
  * A fallback storage used by storage decommissioners.
@@ -59,10 +59,14 @@ private[storage] class FallbackStorage(conf: SparkConf) extends Logging {
     // indicate cleanup thread to shut down once queue is drained
     stopped.set(true)
 
-    // only wait for cleanup thread to finish when configured so
-    // thread is set daemon so JVM can shut down while it is running
-    if (conf.get(STORAGE_DECOMMISSION_FALLBACK_STORAGE_CLEANUP_WAIT_ON_SHUTDOWN)) {
-      cleanupShufflesThread.join()
+    // only wait for cleanups to finish when configured so
+    if (cleanupShufflesThreadPool.isDefined) {
+      if (conf.get(STORAGE_DECOMMISSION_FALLBACK_STORAGE_CLEANUP_WAIT_ON_SHUTDOWN)) {
+        cleanupShufflesThreadPool.get.shutdown()
+        while (!cleanupShufflesThreadPool.get.awaitTermination(1L, TimeUnit.SECONDS)) { }
+      } else {
+        cleanupShufflesThreadPool.get.shutdownNow()
+      }
     }
   }
 
@@ -139,25 +143,23 @@ private[spark] object FallbackStorage extends Logging {
 
   private val stopped = new AtomicBoolean(false)
 
-  // a queue of shuffle cleanup requests and a daemon thread processing them
-  private val cleanupShufflesQueue: BlockingQueue[CleanUp] = new LinkedBlockingQueue[CleanUp]()
-  private val cleanupShufflesThread: Thread = new Thread(new Runnable {
-    override def run(): Unit = {
-      // if stopped and queue is empty, this thread terminates
-      while (!stopped.get() || !cleanupShufflesQueue.isEmpty) {
-        Utils.tryLogNonFatalError {
-          // wait a second for another cleanup request, then check while condition
-          val cleanup = Option(cleanupShufflesQueue.poll(1L, TimeUnit.SECONDS))
-          cleanup.foreach { c => cleanUp(c.conf, c.hadoopConf, c.shuffleId) }
-        }
+  // a daemon thread pool for shuffle cleanups
+  private var cleanupShufflesThreadPool: Option[ThreadPoolExecutor] = None
+  private var cleanupShufflesExecutionContext: Option[ExecutionContextExecutor] = None
+
+  private def getCleanupShufflesExecutionContext(conf: SparkConf): ExecutionContextExecutor = {
+    FallbackStorage.synchronized {
+      if (cleanupShufflesExecutionContext.isEmpty) {
+        val threads = conf.get(STORAGE_DECOMMISSION_FALLBACK_STORAGE_CLEANUP_THREADS)
+        cleanupShufflesThreadPool =
+          Some(ThreadUtils.newDaemonFixedThreadPool(threads, "fallback-storage-cleanup"))
+        cleanupShufflesExecutionContext =
+          cleanupShufflesThreadPool.map(ExecutionContext.fromExecutor)
       }
+
+      cleanupShufflesExecutionContext.get
     }
-  }, "fallback-storage-cleanup")
-  // this is a daemon thread so JVM shutdown is not blocked by this thread running
-  // we block ShutdownHook above if STORAGE_DECOMMISSION_FALLBACK_STORAGE_CLEANUP_WAIT_ON_SHUTDOWN
-  // is true
-  cleanupShufflesThread.setDaemon(true)
-  cleanupShufflesThread.start()
+  }
 
   def getFallbackStorage(conf: SparkConf): Option[FallbackStorage] = {
     if (conf.get(STORAGE_DECOMMISSION_FALLBACK_STORAGE_PATH).isDefined) {
@@ -187,7 +189,7 @@ private[spark] object FallbackStorage extends Logging {
     if (stopped.get()) {
       logInfo("Not queueing cleanup due to shutdown")
     } else {
-      cleanupShufflesQueue.put(CleanUp(conf, hadoopConf, shuffleId))
+      Future { cleanUp(conf, hadoopConf, shuffleId) }(getCleanupShufflesExecutionContext(conf))
     }
   }