HyukjinKwon commented on code in PR #59277:
URL: https://github.com/apache/spark/pull/59277#discussion_r4255049992
##########
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:
**Nit (P3):** `startBlockingCloseSession(key)` defaults
`allowReconnectAfterClose` to `true`, so this close takes the reconnect-allowed
branch of `tryRemoveSessionHolder`. The id is fenced in `closingSessions`, not
tombstoned in `closedSessionsCache`. The other tests that hold a close open use
reconnect-allowed closes too, so no test holds a default (tombstoning) close
open.
The default branch is fine today: it puts the tombstone inside
`computeIfPresent`. Suppose it later moved the put to after `close()`. A
same-id `getOrCreateIsolatedSession` during a default close would then succeed,
the in-flight close's key-based cleanup would tear the new session down, and
every test would still pass. Running this test for both values of the opt-in,
or renaming it to say it covers reconnect-allowed closes, would close the gap.
**Verification:**
- **Behavior:** Holding a default tombstoning close open makes same-id
session creation fail with INVALID_HANDLE.SESSION_CLOSED during the close.
Losing duplicate removals do not close the holder again. After the close, the
tombstone is kept with closedTimeMs set.
- **Regression:** Suppose the default branch of tryRemoveSessionHolder wrote
the tombstone only after close() returned. A same-key creation during the
held-open default close would then succeed and fail the new assertion, so the
PR-rewritten default fence now has a regression signal.
- **Behavior:** The reconnect-allowed variant still fences the id during
close and allows a new, empty session with the same id after close.
--
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]
