venkata91 commented on code in PR #59277:
URL: https://github.com/apache/spark/pull/59277#discussion_r4213257706
##########
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:
Good point. Fixed: `removeSessionHolder` now always tombstones before
removing the holder, and `shutdownSessionHolder` drops the tombstone only after
`close()` completes. So nothing can recreate the key while key-based cleanup is
still running. Added: `session id stays tombstoned until close completes when
reconnect is allowed`.
##########
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:
Fixed. The after-the-fact `invalidate` is gone. `shutdownSessionHolder` is
now the only place that decides: after `close()` it either `invalidate`s
(reconnect allowed) or `put`s the final info, so there is no race between the
two.
##########
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:
Agreed. I removed the client-side `release_session(allow_reconnect)` change
entirely. With the opt-in conf, no client call lifts the fence on a closed
session anymore, so the old client can't touch the recreated one through this
path.
##########
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:
Same as [3]: dropped the client changes. A client that opts in has to handle
`SESSION_CHANGED` and start a fresh `SparkSession` anyway. That's also
documented in the conf doc: only opt in if you don't depend on server-side
state surviving a close.
##########
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:
Makes sense. Reworked to your option (1):
`spark.connect.session.allowReconnectAfterClose` is a session-scoped conf,
default `false`. The session has to opt in while it is alive, and the decision
is read from the holder at shutdown. Releases and expiries of sessions that
didn't opt in keep the `SESSION_CLOSED` fence as before.
##########
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:
Done. Removed the tautological assert. Added a
`SparkConnectReleaseSessionHandler` path test: `ReleaseSession honors
allow_reconnect and the session opt-in`. The recreate test now checks that the
new session starts from a clean conf instead.
##########
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:
Done. All negative cases now go through a helper that asserts `getCondition
== "INVALID_HANDLE.SESSION_CLOSED"`.
##########
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:
Removed along with the after-the-fact invalidate.
--
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]
