venkata91 commented on code in PR #59277:
URL: https://github.com/apache/spark/pull/59277#discussion_r4260950164
##########
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, fixed in f3548e5953. The test now runs for both values of the
opt-in, so a default tombstoning close is also held open: same-id creation
fails with `SESSION_CLOSED` during the close, losing duplicate removals do not
close it again, and after the close the tombstone is kept with `closedTimeMs`
set. I checked that writing the default tombstone only after `close()` fails
this test. I also updated the PR description: without the opt-in, closed ids
are still tombstoned, and the `SESSION_CHANGED` change for requests with an
observed id is listed as a user-facing change.
--
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]
