HyukjinKwon commented on code in PR #59277:
URL: https://github.com/apache/spark/pull/59277#discussion_r4241386214
##########
sql/connect/server/src/main/scala/org/apache/spark/sql/connect/service/SparkConnectSessionManager.scala:
##########
@@ -224,44 +224,58 @@ class SparkConnectSessionManager extends Logging {
session
}
- // Removes session from sessionStore and returns it.
- private def removeSessionHolder(
- key: SessionKey,
- allowReconnect: Boolean = false): Option[SessionHolder] = {
- var sessionHolder: Option[SessionHolder] = None
-
- // The session holder should remain in the session store until it is added
to the closed session
- // cache, because of a subtle data race: a new session with the same key
can be created if the
- // closed session cache does not contain the key right after the key has
been removed from the
- // session store.
- sessionHolder = Option(sessionStore.get(key))
-
- sessionHolder.foreach { s =>
- if (!allowReconnect) {
- // Put into closedSessionsCache to prevent the same session from being
recreated by
- // getOrCreateIsolatedSession when reconnection isn't allowed.
- closedSessionsCache.put(s.key, s.getSessionHolderInfo)
- }
+ // Removes the session with the given key from sessionStore, if any, and
returns it.
+ private[service] def removeSessionHolder(key: SessionKey):
Option[SessionHolder] = {
+ Option(sessionStore.get(key)).filter(tryRemoveSessionHolder)
+ }
- // Then, remove the session holder from the session store.
- sessionStore.remove(key)
- }
- sessionHolder
+ // Removes the given session from sessionStore if it is still stored there.
Returns true only
+ // for the one caller that removed it; that caller must then call
shutdownSessionHolder.
+ //
+ // The session is tombstoned before it is removed, and the tombstone is kept
until
+ // shutdownSessionHolder has finished closing it. Otherwise a session with
the same key could be
+ // created while the old session is still closing, and the key-based cleanup
in
+ // SessionHolder.close() would tear it down.
+ private[service] def tryRemoveSessionHolder(sessionHolder: SessionHolder):
Boolean = {
+ // Tombstone and remove the session in one atomic step for the key,
matched by identity, so
+ // that two callers cannot both remove the same session, a caller that
loses cannot re-put a
+ // tombstone that the owner already dropped, and a later session with the
same key is never
+ // removed.
+ val removed = new AtomicBoolean(false)
+ sessionStore.computeIfPresent(
+ sessionHolder.key,
+ (_, current) => {
+ if (current eq sessionHolder) {
+ closedSessionsCache.put(current.key, current.getSessionHolderInfo)
Review Comment:
**Nit (P3):** `tryRemoveSessionHolder` now puts every closing session into
`closedSessionsCache`, including closes that drop the tombstone again
afterwards: `ReleaseSession(allow_reconnect=true)`, opted-in sessions, and the
foreachBatch cloned-session close in `StreamingForeachBatchHelper.close()`.
That last one passes `allowReconnect = true` precisely because "a tombstone in
closedSessionsCache would only take a slot from a real client session". The
cache is a Guava cache bounded by `closedSessionsTombstonesSize` (default
1000). Once it is full, each of these puts evicts another client's tombstone,
and the `invalidate` after `close()` removes only the transient entry. Under
load, released ids lose their `SESSION_CLOSED` protection sooner, in proportion
to foreachBatch churn and reconnect-allowed closes. Before this PR, these
closes never wrote to the cache.
**Recommended change:** Keep reconnect-allowed closes out of
closedSessionsCache: decide drop-or-keep when the session is removed, fence
reconnect-allowed closing sessions in a separate in-close map, and have the
same-key creation checks consult it. Add a session manager test that a
reconnect-allowed close adds no tombstone-cache entry while still blocking
same-key creation until it finishes.
**Why this works:** tryRemoveSessionHolder takes the release's
allowReconnect and computes dropTombstone = allowReconnect ||
holder.allowReconnectAfterClose inside the identity-matched computeIfPresent. A
close that keeps a tombstone puts it into closedSessionsCache as in the merge
target. A reconnect-allowed close instead records the holder in a
ConcurrentHashMap of closing sessions. validateSessionCreate and
getIsolatedSession treat a key present in either structure as closed. After
close() returns, shutdownSessionHolder removes the closing entry with
remove(key, holder) on the drop path, or refreshes the tombstone on the keep
path. The closing map holds at most one entry per in-flight close, so it needs
no bound and never evicts client tombstones.
**Scope:** Move the in-close fence for reconnect-allowed closes out of the
bounded tombstone cache in the Connect session manager, with a session
lifecycle test.
**Compatibility:** Tombstone-until-close and single-owner close semantics
established in this PR.
**Risks:** Every place that reads closedSessionsCache for same-key
protection during close must also consult the closing map, or a
reconnect-allowed session becomes creatable mid-close.
**Constraints:** Keep review-frame:tombstone-until-close and
review-frame:single-owner-close: same-key creation fails with
INVALID_HANDLE.SESSION_CLOSED until the owning close() returns, and only the
removal winner runs shutdownSessionHolder.
**Success:** A reconnect-allowed close
(ReleaseSession(allow_reconnect=true), an opted-in session's close, or the
foreachBatch cloned-session close) never inserts into closedSessionsCache, so
it cannot evict another session's tombstone. Same-key creation still fails with
INVALID_HANDLE.SESSION_CLOSED while any close is running, and succeeds once a
reconnect-allowed close has returned. Closes that keep a tombstone behave as in
the merge target.
**Verification:**
- **Regression:** While a reconnect-allowed close is held open, the
tombstone cache gains no entry for its key (for example, listClosedSessions
does not contain it), yet same-key creation fails with SESSION_CLOSED until the
close returns.
- **Behavior:** After an opted-in or allow_reconnect close returns, the id
can be used for a new, empty session.
- **Compatibility:** Default releases and expiries still tombstone the id
after close, and listClosedSessions still reports them.
##########
sql/connect/server/src/main/scala/org/apache/spark/sql/connect/service/SparkConnectSessionManager.scala:
##########
@@ -437,13 +454,21 @@ class SparkConnectSessionManager extends Logging {
})
}
+ /**
+ * Used for testing.
+ */
+ private[service] def putSessionHolderForTesting(sessionHolder:
SessionHolder): Unit = {
Review Comment:
**Nit (P3):** This has the same body as the existing `putSessionForTesting`
a few lines below. That hook is `private[connect]`, so this suite's package can
already reach it, and `SparkConnectTestUtils` already uses it. Could
`startBlockingCloseSession` call `putSessionForTesting` instead, so there is
only one test hook?
**Verification:**
- **Inspection:** No putSessionHolderForTesting method remains, and the
blocking-close setup uses putSessionForTesting.
- **Behavior:** The blocking-close tests still pass with the existing hook.
##########
sql/connect/server/src/main/scala/org/apache/spark/sql/connect/config/Connect.scala:
##########
@@ -150,6 +150,18 @@ object Connect {
.intConf
.createWithDefaultString("1000")
+ val CONNECT_SESSION_ALLOW_RECONNECT_AFTER_CLOSE =
+ buildConf("spark.connect.session.allowReconnectAfterClose")
+ .doc(
+ "When true, the session id is not tombstoned when this session is
closed, for example " +
Review Comment:
**Nit (P3):** The doc says the id "is not tombstoned when this session is
closed". In fact `tryRemoveSessionHolder` tombstones it on removal, and
`shutdownSessionHolder` only drops the tombstone after `close()` returns, so
requests during the close still get `INVALID_HANDLE.SESSION_CLOSED`; the new
blocking-close test asserts exactly that. Also, a client object that used the
closed session keeps sending its observed server-side session id and gets
`INVALID_HANDLE.SESSION_CHANGED`; only a new client gets the new session. Could
the doc say that the id is released once the close finishes and that the old
client has to be replaced? Then opted-in clients know to retry `SESSION_CLOSED`
and to start a new `SparkSession`.
See **Shared repair plan 1** in the review body.
##########
sql/connect/server/src/test/scala/org/apache/spark/sql/connect/service/SparkConnectSessionManagerSuite.scala:
##########
@@ -159,6 +164,143 @@ class SparkConnectSessionManagerSuite extends
SharedSparkSession {
assert(closedSessionInfo.get.closedTimeMs.isDefined)
}
+ private def assertSessionClosed(key: SessionKey): Unit = {
+ val ex = intercept[SparkSQLException] {
+ SparkConnectService.sessionManager.getOrCreateIsolatedSession(key, None)
+ }
+ assert(ex.getCondition == "INVALID_HANDLE.SESSION_CLOSED")
+ }
+
+ private def releaseSession(key: SessionKey, allowReconnect: Boolean): Unit =
{
+ val observer = new StreamObserver[proto.ReleaseSessionResponse] {
+ override def onNext(v: proto.ReleaseSessionResponse): Unit = {}
+ override def onError(t: Throwable): Unit = throw t
+ override def onCompleted(): Unit = {}
+ }
+ val request = proto.ReleaseSessionRequest
+ .newBuilder()
+ .setUserContext(proto.UserContext.newBuilder().setUserId(key.userId))
+ .setSessionId(key.sessionId)
+ .setAllowReconnect(allowReconnect)
+ .build()
+ new SparkConnectReleaseSessionHandler(observer).handle(request)
+ }
+
+ test("expired session is tombstoned by default") {
+ val key = SessionKey("user", UUID.randomUUID().toString)
+ SparkConnectService.sessionManager.getOrCreateIsolatedSession(key, None)
+
+
SparkConnectService.sessionManager.periodicMaintenance(defaultInactiveTimeoutMs
= 0L)
+
+ assertSessionClosed(key)
+ }
+
+ test("expired session that allows reconnect after close can be recreated as
a new session") {
+ val key = SessionKey("user", UUID.randomUUID().toString)
+ val oldHolder =
SparkConnectService.sessionManager.getOrCreateIsolatedSession(key, None)
+
oldHolder.session.conf.set(Connect.CONNECT_SESSION_ALLOW_RECONNECT_AFTER_CLOSE.key,
"true")
+ oldHolder.session.conf.set("spark.sql.shuffle.partitions", "7")
+
+
SparkConnectService.sessionManager.periodicMaintenance(defaultInactiveTimeoutMs
= 0L)
+ assert(SparkConnectService.sessionManager.listActiveSessions.isEmpty)
+
+ // The id is reusable, and the recreated session starts from a clean state.
+ val newHolder =
SparkConnectService.sessionManager.getOrCreateIsolatedSession(key, None)
+ assert(newHolder.session.conf.get("spark.sql.shuffle.partitions") != "7")
+ // Opting in is per session: the new session is tombstoned again unless it
opts in too.
+
SparkConnectService.sessionManager.periodicMaintenance(defaultInactiveTimeoutMs
= 0L)
+ assertSessionClosed(key)
+ }
+
+ /** A session holder whose close() blocks until `release` is counted down. */
+ private class BlockingCloseSessionHolder(key: SessionKey)
+ extends SessionHolder(key.userId, key.sessionId, spark.newSession()) {
+ val closeStarted = new CountDownLatch(1)
+ val release = new CountDownLatch(1)
+ val closeCalls = new AtomicInteger(0)
+
+ override private[connect] def close(): Unit = {
+ closeCalls.incrementAndGet()
+ closeStarted.countDown()
+ assert(release.await(1, TimeUnit.MINUTES))
+ super.close()
+ }
+ }
+
+ private def startBlockingCloseSession(key: SessionKey):
BlockingCloseSessionHolder = {
+ val holder = new BlockingCloseSessionHolder(key)
+ holder.initializeSession()
+
holder.session.conf.set(Connect.CONNECT_SESSION_ALLOW_RECONNECT_AFTER_CLOSE.key,
"true")
+ SparkConnectService.sessionManager.putSessionHolderForTesting(holder)
+ holder
+ }
+
+ test("session id stays tombstoned until the owning close completes") {
+ val sessionManager = SparkConnectService.sessionManager
+ val key = SessionKey("user", UUID.randomUUID().toString)
+ val holder = startBlockingCloseSession(key)
+
+ val owner = new Thread(() => sessionManager.closeSession(key))
+ owner.start()
+ try {
+ assert(holder.closeStarted.await(1, TimeUnit.MINUTES))
+ // The old session is removed but still closing: it must not be
recreated in this window.
+ assertSessionClosed(key)
+
+ // Callers that lose the removal, such as a concurrent ReleaseSession or
a maintenance run
+ // that observed the holder before it was removed, must not close it
again or touch its
+ // tombstone.
+ assert(!sessionManager.tryRemoveSessionHolder(holder))
+ sessionManager.closeSession(key)
+ sessionManager.periodicMaintenance(defaultInactiveTimeoutMs = 0L)
Review Comment:
**Nit (P3):** By the time this runs, the owner thread has already removed
the holder from `sessionStore`, since removal happens before `close()` starts.
So `periodicMaintenance`'s `forEach` never sees the holder, and the
`closeSession(key)` above finds nothing either. Only the direct
`tryRemoveSessionHolder(holder)` call exercises a losing removal. If
`periodicMaintenance` went back to shutting down the holder it observed without
checking `tryRemoveSessionHolder`, every test would still pass. A test where
maintenance reads the holder and then loses the removal to a concurrent
`closeSession` would cover the race fixed in this round.
**Verification:**
- **Regression:** When maintenance observes an expired holder and another
caller removes it first, close() runs exactly once and same-key creation keeps
failing with SESSION_CLOSED until the owner's close returns.
##########
sql/connect/server/src/main/scala/org/apache/spark/sql/connect/config/Connect.scala:
##########
@@ -150,6 +150,18 @@ object Connect {
.intConf
.createWithDefaultString("1000")
+ val CONNECT_SESSION_ALLOW_RECONNECT_AFTER_CLOSE =
+ buildConf("spark.connect.session.allowReconnectAfterClose")
Review Comment:
**Nit (P3):** This conf is user-facing (not `.internal()`) and meant to be
set by clients, but it isn't in the Spark Connect table in
`docs/configuration.md`. Every other non-internal Connect conf except one is
listed there, including session-level ones like
`spark.sql.connect.enrichError.enabled`. Could you add a row so users can find
the opt-in and its caveats?
See **Shared repair plan 1** in the review body.
##########
sql/connect/server/src/main/scala/org/apache/spark/sql/connect/service/SparkConnectSessionManager.scala:
##########
@@ -224,44 +224,58 @@ class SparkConnectSessionManager extends Logging {
session
}
- // Removes session from sessionStore and returns it.
- private def removeSessionHolder(
- key: SessionKey,
- allowReconnect: Boolean = false): Option[SessionHolder] = {
- var sessionHolder: Option[SessionHolder] = None
-
- // The session holder should remain in the session store until it is added
to the closed session
- // cache, because of a subtle data race: a new session with the same key
can be created if the
- // closed session cache does not contain the key right after the key has
been removed from the
- // session store.
- sessionHolder = Option(sessionStore.get(key))
-
- sessionHolder.foreach { s =>
- if (!allowReconnect) {
- // Put into closedSessionsCache to prevent the same session from being
recreated by
- // getOrCreateIsolatedSession when reconnection isn't allowed.
- closedSessionsCache.put(s.key, s.getSessionHolderInfo)
- }
+ // Removes the session with the given key from sessionStore, if any, and
returns it.
+ private[service] def removeSessionHolder(key: SessionKey):
Option[SessionHolder] = {
+ Option(sessionStore.get(key)).filter(tryRemoveSessionHolder)
+ }
- // Then, remove the session holder from the session store.
- sessionStore.remove(key)
- }
- sessionHolder
+ // Removes the given session from sessionStore if it is still stored there.
Returns true only
+ // for the one caller that removed it; that caller must then call
shutdownSessionHolder.
+ //
+ // The session is tombstoned before it is removed, and the tombstone is kept
until
+ // shutdownSessionHolder has finished closing it. Otherwise a session with
the same key could be
+ // created while the old session is still closing, and the key-based cleanup
in
+ // SessionHolder.close() would tear it down.
+ private[service] def tryRemoveSessionHolder(sessionHolder: SessionHolder):
Boolean = {
+ // Tombstone and remove the session in one atomic step for the key,
matched by identity, so
+ // that two callers cannot both remove the same session, a caller that
loses cannot re-put a
+ // tombstone that the owner already dropped, and a later session with the
same key is never
+ // removed.
+ val removed = new AtomicBoolean(false)
+ sessionStore.computeIfPresent(
+ sessionHolder.key,
+ (_, current) => {
+ if (current eq sessionHolder) {
+ closedSessionsCache.put(current.key, current.getSessionHolderInfo)
+ removed.set(true)
+ null
+ } else {
+ current
+ }
+ })
+ removed.get()
}
// Shuts down the session after removing.
- private def shutdownSessionHolder(
+ private[service] def shutdownSessionHolder(
Review Comment:
**Nit (P3):** `removeSessionHolder` and `shutdownSessionHolder` were
`private` before, and nothing outside this class calls them now; the suite uses
`tryRemoveSessionHolder`. Widening `shutdownSessionHolder` lets other code in
the package call it without winning `tryRemoveSessionHolder`, which breaks the
single-owner rule the new comment describes. Could both stay `private`?
```suggestion
private def shutdownSessionHolder(
```
**Verification:**
- **Compile:** The Connect server module and its tests compile with both
methods private.
--
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]
