diff --git a/sql/connect/server/src/main/scala/org/apache/spark/sql/connect/config/Connect.scala b/sql/connect/server/src/main/scala/org/apache/spark/sql/connect/config/Connect.scala index 10531e063224d..62f587c0ab01b 100644 --- a/sql/connect/server/src/main/scala/org/apache/spark/sql/connect/config/Connect.scala +++ b/sql/connect/server/src/main/scala/org/apache/spark/sql/connect/config/Connect.scala @@ -156,6 +156,14 @@ object Connect { .timeConf(TimeUnit.MILLISECONDS) .createWithDefaultString("30s") + val CONNECT_SESSION_MANAGER_CLEANUP_CACHED_DATA_ENABLED = + buildStaticConf("spark.connect.session.manager.cleanupCachedData.enabled") + .doc("When true, cached data persisted by an isolated session is removed when the session " + + "is closed. Cached data that is also persisted by another session is preserved.") + .version("4.4.0") + .booleanConf + .createWithDefault(false) + val CONNECT_EXECUTE_MANAGER_DETACHED_TIMEOUT = buildStaticConf("spark.connect.execute.manager.detachedTimeout") .internal() diff --git a/sql/connect/server/src/main/scala/org/apache/spark/sql/connect/service/SessionHolder.scala b/sql/connect/server/src/main/scala/org/apache/spark/sql/connect/service/SessionHolder.scala index e97bcb0c786e3..0c76cafc40f51 100644 --- a/sql/connect/server/src/main/scala/org/apache/spark/sql/connect/service/SessionHolder.scala +++ b/sql/connect/server/src/main/scala/org/apache/spark/sql/connect/service/SessionHolder.scala @@ -434,6 +434,10 @@ case class SessionHolder(userId: String, sessionId: String, session: SparkSessio // Clean up ML cache (only if ML models were created) mlCache.close() + if (SparkEnv.get.conf.get(Connect.CONNECT_SESSION_MANAGER_CLEANUP_CACHED_DATA_ENABLED)) { + session.sharedState.cacheManager.clearCache(session) + } + session.cleanupPythonWorkerLogs() eventManager.postClosed() diff --git a/sql/connect/server/src/test/scala/org/apache/spark/sql/connect/service/SparkConnectSessionManagerSuite.scala b/sql/connect/server/src/test/scala/org/apache/spark/sql/connect/service/SparkConnectSessionManagerSuite.scala index 5fae135bcf18a..e5bbc683d9084 100644 --- a/sql/connect/server/src/test/scala/org/apache/spark/sql/connect/service/SparkConnectSessionManagerSuite.scala +++ b/sql/connect/server/src/test/scala/org/apache/spark/sql/connect/service/SparkConnectSessionManagerSuite.scala @@ -21,14 +21,28 @@ import java.util.UUID import org.scalatest.time.SpanSugar._ -import org.apache.spark.SparkSQLException +import org.apache.spark.{SparkEnv, SparkSQLException} import org.apache.spark.sql.SparkSession +import org.apache.spark.sql.connect.config.Connect import org.apache.spark.sql.pipelines.graph.{DataflowGraph, PipelineUpdateContextImpl} import org.apache.spark.sql.pipelines.logging.PipelineEvent import org.apache.spark.sql.test.SharedSparkSession class SparkConnectSessionManagerSuite extends SharedSparkSession { + private def withSparkConf(pairs: (String, String)*)(f: => Unit): Unit = { + val conf = SparkEnv.get.conf + val previousValues = pairs.map { case (key, _) => key -> conf.getOption(key) } + pairs.foreach { case (key, value) => conf.set(key, value) } + try f + finally { + previousValues.foreach { + case (key, Some(value)) => conf.set(key, value) + case (key, None) => conf.remove(key) + } + } + } + override def beforeEach(): Unit = { super.beforeEach() SparkConnectService.sessionManager.invalidateAllSessions() @@ -177,6 +191,60 @@ class SparkConnectSessionManagerSuite extends SharedSparkSession { "pipeline execution was not removed") } + test("SPARK-50569: cached data cleanup on session close is configurable and isolated") { + Seq(false, true).foreach { cleanupCachedData => + withClue(s"cleanupCachedData=$cleanupCachedData") { + withSparkConf( + Connect.CONNECT_SESSION_MANAGER_CLEANUP_CACHED_DATA_ENABLED.key -> + cleanupCachedData.toString) { + val first = SparkConnectService.sessionManager.getOrCreateIsolatedSession( + SessionKey("user", UUID.randomUUID().toString), None) + val second = SparkConnectService.sessionManager.getOrCreateIsolatedSession( + SessionKey("user", UUID.randomUUID().toString), None) + val firstDataFrame = first.session.range(1) + second.session.range(1, 2).createTempView("second_view") + second.session.catalog.cacheTable("second_view") + val secondDataFrame = second.session.table("second_view") + + firstDataFrame.persist() + SparkConnectService.sessionManager.closeSession(first.key) + + assert( + first.session.sharedState.cacheManager.lookupCachedData(firstDataFrame).isDefined === + !cleanupCachedData) + assert( + second.session.sharedState.cacheManager.lookupCachedData(secondDataFrame).isDefined) + + SparkConnectService.sessionManager.closeSession(second.key) + assert( + second.session.sharedState.cacheManager.lookupCachedData(secondDataFrame).isDefined === + !cleanupCachedData) + spark.catalog.clearCache() + } + } + } + } + + test("SPARK-50569: cached data cleanup preserves entries persisted by another session") { + withSparkConf(Connect.CONNECT_SESSION_MANAGER_CLEANUP_CACHED_DATA_ENABLED.key -> "true") { + val first = SparkConnectService.sessionManager.getOrCreateIsolatedSession( + SessionKey("user", UUID.randomUUID().toString), None) + val second = SparkConnectService.sessionManager.getOrCreateIsolatedSession( + SessionKey("user", UUID.randomUUID().toString), None) + val firstDataFrame = first.session.range(1) + val secondDataFrame = second.session.range(1) + + firstDataFrame.persist() + secondDataFrame.persist() + SparkConnectService.sessionManager.closeSession(first.key) + + assert(second.session.sharedState.cacheManager.lookupCachedData(secondDataFrame).isDefined) + + SparkConnectService.sessionManager.closeSession(second.key) + assert(second.session.sharedState.cacheManager.lookupCachedData(secondDataFrame).isEmpty) + } + } + test("baseSession allows creating sessions after default session is cleared") { // Create a new session manager to test initialization val sessionManager = new SparkConnectSessionManager() diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/CacheManager.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/CacheManager.scala index a7adfdf4049e1..a5780e62ee3dc 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/CacheManager.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/CacheManager.scala @@ -17,6 +17,8 @@ package org.apache.spark.sql.execution +import java.util.IdentityHashMap + import scala.util.control.NonFatal import org.apache.hadoop.fs.{FileSystem, Path} @@ -79,13 +81,47 @@ class CacheManager extends Logging with AdaptiveSparkPlanHelper { @transient @volatile private var cachedData = IndexedSeq[CachedData]() + /** + * Tracks the sessions that explicitly cached each entry. A cache entry can be shared by + * multiple sessions because cache lookup uses plan semantics rather than session identity. + */ + @transient + private val cacheOwners = new IdentityHashMap[CachedData, Set[String]] + /** Clears all cached tables. */ def clearCache(): Unit = this.synchronized { cachedData.foreach(_.cachedRepresentation.cacheBuilder.clearCache()) cachedData = IndexedSeq[CachedData]() + cacheOwners.clear() CacheManager.logCacheOperation(log"Cleared all Dataframe cache entries") } + /** Clears cached data that is owned only by the given session. */ + private[sql] def clearCache(session: SparkSession): Unit = { + val sessionUUID = session.sessionUUID + val plansToUncache = this.synchronized { + val plans = cachedData.filter { cd => + Option(cacheOwners.get(cd)).contains(Set(sessionUUID)) + } + cachedData = cachedData.filterNot(cd => plans.exists(_ eq cd)) + val ownerEntries = cacheOwners.entrySet().iterator() + while (ownerEntries.hasNext) { + val entry = ownerEntries.next() + val remainingOwners = entry.getValue - sessionUUID + if (remainingOwners.nonEmpty) { + entry.setValue(remainingOwners) + } else { + ownerEntries.remove() + } + } + plans + } + plansToUncache.foreach(_.cachedRepresentation.cacheBuilder.clearCache()) + CacheManager.logCacheOperation( + log"Cleared ${MDC(SIZE, plansToUncache.size)} Dataframe cache entries for session " + + log"${MDC(SESSION_ID, sessionUUID)}") + } + /** Checks if the cache is empty. */ def isEmpty: Boolean = { cachedData.isEmpty @@ -144,7 +180,7 @@ class CacheManager extends Logging with AdaptiveSparkPlanHelper { log"Asked to cache a plan that is inapplicable for caching: " + log"${MDC(LOGICAL_PLAN, unnormalizedPlan)}" ) - } else if (lookupCachedDataInternal(normalizedPlan).nonEmpty) { + } else if (registerCacheOwner(normalizedPlan, spark.sessionUUID)) { logWarning("Asked to cache already cached data.") } else { val sessionWithConfigsOff = getOrCloneSessionWithConfigsOff(spark) @@ -158,12 +194,13 @@ class CacheManager extends Logging with AdaptiveSparkPlanHelper { } this.synchronized { - if (lookupCachedDataInternal(normalizedPlan).nonEmpty) { + if (registerCacheOwner(normalizedPlan, spark.sessionUUID)) { logWarning("Data has already been cached.") } else { // the cache key is the normalized plan val cd = CachedData(normalizedPlan, inMemoryRelation) cachedData = cd +: cachedData + cacheOwners.put(cd, Set(spark.sessionUUID)) CacheManager.logCacheOperation(log"Added Dataframe cache entry:" + log"${MDC(DATAFRAME_CACHE_ENTRY, cd)}") } @@ -171,6 +208,17 @@ class CacheManager extends Logging with AdaptiveSparkPlanHelper { } } + private def registerCacheOwner(plan: LogicalPlan, sessionUUID: String): Boolean = + this.synchronized { + lookupCachedDataInternal(plan) match { + case Some(cd) => + val owners = Option(cacheOwners.get(cd)).getOrElse(Set.empty) + cacheOwners.put(cd, owners + sessionUUID) + true + case None => false + } + } + /** * Un-cache the given plan or all the cache entries that refer to the given plan. * @@ -301,6 +349,7 @@ class CacheManager extends Logging with AdaptiveSparkPlanHelper { val plansToUncache = cachedData.filter(cd => shouldRemove(cd.plan)) this.synchronized { cachedData = cachedData.filterNot(cd => plansToUncache.exists(_ eq cd)) + plansToUncache.foreach(cd => cacheOwners.remove(cd)) } plansToUncache.foreach { _.cachedRepresentation.cacheBuilder.clearCache(blocking) } CacheManager.logCacheOperation(log"Removed ${MDC(SIZE, plansToUncache.size)} Dataframe " + @@ -374,14 +423,22 @@ class CacheManager extends Logging with AdaptiveSparkPlanHelper { } needToRecache.foreach { cd => cd.cachedRepresentation.cacheBuilder.clearCache() - tryRebuildCacheEntry(spark, cd).foreach { entry => - this.synchronized { - if (lookupCachedDataInternal(entry.plan).nonEmpty) { - logWarning("While recaching, data was already added to cache.") - } else { - cachedData = entry +: cachedData - CacheManager.logCacheOperation(log"Re-cached Dataframe cache entry:" + - log"${MDC(DATAFRAME_CACHE_ENTRY, entry)}") + val rebuiltEntry = tryRebuildCacheEntry(spark, cd) + this.synchronized { + val previousOwners = Option(cacheOwners.remove(cd)).getOrElse(Set.empty) + rebuiltEntry.foreach { entry => + if (previousOwners.nonEmpty) { + lookupCachedDataInternal(entry.plan) match { + case Some(existing) => + val existingOwners = Option(cacheOwners.get(existing)).getOrElse(Set.empty) + cacheOwners.put(existing, existingOwners ++ previousOwners) + logWarning("While recaching, data was already added to cache.") + case None => + cachedData = entry +: cachedData + cacheOwners.put(entry, previousOwners) + CacheManager.logCacheOperation(log"Re-cached Dataframe cache entry:" + + log"${MDC(DATAFRAME_CACHE_ENTRY, entry)}") + } } } }