HyukjinKwon commented on code in PR #59277:
URL: https://github.com/apache/spark/pull/59277#discussion_r4242765483
##########
sql/connect/server/src/main/scala/org/apache/spark/sql/connect/service/SparkConnectSessionManager.scala:
##########
@@ -224,46 +241,75 @@ class SparkConnectSessionManager extends Logging {
session
}
- // Removes session from sessionStore and returns it.
+ // Removes the session with the given key from sessionStore, if any. Returns
the removed
+ // session and whether its tombstone is dropped after the close, see
tryRemoveSessionHolder.
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)
- }
-
- // Then, remove the session holder from the session store.
- sessionStore.remove(key)
+ allowReconnect: Boolean): Option[(SessionHolder, Boolean)] = {
+ Option(sessionStore.get(key)).flatMap { sessionHolder =>
+ tryRemoveSessionHolder(sessionHolder, allowReconnect).map(sessionHolder
-> _)
}
- sessionHolder
+ }
+
+ // Removes the given session from sessionStore if it is still stored there.
Returns
+ // Some(dropTombstone) only for the one caller that removed it; that caller
must then call
+ // shutdownSessionHolder with it. Returns None if the session was not
removed.
+ //
+ // The session id is fenced before the session is removed, and the fence 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. dropTombstone is decided here,
once: if true, the
+ // fence is an entry in closingSessions that is dropped after the close, so
the session id can
+ // be reused; otherwise it is a tombstone in closedSessionsCache that is
kept.
+ private[service] def tryRemoveSessionHolder(
+ sessionHolder: SessionHolder,
+ allowReconnect: Boolean = false): Option[Boolean] = {
+ // Fence 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 fence
+ // that the owner already dropped, and a later session with the same key
is never removed.
+ var removed: Option[Boolean] = None
+ sessionStore.computeIfPresent(
+ sessionHolder.key,
+ (_, current) => {
+ if (current eq sessionHolder) {
+ val dropTombstone = allowReconnect ||
current.allowReconnectAfterClose
+ if (dropTombstone) {
+ closingSessions.put(current.key, current)
+ } else {
+ closedSessionsCache.put(current.key, current.getSessionHolderInfo)
+ }
+ removed = Some(dropTombstone)
+ null
+ } else {
+ current
+ }
+ })
+ removed
}
// Shuts down the session after removing.
private 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)
+ dropTombstone: Boolean): Unit = {
+ try {
+ sessionHolder.close()
+ } finally {
Review Comment:
**Nit (P3):** Two guarantees added here would regress without any test
failing:
- If `close()` throws, this `finally` still releases a reconnect-allowed id
and refreshes a default tombstone. `BlockingCloseSessionHolder` always calls
`super.close()`, so no test makes `close()` throw. If the release moved back
after `close()`, a throwing close would leave a permanent `closingSessions`
entry, because that map is never evicted, and the id would get `SESSION_CLOSED`
until the server restarts.
- `validateCloneTargetSession` now checks `isSessionClosed`, so cloning into
an id whose reconnect-allowed close is still running fails with
`TARGET_SESSION_ID_ALREADY_CLOSED`. No test clones into an id that is still
closing, so reverting line 212 to the cache-only check would keep every test
green.
A throwing variant of the blocking holder and one `cloneSession` call while
the close is held open would cover both.
**Verification:**
- **Regression:** With a close() that throws, a reconnect-allowed close
leaves the id reusable for a new session, and a default close leaves it
tombstoned with a closed time.
- **Regression:** While a reconnect-allowed close is held open, cloning into
its id fails with TARGET_SESSION_ID_ALREADY_CLOSED. After the close returns,
the clone succeeds.
##########
sql/connect/server/src/main/scala/org/apache/spark/sql/connect/config/Connect.scala:
##########
@@ -150,6 +150,23 @@ object Connect {
.intConf
.createWithDefaultString("1000")
+ val CONNECT_SESSION_ALLOW_RECONNECT_AFTER_CLOSE =
+ buildConf("spark.connect.session.allowReconnectAfterClose")
+ .doc(
+ "When true, the session id can be reused after this session is closed,
for example " +
+ "on inactivity expiry. The id stays reserved, and requests fail with
" +
+ "INVALID_HANDLE.SESSION_CLOSED, until the close finishes. After
that, a client that " +
+ "used the closed session gets INVALID_HANDLE.SESSION_CHANGED and
must create a new " +
Review Comment:
**Nit (P3):** This holds for session-creating RPCs that carry
`client_observed_server_side_session_id` (ExecutePlan, AnalyzePlan, Config,
AddArtifacts and so on), but not in general:
- ReattachExecute, GetStatus, ReleaseExecute, and the CloneSession source go
through `getIsolatedSession`. Its not-found branch throws `SESSION_NOT_FOUND`
before `validateSessionId` runs. A PySpark client that is mid-stream when its
opted-in session expires gets `RESPONSE_ALREADY_RECEIVED` from `reattach.py`
and is not marked closed.
- A client that never sends the observed id, which is optional in
`base.proto`, gets the new, empty session without any error.
Two changes would make the doc hold. First, `getIsolatedSession` could throw
`sessionChangedError(key)` from its not-found branch when an observed id is
supplied, after the `isSessionClosed` check. Second, the doc here and in
`docs/configuration.md` could say that the signal relies on the client sending
its observed server-side session id, as the built-in Python and Scala clients
do.
**Verification:**
- **Behavior:** After an opted-in session expires, a lookup through
getIsolatedSession that carries the old server-side session id reports
SESSION_CHANGED, and one without an observed id reports SESSION_NOT_FOUND.
- **Compatibility:** A closed or still-closing id keeps reporting
SESSION_CLOSED through getIsolatedSession whether or not an observed id is
supplied, and a live session with a matching observed id is still returned.
- **Inspection:** The conf doc and the docs/configuration.md row describe
the observed-id precondition for SESSION_CHANGED and contain the same text.
##########
sql/connect/server/src/main/scala/org/apache/spark/sql/connect/service/SparkConnectSessionManager.scala:
##########
@@ -97,6 +109,11 @@ class SparkConnectSessionManager extends Logging {
key,
Some(() => {
validateSessionCreate(key)
+ // A client that observed a previous session expects that session, so
do not create a
+ // new one that it cannot use. Throwing here leaves sessionStore
unchanged.
+ if (previouslyObservedSesssionId.isDefined) {
+ throw sessionChangedError(key)
Review Comment:
**Nit (P3):** This error is now the documented signal for a client whose
opted-in session was closed on a server that is still running. The
`INVALID_HANDLE.SESSION_CHANGED` message in `error-conditions.json` says "The
existing Spark server driver instance has restarted. Please reconnect.", so
users who hit it after an inactivity expiry are sent looking for a restart that
never happened. The message also doesn't say that a new `SparkSession` is
needed. PySpark and the Scala `ResponseValidator` key on the condition name, so
a more general message would not change client handling, for example: "The
server-side session has changed, for example because the server restarted or
the session was closed. Please create a new session."
**Verification:**
- **Inspection:** The SESSION_CHANGED message names both a server restart
and a closed session whose id was reused as possible causes, and asks the user
to create a new session.
- **Compatibility:** Error-condition registry format checks still pass, and
client detection by the INVALID_HANDLE.SESSION_CHANGED condition is unchanged.
--
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]
