HyukjinKwon commented on code in PR #59277:
URL: https://github.com/apache/spark/pull/59277#discussion_r4228193108
##########
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:
**Nit (P3):** `periodicMaintenance` (lines 353-355) still calls
`shutdownSessionHolder(sessionHolder)` for the holder it saw during
`sessionStore.forEach`, even when `removeSessionHolder(info.key)` returned
`None` because a concurrent `closeSession` already removed it. That second
`close()` throws `sessionAlreadyClosed`, but this `finally` block still runs:
1. An opted-in session expires, and at the same moment the client calls
`ReleaseSession`.
2. The release thread wins `removeSessionHolder`, which writes the
tombstone, and is now inside `close()` (for example, blocked in
`cleanupRunningQueries(this, blocking = true)`).
3. The maintenance thread's duplicate `shutdownSessionHolder(H)` hits
`sessionAlreadyClosed`, and the `finally` block calls
`closedSessionsCache.invalidate(key)`.
4. The tombstone is gone while the first `close()` is still doing key-based
cleanup. A same-key request can create a new session, and
`removeAllExecutionsForSession(this.key)` and the streaming cleanup then tear
it down.
For a session that did not opt in, the duplicate call instead re-puts the
tombstone, which can undo an `allow_reconnect=true` release that finished
first. In the merge target the throwing `close()` skipped the cache update, so
neither case could happen. Two concurrent `closeSession` calls can reach the
same path, because `removeSessionHolder` does a non-atomic `get` followed by
`remove(key)`. The window is narrow, but this is the exact guarantee the new
comment in `removeSessionHolder` describes.
See **Shared repair plan 1** in the review body.
##########
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:
**Nit (P3):** The new accessor sits between the existing `/** Get
SessionInfo with information about this SessionHolder. */` comment (line 460)
and `getSessionHolderInfo`. That stacks two doc comments on
`allowReconnectAfterClose` and leaves `getSessionHolderInfo` undocumented.
Moving the new member and its comment above line 460, or below
`getSessionHolderInfo`, keeps each comment on its own member.
**Verification:**
- **Inspection:** Each of the two members has exactly one doc comment
immediately above it.
##########
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:
**Nit (P3):** Per the versioning policy in `AGENTS.md`, a PR against
`master` that isn't master-only ships first in the latest `branch-4.x`, which
is currently `4.4.0-SNAPSHOT`. Master-only means a breaking change or a
non-critical dependency upgrade. This is an opt-in, non-breaking conf, so its
since-version would be `4.4.0`. Otherwise a 4.4.0 release documents it as added
in 5.0.0.
```suggestion
.version("4.4.0")
```
**Verification:**
- **Inspection:** The entry's declared version equals the branch-4.x pom
version with -SNAPSHOT stripped.
##########
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:
**Nit (P3):** This test checks `SESSION_CLOSED` after `removeSessionHolder`
(line 220) and checks reuse after `shutdownSessionHolder` returns. Nothing
observes the cache while `SessionHolder.close()` is running. If
`shutdownSessionHolder` invalidated the tombstone *before* calling `close()`,
which is the race this PR fixes, the test would still pass. The
duplicate-shutdown path in `shutdownSessionHolder` isn't covered either. To
match the test name, it could hold the old holder's `close()` open and assert
that same-key creation keeps failing until `close()` returns.
See **Shared repair plan 1** in the review body.
--
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]
