[CARMEL-7330][CARMEL-4598] Refine Context Cleaner to improve the clean performance (apache#93)

Author: wangyum

core/src/main/scala/org/apache/spark/ContextCleaner.scala

@@ -1780,6 +1780,12 @@ package object config {
       .booleanConf
       .createWithDefault(false)
 
+  private[spark] val CLEANER_REFERENCE_TRACKING_BLOCKING_BROADCAST =
+    ConfigBuilder("spark.cleaner.referenceTracking.blocking.broadcast")
+      .version("3.5.0")
+      .booleanConf
+      .createWithDefault(false)
+
   private[spark] val CLEANER_REFERENCE_TRACKING_CLEAN_CHECKPOINTS =
     ConfigBuilder("spark.cleaner.referenceTracking.cleanCheckpoints")
       .version("1.4.0")

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

@@ -1847,7 +1847,9 @@ abstract class RDD[T: ClassTag](
         dep match {
           case dependency: ShuffleDependency[_, _, _] =>
             val shuffleId = dependency.shuffleId
-            cleaner.doCleanupShuffle(shuffleId, blocking)
+            cleaner.getContextCleanupWorker(classOf[CleanShuffle].getName).
+              asInstanceOf[ShuffleCleanupWorker].
+              doCleanupShuffle(shuffleId, blocking)
           case _ => // do nothing
         }
         val rdd = dep.rdd

core/src/main/scala/org/apache/spark/rdd/RDD.scala

@@ -122,7 +122,9 @@ class ContextCleanerSuite extends ContextCleanerSuiteBase {
     val tester = new CleanerTester(sc, rddIds = Seq(rdd.id))
 
     // Explicit cleanup
-    cleaner.doCleanupRDD(rdd.id, blocking = true)
+    cleaner.getContextCleanupWorker(classOf[CleanRDD].getName).
+      asInstanceOf[RDDCleanupWorker].
+      doCleanupRDD(rdd.id, blocking = true)
     tester.assertCleanup()
 
     // Verify that RDDs can be re-executed after cleaning up
@@ -135,7 +137,8 @@ class ContextCleanerSuite extends ContextCleanerSuiteBase {
     val tester = new CleanerTester(sc, shuffleIds = shuffleDeps.map(_.shuffleId))
 
     // Explicit cleanup
-    shuffleDeps.foreach(s => cleaner.doCleanupShuffle(s.shuffleId, blocking = true))
+    shuffleDeps.foreach(s => cleaner.getContextCleanupWorker(classOf[CleanShuffle].getName).
+      asInstanceOf[ShuffleCleanupWorker].doCleanupShuffle(s.shuffleId, blocking = true))
     tester.assertCleanup()
 
     // Verify that shuffles can be re-executed after cleaning up
@@ -147,7 +150,8 @@ class ContextCleanerSuite extends ContextCleanerSuiteBase {
     val tester = new CleanerTester(sc, broadcastIds = Seq(broadcast.id))
 
     // Explicit cleanup
-    cleaner.doCleanupBroadcast(broadcast.id, blocking = true)
+    cleaner.getContextCleanupWorker(classOf[CleanBroadcast].getName).
+      asInstanceOf[BroadcastCleanupWorker].doCleanupBroadcast(broadcast.id, blocking = true)
     tester.assertCleanup()
   }
 

core/src/test/scala/org/apache/spark/ContextCleanerSuite.scala

@@ -245,12 +245,17 @@ class ExternalShuffleServiceSuite extends ShuffleSuite with BeforeAndAfterAll wi
           assert(sc.env.blockManager.master.getExecutorEndpointRef("0").isEmpty)
         }
 
-        sc.cleaner.foreach(_.doCleanupShuffle(0, true))
+        sc.cleaner.foreach(
+          _.getContextCleanupWorker(classOf[CleanShuffle].getName).doCleanup(CleanShuffle(0)))
 
         if (enabled) {
-          assert(filesToCheck.forall(!_.exists()))
+          eventually(timeout(2.seconds), interval(100.milliseconds)) {
+            assert(filesToCheck.forall(!_.exists()))
+          }
         } else {
-          assert(filesToCheck.forall(_.exists()))
+          eventually(timeout(2.seconds), interval(100.milliseconds)) {
+            assert(filesToCheck.forall(_.exists()))
+          }
         }
       } finally {
         rpcHandler.applicationRemoved(sc.conf.getAppId, true)

core/src/test/scala/org/apache/spark/ExternalShuffleServiceSuite.scala

@@ -182,7 +182,10 @@ class InternalAccumulatorSuite extends SparkFunSuite with LocalSparkContext {
   test("internal accumulators are registered for cleanups") {
     sc = new SparkContext("local", "test") {
       private val myCleaner = new SaveAccumContextCleaner(this)
-      override def cleaner: Option[ContextCleaner] = Some(myCleaner)
+      override def cleaner: Option[ContextCleaner] = {
+        myCleaner.start()
+        Some(myCleaner)
+      }
     }
     assert(AccumulatorContext.numAccums == 0)
     sc.parallelize(1 to 100).map { i => (i, i) }.reduceByKey { _ + _ }.count()

core/src/test/scala/org/apache/spark/InternalAccumulatorSuite.scala

@@ -27,6 +27,9 @@ import org.mockito.ArgumentMatchers.any
 import org.mockito.Mockito._
 import org.mockito.invocation.InvocationOnMock
 import org.roaringbitmap.RoaringBitmap
+import org.scalatest.concurrent.Eventually.{eventually, interval}
+import org.scalatest.concurrent.Eventually
+import org.scalatest.time.SpanSugar._
 
 import org.apache.spark.LocalSparkContext._
 import org.apache.spark.broadcast.BroadcastManager
@@ -966,8 +969,12 @@ class MapOutputTrackerSuite extends SparkFunSuite with LocalSparkContext {
           true
         })
 
-      sc.cleaner.get.doCleanupShuffle(SHUFFLE_ID, blocking = true)
-      assert(foundHosts.asScala == mergerLocs.map(_.host).toSet)
+      sc.cleaner.foreach(
+        _.getContextCleanupWorker(classOf[CleanShuffle].getName)
+          .doCleanup(CleanShuffle(SHUFFLE_ID)))
+      eventually(Eventually.timeout(2.seconds), interval(100.milliseconds)) {
+        assert(foundHosts.asScala == mergerLocs.map(_.host).toSet)
+      }
     }
   }
 

core/src/test/scala/org/apache/spark/MapOutputTrackerSuite.scala

@@ -118,6 +118,14 @@ private[v1] class SqlResource extends BaseAppResource {
     }
   }
 
+  @GET
+  @Path("context-cleaner")
+  def contextCleaner(): Seq[(String, Int)] = {
+    withUI { ui =>
+      ui.sc.get.cleaner.get.contextCleanupWorkerSummary()
+    }
+  }
+
   private def printableMetrics(allNodes: collection.Seq[SparkPlanGraphNode],
     metricValues: Map[Long, String]): collection.Seq[Node] = {