HyukjinKwon commented on code in PR #59277:
URL: https://github.com/apache/spark/pull/59277#discussion_r4262598365
##########
sql/connect/server/src/main/scala/org/apache/spark/sql/connect/config/Connect.scala:
##########
@@ -150,6 +150,28 @@ 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 if it
sends its " +
+ "observed server-side session id, as the built-in Python and Scala
clients do. " +
Review Comment:
**Nit (P3):** The doc says a client that used the closed session gets
`INVALID_HANDLE.SESSION_CHANGED` if it sends its observed server-side session
id, "as the built-in Python and Scala clients do". Neither built-in client sets
`client_observed_server_side_session_id` on `AddArtifactsRequest` or
`ArtifactStatusesRequest`: see the request builders in
`python/pyspark/sql/connect/client/artifact.py` and in the Scala client
`ArtifactManager`. `SparkConnectAddArtifactsHandler` and
`SparkConnectArtifactStatusesHandler` then call
`getOrCreateIsolatedSession(key, None)`, so the new creation-path check never
runs for them.
These are often the first RPCs a stale client sends. The Scala client's
`execute` and `analyze` call `uploadAllClassFileArtifacts()` before the plan
RPC, so a REPL cell that defines classes uploads them first. PySpark's
`createDataFrame` uses `ArtifactStatuses` and `AddArtifacts` to cache a large
local relation. After an opted-in session expires, that first upload silently
creates a new session under the id. That session has default confs, so it has
not opted in, and it already holds the stale client's artifacts. The
replacement `SparkSession` the doc tells users to create then attaches to that
session rather than to an empty one. If nobody reattaches and sets the conf
again before that session times out, maintenance tombstones the id again. That
is the outcome the observed-id check in `getOrCreateIsolatedSession` was added
to prevent.
**Recommended change:** Set client_observed_server_side_session_id on the
AddArtifacts and ArtifactStatuses requests built by the Python and Scala client
artifact managers once the client has observed a server-side session id, as the
clients already do for their other session RPCs. Add Spark Connect client
artifact tests and a session manager test for an artifact RPC from a stale
client after an opted-in expiry.
**Why this works:** With the observed id present, both artifact handlers
reach getOrCreateIsolatedSession's new check, which throws SESSION_CHANGED
before creating a session. A stale client's first artifact upload after an
opted-in expiry then gets the same signal as its other RPCs, leaves no
non-opted-in session behind, and the doc's built-in-client claim holds.
**Scope:** Make the built-in clients' artifact requests carry the client's
observed server-side session id.
**Compatibility:** A fresh client with no observed server-side session id
still creates its session on its first artifact upload.
**Risks:** The Python ArtifactManager does not hold the client's observed
server-side session id today, so it needs the client's current value rather
than a copy taken at construction. Older servers already read the optional
field on artifact requests, so sending it only changes the outcome when the
observed session no longer exists.
**Constraints:** Preserve review-frame:observed-id-session-changed: a closed
or still-closing id still reports SESSION_CLOSED before SESSION_CHANGED. Leave
ReleaseSession behavior unchanged under review-frame:stale-release-limitation.
**Success:** After an opted-in session expires, an AddArtifacts or
ArtifactStatuses request from a stale built-in Python or Scala client fails
with INVALID_HANDLE.SESSION_CHANGED, and no session exists for the id
afterwards. A client that has not yet observed a server-side session id can
still upload artifacts before its first query and gets a new session. The conf
doc's statement that the built-in clients send their observed server-side
session id holds for every session-creating RPC they issue.
**Verification:**
- **Behavior:** A built-in client that has observed a server-side session id
sends it on artifact upload and artifact status requests, and a client that has
not observed one sends none.
- **Regression:** After an opted-in session expires, an artifact upload or
status request carrying the old observed id fails with SESSION_CHANGED and
leaves no session for the id, while a new client without an observed id can
then upload artifacts and create the session.
- **Compatibility:** Artifact uploads and status checks on a live session
whose observed id matches keep succeeding.
##########
sql/connect/server/src/test/scala/org/apache/spark/sql/connect/service/SparkConnectSessionManagerSuite.scala:
##########
@@ -159,6 +235,295 @@ 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)
+ }
+
+ /**
+ * A session holder whose close() blocks until `release` is counted down. If
`blockInfoOn` is
+ * set, the first getSessionHolderInfo call on that thread blocks until
`releaseInfo` is counted
+ * down.
+ */
+ private class BlockingCloseSessionHolder(key: SessionKey)
+ extends SessionHolder(key.userId, key.sessionId, spark.newSession()) {
+ val closeStarted = new CountDownLatch(1)
+ val release = new CountDownLatch(1)
+ val closeCalls = new AtomicInteger(0)
+
+ @volatile var blockInfoOn: Thread = null
+ val infoBlocked = new CountDownLatch(1)
+ val releaseInfo = new CountDownLatch(1)
+
+ override private[connect] def close(): Unit = {
+ closeCalls.incrementAndGet()
+ closeStarted.countDown()
+ assert(release.await(1, TimeUnit.MINUTES))
+ super.close()
+ }
+
+ override def getSessionHolderInfo: SessionHolderInfo = {
+ if (blockInfoOn eq Thread.currentThread()) {
+ blockInfoOn = null
+ infoBlocked.countDown()
+ assert(releaseInfo.await(1, TimeUnit.MINUTES))
+ }
+ super.getSessionHolderInfo
+ }
+ }
+
+ private def startBlockingCloseSession(
+ key: SessionKey,
+ allowReconnectAfterClose: Boolean = true): BlockingCloseSessionHolder = {
+ val holder = new BlockingCloseSessionHolder(key)
+ holder.initializeSession()
+ holder.session.conf.set(
+ Connect.CONNECT_SESSION_ALLOW_RECONNECT_AFTER_CLOSE.key,
+ allowReconnectAfterClose.toString)
+ SparkConnectService.sessionManager.putSessionForTesting(holder)
+ holder
+ }
+
+ test("session id stays tombstoned until the owning close completes") {
Review Comment:
Thanks @venkata91, I checked f3548e5953. `session id stays tombstoned until
the owning close completes` now runs for both values of the opt-in. For a
default close, same-id creation fails with `SESSION_CLOSED` while the close is
held open, the losing removals don't close it again, and the tombstone is kept
afterwards with `closedTimeMs` set. The PR description also now lists the
`SESSION_CHANGED` change for requests with an observed id as a user-facing
change that applies without the opt-in.
<!-- SPARK_DEV_REVIEW_REPLY
{"feedback_id":"inline:4255049992","thread_id":"inline:4255049992","verdict_sha256":"c3e71497203cd96fb8e733fa958df1d21bdb50af02c780290ea88c200dc241be"}
-->
--
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]
