You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@spark.apache.org by db...@apache.org on 2026/06/10 18:58:50 UTC
(spark) branch branch-4.x updated: [SPARK-57325][CONNECT] Stop streaming queries registered while the Connect session is closing
This is an automated email from the ASF dual-hosted git repository.
dbtsai pushed a commit to branch branch-4.x
in repository https://gitbox.apache.org/repos/asf/spark.git
The following commit(s) were added to refs/heads/branch-4.x by this push:
new 7fa2bf19a338 [SPARK-57325][CONNECT] Stop streaming queries registered while the Connect session is closing
7fa2bf19a338 is described below
commit 7fa2bf19a338cdbcbde47154e26f64469671f08d
Author: DB Tsai <db...@dbtsai.com>
AuthorDate: Wed Jun 10 11:58:17 2026 -0700
[SPARK-57325][CONNECT] Stop streaming queries registered while the Connect session is closing
### What changes were proposed in this pull request?
This PR fixes a race condition between registering a newly started streaming query and the owning Connect session being closed. The race strands two kinds of per-session resources, so the fix touches two server-side caches with the same publish-then-recheck pattern.
#### 1. `SparkConnectStreamingQueryCache` (the query cache)
When a Connect session is closed (`SessionHolder.close()`), it stops all of the session's streaming queries via `SparkConnectStreamingQueryCache.cleanupRunningQueries()`, which iterates over the query cache. As the existing code comment in `close()` notes, *"there can be concurrent streaming queries being started."* A query that finishes `DataStreamWriter.start()` on another thread and is registered **after** that iteration is missed by the cleanup. It is left in the cache as an activ [...]
The fix makes `registerNewStreamingQuery` coordinate with session shutdown without introducing additional locking:
1. **`SessionHolder.isClosing`** — a new `private[connect]` accessor exposing `closedTimeMs.isDefined`. `closedTimeMs` is set at the very start of `close()`, before any session resources (including streaming queries) are cleaned up, so it is a reliable "session is shutting down" signal.
2. **Publish-then-recheck in `registerNewStreamingQuery`** — after the query is inserted into the cache, we re-check `sessionHolder.isClosing`. If the session is closing, we stop the query (asynchronously, so we don't block the caller) and drop the entry. Because `closedTimeMs` is set before `cleanupRunningQueries()` runs, and both the cache publish and `closedTimeMs` are volatile, every interleaving is covered:
- if we observe the session as closing, we stop and drop the query here;
- otherwise `close()` has not set `closedTimeMs` yet, so its `cleanupRunningQueries()` runs after our cache insertion and observes the entry we just published.
`StreamingQuery.stop()` is idempotent and `isActive`-guarded, so both sides firing is harmless.
3. **Identity-based, stop-then-remove async cleanup** — when the recheck fires, we stop the query and remove the cache entry on a `Future` (since `stop()` may block). Two subtleties:
- **Stop before remove.** We drop the entry only *after* the query has actually stopped. Removing it first would discard the only server-side handle to a query that might still be running, re-introducing the very leak this guards against. If `stop()` throws, we leave the entry cached so `periodicMaintenance` (and `cleanupRunningQueries`) can still find and reap it once the query becomes inactive.
- **Match by query identity, not value equality.** Removal uses `queryCache.computeIfPresent` and only nulls the entry when `current.query eq query`. A plain `queryCache.remove(queryKey, value)` by `QueryCacheValue` case-class equality would *fail to remove* the entry if the maintenance thread had concurrently rewritten its `expiresAtMs` after observing the just-stopped query (the value no longer equals the one we inserted). Identity matching removes our entry regardless of such re [...]
4. **`isActive` check at insertion time** — the new entry's `expiresAtMs` is now derived from `query.isActive`. A query that is already inactive at registration time (e.g. a `Trigger.AvailableNow` query that already finished, or one stopped right after `start()`) gets an expiry time immediately, instead of lingering as a falsely "active" entry until a later maintenance cycle notices it stopped.
#### 2. `StreamingForeachBatchHelper.CleanerCache` (the foreachBatch runner cache)
The same shutdown window applies to the Python `foreachBatch` runner cleaner, which is registered (via `registerCleanerForQuery`) immediately after the query is registered. `close()` reaps runners through `CleanerCache.cleanUpAll()`, and the `onQueryTerminated` listener is the other reaper. A cleaner registered for a query started concurrently with `close()` can be missed by `cleanUpAll()`, and if its query already terminated it can also be missed by the listener — stranding a Python [...]
1. **Pre-insert fast path** — if `sessionHolder.isClosing` is already true on entry, we close the runner immediately and return without registering anything (in particular without adding the listener; see below).
2. **Post-insert recheck** — after inserting the cleaner, we re-check `isClosing` and, if the session started closing in the meantime, clean the runner up here and drop the listener.
3. **Listener-lifecycle rework** — the runner clean-up listener was a `lazy val` initialized on first query and never removed. Crucially, `SessionHolder.close()` does **not** remove this listener (it is not tracked in the session's `listenerCache`), so a listener left registered keeps the `CleanerCache` / `SessionHolder` reachable after the session is closed. It is now a `var` managed by `ensureListenerRegistered()` / `removeListenerIfRegistered()`, both guarded by `this`:
- `cleanUpAll()` removes the listener after reaping the runners;
- the post-insert recheck removes it too (the recheck may have just re-added it after `cleanUpAll()` ran);
- the field is **recoverable** — a later registration re-adds a fresh listener, so `cleanUpAll()` does not permanently disable the cache if it is ever reused. Today `cleanUpAll()` is only called on the close path (after which registration fast-paths on `isClosing`), but correctness no longer depends on that.
- `listenerForTesting` now reads the field under the same lock, so concurrent tests see a consistent value rather than a stale/torn read.
### Why are the changes needed?
The race strands streaming queries: they keep running, hold a reference to a closed session, and are never reaped. In production this manifested as Connect structured-streaming jobs in multi-task workflows that complete their computation but never exit — the driver sits idle (0% CPU) for many hours until a manual cancellation or a much longer infra-level cleanup timeout, incurring significant unnecessary cost. The foreachBatch-cleaner guard closes the symmetric leak for Python `foreac [...]
### Does this PR introduce _any_ user-facing change?
No. This fixes an internal resource-cleanup race in the Spark Connect server. There is no API or behavior change visible to users other than streaming queries (and their Python `foreachBatch` workers) no longer being stranded when a session is closed concurrently with a query starting.
### How was this patch tested?
Added 7 unit tests across two suites.
`SparkConnectStreamingQueryCacheSuite`:
- `"Query registered when the session is already closing is stopped and dropped"`
- `"Query registered for a closing session is retained when stopping it fails"`
- `"Query registration racing with session shutdown leaves no query running"` (200-iteration race test)
`StreamingForeachBatchHelperSuite`:
- `"CleanerCache: a runner registered for a closing session is cleaned up immediately"`
- `"CleanerCache.cleanUpAll unregisters the streaming listener"`
- `"CleanerCache: listener is recoverable -- re-registered after cleanUpAll"`
- `"CleanerCache: registration racing with session shutdown strands no runner or listener"` (200-iteration race test)
The existing happy-path tests continue to pass with the `isActive`-based expiry change (an active query is still cached with no expiry).
### Was this patch authored or co-authored using generative AI tooling?
Generated-by: Claude Code (Claude Opus 4.8)
Closes #56377 from dbtsai/spark-connect-fix.
Lead-authored-by: DB Tsai <db...@dbtsai.com>
Co-authored-by: Wenchen Fan <we...@databricks.com>
Signed-off-by: DB Tsai <db...@dbtsai.com>
(cherry picked from commit a852aa3cd81dc23080f714ace207aeb5f7ca6dc8)
Signed-off-by: DB Tsai <db...@dbtsai.com>
---
.../planner/StreamingForeachBatchHelper.scala | 67 +++++++++--
.../spark/sql/connect/service/SessionHolder.scala | 7 ++
.../service/SparkConnectStreamingQueryCache.scala | 80 +++++++++++-
.../planner/StreamingForeachBatchHelperSuite.scala | 114 ++++++++++++++++++
.../SparkConnectStreamingQueryCacheSuite.scala | 134 ++++++++++++++++++++-
5 files changed, 389 insertions(+), 13 deletions(-)
diff --git a/sql/connect/server/src/main/scala/org/apache/spark/sql/connect/planner/StreamingForeachBatchHelper.scala b/sql/connect/server/src/main/scala/org/apache/spark/sql/connect/planner/StreamingForeachBatchHelper.scala
index 97da96894e2a..56973304f41e 100644
--- a/sql/connect/server/src/main/scala/org/apache/spark/sql/connect/planner/StreamingForeachBatchHelper.scala
+++ b/sql/connect/server/src/main/scala/org/apache/spark/sql/connect/planner/StreamingForeachBatchHelper.scala
@@ -215,21 +215,49 @@ object StreamingForeachBatchHelper extends Logging {
// Mapping from streaming (queryId, runId) to runner cleaner. Used for Python foreachBatch.
private val cleanerCache: ConcurrentMap[CacheKey, AutoCloseable] = new ConcurrentHashMap()
- private lazy val streamingListener = { // Initialized on first registered query
- val listener = new StreamingRunnerCleanerListener
- sessionHolder.session.streams.addListener(listener)
- logInfo(
- log"[session: ${MDC(SESSION_ID, sessionHolder.sessionId)}] " +
- log"[userId: ${MDC(USER_ID, sessionHolder.userId)}] " +
- log"Registered runner clean up listener.")
- listener
+ // The runner clean-up listener, registered on first use and removed on cleanup. Both
+ // operations are guarded by `this`, and the field is recoverable: after removal a later
+ // registration re-adds a fresh listener, so cleanUpAll() does not permanently disable the cache
+ // if it is reused. (Today cleanUpAll() is only called on the close path, after which
+ // registration fast-paths on isClosing -- but correctness no longer depends on that.)
+ private var streamingListener: StreamingRunnerCleanerListener = _
+
+ private def ensureListenerRegistered(): Unit = synchronized {
+ if (streamingListener == null) {
+ val listener = new StreamingRunnerCleanerListener
+ sessionHolder.session.streams.addListener(listener)
+ streamingListener = listener
+ logInfo(
+ log"[session: ${MDC(SESSION_ID, sessionHolder.sessionId)}] " +
+ log"[userId: ${MDC(USER_ID, sessionHolder.userId)}] " +
+ log"Registered runner clean up listener.")
+ }
+ }
+
+ // Removes the listener from session.streams if it is currently registered.
+ // SessionHolder.close() does not remove this listener (it is not tracked in the session's
+ // listenerCache), so the cache must drop it on cleanup; otherwise the listener keeps this
+ // CleanerCache / SessionHolder reachable after the session is closed.
+ private def removeListenerIfRegistered(): Unit = synchronized {
+ if (streamingListener != null) {
+ sessionHolder.session.streams.removeListener(streamingListener)
+ streamingListener = null
+ }
}
private[connect] def registerCleanerForQuery(
query: StreamingQuery,
cleaner: AutoCloseable): Unit = {
- streamingListener // Access to initialize
+ // Fast path: if the session is already closing, do not even register the listener (it is
+ // added to session.streams and is not removed by SessionHolder.close(), so it would leak and
+ // keep the closed session reachable). Just close the runner and return.
+ if (sessionHolder.isClosing) {
+ cleaner.close()
+ return
+ }
+
+ ensureListenerRegistered() // Register the runner clean-up listener if not already.
val key = CacheKey(query.id.toString, query.runId.toString)
Option(cleanerCache.putIfAbsent(key, cleaner)) match {
@@ -237,12 +265,27 @@ object StreamingForeachBatchHelper extends Logging {
throw IllegalStateErrors.cleanerAlreadySet(sessionHolder.key.toString, key.toString)
case None => // Inserted. Normal.
}
+
+ // Guard against the same shutdown race that SparkConnectStreamingQueryCache handles for
+ // queries: SessionHolder.close() reaps runners via cleanUpAll(), and a cleaner registered
+ // after that pass (for a query started concurrently with close()) would be missed by it. The
+ // onQueryTerminated listener is the other reaper, but it too can miss a cleaner whose query
+ // already terminated before this registration. So after inserting we re-check isClosing; if
+ // the session started closing in the meantime we clean the runner up here to avoid stranding
+ // a Python worker. We also drop the listener: it was just added above (possibly after close()
+ // ran cleanUpAll()), and close() does not remove it, so it would otherwise leak.
+ if (sessionHolder.isClosing) {
+ cleanupStreamingRunner(key)
+ removeListenerIfRegistered()
+ }
}
/** Cleans up all the registered runners. */
private[connect] def cleanUpAll(): Unit = {
// Clean up all remaining registered runners.
cleanerCache.keySet().asScala.foreach(cleanupStreamingRunner(_))
+ // Drop the listener as well; close() does not remove it otherwise.
+ removeListenerIfRegistered()
}
private def cleanupStreamingRunner(key: CacheKey): Unit = {
@@ -279,6 +322,10 @@ object StreamingForeachBatchHelper extends Logging {
.toMap
}
- private[connect] def listenerForTesting: StreamingQueryListener = streamingListener
+ // Reads the listener under the same lock that guards registration/removal so concurrent tests
+ // see a consistent value rather than a stale/torn read of the field.
+ private[connect] def listenerForTesting: StreamingQueryListener = synchronized {
+ streamingListener
+ }
}
}
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 307416a659f7..2276230545e6 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
@@ -89,6 +89,13 @@ case class SessionHolder(userId: String, sessionId: String, session: SparkSessio
// Set only by close(), and only once.
@volatile private var closedTimeMs: Option[Long] = None
+ /**
+ * Whether this session is closing or already closed. closedTimeMs is set at the very beginning
+ * of [[close]], before any session resources (e.g. running streaming queries) are cleaned up,
+ * so this can be used to detect a session shutdown that races with newly started operations.
+ */
+ private[connect] def isClosing: Boolean = closedTimeMs.isDefined
+
// Custom timeout after a session expires due to inactivity.
// Used by SparkConnectSessionManager instead of default timeout if set.
// Setting it to -1 indicated forever.
diff --git a/sql/connect/server/src/main/scala/org/apache/spark/sql/connect/service/SparkConnectStreamingQueryCache.scala b/sql/connect/server/src/main/scala/org/apache/spark/sql/connect/service/SparkConnectStreamingQueryCache.scala
index 8ff13f5afe19..17f94f56f11c 100644
--- a/sql/connect/server/src/main/scala/org/apache/spark/sql/connect/service/SparkConnectStreamingQueryCache.scala
+++ b/sql/connect/server/src/main/scala/org/apache/spark/sql/connect/service/SparkConnectStreamingQueryCache.scala
@@ -18,7 +18,7 @@
package org.apache.spark.sql.connect.service
import java.util.concurrent.{ConcurrentHashMap, ConcurrentMap, Executors, ScheduledExecutorService, TimeUnit}
-import java.util.concurrent.atomic.AtomicReference
+import java.util.concurrent.atomic.{AtomicBoolean, AtomicReference}
import scala.collection.mutable
import scala.concurrent.{ExecutionContext, Future}
@@ -59,13 +59,20 @@ private[connect] class SparkConnectStreamingQueryCache(
query: StreamingQuery,
tags: Set[String],
operationId: String): Unit = {
+ // If the query is already inactive by the time it is registered (e.g. a Trigger.AvailableNow
+ // query that already finished, or a query that was stopped right after start()), set its
+ // expiry time immediately so that it is reaped on the regular schedule, instead of lingering
+ // in the cache as a falsely "active" entry until a later maintenance cycle notices it stopped.
+ val expiresAtMs =
+ if (query.isActive) None
+ else Some(clock.getTimeMillis() + stoppedQueryInactivityTimeout.toMillis)
val value = QueryCacheValue(
userId = sessionHolder.userId,
sessionId = sessionHolder.sessionId,
session = sessionHolder.session,
query = query,
operationId = operationId,
- expiresAtMs = None)
+ expiresAtMs = expiresAtMs)
val queryKey = QueryCacheKey(query.id.toString, query.runId.toString)
tags.foreach { tag => addTaggedQuery(tag, queryKey) }
@@ -86,6 +93,66 @@ private[connect] class SparkConnectStreamingQueryCache(
})
schedulePeriodicChecks() // Start the scheduler thread if it has not been started.
+
+ // Guard against a race with session shutdown. SessionHolder.close() stops all of a session's
+ // streaming queries through cleanupRunningQueries(), which iterates over this cache. A query
+ // registered *after* that iteration would otherwise be missed and left running, holding a
+ // strong reference to the now-closed session so that the driver can never exit.
+ //
+ // We publish the entry with queryCache.compute() *first* and only then read
+ // sessionHolder.isClosing; close() writes the volatile closedTimeMs *before* it iterates this
+ // cache in cleanupRunningQueries(). Because both the cache publish and closedTimeMs are
+ // volatile, at least one side always observes the other:
+ // - if the cleanup forEach observes our entry, cleanupRunningQueries() stops the query;
+ // - otherwise the forEach missed it, which means our compute() linearized *after* close()'s
+ // read of that key's bin. Since closedTimeMs was written before that bin read (program
+ // order in close()) and our compute() precedes the isClosing read (program order here),
+ // transitivity guarantees we observe isClosing == true and stop the query ourselves.
+ // Either one or both sides stop the query. StreamingQuery.stop() and the cache removal below
+ // are idempotent and isActive-guarded, so both sides firing is harmless.
+ if (sessionHolder.isClosing) {
+ logWarning(
+ log"Stopping streaming query registered for a closing session. " +
+ log"Query Id: ${MDC(QUERY_ID, query.id)}, " +
+ log"runId: ${MDC(QUERY_RUN_ID, query.runId)}, " +
+ log"session ${MDC(SESSION_ID, sessionHolder.sessionId)}.")
+ // Stop asynchronously (stop() may block) and drop the cache entry only after the query has
+ // actually been stopped. Removing it before the stop succeeds would discard the only
+ // server-side handle to a query that might still be running, re-introducing the leak this
+ // guards against. If the stop fails we keep the entry cached so periodicMaintenance can reap
+ // it once the query becomes inactive (and so cleanupRunningQueries can still find it).
+ Future {
+ try {
+ if (query.isActive) query.stop()
+ // Drop only the entry we inserted, matched by query identity. Identity (rather than
+ // queryCache.remove(queryKey, value) by case-class equality) ensures we still remove the
+ // entry if the maintenance thread concurrently rewrote its expiresAtMs after observing
+ // the just-stopped query, while still never evicting a later replacement for the same key
+ // (queryCache.compute allows replacement, though it is not expected).
+ val removed = new AtomicBoolean(false)
+ queryCache.computeIfPresent(
+ queryKey,
+ (_, current) => {
+ if (current.query eq query) {
+ removed.set(true)
+ null
+ } else {
+ current
+ }
+ })
+ if (removed.get()) {
+ tags.foreach { tag => removeTaggedQuery(tag, queryKey) }
+ }
+ } catch {
+ case NonFatal(ex) =>
+ logWarning(
+ log"Failed to stop streaming query ${MDC(QUERY_ID, query.id)} " +
+ log"runId: ${MDC(QUERY_RUN_ID, query.runId)} " +
+ log"for a closing session; leaving it cached for later cleanup.",
+ ex)
+ }
+ }(ExecutionContext.global)
+ }
}
/**
@@ -227,6 +294,15 @@ private[connect] class SparkConnectStreamingQueryCache(
})
}
+ private def removeTaggedQuery(tag: String, queryKey: QueryCacheKey): Unit = {
+ taggedQueries.computeIfPresent(
+ tag,
+ (k, v) => {
+ // removeKey returns true once the set is empty; drop the tag entry in that case.
+ if (v.removeKey(queryKey)) null else v
+ })
+ }
+
/**
* Periodic maintenance task to do the following:
* - Update status of query if it is inactive. Sets an expiry time for such queries
diff --git a/sql/connect/server/src/test/scala/org/apache/spark/sql/connect/planner/StreamingForeachBatchHelperSuite.scala b/sql/connect/server/src/test/scala/org/apache/spark/sql/connect/planner/StreamingForeachBatchHelperSuite.scala
index 657be7b7954b..6a9656e5d6f8 100644
--- a/sql/connect/server/src/test/scala/org/apache/spark/sql/connect/planner/StreamingForeachBatchHelperSuite.scala
+++ b/sql/connect/server/src/test/scala/org/apache/spark/sql/connect/planner/StreamingForeachBatchHelperSuite.scala
@@ -17,19 +17,31 @@
package org.apache.spark.sql.connect.planner
import java.util.UUID
+import java.util.concurrent.CountDownLatch
+import org.mockito.Mockito.atLeastOnce
import org.mockito.Mockito.times
import org.mockito.Mockito.verify
import org.mockito.Mockito.when
import org.scalatestplus.mockito.MockitoSugar
import org.apache.spark.sql.connect.SparkConnectTestUtils
+import org.apache.spark.sql.connect.service.SessionHolder
import org.apache.spark.sql.streaming.StreamingQuery
import org.apache.spark.sql.streaming.StreamingQueryListener
import org.apache.spark.sql.test.SharedSparkSession
class StreamingForeachBatchHelperSuite extends SharedSparkSession with MockitoSugar {
+ // A session holder that is NOT registered in the global SparkConnectService.sessionManager
+ // (unlike SparkConnectTestUtils.createDummySessionHolder). The closing-session tests below call
+ // sessionHolder.close() directly, bypassing the manager; a holder registered in the manager's
+ // store at that point would stay there as a closed entry, and a later suite's
+ // invalidateAllSessions() would fail closing it a second time (SESSION_ALREADY_CLOSED).
+ private def unregisteredSessionHolder(): SessionHolder = {
+ SessionHolder(userId = "testUser", sessionId = UUID.randomUUID().toString, session = spark)
+ }
+
private def mockQuery(): StreamingQuery = {
val query = mock[StreamingQuery]
val (queryId, runId) = (UUID.randomUUID(), UUID.randomUUID())
@@ -78,4 +90,106 @@ class StreamingForeachBatchHelperSuite extends SharedSparkSession with MockitoSu
// No more entries left in it now.
assert(cache.listEntriesForTesting().isEmpty)
}
+
+ test("CleanerCache: a runner registered for a closing session is cleaned up immediately") {
+ // Mirrors the SparkConnectStreamingQueryCache shutdown-race guard. A runner registered after
+ // SessionHolder.close() has already run cleanUpAll() (for a query started concurrently with
+ // close()) would otherwise be missed by both reapers and strand a Python worker.
+ val cleaner = mock[AutoCloseable]
+ val query = mockQuery()
+ val sessionHolder = unregisteredSessionHolder()
+ val cache = new StreamingForeachBatchHelper.CleanerCache(sessionHolder)
+
+ // Mark the session as closing before the runner is registered.
+ sessionHolder.close()
+ assert(sessionHolder.isClosing)
+
+ val listenersBefore = spark.streams.listListeners().length
+ cache.registerCleanerForQuery(query, cleaner)
+
+ // The runner must not be stranded: it is closed and never cached.
+ verify(cleaner, times(1)).close()
+ assert(cache.listEntriesForTesting().isEmpty)
+ // The fast path must not register a (leaking) listener on the closing session's streams.
+ assert(spark.streams.listListeners().length == listenersBefore)
+ }
+
+ test("CleanerCache.cleanUpAll unregisters the streaming listener") {
+ // close() does not remove the StreamingRunnerCleanerListener (it is not tracked in the
+ // session's listenerCache), so cleanUpAll() must drop it; otherwise the listener keeps the
+ // cache / session reachable after the session is closed.
+ val cleaner = mock[AutoCloseable]
+ val query = mockQuery()
+ val cache = new StreamingForeachBatchHelper.CleanerCache(
+ SparkConnectTestUtils.createDummySessionHolder(spark))
+
+ cache.registerCleanerForQuery(query, cleaner)
+ val listener = cache.listenerForTesting
+ assert(spark.streams.listListeners().contains(listener))
+
+ cache.cleanUpAll()
+
+ verify(cleaner, times(1)).close()
+ assert(!spark.streams.listListeners().contains(listener))
+ }
+
+ test("CleanerCache: listener is recoverable -- re-registered after cleanUpAll") {
+ // streamingListener is no longer a one-shot lazy val: after cleanUpAll() removes it, a later
+ // registration must re-add a working listener so the cache is safe to reuse.
+ val cache = new StreamingForeachBatchHelper.CleanerCache(
+ SparkConnectTestUtils.createDummySessionHolder(spark))
+
+ cache.registerCleanerForQuery(mockQuery(), mock[AutoCloseable])
+ val firstListener = cache.listenerForTesting
+ assert(spark.streams.listListeners().contains(firstListener))
+
+ cache.cleanUpAll()
+ assert(!spark.streams.listListeners().contains(firstListener))
+
+ // Reuse: a new registration re-registers a listener on session.streams.
+ cache.registerCleanerForQuery(mockQuery(), mock[AutoCloseable])
+ val secondListener = cache.listenerForTesting
+ assert(spark.streams.listListeners().contains(secondListener))
+
+ cache.cleanUpAll()
+ assert(!spark.streams.listListeners().contains(secondListener))
+ }
+
+ test("CleanerCache: registration racing with session shutdown strands no runner or listener") {
+ // Mirrors the SparkConnectStreamingQueryCache race test for the foreachBatch cleaner:
+ // registration runs concurrently with the shutdown sequence (close() marks the session closing,
+ // then cleanUpAll() reaps runners and the listener). Whatever the interleaving, the runner must
+ // be closed and no listener may be left registered on session.streams.
+ val baselineListeners = spark.streams.listListeners().length
+ val numIterations = 200
+ (1 to numIterations).foreach { _ =>
+ val cleaner = mock[AutoCloseable]
+ val query = mockQuery()
+ val sessionHolder = unregisteredSessionHolder()
+ val cache = new StreamingForeachBatchHelper.CleanerCache(sessionHolder)
+
+ val startLatch = new CountDownLatch(1)
+ val closeThread = new Thread(() => {
+ startLatch.await()
+ sessionHolder.close() // Marks the session closing.
+ cache.cleanUpAll() // Mirrors close()'s runner + listener reaping.
+ })
+ val registerThread = new Thread(() => {
+ startLatch.await()
+ cache.registerCleanerForQuery(query, cleaner)
+ })
+ closeThread.start()
+ registerThread.start()
+ startLatch.countDown()
+ closeThread.join()
+ registerThread.join()
+
+ // The runner must be closed by one of the paths: the fast path, the post-insert guard, or
+ // cleanUpAll(). registerCleanerForQuery and cleanUpAll are synchronous, so this is settled
+ // once both threads have joined.
+ verify(cleaner, atLeastOnce()).close()
+ }
+ // No iteration may leave a listener registered on the shared streams manager.
+ assert(spark.streams.listListeners().length == baselineListeners)
+ }
}
diff --git a/sql/connect/server/src/test/scala/org/apache/spark/sql/connect/service/SparkConnectStreamingQueryCacheSuite.scala b/sql/connect/server/src/test/scala/org/apache/spark/sql/connect/service/SparkConnectStreamingQueryCacheSuite.scala
index 6a133f87c1b1..1a335c879a8c 100644
--- a/sql/connect/server/src/test/scala/org/apache/spark/sql/connect/service/SparkConnectStreamingQueryCacheSuite.scala
+++ b/sql/connect/server/src/test/scala/org/apache/spark/sql/connect/service/SparkConnectStreamingQueryCacheSuite.scala
@@ -18,10 +18,11 @@
package org.apache.spark.sql.connect.service
import java.util.UUID
+import java.util.concurrent.CountDownLatch
import scala.concurrent.duration.DurationInt
-import org.mockito.Mockito.when
+import org.mockito.Mockito.{atLeastOnce, doThrow, verify, when}
import org.scalatest.concurrent.Eventually.eventually
import org.scalatest.concurrent.Futures.timeout
import org.scalatestplus.mockito.MockitoSugar
@@ -156,4 +157,135 @@ class SparkConnectStreamingQueryCacheSuite extends SparkFunSuite with MockitoSug
}
sessionMgr.shutdown()
}
+
+ test("Query registered when the session is already closing is stopped and dropped") {
+ // Tests the closing-session guard in registerNewStreamingQuery in isolation (the session is
+ // already marked closing before the query is registered). The concurrent race is covered by the
+ // next test.
+
+ val queryId = UUID.randomUUID().toString
+ val runId = UUID.randomUUID().toString
+ val tag = "test_tag"
+ val mockSession = mock[SparkSession]
+ val mockQuery = mock[StreamingQuery]
+ val mockStreamingQueryManager = mock[StreamingQueryManager]
+
+ val sessionHolder =
+ SessionHolder(userId = "test_user_1", sessionId = "test_session_1", session = mockSession)
+
+ val sessionMgr = createSessionManager()
+
+ when(mockQuery.id).thenReturn(UUID.fromString(queryId))
+ when(mockQuery.runId).thenReturn(UUID.fromString(runId))
+ when(mockQuery.isActive).thenReturn(true)
+ when(mockSession.streams).thenReturn(mockStreamingQueryManager)
+ when(mockStreamingQueryManager.get(queryId)).thenReturn(mockQuery)
+
+ // Mark the session as closing. close() sets closedTimeMs and, for a not-yet-started session in
+ // tests, returns early without running the rest of the cleanup.
+ sessionHolder.close()
+ assert(sessionHolder.isClosing)
+
+ sessionMgr.registerNewStreamingQuery(sessionHolder, mockQuery, Set(tag), "")
+
+ // The query must be stopped and, once stopped, dropped from the cache (directly and by tag).
+ // The cleanup removes the entry by query identity (computeIfPresent matching current.query eq
+ // query), not by case-class value equality, so removal still succeeds even if the maintenance
+ // thread concurrently rewrites the entry's expiresAtMs after seeing the just-stopped query.
+ // That specific maintenance-vs-cleanup interleaving is not exercised as a separate test
+ // because it is not deterministically reproducible; the identity match makes it correct.
+ eventually(timeout(1.minute)) {
+ verify(mockQuery).stop()
+ assert(sessionMgr.getCachedValue(queryId, runId).isEmpty)
+ assert(!sessionMgr.taggedQueries.containsKey(tag))
+ }
+ sessionMgr.shutdown()
+ }
+
+ test("Query registered for a closing session is retained when stopping it fails") {
+ // If stopping the query fails, the cache entry must NOT be dropped: removing it would discard
+ // the only server-side handle to a possibly still-running query, re-creating the leak. The
+ // entry is kept so a later cleanup / maintenance pass can reap it.
+ val queryId = UUID.randomUUID().toString
+ val runId = UUID.randomUUID().toString
+ val mockSession = mock[SparkSession]
+ val mockQuery = mock[StreamingQuery]
+ val mockStreamingQueryManager = mock[StreamingQueryManager]
+
+ val sessionHolder =
+ SessionHolder(userId = "test_user_1", sessionId = "test_session_1", session = mockSession)
+
+ val sessionMgr = createSessionManager()
+
+ when(mockQuery.id).thenReturn(UUID.fromString(queryId))
+ when(mockQuery.runId).thenReturn(UUID.fromString(runId))
+ when(mockQuery.isActive).thenReturn(true)
+ when(mockSession.streams).thenReturn(mockStreamingQueryManager)
+ when(mockStreamingQueryManager.get(queryId)).thenReturn(mockQuery)
+ doThrow(new RuntimeException("stop failed")).when(mockQuery).stop()
+
+ sessionHolder.close()
+ sessionMgr.registerNewStreamingQuery(sessionHolder, mockQuery, Set("test_tag"), "")
+
+ // The stop is attempted...
+ eventually(timeout(1.minute)) {
+ verify(mockQuery).stop()
+ }
+ // ...but because it failed, the entry is retained rather than dropped.
+ assert(sessionMgr.getCachedValue(queryId, runId).nonEmpty)
+ sessionMgr.shutdown()
+ }
+
+ test("Query registration racing with session shutdown leaves no query running") {
+ // Exercises the actual race: registerNewStreamingQuery runs concurrently with the session
+ // shutdown sequence (close() sets closedTimeMs, then cleanupRunningQueries() stops the
+ // session's queries by iterating the cache). Whatever the interleaving, the query must end up
+ // stopped and never stranded (left running while holding a reference to the closed session).
+ val sessionMgr = createSessionManager()
+ val numIterations = 200
+ try {
+ (1 to numIterations).foreach { i =>
+ val queryId = UUID.randomUUID().toString
+ val runId = UUID.randomUUID().toString
+ val mockSession = mock[SparkSession]
+ val mockQuery = mock[StreamingQuery]
+ val mockStreamingQueryManager = mock[StreamingQueryManager]
+ when(mockQuery.id).thenReturn(UUID.fromString(queryId))
+ when(mockQuery.runId).thenReturn(UUID.fromString(runId))
+ when(mockQuery.isActive).thenReturn(true)
+ when(mockSession.streams).thenReturn(mockStreamingQueryManager)
+ when(mockStreamingQueryManager.get(queryId)).thenReturn(mockQuery)
+
+ val sessionHolder = SessionHolder(
+ userId = "test_user",
+ sessionId = s"test_session_$i",
+ session = mockSession)
+
+ // Release both threads together to maximize the chance of interleaving.
+ val startLatch = new CountDownLatch(1)
+ val closeThread = new Thread(() => {
+ startLatch.await()
+ sessionHolder.close() // Sets closedTimeMs, i.e. isClosing.
+ sessionMgr.cleanupRunningQueries(sessionHolder) // Mirrors close()'s query cleanup.
+ })
+ val registerThread = new Thread(() => {
+ startLatch.await()
+ sessionMgr.registerNewStreamingQuery(sessionHolder, mockQuery, Set.empty[String], "")
+ })
+ closeThread.start()
+ registerThread.start()
+ startLatch.countDown()
+ closeThread.join()
+ registerThread.join()
+
+ // Whatever the interleaving, the query must have been stopped by one of the two paths
+ // (registration's closing-session branch stops asynchronously, hence eventually()).
+ eventually(timeout(10.seconds)) {
+ verify(mockQuery, atLeastOnce()).stop()
+ }
+ }
+ } finally {
+ sessionMgr.shutdown()
+ }
+ }
}
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@spark.apache.org
For additional commands, e-mail: commits-help@spark.apache.org