dongjoon-hyun commented on code in PR #59277:
URL: https://github.com/apache/spark/pull/59277#discussion_r4212641691
##########
sql/connect/server/src/main/scala/org/apache/spark/sql/connect/service/SparkConnectSessionManager.scala:
##########
@@ -262,6 +262,12 @@ class SparkConnectSessionManager extends Logging {
def closeSession(key: SessionKey, allowReconnect: Boolean = false): Unit = {
val sessionHolder = removeSessionHolder(key, allowReconnect)
+ if (sessionHolder.isEmpty && allowReconnect) {
+ // The session is already gone (e.g. expired) and the client explicitly
asked to be able to
+ // reconnect, accepting that its state is lost: drop the tombstone so
the same session id
+ // can be recreated by getOrCreateIsolatedSession.
+ closedSessionsCache.invalidate(key)
Review Comment:
**[1] The tombstone can be dropped while the old session is still being
closed.**
`sessionHolder.isEmpty` does not mean the old session is fully gone.
`periodicMaintenance` (or a concurrent `ReleaseSession`) first runs
`removeSessionHolder`, which puts a tombstone with `closedTimeMs = None` and
removes the store entry, and only then calls `SessionHolder.close()`. `close()`
can take a long time, e.g. `cleanupRunningQueries(this, blocking = true)` stops
streaming queries synchronously.
If `ReleaseSession(allow_reconnect=true)` arrives in that window:
1. `sessionStore.get(key)` is null, so the tombstone is invalidated and the
RPC returns immediately.
2. The client recreates a session with the same `SessionKey` and starts
running work.
3. The old `close()` continues with
`SparkConnectExecutionManager.removeAllExecutionsForSession(this.key)` and
`cleanupRunningQueries`, both of which match on `userId`/`sessionId`, so the
**new** session's executions and streaming queries are abandoned or stopped.
Note that the existing `allowReconnect = true` path on an active session
does not have this problem, because `shutdownSessionHolder` runs synchronously
in the handler thread before the response is sent.
Could we invalidate only when the shutdown has finished? For example, only
invalidate when the existing tombstone has `closedTimeMs.isDefined` (that is
the snapshot written by `shutdownSessionHolder` after `close()`), and otherwise
keep it or fail the request.
##########
sql/connect/server/src/main/scala/org/apache/spark/sql/connect/service/SparkConnectSessionManager.scala:
##########
@@ -262,6 +262,12 @@ class SparkConnectSessionManager extends Logging {
def closeSession(key: SessionKey, allowReconnect: Boolean = false): Unit = {
val sessionHolder = removeSessionHolder(key, allowReconnect)
+ if (sessionHolder.isEmpty && allowReconnect) {
+ // The session is already gone (e.g. expired) and the client explicitly
asked to be able to
+ // reconnect, accepting that its state is lost: drop the tombstone so
the same session id
+ // can be recreated by getOrCreateIsolatedSession.
+ closedSessionsCache.invalidate(key)
Review Comment:
**[2] A concurrent shutdown re-puts the tombstone that was just
invalidated.**
In the same window as [1], `shutdownSessionHolder` unconditionally calls
`closedSessionsCache.put(sessionHolder.key,
sessionHolder.getSessionHolderInfo)` after `close()` returns. That silently
undoes the `invalidate(key)` here:
- `ReleaseSession(allow_reconnect=true)` succeeds, but the next
`getOrCreateIsolatedSession` still fails with `INVALID_HANDLE.SESSION_CLOSED`,
or
- if the client already recreated the session, the same key is now both
active in `sessionStore` and listed in `listClosedSessions`.
The guarantee described in the proto comment is therefore not reliable under
concurrency.
##########
python/pyspark/sql/connect/client/core.py:
##########
@@ -2182,10 +2182,22 @@ def interrupt_operation(self, op_id: str) ->
Optional[List[str]]:
except Exception as error:
self._handle_error(error)
- def release_session(self) -> None:
+ def release_session(self, allow_reconnect: bool = False) -> None:
Review Comment:
**[3] The old client can later close and tombstone the recreated session.**
After `release_session(allow_reconnect=True)`, the old
`SparkConnectClient`/`SparkSession` is still fully usable and is not marked
closed. With the flow in the PR description:
```python
spark.client.release_session(allow_reconnect=True)
spark2 =
SparkSession.builder.remote(f"sc://...;session_id={same_id}").create()
...
spark.stop() # or process exit with
SPARK_CONNECT_RELEASE_SESSION_ON_EXIT=true
```
`spark.stop()` (and `_on_exit`) sends `ReleaseSession` for the same
`session_id` with `allow_reconnect=false`. `SparkConnectReleaseSessionHandler`
does not validate the server-side session id, so it closes `spark2`'s live
session and puts the tombstone back. Every later call on `spark2` then fails
with `SESSION_CLOSED`.
Should the client be closed (or at least have release-on-stop disabled)
after a reconnect-allowing release?
##########
sql/connect/server/src/main/scala/org/apache/spark/sql/connect/service/SparkConnectSessionManager.scala:
##########
@@ -262,6 +262,12 @@ class SparkConnectSessionManager extends Logging {
def closeSession(key: SessionKey, allowReconnect: Boolean = false): Unit = {
val sessionHolder = removeSessionHolder(key, allowReconnect)
+ if (sessionHolder.isEmpty && allowReconnect) {
Review Comment:
**[8] The `sessionHolder.isEmpty` guard is redundant.**
With `allowReconnect = true`, `removeSessionHolder` never writes a
tombstone, and an active holder cannot coexist with a tombstone for the same
key. So `if (allowReconnect) closedSessionsCache.invalidate(key)` (or doing it
inside `removeSessionHolder`) behaves the same and keeps the tombstone handling
in fewer places. Currently it is split across `removeSessionHolder`,
`shutdownSessionHolder`, and `closeSession`.
##########
python/pyspark/sql/connect/client/core.py:
##########
@@ -2182,10 +2182,22 @@ def interrupt_operation(self, op_id: str) ->
Optional[List[str]]:
except Exception as error:
self._handle_error(error)
- def release_session(self) -> None:
+ def release_session(self, allow_reconnect: bool = False) -> None:
+ """
+ Release the session on the server.
+
+ Parameters
+ ----------
+ allow_reconnect : bool, optional
+ If True, the session id can be reused afterwards to create a new,
empty session,
+ including when the session has already been closed on the server
(for example,
+ after it expired). All server-side state of the old session
(temporary views,
+ temporary functions, cached data, the current catalog) is lost.
+ """
req = pb2.ReleaseSessionRequest()
req.session_id = self._session_id
req.client_type = self._builder.userAgent
+ req.allow_reconnect = allow_reconnect
Review Comment:
**[4] The same client keeps the stale `_server_session_id`.**
The docstring says the session id "can be reused afterwards", but this
client still sends `client_observed_server_side_session_id =
self._server_session_id` on subsequent requests. On the server,
`getOrCreateIsolatedSession` first creates a fresh holder via `computeIfAbsent`
and only then `validateSessionId` throws `INVALID_HANDLE.SESSION_CHANGED`. As a
result:
- the client sets `_closed = True` on the first follow-up call, and
- an empty orphan session is left on the server until it expires.
Please either reset the client state accordingly or document clearly that a
new `SparkSession` object is required.
##########
sql/connect/server/src/main/scala/org/apache/spark/sql/connect/service/SparkConnectSessionManager.scala:
##########
@@ -262,6 +262,12 @@ class SparkConnectSessionManager extends Logging {
def closeSession(key: SessionKey, allowReconnect: Boolean = false): Unit = {
val sessionHolder = removeSessionHolder(key, allowReconnect)
+ if (sessionHolder.isEmpty && allowReconnect) {
Review Comment:
**[5] This lifts the `SESSION_CLOSED` fence for any caller, with no
server-side opt-out.**
Previously, the tombstone guaranteed that a released or expired session id
never silently came back as an empty session, so stale clients fail fast
instead of running against missing temp views / current catalog. With this
change, any caller that knows the `user_id` (client-supplied, unauthenticated
by default) and `session_id` can remove that fence.
Since this changes the server's semantics for existing clients that already
send `allow_reconnect = true`, could we gate it with a server-side
configuration? Alternatively, the root cause described in the PR (tombstones
only disappear by size eviction) could be addressed directly, e.g. with a
time-based `expireAfterWrite` on `closedSessionsCache`.
##########
sql/connect/server/src/test/scala/org/apache/spark/sql/connect/service/SparkConnectSessionManagerSuite.scala:
##########
@@ -159,6 +159,29 @@ class SparkConnectSessionManagerSuite extends
SharedSparkSession {
assert(closedSessionInfo.get.closedTimeMs.isDefined)
}
+ test("releasing an already closed session with allowReconnect allows
recreating it") {
+ val key = SessionKey("user", UUID.randomUUID().toString)
+ val sessionHolder =
SparkConnectService.sessionManager.getOrCreateIsolatedSession(key, None)
+
SparkConnectService.sessionManager.periodicMaintenance(defaultInactiveTimeoutMs
= 0L)
+
+ val ex = intercept[SparkSQLException] {
+ SparkConnectService.sessionManager.getOrCreateIsolatedSession(key, None)
+ }
+ assert(ex.getCondition == "INVALID_HANDLE.SESSION_CLOSED")
+
+ // Without allowReconnect the tombstone stays.
+ SparkConnectService.sessionManager.closeSession(key)
Review Comment:
**[7] Please check the error condition here.**
This `intercept[SparkSQLException]` accepts any `SparkSQLException`. If this
ever started failing for a different reason (e.g. `INVALID_HANDLE.FORMAT` or
`SESSION_NOT_FOUND`), the test would still pass while the behavior it claims to
verify ("the tombstone stays") regressed. Could you assert `getCondition ==
"INVALID_HANDLE.SESSION_CLOSED"` like the earlier block?
##########
sql/connect/server/src/test/scala/org/apache/spark/sql/connect/service/SparkConnectSessionManagerSuite.scala:
##########
@@ -159,6 +159,29 @@ class SparkConnectSessionManagerSuite extends
SharedSparkSession {
assert(closedSessionInfo.get.closedTimeMs.isDefined)
}
+ test("releasing an already closed session with allowReconnect allows
recreating it") {
+ val key = SessionKey("user", UUID.randomUUID().toString)
+ val sessionHolder =
SparkConnectService.sessionManager.getOrCreateIsolatedSession(key, None)
+
SparkConnectService.sessionManager.periodicMaintenance(defaultInactiveTimeoutMs
= 0L)
+
+ val ex = intercept[SparkSQLException] {
+ SparkConnectService.sessionManager.getOrCreateIsolatedSession(key, None)
+ }
+ assert(ex.getCondition == "INVALID_HANDLE.SESSION_CLOSED")
+
+ // Without allowReconnect the tombstone stays.
+ SparkConnectService.sessionManager.closeSession(key)
+ intercept[SparkSQLException] {
+ SparkConnectService.sessionManager.getOrCreateIsolatedSession(key, None)
+ }
+
+ SparkConnectService.sessionManager.closeSession(key, allowReconnect = true)
+ val recreated =
SparkConnectService.sessionManager.getOrCreateIsolatedSession(key, None)
+ assert(recreated.serverSessionId != sessionHolder.serverSessionId)
Review Comment:
**[6] This assertion is always true, and the test bypasses the RPC path.**
`serverSessionId` is a fresh random UUID for every new `SessionHolder`, so
`recreated.serverSessionId != sessionHolder.serverSessionId` can never fail.
Also, the test calls `closeSession` directly, so it does not cover:
- `SparkConnectReleaseSessionHandler` forwarding `allow_reconnect`,
- the Scala client's `releaseSession(allowReconnect = true)`,
- the concurrent-shutdown window in [1] and [2].
An end-to-end test through the handler (or the Scala client) would make this
much stronger.
--
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]
