Skip to content

[SPARK-54367][Connect]Propagate close() from SessionState to CatalogPlugins to prevent leak #53078

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
MonkeyCanCode wants to merge 5 commits into
base: master
Choose a base branch
from
Open

[SPARK-54367][Connect]Propagate close() from SessionState to CatalogPlugins to prevent leak #53078

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
5 commits
Select commit Hold shift + click to select a range
f135cc1
Propagate close() from SessionState to CatalogPlugins to prevent reso…
MonkeyCanCode Nov 15, 2025
0cd632e
Trigger build
MonkeyCanCode Nov 15, 2025
e769061
Fix style
MonkeyCanCode Nov 15, 2025
5cfed61
Fix test cases
MonkeyCanCode Nov 15, 2025
461e7bb
Fix test cases
MonkeyCanCode Nov 16, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
7 changes: 6 additions & 1 deletion sql/catalyst/src/main/java/org/apache/spark/sql/connector/catalog/CatalogPlugin.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -17,6 +17,8 @@

package org.apache.spark.sql.connector.catalog;

import java.io.Closeable;

import org.apache.spark.annotation.Evolving;
import org.apache.spark.sql.internal.SQLConf;
import org.apache.spark.sql.util.CaseInsensitiveStringMap;
Expand All @@ -42,7 +44,7 @@
* @since 3.0.0
*/
@Evolving
public interface CatalogPlugin {
public interface CatalogPlugin extends Closeable {
/**
* Called to initialize configuration.
* <p>
Expand Down Expand Up @@ -74,4 +76,7 @@ public interface CatalogPlugin {
default String[] defaultNamespace() {
return new String[0];
}

@Override
default void close() {}
}
2 changes: 2 additions & 0 deletions sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/v2ResolutionPlans.scala
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -265,6 +265,7 @@ object ResolvedIdentifier {
object FakeSystemCatalog extends CatalogPlugin {
override def initialize(name: String, options: CaseInsensitiveStringMap): Unit = {}
override def name(): String = "system"
override def close(): Unit = {}
}

/**
Expand All @@ -273,4 +274,5 @@ object FakeSystemCatalog extends CatalogPlugin {
object FakeLocalCatalog extends CatalogPlugin {
override def initialize(name: String, options: CaseInsensitiveStringMap): Unit = {}
override def name(): String = "local"
override def close(): Unit = {}
}
19 changes: 18 additions & 1 deletion sql/catalyst/src/main/scala/org/apache/spark/sql/connector/catalog/CatalogManager.scala
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -17,6 +17,8 @@

package org.apache.spark.sql.connector.catalog

import java.io.Closeable

import scala.collection.mutable

import org.apache.spark.internal.Logging
Expand All @@ -40,7 +42,7 @@ import org.apache.spark.sql.internal.SQLConf
private[sql]
class CatalogManager(
defaultSessionCatalog: CatalogPlugin,
val v1SessionCatalog: SessionCatalog) extends SQLConfHelper with Logging {
val v1SessionCatalog: SessionCatalog) extends SQLConfHelper with Logging with Closeable {
import CatalogManager.SESSION_CATALOG_NAME
import CatalogV2Util._

Expand All @@ -57,6 +59,21 @@ class CatalogManager(
}
}

override def close(): Unit = synchronized {
val allCatalogs = (catalogs.values.toSet + defaultSessionCatalog).toSeq
allCatalogs.foreach {
case c: Closeable =>
try {
c.close()
} catch {
case e: Throwable =>
logWarning(s"Failed to close catalog of class ${c.getClass.getName}", e)
}
case _ =>
}
catalogs.clear()
}

def isCatalogRegistered(name: String): Boolean = {
try {
catalog(name)
Expand Down
21 changes: 21 additions & 0 deletions sql/catalyst/src/test/scala/org/apache/spark/sql/connector/catalog/CatalogManagerSuite.scala
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -127,6 +127,18 @@ class CatalogManagerSuite extends SparkFunSuite with SQLHelper {
}
}
}

test("CatalogManager.close should close all closeable catalogs") {
val catalogManager = new CatalogManager(FakeV2SessionCatalog, createSessionCatalog())
withSQLConf("spark.sql.catalog.dummy" -> classOf[DummyCatalog].getName,
"spark.sql.catalog.closeable" -> classOf[CloseableCatalog].getName) {
catalogManager.setCurrentCatalog("dummy")
val closeable = catalogManager.catalog("closeable").asInstanceOf[CloseableCatalog]
assert(!closeable.isClosed)
catalogManager.close()
assert(closeable.isClosed)
}
}
}

class DummyCatalog extends CatalogPlugin {
Expand All @@ -136,4 +148,13 @@ class DummyCatalog extends CatalogPlugin {
private var _name: String = null
override def name(): String = _name
override def defaultNamespace(): Array[String] = Array("a", "b")
override def close(): Unit = {}
}

class CloseableCatalog extends DummyCatalog with java.io.Closeable {
private var closed = false
override def close(): Unit = {
closed = true
}
def isClosed: Boolean = closed
}
1 change: 1 addition & 0 deletions sql/catalyst/src/test/scala/org/apache/spark/sql/connector/catalog/LookupCatalogSuite.scala
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -33,6 +33,7 @@ import org.apache.spark.sql.util.CaseInsensitiveStringMap
private case class DummyCatalogPlugin(override val name: String) extends CatalogPlugin {

override def initialize(name: String, options: CaseInsensitiveStringMap): Unit = ()
override def close(): Unit = {}
}

class LookupCatalogSuite extends SparkFunSuite with LookupCatalog with Inside {
Expand Down
1 change: 1 addition & 0 deletions sql/connect/server/src/main/scala/org/apache/spark/sql/connect/service/SessionHolder.scala
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -380,6 +380,7 @@ case class SessionHolder(userId: String, sessionId: String, session: SparkSessio

mlCache.clear()

session.sessionState.close()
session.cleanupPythonWorkerLogs()

eventManager.postClosed()
Expand Down
17 changes: 10 additions & 7 deletions ...src/main/scala/org/apache/spark/sql/connect/service/SparkConnectListenerBusListener.scala
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -137,13 +137,16 @@ private[sql] class SparkConnectListenerBusListener(
}

def sendResultComplete(): Unit = {
responseObserver
.asInstanceOf[ExecuteResponseObserver[ExecutePlanResponse]]
.onNextComplete(
ExecutePlanResponse
.newBuilder()
.setResultComplete(ExecutePlanResponse.ResultComplete.newBuilder().build())
.build())
responseObserver match {
Copy link
Contributor

@hvanhovell hvanhovell Dec 4, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

How is this a problem now?

Copy link
Author

@MonkeyCanCode MonkeyCanCode Dec 4, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This was added due to failed unit tests with mock object where i was getting class cast exception earlier (sample error can be found in https://github.com/MonkeyCanCode/spark/actions/runs/19394233457/job/55492018508).

Here is the snippet:

[info] org.apache.spark.sql.connect.pipelines.EndToEndAPISuite *** ABORTED *** (81 milliseconds)
[info]   java.lang.ClassCastException: class io.grpc.stub.StreamObserver$MockitoMock$lemjBfIA cannot be cast to class org.apache.spark.sql.connect.execution.ExecuteResponseObserver (io.grpc.stub.StreamObserver$MockitoMock$lemjBfIA and org.apache.spark.sql.connect.execution.ExecuteResponseObserver are in unnamed module of loader 'app')
[info]   at org.apache.spark.sql.connect.service.SparkConnectListenerBusListener.sendResultComplete(SparkConnectListenerBusListener.scala:140)
[info]   at org.apache.spark.sql.connect.service.ServerSideListenerHolder.cleanUp(SparkConnectListenerBusListener.scala:84)
[info]   at org.apache.spark.sql.connect.service.SessionHolder.close(SessionHolder.scala:370)
[info]   at org.apache.spark.sql.connect.service.SparkConnectSessionManager.shutdownSessionHolder(SparkConnectSessionManager.scala:232)
[info]   at org.apache.spark.sql.connect.service.SparkConnectSessionManager.$anonfun$invalidateAllSessions$1(SparkConnectSessionManager.scala:369)
[info]   at java.base/java.util.concurrent.ConcurrentHashMap.forEach(ConcurrentHashMap.java:1603)
[info]   at org.apache.spark.sql.connect.service.SparkConnectSessionManager.invalidateAllSessions(SparkConnectSessionManager.scala:367)
[info]   at org.apache.spark.sql.connect.SparkConnectServerTest.clearAllExecutions(SparkConnectServerTest.scala:83)
[info]   at org.apache.spark.sql.connect.SparkConnectServerTest.clearAllExecutions$(SparkConnectServerTest.scala:80)
[info]   at org.apache.spark.sql.connect.pipelines.EndToEndAPISuite.clearAllExecutions(EndToEndAPISuite.scala:37)
[info]   at org.apache.spark.sql.connect.SparkConnectServerTest.beforeEach(SparkConnectServerTest.scala:72)
[info]   at org.apache.spark.sql.connect.SparkConnectServerTest.beforeEach$(SparkConnectServerTest.scala:70)
[info]   at org.apache.spark.sql.connect.pipelines.EndToEndAPISuite.beforeEach(EndToEndAPISuite.scala:37)
[info]   at org.scalatest.BeforeAndAfterEach.runTest(BeforeAndAfterEach.scala:233)
[info]   at org.scalatest.BeforeAndAfterEach.runTest$(BeforeAndAfterEach.scala:227)
[info]   at org.apache.spark.SparkFunSuite.runTest(SparkFunSuite.scala:68)
[info]   at org.scalatest.funsuite.AnyFunSuiteLike.$anonfun$runTests$1(AnyFunSuiteLike.scala:269)
[info]   at org.scalatest.SuperEngine.$anonfun$runTestsInBranch$1(Engine.scala:413)
[info]   at scala.collection.immutable.List.foreach(List.scala:323)
[info]   at org.scalatest.SuperEngine.traverseSubNodes$1(Engine.scala:401)
[info]   at org.scalatest.SuperEngine.runTestsInBranch(Engine.scala:396)
[info]   at org.scalatest.SuperEngine.runTestsImpl(Engine.scala:475)
[info]   at org.scalatest.funsuite.AnyFunSuiteLike.runTests(AnyFunSuiteLike.scala:269)
[info]   at org.scalatest.funsuite.AnyFunSuiteLike.runTests$(AnyFunSuiteLike.scala:268)
[info]   at org.scalatest.funsuite.AnyFunSuite.runTests(AnyFunSuite.scala:1564)
[info]   at org.scalatest.Suite.run(Suite.scala:1114)
[info]   at org.scalatest.Suite.run$(Suite.scala:1096)
[info]   at org.scalatest.funsuite.AnyFunSuite.org$scalatest$funsuite$AnyFunSuiteLike$$super$run(AnyFunSuite.scala:1564)
[info]   at org.scalatest.funsuite.AnyFunSuiteLike.$anonfun$run$1(AnyFunSuiteLike.scala:273)
[info]   at org.scalatest.SuperEngine.runImpl(Engine.scala:535)
[info]   at org.scalatest.funsuite.AnyFunSuiteLike.run(AnyFunSuiteLike.scala:273)
[info]   at org.scalatest.funsuite.AnyFunSuiteLike.run$(AnyFunSuiteLike.scala:272)
[info]   at org.apache.spark.SparkFunSuite.org$scalatest$BeforeAndAfterAll$$super$run(SparkFunSuite.scala:68)
[info]   at org.scalatest.BeforeAndAfterAll.liftedTree1$1(BeforeAndAfterAll.scala:213)
[info]   at org.scalatest.BeforeAndAfterAll.run(BeforeAndAfterAll.scala:210)
[info]   at org.scalatest.BeforeAndAfterAll.run$(BeforeAndAfterAll.scala:208)
[info]   at org.apache.spark.SparkFunSuite.run(SparkFunSuite.scala:68)
[info]   at org.scalatest.tools.Framework.org$scalatest$tools$Framework$$runSuite(Framework.scala:321)
[info]   at org.scalatest.tools.Framework$ScalaTestTask.execute(Framework.scala:517)
[info]   at sbt.ForkMain$Run.lambda$runTest$1(ForkMain.java:414)
[info]   at java.base/java.util.concurrent.FutureTask.run(FutureTask.java:264)
[info]   at java.base/java.util.concurrent.ThreadPoolExecutor.runWorker(ThreadPoolExecutor.java:1136)
[info]   at java.base/java.util.concurrent.ThreadPoolExecutor$Worker.run(ThreadPoolExecutor.java:635)
[info]   at java.base/java.lang.Thread.run(Thread.java:840)

Before this PR, the session cleanup process in the tests was less explicit. Now with proper resources closed, the generic test StreamObserver is causing issue during casting with existed code.

case obs: ExecuteResponseObserver[ExecutePlanResponse] =>
obs.onNextComplete(
ExecutePlanResponse
.newBuilder()
.setResultComplete(ExecutePlanResponse.ResultComplete.newBuilder().build())
.build())
case _ =>
responseObserver.onCompleted()
}
}

// QueryStartedEvent is sent to client along with WriteStreamOperationStartResult
Expand Down
3 changes: 3 additions & 0 deletions sql/connect/server/src/test/scala/org/apache/spark/sql/connect/SparkConnectTestUtils.scala
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -31,6 +31,9 @@ object SparkConnectTestUtils {
sessionId = UUID.randomUUID().toString,
session = session)
SparkConnectService.sessionManager.putSessionForTesting(ret)
if (session != null) {
ret.initializeSession()
Copy link
Contributor

@hvanhovell hvanhovell Dec 4, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Why is this needed now?

Copy link
Author

@MonkeyCanCode MonkeyCanCode Dec 4, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Earlier tests create dummy session that are not yet fully initialized. Now as all sessions got properly closed, calling closed on not initialized sessions was causing failed tests as well.

}
ret
}
}
32 changes: 32 additions & 0 deletions .../src/test/scala/org/apache/spark/sql/connect/service/SparkConnectSessionHolderSuite.scala
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -40,10 +40,24 @@ import org.apache.spark.sql.connect.planner.StreamingForeachBatchHelper.RunnerCl
import org.apache.spark.sql.pipelines.graph.{DataflowGraph, PipelineUpdateContextImpl}
import org.apache.spark.sql.pipelines.logging.PipelineEvent
import org.apache.spark.sql.test.SharedSparkSession
import org.apache.spark.sql.util.CaseInsensitiveStringMap
import org.apache.spark.util.ArrayImplicits._

class SparkConnectSessionHolderSuite extends SharedSparkSession {

test("SessionHolder.close should close catalogs") {
val sessionHolder = SparkConnectTestUtils.createDummySessionHolder(spark)
val catalogName = "my_closeable_catalog"
sessionHolder.session.conf
.set(s"spark.sql.catalog.$catalogName", classOf[CloseableCatalog].getName)

val catalog = sessionHolder.session.sessionState.catalogManager.catalog(catalogName)
val closeableCatalog = catalog.asInstanceOf[CloseableCatalog]
assert(!closeableCatalog.isClosed)
sessionHolder.close()
assert(closeableCatalog.isClosed)
}

test("DataFrame cache: Successful put and get") {
val sessionHolder = SparkConnectTestUtils.createDummySessionHolder(spark)
import sessionHolder.session.implicits._
Expand Down Expand Up @@ -484,3 +498,21 @@ class SparkConnectSessionHolderSuite extends SharedSparkSession {
assertPlanCache(sessionHolder, Some(expected))
}
}

class CloseableCatalog
extends org.apache.spark.sql.connector.catalog.CatalogPlugin
with java.io.Closeable {
private var _name: String = _
private var closed = false

override def initialize(name: String, options: CaseInsensitiveStringMap): Unit = {
_name = name
}

override def name(): String = _name
override def close(): Unit = {
closed = true
}

def isClosed: Boolean = closed
}
2 changes: 2 additions & 0 deletions sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/V2SessionCatalog.scala
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -55,6 +55,8 @@ class V2SessionCatalog(catalog: SessionCatalog)
// This class is instantiated by Spark, so `initialize` method will not be called.
override def initialize(name: String, options: CaseInsensitiveStringMap): Unit = {}

override def close(): Unit = {}

override def capabilities(): util.Set[TableCatalogCapability] = {
Set(
TableCatalogCapability.SUPPORT_COLUMN_DEFAULT_VALUE
Expand Down
8 changes: 6 additions & 2 deletions sql/core/src/main/scala/org/apache/spark/sql/internal/SessionState.scala
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -17,7 +17,7 @@

package org.apache.spark.sql.internal

import java.io.File
import java.io.{Closeable, File}
import java.net.URI

import org.apache.hadoop.conf.Configuration
Expand Down Expand Up @@ -91,7 +91,7 @@ private[sql] class SessionState(
val columnarRules: Seq[ColumnarRule],
val adaptiveRulesHolder: AdaptiveRulesHolder,
val planNormalizationRules: Seq[Rule[LogicalPlan]],
val artifactManagerBuilder: () => ArtifactManager) {
val artifactManagerBuilder: () => ArtifactManager) extends Closeable {

// The following fields are lazy to avoid creating the Hive client when creating SessionState.
lazy val catalog: SessionCatalog = catalogBuilder()
Expand All @@ -110,6 +110,10 @@ private[sql] class SessionState(

def catalogManager: CatalogManager = analyzer.catalogManager

override def close(): Unit = {
catalogManager.close()
Copy link
Contributor

@hvanhovell hvanhovell Dec 4, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

We need to make sure that the session manager supports cloning.

Copy link
Author

@MonkeyCanCode MonkeyCanCode Dec 4, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thanks for the review and feedback. I just want to make sure I fully understand the requirement here. When you mention this, are you thinking of the spark.newSession() case? My main concern is avoid issues where cloning one session could affected a cloned session as now the catalogs may soon be closable as well (https://github.com/apache/iceberg/pull/14590/files#diff-e263afc81b5b90944281c544005466e8671e7756f52b5b59b67a77e7c63aaf13).

Do you mind point me in the right direction for the implementation you have in mind?

Again, thanks so much for your review and time.

}

def newHadoopConf(): Configuration = SessionState.newHadoopConf(
sharedState.sparkContext.hadoopConfiguration,
conf)
Expand Down