venkata91 commented on code in PR #59277:
URL: https://github.com/apache/spark/pull/59277#discussion_r4240563171
##########
sql/connect/server/src/main/scala/org/apache/spark/sql/connect/service/SparkConnectSessionManager.scala:
##########
@@ -225,43 +225,39 @@ class SparkConnectSessionManager extends Logging {
}
// 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))
-
+ private[service] def removeSessionHolder(key: SessionKey):
Option[SessionHolder] = {
+ // The session is always tombstoned before it is removed from the session
store, and the
+ // tombstone is kept until shutdownSessionHolder has finished closing it.
Otherwise a session
+ // with the same key could be created right after removal, or while the
old session is still
+ // closing, and the key-based cleanup in SessionHolder.close() would tear
it down.
+ val 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)
- }
-
- // Then, remove the session holder from the session store.
+ closedSessionsCache.put(s.key, s.getSessionHolderInfo)
sessionStore.remove(key)
}
sessionHolder
}
// Shuts down the session after removing.
- private def shutdownSessionHolder(
+ private[service] def shutdownSessionHolder(
sessionHolder: SessionHolder,
allowReconnect: Boolean = false): Unit = {
- sessionHolder.close()
- if (!allowReconnect) {
- // Update in closedSessionsCache: above it wasn't updated with
closedTime etc. yet.
- closedSessionsCache.put(sessionHolder.key,
sessionHolder.getSessionHolderInfo)
+ val dropTombstone = allowReconnect ||
sessionHolder.allowReconnectAfterClose
+ try {
+ sessionHolder.close()
+ } finally {
Review Comment:
Thanks, good catch. Fixed in 7853a51d06 and 710d9bb5e2, following Shared
repair plan 1. My first attempt in 0a036296059 only covered the
`periodicMaintenance` path; it didn't make the removal itself atomic.
The new `tryRemoveSessionHolder(holder)` writes the tombstone and removes
the session in one `computeIfPresent` step, only if the stored holder is still
that exact holder (identity check, as in `SparkConnectStreamingQueryCache`). It
returns true only to the caller that removed it. `removeSessionHolder(key)` is
now `Option(sessionStore.get(key)).filter(tryRemoveSessionHolder)`, and
`periodicMaintenance` and `invalidateAllSessions` call `shutdownSessionHolder`
only when `tryRemoveSessionHolder` returns true.
So exactly one caller owns `close()` and the post-close drop-or-put of the
tombstone. A losing `closeSession` or maintenance run never reaches the
`finally` block. Because the tombstone is written inside the atomic step, a
losing caller can't re-put a tombstone the owner already dropped. Maintenance
also can't remove a later session with the same key. Tombstone-before-removal
and the drop-or-put in `shutdownSessionHolder` are unchanged, and maintenance
still logs and continues on a per-session shutdown failure.
##########
sql/connect/server/src/test/scala/org/apache/spark/sql/connect/service/SparkConnectSessionManagerSuite.scala:
##########
@@ -159,6 +162,87 @@ 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)
+ }
+
+ test("session id stays tombstoned until close completes when reconnect is
allowed") {
+ val key = SessionKey("user", UUID.randomUUID().toString)
+ val holder =
SparkConnectService.sessionManager.getOrCreateIsolatedSession(key, None)
+
holder.session.conf.set(Connect.CONNECT_SESSION_ALLOW_RECONNECT_AFTER_CLOSE.key,
"true")
+
+ SparkConnectService.sessionManager.removeSessionHolder(key)
+ // The old session is removed but not closed yet: it must not be recreated
in this window.
+ assertSessionClosed(key)
+
+ SparkConnectService.sessionManager.shutdownSessionHolder(holder)
Review Comment:
Agreed. Replaced in 7853a51d06. A test-only holder subclass blocks inside
`close()`. The test `session id stays tombstoned until the owning close
completes` holds the owning `close()` open and, during that time:
- checks that same-key creation fails with `SESSION_CLOSED`;
- runs a duplicate removal of the same holder, a second `closeSession`, and
`periodicMaintenance`, and checks the tombstone is still there;
- after releasing, checks `close()` ran exactly once and the id is reusable.
`a late duplicate shutdown does not re-tombstone or remove a recreated
session` covers the `allow_reconnect=true` case. A late removal of the old
holder doesn't re-tombstone the id or remove the session recreated with that
key.
I checked that the first test fails if the tombstone is dropped before
`close()`, and the second fails if removal ignores holder identity.
##########
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 " +
+ "on inactivity expiry, so a client can create a new, empty session
with the same " +
+ "session id afterwards. Only set this when the client does not
depend on server-side " +
+ "session state, such as temporary views, temporary functions, cached
data or the " +
+ "current catalog, surviving a close.")
+ .version("5.0.0")
Review Comment:
Done in 0a036296059: `.version("4.4.0")`, which matches
`dev/next_version_candidates.py` (branch-4.x 4.4.0).
##########
sql/connect/server/src/main/scala/org/apache/spark/sql/connect/service/SessionHolder.scala:
##########
@@ -458,6 +458,10 @@ case class SessionHolder(userId: String, sessionId:
String, session: SparkSessio
}
/** Get SessionInfo with information about this SessionHolder. */
+ /** Whether the client opted in to reusing this session id after the session
is closed. */
Review Comment:
Done in 0a036296059. The new accessor now sits above the
`getSessionHolderInfo` scaladoc, so each member has exactly one doc comment.
--
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]
