spark-test commented on code in PR #59952:
URL: https://github.com/apache/spark/pull/59952#discussion_r4429037137
##########
sql/connect/server/src/main/scala/org/apache/spark/sql/connect/service/SessionHolder.scala:
##########
@@ -439,6 +435,15 @@ case class SessionHolder(userId: String, sessionId:
String, session: SparkSessio
// Clean up ML cache (only if ML models were created)
mlCache.close()
+ // Close the catalogs loaded by this session, then clean up all artifacts.
Catalogs are closed
+ // first because they may have been loaded by the session classloader,
which is closed with
+ // the artifacts. Otherwise resources held by the catalogs (e.g. thread
pools) keep the
+ // session classloader alive. Cloned sessions own their catalogs and
artifacts, so this does
+ // not affect them.
+ // Note: interrupted executions may still be running; they fail if they
try to load a catalog
+ // after this point. There can also be concurrent AddArtifact calls still
adding something.
+ session.sessionState.close()
Review Comment:
**Non-blocking (P2):** Artifacts used to be released first in `close()`, so
a failure in a later step could not skip them. Now they are released only here,
after `removeAllPipelineExecutions()` and the other teardown steps, and
`removeAllPipelineExecutions()` throws when a pipeline flow fails to stop
(`RunTerminationException`) or its control thread does not stop in time
(`TimeoutException`). When any earlier step throws, neither the session's
catalogs nor its artifacts are released. The cached artifact blocks stay in the
driver `BlockManager`, which the garbage-collection cleaner never cleans, and
the session classloader stays open. The holder has already been removed and
refuses a second close. Running `session.sessionState.close()` in a `finally`
around the earlier steps, with `Utils.tryWithSafeFinally`, keeps the
catalogs-before-artifacts order and restores the old guarantee.
See **Shared repair plan 1** in the review body.
##########
sql/core/src/main/scala/org/apache/spark/sql/internal/SessionState.scala:
##########
@@ -133,6 +133,21 @@ private[sql] class SessionState(
*/
def clone(newSparkSession: SparkSession): SessionState =
createClone(newSparkSession, this)
+ /**
+ * Release the resources owned by this `SessionState`: the catalogs loaded
by its
+ * [[CatalogManager]], then the session artifacts. Catalogs are closed first
since they may have
+ * been loaded by the session classloader, which is closed with the
artifacts. Clones have their
+ * own catalogs and artifacts, so closing a `SessionState` does not affect
its clones. No catalog
+ * can be loaded afterwards.
+ */
+ def close(): Unit = {
+ try {
+ catalogManager.close()
Review Comment:
**Blocking (P1):** `catalogManager` is `analyzer.catalogManager`, and
`analyzer` is a lazy val, so this line builds the analyzer for any Spark
Connect session that never analyzed a plan, such as one that only sent Config
requests before it was released or expired. That creates the `SessionCatalog`
(copying the Hadoop and SQL configuration), the `V2SessionCatalog`, the
`DefaultCatalogManager` and the `Analyzer`, and runs the
`SparkSessionExtensions` rule builders on the ReleaseSession or maintenance
thread, only to close an empty catalog map. Closing such a session did none of
this before. If the construction throws, `SessionHolder.close()` fails after
the artifacts are released and skips `cleanupPythonWorkerLogs()` and
`postClosed()`, so the session manager never records the session as closed.
**Recommended change:** Make SessionState.close() close the catalog manager
only when the analyzer has already been built. If the analyzer is first built
after close(), close its catalog manager before the analyzer is handed out, so
either way no catalog can be loaded through a released session.
**Why this works:** SessionState records that close() ran and whether its
analyzer has been initialized. close() sets the closed flag and reads the
initialized flag under the lock that guards the lazy analyzer's initialization,
then closes the existing catalog manager if there is one, and then the
artifacts. The analyzer initializer marks itself initialized and, when the
closed flag is already set, closes the new catalog manager before returning. A
session that never analyzed a plan therefore releases without building
anything, while a racing or later first analysis still finds a terminal manager.
**Scope:** Track analyzer initialization and closed state in SessionState so
that releasing a session never builds its analyzer stack.
**Compatibility:** Sessions that already analyzed plans release exactly as
before: their loaded catalogs are closed first, then their artifacts.
**Risks:** The analyzer initializer closes an empty catalog manager while
holding the SessionState lazy-initialization lock. That adds a
SessionState-then-catalog-manager lock order, which no current code path
reverses.
**Constraints:** Catalog loads after close() must keep failing with the
closed-manager error whether the analyzer was built before or after close().
Catalogs must still be closed before artifacts.
**Success:** Releasing a Spark Connect session that never analyzed a plan
does not construct its analyzer, SessionCatalog, V2SessionCatalog or catalog
manager, and the session is still recorded as closed. Releasing a session whose
analyzer exists still closes its loaded catalogs before its artifacts. If a
released session's analyzer is first built after release, its catalog manager
is already closed, so catalog loads fail with the closed-manager error.
**Verification:**
- **Regression:** Releasing a Spark Connect session that only issued
configuration requests leaves its analyzer unbuilt (for example, an injected
resolution-rule builder is never invoked), and the session is still recorded as
closed.
- **Behavior:** After such a release, the first analysis on the released
session's state finds a closed catalog manager: loading a configured catalog
fails with the closed-manager error.
- **Behavior:** Releasing a session that loaded a closeable catalog still
closes that catalog.
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/connector/catalog/CatalogManager.scala:
##########
@@ -156,7 +174,10 @@ private[sql] class DefaultCatalogManager(
override private[sql] def v2SessionCatalog: CatalogPlugin = {
conf.getConf(SQLConf.V2_SESSION_CATALOG_IMPLEMENTATION) match {
case "builtin" => defaultSessionCatalog
- case _ => catalogs.getOrElseUpdate(SESSION_CATALOG_NAME,
loadV2SessionCatalog())
+ case _ => catalogs.getOrElseUpdate(SESSION_CATALOG_NAME, {
+ checkNotClosed(SESSION_CATALOG_NAME)
Review Comment:
**Non-blocking (P2):** `catalog()` is `synchronized`, so its check, load and
insert are atomic with `close()`. This branch is not, and `v2SessionCatalog` is
called directly without the lock, for example by `LookupCatalog` for
`global_temp` names and by `CatalogV2Util.getTableProviderCatalog`. With a
custom `spark_catalog` that the session has not loaded yet, a lookup can pass
`checkNotClosed` before `close()` runs, finish `loadV2SessionCatalog()` after
`close()` has cleared the map, and register an instance that is never closed.
Later lookups then return it from the map without reaching the check. Running
this branch under the manager lock, as `catalog()` already does when it loads
the session catalog, closes the window without adding a new lock order.
**Verification:**
- **Regression:** With a custom spark_catalog whose initialization is held
on a latch, a direct v2SessionCatalog load that starts before close() and
finishes after it ends with the loaded instance closed or with the
closed-manager error, and the manager serves no unclosed session catalog
afterwards.
- **Behavior:** A custom v2 session catalog loaded before close() is still
closed by close(), later access fails with the closed-manager error, and the
builtin session catalog stays available.
##########
sql/core/src/main/scala/org/apache/spark/sql/internal/SessionState.scala:
##########
@@ -133,6 +133,21 @@ private[sql] class SessionState(
*/
def clone(newSparkSession: SparkSession): SessionState =
createClone(newSparkSession, this)
+ /**
+ * Release the resources owned by this `SessionState`: the catalogs loaded
by its
+ * [[CatalogManager]], then the session artifacts. Catalogs are closed first
since they may have
+ * been loaded by the session classloader, which is closed with the
artifacts. Clones have their
+ * own catalogs and artifacts, so closing a `SessionState` does not affect
its clones. No catalog
+ * can be loaded afterwards.
+ */
+ def close(): Unit = {
Review Comment:
**Non-blocking (P2):** This is the new release point, but only Spark
Connect's `SessionHolder.close()` calls it. Two other places discard sessions
that own their own catalog manager. In multi-session mode, the Thrift server's
`SparkSQLSessionManager.closeSession` drops each connection's `newSession()`
after uncaching its temp views. Each streaming query also runs in
`StreamExecution.sparkSessionForStream`, which nothing releases when the query
ends. That per-query session loads its own catalog instance without any user
code when the sink is a `TransactionalCatalogPlugin` table, and it loads more
when Spark Connect Scala `foreachBatch` code resolves tables through the batch
DataFrame's session. Catalogs loaded on those paths keep leaking their threads
and clients. That includes Spark Connect sessions that ran such a query, which
is the leak this change sets out to fix.
**Recommended change:** Release the connection's own session in the Thrift
server's closeSession, and only when it is the per-connection session
registered for that connection, never the server's root session. Release a
streaming query's per-query session once the query has terminated and stopped
its sources and sinks, using the same SessionState.close(). Keep a repeated
release of the same session state harmless, so the Python foreachBatch holder
and the streaming engine can both release it. Then update the CatalogPlugin
Javadoc's description of when Spark releases sessions and the catalog manager's
caller comments to match.
**Why this works:** Each owner that discards a session now calls the shared
release at its discard point, so catalogs loaded through Thrift and streaming
per-query sessions are closed the same way as those of Spark Connect sessions.
**Scope:** Extend the session-state release to Thrift server and streaming
per-query sessions.
**Compatibility:** Spark Connect session release, the Python foreachBatch
cloned-holder release and the Thrift server's single-session mode behave as
before.
**Risks:** Releasing a streaming query's per-query session at termination
closes its catalogs and artifact copies, so user code that kept a batch
DataFrame's session beyond the query's lifetime would see a closed catalog
manager.
**Constraints:** Never release the Thrift server's root session: not in
single-session mode, where every connection shares it, and not when
closeSession falls back to it because no session was registered for the
connection, as happens when opening the session fails. Release a streaming
query's session only after the query has stopped its sources and sinks.
Releasing a session that never analyzed a plan must stay cheap; this relies on
the repair that stops session release from building the analyzer.
**Exclusions:**
- Classic sessions that applications create and drop themselves, which have
no release point.
**Success:** Closing a multi-session Thrift server connection closes the
closeable catalogs its session loaded. When a streaming query terminates,
closeable catalogs loaded through its per-query session are closed. This
includes the catalog of a transactional streaming sink, Spark Connect Scala
foreachBatch queries and queries stopped by a Spark Connect session release.
Releasing the same session state twice has no further effect.
**Verification:**
- **Behavior:** A closeable catalog loaded by a Thrift server connection in
multi-session mode is closed when that connection's session is closed.
- **Behavior:** A closeable catalog loaded through a streaming query's
per-query session, both for a transactional catalog sink and by foreachBatch
code through the batch DataFrame's session, is closed after the query
terminates, both for a query stopped directly and for one stopped by a Spark
Connect session release.
- **Compatibility:** In Thrift server single-session mode, closing a
connection does not release the shared session or its catalogs.
- **Compatibility:** When opening a Thrift server session fails after the
session was created, for example because the requested database does not exist,
the cleanup releases neither the server's root session nor its catalogs.
- **Regression:** A Python foreachBatch query whose cloned session holder
and streaming termination both release the same session state closes each
catalog once and raises no error.
##########
sql/connect/server/src/test/scala/org/apache/spark/sql/connect/service/SparkConnectSessionManagerSuite.scala:
##########
@@ -177,6 +178,21 @@ class SparkConnectSessionManagerSuite extends
SharedSparkSession {
"pipeline execution was not removed")
}
+ test("Catalogs loaded by the session are closed when the session holder is
closed") {
+ val key = SessionKey("user", UUID.randomUUID().toString)
+ val sessionHolder =
SparkConnectService.sessionManager.getOrCreateIsolatedSession(key, None)
+ sessionHolder.session.conf
+ .set("spark.sql.catalog.closeable", classOf[CloseableCatalog].getName)
+ // Catalogs are loaded with the active session's conf.
+ val catalog = sessionHolder.withSession {
+
_.sessionState.catalogManager.catalog("closeable").asInstanceOf[CloseableCatalog]
+ }
+ assert(!catalog.closed)
+
+ SparkConnectService.sessionManager.closeSession(sessionHolder.key)
+ assert(catalog.closed)
Review Comment:
**Nit (P3):** This only checks the closed flag of a catalog loaded from the
parent classpath, so nothing fails if `SessionState.close()` releases the
artifacts first. The order matters for catalogs loaded from session jars. If
the session classloader is closed first, their `close()` can fail with
`NoClassDefFoundError`, which `DefaultCatalogManager.close()` only logs, and
they leak again. A closeable test catalog whose `close()` records whether a
previously added session artifact is still present would catch that.
See **Shared repair plan 1** in the review body.
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/connector/catalog/CatalogManager.scala:
##########
@@ -526,6 +547,34 @@ private[sql] class DefaultCatalogManager(
}
v1SessionCatalog.setCurrentDatabase(conf.defaultDatabase)
}
+
+ // Close and unregister every catalog loaded by this manager that implements
`AutoCloseable`.
Review Comment:
**Nit (P3):** As written, this limits both closing and unregistering to
`AutoCloseable` catalogs. `close()` clears every entry and closes only the
`AutoCloseable` ones, so a non-closeable catalog such as the test's
`DummyCatalog` is unregistered too, and a later `catalog(name)` throws.
```suggestion
// Unregister every catalog loaded by this manager; close the ones that
are `AutoCloseable`.
```
**Verification:**
- **Inspection:** The comment above DefaultCatalogManager.close() matches
its behavior: every loaded catalog is unregistered, the AutoCloseable ones are
closed, and later loads fail.
##########
sql/catalyst/src/main/java/org/apache/spark/sql/connector/catalog/CatalogPlugin.java:
##########
@@ -39,6 +39,12 @@
* {@code spark.sql.catalog.catalog-name.(key)=(value)} will be passed in the
case insensitive
* string map of options in initialization with the prefix removed.
* {@code name}, is also passed and is the catalog's name; in this case,
"catalog-name".
+ * <p>
+ * Catalog instances are loaded per session. Implementations that hold
resources such as clients,
+ * connections or thread pools can also implement {@link java.io.Closeable} (or
+ * {@link AutoCloseable}). Spark calls {@code close()} when it releases the
session that loaded
Review Comment:
**Non-blocking (P2):** This paragraph reads as if `close()` ends the
instance's use, so an implementation that follows it will shut down its clients
and thread pools there. Spark can still call the instance afterwards. Session
release interrupts executions without waiting for them, so an interrupted CTAS
can still run `dropTable` or `abortStagedChanges` on it. A table that the
released session cached is refreshed through this instance when another session
writes it, and a failed refresh silently drops the cache entry. Temp views over
time-travel reads that were copied into a cloned session also keep using it.
The paragraph also leaves out that `close()` runs at most once, on a
Spark-internal thread outside the session's classloader and active session,
that its exceptions are only logged, that it should not block for long, and
that existing `AutoCloseable` implementations start receiving it with this
change. The new comments on `SessionState.close()` and `SessionHolder.close()`,
wh
ich say that closing a session does not affect its clones, need the same
correction.
**Verification:**
- **Inspection:** The CatalogPlugin class Javadoc describes the call count,
the calling thread and context, possible overlap with and later calls on the
instance, failure handling, and the first release that calls close(),
consistently with the catalog manager's close() and Spark Connect session
release.
- **Inspection:** The SessionState.close() scaladoc and the
SessionHolder.close() comment say that temp views copied into a cloned session
can still use the closed session's catalog instances.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
