This is an automated email from the ASF dual-hosted git repository.
dbtsai pushed a commit to branch branch-4.2
in repository https://gitbox.apache.org/repos/asf/spark.git
The following commit(s) were added to refs/heads/branch-4.2 by this push:
new eca044702155 [SPARK-57325][CONNECT] Stop streaming queries registered
while the Connect session is closing
eca044702155 is described below
commit eca044702155cc146ea6673f5713b2c88bd679db
Author: DB Tsai <[email protected]>
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 <[email protected]>
Co-authored-by: Wenchen Fan <[email protected]>
Signed-off-by: DB Tsai <[email protected]>
(cherry picked from commit a852aa3cd81dc23080f714ace207aeb5f7ca6dc8)
Signed-off-by: DB Tsai <[email protected]>
---
.../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: [email protected]
For additional commands, e-mail: [email protected]
