[CORE][SPARK-49788] Add optional TTL based cleaning for blocks to better support notebooks

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

@@ -0,0 +1,96 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.spark
+
+import java.util.concurrent.ConcurrentHashMap
+
+import scala.jdk.CollectionConverters._
+import scala.util.control.NonFatal
+
+import org.apache.spark.internal.Logging
+
+/**
+ * Shared periodic TTL sweep for a block/shuffle access-time map. Finds ids whose recorded access
+ * time is older than `ttlMillis`, reaps them, and sleeps until the next possible expiry. Used by
+ * both the RDD-cache TTL cleaner (in `BlockManagerMasterEndpoint`) and the shuffle TTL cleaner (in
+ * `MapOutputTrackerMaster`); they differ only in the map, the `shouldReap` gate, and the `reap`
+ * action, so the loop lives here once.
+ *
+ * The runnable loops until interrupted (its owner interrupts it via `shutdownNow` on stop).
+ *
+ * @param name        label used in log messages (e.g. "RDD" / "shuffle")
+ * @param ttlMillis   the TTL; only constructed when the corresponding config is set
+ * @param accessTimes id -> last-access-time (millis). Must be a `ConcurrentHashMap`: this sweep
+ *                    iterates it (weakly consistent) while other threads `put` to it, and plain
+ *                    HashMap structural mutation from multiple threads can corrupt the map.
+ * @param shouldReap  gate checked before removal; return false to leave an id tracked this pass
+ *                    (e.g. a shuffle that has not produced output yet has nothing to reclaim, and
+ *                    removing its state would break a later registration).
+ * @param reap        performs the actual removal for an id whose atime was still stale.
+ */
+private[spark] class BlockTtlCleaner(
+    name: String,
+    ttlMillis: Long,
+    accessTimes: ConcurrentHashMap[Int, Long],
+    shouldReap: Int => Boolean,
+    reap: Int => Unit) extends Runnable with Logging {
+
+  override def run(): Unit = {
+    try {
+      while (!Thread.currentThread().isInterrupted) {
+        val maxAge = System.currentTimeMillis() - ttlMillis
+        // Track the oldest still-live atime so we can sleep until the next possible expiry.
+        var oldest = System.currentTimeMillis()
+        val toBeRemoved = accessTimes.asScala.toList.flatMap { case (id, atime) =>
+          if (atime < maxAge) {
+            Some((id, atime))
+          } else {
+            if (atime < oldest) {
+              oldest = atime
+            }
+            None
+          }
+        }
+        toBeRemoved.foreach { case (id, atime) =>
+          try {
+            // `shouldReap` is checked before the removal so a skipped id stays tracked (its atime
+            // is unchanged). `remove(key, value)` only succeeds if the atime is unchanged since the
+            // snapshot, so a concurrent access in the window leaves the entry: it is back in use.
+            if (shouldReap(id) && accessTimes.remove(id, atime)) {
+              reap(id)
+            }
+          } catch {
+            // Warn, not debug: this loop's whole value is reclaiming space, so a reap that always
+            // fails (e.g. an unwired remover, or an RPC failure) must not be invisible at the
+            // default log level. The id has already been dropped from `accessTimes`, so it is not
+            // retried -- a persistent failure means that id is simply never reclaimed.
+            case NonFatal(e) =>
+              logWarning(s"Error reaping $id in the $name TTL cleaner", e)
+          }
+        }
+        // Wait until the next possible element to be removed.
+        val delay = math.max((oldest + ttlMillis) - System.currentTimeMillis(), 100)
+        Thread.sleep(delay)
+      }
+      logInfo(s"$name TTL cleaner thread interrupted, exiting.")
+    } catch {
+      case _: InterruptedException =>
+        logInfo(s"$name TTL cleaner thread interrupted, exiting.")
+    }
+  }
+}

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

@@ -124,6 +124,18 @@ private[spark] class ContextCleaner(
     listeners.add(listener)
   }
 
+  /**
+   * Notify listeners that a shuffle's data was reclaimed by something other than the
+   * reference-tracking path -- specifically the shuffle TTL cleaner, which removes the shuffle
+   * itself and so must issue this fan-out on its own. Without it `ExecutorMonitor` never sees the
+   * shuffle go away, so under `spark.dynamicAllocation.shuffleTracking.enabled` an executor holding
+   * only that shuffle would never become idle and the reclaimed space would not release the
+   * executor.
+   */
+  private[spark] def notifyShuffleCleaned(shuffleId: Int): Unit = {
+    listeners.asScala.foreach(_.shuffleCleaned(shuffleId))
+  }
+
   /** Start the cleaner. */
   def start(): Unit = {
     cleaningThread.setDaemon(true)

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

@@ -801,6 +801,24 @@ private[spark] class MapOutputTrackerMaster(
     private[spark] val isLocal: Boolean)
   extends MapOutputTracker(conf) with ShuffleOutputTrackerMaster {
 
+  // Keep track of last access times for shuffle based TTL. We don't care about overwriting times
+  // that are "close", but this is written concurrently by the (multi-threaded) map-output
+  // dispatcher, the DAGScheduler, and the TTL cleaner thread, so it must be a concurrent map:
+  // plain HashMap structural mutation from multiple threads can corrupt the map, not merely skew
+  // a timestamp.
+  private[spark] val shuffleAccessTime = new ConcurrentHashMap[Int, Long]
+
+  // Hook used by the shuffle TTL cleaner to reclaim a shuffle's on-disk blocks through the normal
+  // removal path (ShuffleDriverComponents.removeShuffle -> RemoveShuffle RPC to executors/ESS).
+  // Wired by SparkContext, which owns the ShuffleDriverComponents; the tracker itself is created in
+  // SparkEnv before those exist. When unset the cleaner still unregisters the driver-side status
+  // but cannot delete executor disk, so this must be wired for the shuffle TTL to reclaim space.
+  @volatile private[spark] var shuffleFileRemover: Option[Int => Unit] = None
+
+  // Cache the (immutable-after-start) shuffle TTL config once rather than re-parsing the time
+  // string on every updateShuffleAtime, which is on the hot path of serving map-output requests.
+  private val shuffleTtl: Option[Long] = conf.get(SPARK_TTL_SHUFFLE_BLOCK_CLEANER)
+
   // The size at which we use Broadcast to send the map output statuses to the executors
   private val minSizeForBroadcast = conf.get(SHUFFLE_MAPOUTPUT_MIN_SIZE_FOR_BROADCAST).toInt
 
@@ -841,6 +859,16 @@ private[spark] class MapOutputTrackerMaster(
 
   private val pushBasedShuffleEnabled = Utils.isPushBasedShuffleEnabled(conf, isDriver = true)
 
+  // The cleaner daemon is started at the end of construction (after the fail-fast broadcast-size
+  // check below), not here, so a failed construction can't orphan the thread.
+  private[spark] val cleanerThreadpool: Option[ThreadPoolExecutor] = {
+    if (shuffleTtl.isDefined) {
+      Some(ThreadUtils.newDaemonFixedThreadPool(1, "map-output-ttl-cleaner"))
+    } else {
+      None
+    }
+  }
+
   // Thread pool used for handling map output status requests. This is a separate thread pool
   // to ensure we don't block the normal dispatcher threads.
   private val threadpool: ThreadPoolExecutor = {
@@ -854,6 +882,12 @@ private[spark] class MapOutputTrackerMaster(
 
   private val availableProcessors = Runtime.getRuntime.availableProcessors()
 
+  def updateShuffleAtime(shuffleId: Int): Unit = {
+    if (shuffleTtl.isDefined) {
+      shuffleAccessTime.put(shuffleId, System.currentTimeMillis())
+    }
+  }
+
   // Make sure that we aren't going to exceed the max RPC message size by making sure
   // we use broadcast to send large map output statuses.
   if (minSizeForBroadcast > maxRpcMessageSize) {
@@ -865,6 +899,40 @@ private[spark] class MapOutputTrackerMaster(
     throw new IllegalArgumentException(logEntry.message)
   }
 
+  // Start the shuffle TTL cleaner only after the fail-fast check above, so a failed construction
+  // can't leave the daemon running.
+  //
+  // Reap only a shuffle that has actually produced output: a registered-but-not-yet-produced
+  // shuffle (e.g. a map stage still waiting on parents) has no files to reclaim.
+  //
+  // Reaping reclaims the on-disk blocks on executors/ESS (shuffleFileRemover, i.e. the normal
+  // ShuffleDriverComponents.removeShuffle path) and then clears the driver-side outputs via
+  // unregisterAllMapAndMergeOutput. Deliberately NOT unregisterShuffle: keeping the (now empty)
+  // ShuffleStatus registered is what makes this safe against a shuffle that is still referenced.
+  //   - unregisterAllMapAndMergeOutput calls incrementEpoch, so executors drop their cached
+  //     statuses and re-ask rather than fetching files we just deleted. Re-asking yields an empty
+  //     status -> MetadataFetchFailedException, which is a FetchFailed, so the DAGScheduler
+  //     recomputes the map stage. unregisterShuffle bumps no epoch, so executors would instead
+  //     fetch deleted files.
+  //   - The DAGScheduler's own FetchFailed recovery calls unregisterAllMapAndMergeOutput /
+  //     unregisterMapOutput, which go through getShuffleStatusOrError. With the status removed
+  //     those throw ShuffleStatusNotFoundException on the event-loop thread, and
+  //     DAGSchedulerEventProcessLoop.onError responds by cancelling all jobs and stopping the
+  //     SparkContext -- i.e. reaping a still-referenced shuffle would kill the application.
+  // The cost is that an emptied ShuffleStatus stays in shuffleStatuses until the ContextCleaner
+  // collects it; emptying it also makes numAvailableMapOutputs 0, so it is not reaped again.
+  cleanerThreadpool.foreach { pool =>
+    pool.execute(new BlockTtlCleaner(
+      name = "shuffle",
+      ttlMillis = shuffleTtl.get,
+      accessTimes = shuffleAccessTime,
+      shouldReap = shuffleId => shuffleStatuses.get(shuffleId).exists(_.numAvailableMapOutputs > 0),
+      reap = shuffleId => {
+        shuffleFileRemover.foreach(_(shuffleId))
+        unregisterAllMapAndMergeOutput(shuffleId)
+      }))
+  }
+
   def post(message: MapOutputTrackerMasterMessage): Unit = {
     mapOutputTrackerMasterMessages.offer(message)
   }
@@ -879,6 +947,7 @@ private[spark] class MapOutputTrackerMaster(
       val shuffleStatus = shuffleStatuses.get(shuffleId).head
       logDebug(s"Handling request to send ${if (needMergeOutput) "map/merge" else "map"}" +
         s" output locations for shuffle $shuffleId to $hostPort")
+      updateShuffleAtime(shuffleId)
       if (needMergeOutput) {
         context.reply(
           shuffleStatus.
@@ -930,6 +999,7 @@ private[spark] class MapOutputTrackerMaster(
   }
 
   def registerShuffle(shuffleId: Int, numMaps: Int, numReduces: Int): Unit = {
+    updateShuffleAtime(shuffleId)
     if (pushBasedShuffleEnabled) {
       if (shuffleStatuses.put(shuffleId,
         new ShuffleStatus(numMaps, numReduces, bufferRacingMigrations)).isDefined) {
@@ -967,6 +1037,11 @@ private[spark] class MapOutputTrackerMaster(
   }
 
   def registerMapOutput(shuffleId: Int, mapIndex: Int, status: MapStatus): Boolean = {
+    // A map task completing output for this shuffle is an active use of it: refresh the atime so a
+    // map stage that runs longer than the TTL is not reaped mid-production. Reduce-side fetches
+    // refresh via handleStatusMessage; this covers the produce side (registerShuffle only stamps
+    // the atime once, at stage submission).
+    updateShuffleAtime(shuffleId)
     getShuffleStatusOrError(shuffleId, "registerMapOutput").addMapOutput(mapIndex, status)
   }
 
@@ -978,6 +1053,8 @@ private[spark] class MapOutputTrackerMaster(
 
   /** Unregister all map and merge output information of the given shuffle. */
   def unregisterAllMapAndMergeOutput(shuffleId: Int): Unit = {
+    // Drop any TTL tracking (a bare no-op when the map is empty / TTL disabled).
+    shuffleAccessTime.remove(shuffleId)
     val shuffleStatus = getShuffleStatusOrError(shuffleId, "unregisterAllMapAndMergeOutput")
     shuffleStatus.removeOutputsByFilter(x => true)
     shuffleStatus.removeMergeResultsByFilter(x => true)
@@ -1035,6 +1112,10 @@ private[spark] class MapOutputTrackerMaster(
 
   /** Unregister shuffle data */
   override def unregisterShuffle(shuffleId: Int): Unit = {
+    // Drop any TTL tracking so the cleaner doesn't later wake up and try to unregister a shuffle
+    // that has already been GC-cleaned (which would throw ShuffleStatusNotFoundException). This
+    // mirrors removeRdd dropping rddAccessTime. A no-op when TTL tracking is disabled (map empty).
+    shuffleAccessTime.remove(shuffleId)
     shuffleStatuses.remove(shuffleId).foreach { shuffleStatus =>
       shuffleStatus.invalidateSerializedMapOutputStatusCache()
       shuffleStatus.invalidateSerializedMergeOutputStatusCache()
@@ -1379,6 +1460,7 @@ private[spark] class MapOutputTrackerMaster(
 
   // This method is only called in local-mode.
   override def getShufflePushMergerLocations(shuffleId: Int): Seq[BlockManagerId] = {
+    updateShuffleAtime(shuffleId)
     shuffleStatuses.get(shuffleId).map(_.getShufflePushMergerLocations).getOrElse(Seq.empty)
   }
 
@@ -1389,6 +1471,7 @@ private[spark] class MapOutputTrackerMaster(
   override def stop(): Unit = {
     mapOutputTrackerMasterMessages.offer(PoisonPill)
     threadpool.shutdown()
+    cleanerThreadpool.foreach(_.shutdownNow())
     try {
       sendTracker(StopMapOutputTracker)
     } catch {

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

@@ -314,6 +314,19 @@ class SparkContext(config: SparkConf) extends Logging {
     val map: ConcurrentMap[Int, RDD[_]] = new MapMaker().weakValues().makeMap[Int, RDD[_]]()
     map.asScala
   }
+
+  // Ids of RDDs that have been locally checkpointed. A local checkpoint truncates lineage, so its
+  // cache blocks are the only copy of the data and the RDD-cache TTL cleaner must never reap it.
+  // Kept here rather than derived from the live RDD because the cleaner runs on its own thread:
+  // this set gives it a safely-published answer, whereas RDD.checkpointData is not volatile. Not
+  // derived from `persistentRdds` either, since unpersisting drops that entry permanently
+  // (RDD.persist only registers on the first transition out of StorageLevel.NONE).
+  private[spark] val locallyCheckpointedRddIds = ConcurrentHashMap.newKeySet[Int]()
+
+  private[spark] def registerLocallyCheckpointedRdd(rddId: Int): Unit = {
+    locallyCheckpointedRddIds.add(rddId)
+  }
+
   def statusTracker: SparkStatusTracker = _statusTracker
 
   private[spark] def progressBar: Option[ConsoleProgressBar] = _progressBar
@@ -645,6 +658,42 @@ class SparkContext(config: SparkConf) extends Logging {
       _conf.set(ShuffleDataIOUtils.SHUFFLE_SPARK_CONF_PREFIX + k, v)
     }
 
+    // Wire the TTL cleaners, which live in SparkEnv components created before the pieces they need
+    // exist (ShuffleDriverComponents, the ContextCleaner's listeners, and this SparkContext).
+    // Without this wiring a reap still frees blocks, but reclaims no executor disk for shuffles and
+    // leaves the unpersist bookkeeping stale.
+    _env.mapOutputTracker match {
+      case mapOutputTrackerMaster: MapOutputTrackerMaster =>
+        mapOutputTrackerMaster.shuffleFileRemover = Some { shuffleId =>
+          // Reclaim the on-disk blocks on executors/ESS, then tell the CleanerListeners, which is
+          // what lets dynamic allocation's shuffle tracking release an executor that only held this
+          // shuffle. Mirrors the tail of ContextCleaner.doCleanupShuffle.
+          _shuffleDriverComponents.removeShuffle(shuffleId, false)
+          _cleaner.foreach(_.notifyShuffleCleaned(shuffleId))
+        }
+      case _ =>
+    }
+    _env.blockManagerMasterEndpoint.foreach { endpoint =>
+      // Never TTL-reap a locally-checkpointed RDD: localCheckpoint truncates lineage, so its cache
+      // blocks are the only copy of the data and losing them is unrecoverable. Read from the
+      // dedicated id set, not from persistentRdds: the reap below unpersists, which drops the
+      // persistentRdds entry permanently (RDD.persist only registers on the first transition out of
+      // StorageLevel.NONE), so a persistentRdds-based gate would stop protecting the RDD after the
+      // first reap.
+      endpoint.rddReapable = rddId => !locallyCheckpointedRddIds.contains(rddId)
+      // Reap through ContextCleaner.doCleanupRDD, the same entry point the GC-driven cleanup uses,
+      // rather than re-implementing it: it unpersists the RDD's blocks and notifies the
+      // CleanerListeners. Note this frees the blocks but does not reset the RDD's storage level
+      // (the cleaner only has an id), so a later action re-caches it.
+      // Falls back to unpersistRDD when reference tracking is disabled and there is no cleaner.
+      endpoint.rddReaper = Some { rddId =>
+        _cleaner match {
+          case Some(contextCleaner) => contextCleaner.doCleanupRDD(rddId, blocking = false)
+          case None => unpersistRDD(rddId, blocking = false)
+        }
+      }
+    }
+
     if (_conf.get(UI_REVERSE_PROXY)) {
       val proxyUrl = _conf.get(UI_REVERSE_PROXY_URL).getOrElse("").stripSuffix("/")
       System.setProperty("spark.ui.proxyBase", proxyUrl + "/proxy/" + _applicationId)
@@ -2162,6 +2211,7 @@ class SparkContext(config: SparkConf) extends Logging {
   private[spark] def unpersistRDD(rddId: Int, blocking: Boolean): Unit = {
     env.blockManager.master.removeRdd(rddId, blocking)
     persistentRdds.remove(rddId)
+    locallyCheckpointedRddIds.remove(rddId)
     listenerBus.post(SparkListenerUnpersistRDD(rddId))
   }