[SPARK-59776][CORE][SQL] Notify ExecutorMonitor when SQL shuffle cleanup removes shuffle files

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

@@ -306,6 +306,21 @@ private[spark] class ContextCleaner(
     mgr != null && !mgr.usesStreamingShuffleOutputTracker && mgr.holdsShuffle(shuffleId)
   }
 
+  /**
+   * Notify listeners that the files of `shuffleId` were removed outside of this cleaner (e.g. by
+   * SQL shuffle dependency cleanup) while the shuffle stays registered on the MapOutputTracker.
+   * Unlike `doCleanupShuffle`, this fires `shuffleFilesRemoved` rather than `shuffleCleaned`, since
+   * the shuffle's map outputs are not unregistered.
+   */
+  def notifyShuffleFilesRemoved(shuffleId: Int): Unit = {
+    try {
+      listeners.asScala.foreach(_.shuffleFilesRemoved(shuffleId))
+    } catch {
+      case e: Exception =>
+        logError(log"Error notifying removal of shuffle ${MDC(SHUFFLE_ID, shuffleId)}", e)
+    }
+  }
+
   /** Perform broadcast cleanup. */
   def doCleanupBroadcast(broadcastId: Long, blocking: Boolean): Unit = {
     try {
@@ -376,6 +391,11 @@ private object ContextCleaner {
 private[spark] trait CleanerListener {
   def rddCleaned(rddId: Int): Unit
   def shuffleCleaned(shuffleId: Int): Unit
+  /**
+   * Called when a shuffle's files were removed but the shuffle is still registered on the
+   * MapOutputTracker (see `ContextCleaner.notifyShuffleFilesRemoved`).
+   */
+  def shuffleFilesRemoved(shuffleId: Int): Unit = {}
   def broadcastCleaned(broadcastId: Long): Unit
   def accumCleaned(accId: Long): Unit
   def checkpointCleaned(rddId: Long): Unit

core/src/main/scala/org/apache/spark/scheduler/dynalloc/ExecutorMonitor.scala

@@ -460,6 +460,9 @@ private[spark] class ExecutorMonitor(
     }
   }
 
+  // Executors no longer hold the shuffle's files, so they need not be kept around for it.
+  override def shuffleFilesRemoved(shuffleId: Int): Unit = shuffleCleaned(shuffleId)
+
   override def broadcastCleaned(broadcastId: Long): Unit = { }
 
   override def accumCleaned(accId: Long): Unit = { }

core/src/test/scala/org/apache/spark/scheduler/dynalloc/ExecutorMonitorSuite.scala

@@ -385,6 +385,29 @@ class ExecutorMonitorSuite extends SparkFunSuite {
   }
 
 
+  test("SPARK-59776: executor with removed shuffle files times out at the idle deadline") {
+    val bus = mockListenerBus()
+    conf.set(DYN_ALLOCATION_SHUFFLE_TRACKING_ENABLED, true).set(SHUFFLE_SERVICE_ENABLED, false)
+    monitor = new ExecutorMonitor(conf, client, bus, clock, allocationManagerSource())
+
+    val stage1 = stageInfo(1, shuffleId = 0)
+    val stage2 = stageInfo(2)
+    monitor.onExecutorAdded(SparkListenerExecutorAdded(clock.getTimeMillis(), "1", execInfo))
+    monitor.onJobStart(SparkListenerJobStart(1, clock.getTimeMillis(), Seq(stage1, stage2)))
+    monitor.onTaskStart(SparkListenerTaskStart(1, 0, taskInfo("1", 1)))
+    monitor.onTaskEnd(SparkListenerTaskEnd(1, 0, "foo", Success, taskInfo("1", 1),
+      new ExecutorMetrics, null))
+    monitor.onJobEnd(SparkListenerJobEnd(1, clock.getTimeMillis(), JobSucceeded))
+
+    // The executor holds shuffle 0's files, so only the shuffle timeout applies.
+    assert(monitor.timedOutExecutors(idleDeadline).isEmpty)
+    assert(monitor.timedOutExecutors(shuffleDeadline) === Seq("1"))
+
+    // Files removed while the shuffle stays registered (SQL shuffle cleanup).
+    monitor.shuffleFilesRemoved(0)
+    assert(monitor.timedOutExecutors(idleDeadline) === Seq("1"))
+  }
+
   test("SPARK-28839: Avoids NPE in context cleaner when shuffle service is on") {
     val bus = mockListenerBus()
     conf.set(DYN_ALLOCATION_SHUFFLE_TRACKING_ENABLED, true).set(SHUFFLE_SERVICE_ENABLED, true)

sql/core/src/main/scala/org/apache/spark/sql/execution/SQLExecution.scala

@@ -152,6 +152,9 @@ object SQLExecution extends Logging {
               // MapOutputTracker so that stage retries would be triggered. Blocking is
               // Utils.isTesting to deflake unit tests.
               sc.shuffleDriverComponents.removeShuffle(shuffleId, Utils.isTesting)
+              // Let listeners such as the dynamic allocation ExecutorMonitor stop tracking the
+              // shuffle, as ContextCleaner will not report it cleaned while it stays registered.
+              sc.cleaner.foreach(_.notifyShuffleFilesRemoved(shuffleId))
             } catch {
               case NonFatal(e) =>
                 logWarning(log"Failed to remove shuffle ${MDC(SHUFFLE_ID, shuffleId)} for " +

sql/core/src/test/scala/org/apache/spark/sql/execution/SQLExecutionSuite.scala

@@ -18,14 +18,15 @@
 package org.apache.spark.sql.execution
 
 import java.util.Locale
-import java.util.concurrent.Executors
+import java.util.concurrent.{ConcurrentLinkedQueue, Executors}
 import java.util.concurrent.atomic.{AtomicInteger, AtomicReference}
 
 import scala.collection.parallel.immutable.ParRange
 import scala.concurrent.{ExecutionContext, Future}
 import scala.concurrent.duration._
+import scala.jdk.CollectionConverters._
 
-import org.apache.spark.{SparkConf, SparkContext, SparkEnv, SparkFunSuite}
+import org.apache.spark.{CleanerListener, SparkConf, SparkContext, SparkEnv, SparkFunSuite}
 import org.apache.spark.launcher.SparkLauncher
 import org.apache.spark.scheduler.{SparkListener, SparkListenerEvent, SparkListenerJobStart}
 import org.apache.spark.sql.{Observation, Row, SparkSession}
@@ -603,6 +604,42 @@ class SQLExecutionSuite extends SparkFunSuite with SQLConfHelper {
     }
   }
 
+  test("SPARK-59776: RemoveShuffleFiles cleanup notifies cleaner listeners") {
+    val spark = SparkSession.builder().master("local[*]").appName("test").getOrCreate()
+    try {
+      withSQLConf(
+        SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false",
+        SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "-1",
+        SQLConf.CLASSIC_SHUFFLE_DEPENDENCY_FILE_CLEANUP_ENABLED.key -> "true") {
+        val filesRemoved = new ConcurrentLinkedQueue[Int]()
+        val cleaned = new ConcurrentLinkedQueue[Int]()
+        spark.sparkContext.cleaner.get.attachListener(new CleanerListener {
+          override def rddCleaned(rddId: Int): Unit = {}
+          override def shuffleCleaned(shuffleId: Int): Unit = cleaned.add(shuffleId)
+          override def shuffleFilesRemoved(shuffleId: Int): Unit = filesRemoved.add(shuffleId)
+          override def broadcastCleaned(broadcastId: Long): Unit = {}
+          override def accumCleaned(accId: Long): Unit = {}
+          override def checkpointCleaned(rddId: Long): Unit = {}
+        })
+
+        val df = spark.range(0, 10).join(spark.range(0, 20), "id")
+        val qe = df.queryExecution
+        assert(qe.shuffleCleanupMode == RemoveShuffleFiles)
+        val shuffleIds = qe.executedPlan.collect { case e: ShuffleExchangeLike => e.shuffleId }
+        assert(shuffleIds.size > 1)
+        assert(df.collect().length == 10)
+
+        // Every shuffle's files removal is reported, so e.g. the dynamic allocation
+        // ExecutorMonitor stops keeping executors alive for them. The shuffles are not cleaned
+        // through ContextCleaner, so they must not be reported as fully cleaned.
+        assert(filesRemoved.asScala.toSet == shuffleIds.toSet)
+        assert(cleaned.asScala.toSet.intersect(shuffleIds.toSet).isEmpty)
+      }
+    } finally {
+      spark.stop()
+    }
+  }
+
   test("SPARK-59242: SkipMigration cleanup failure has mode-specific warning") {
     val spark = SparkSession.builder().master("local[*]").appName("test").getOrCreate()
     try {