venkata91 commented on code in PR #59277:
URL: https://github.com/apache/spark/pull/59277#discussion_r4242350497
##########
sql/connect/server/src/main/scala/org/apache/spark/sql/connect/service/SparkConnectSessionManager.scala:
##########
@@ -224,44 +224,58 @@ class SparkConnectSessionManager extends Logging {
session
}
- // Removes session from sessionStore and returns it.
- private def removeSessionHolder(
- key: SessionKey,
- allowReconnect: Boolean = false): Option[SessionHolder] = {
- var sessionHolder: Option[SessionHolder] = None
-
- // The session holder should remain in the session store until it is added
to the closed session
- // cache, because of a subtle data race: a new session with the same key
can be created if the
- // closed session cache does not contain the key right after the key has
been removed from the
- // session store.
- sessionHolder = Option(sessionStore.get(key))
-
- sessionHolder.foreach { s =>
- if (!allowReconnect) {
- // Put into closedSessionsCache to prevent the same session from being
recreated by
- // getOrCreateIsolatedSession when reconnection isn't allowed.
- closedSessionsCache.put(s.key, s.getSessionHolderInfo)
- }
+ // Removes the session with the given key from sessionStore, if any, and
returns it.
+ private[service] def removeSessionHolder(key: SessionKey):
Option[SessionHolder] = {
+ Option(sessionStore.get(key)).filter(tryRemoveSessionHolder)
+ }
- // Then, remove the session holder from the session store.
- sessionStore.remove(key)
- }
- sessionHolder
+ // Removes the given session from sessionStore if it is still stored there.
Returns true only
+ // for the one caller that removed it; that caller must then call
shutdownSessionHolder.
+ //
+ // The session is tombstoned before it is removed, and the tombstone is kept
until
+ // shutdownSessionHolder has finished closing it. Otherwise a session with
the same key could be
+ // created while the old session is still closing, and the key-based cleanup
in
+ // SessionHolder.close() would tear it down.
+ private[service] def tryRemoveSessionHolder(sessionHolder: SessionHolder):
Boolean = {
+ // Tombstone and remove the session in one atomic step for the key,
matched by identity, so
+ // that two callers cannot both remove the same session, a caller that
loses cannot re-put a
+ // tombstone that the owner already dropped, and a later session with the
same key is never
+ // removed.
+ val removed = new AtomicBoolean(false)
+ sessionStore.computeIfPresent(
+ sessionHolder.key,
+ (_, current) => {
+ if (current eq sessionHolder) {
+ closedSessionsCache.put(current.key, current.getSessionHolderInfo)
Review Comment:
Fixed in f2ebb5a9e5. Reconnect-allowed closes (`allow_reconnect=true` or the
session opt-in) now fence the id in a separate, unbounded `closingSessions` map
that is dropped after `close()` finishes. They never go into the bounded
`closedSessionsCache`, so they can no longer evict other sessions' tombstones.
Whether to keep or release the id is decided once, inside the identity-checked
removal. Same-key checks look at both structures. A new test, `a
reconnect-allowed close fences its id without taking a tombstone slot`, covers
this.
##########
sql/connect/server/src/main/scala/org/apache/spark/sql/connect/service/SparkConnectSessionManager.scala:
##########
@@ -437,13 +454,21 @@ class SparkConnectSessionManager extends Logging {
})
}
+ /**
+ * Used for testing.
+ */
+ private[service] def putSessionHolderForTesting(sessionHolder:
SessionHolder): Unit = {
Review Comment:
Removed `putSessionHolderForTesting` in f2ebb5a9e5; the tests use the
existing `putSessionForTesting`.
##########
sql/connect/server/src/main/scala/org/apache/spark/sql/connect/config/Connect.scala:
##########
@@ -150,6 +150,18 @@ object Connect {
.intConf
.createWithDefaultString("1000")
+ val CONNECT_SESSION_ALLOW_RECONNECT_AFTER_CLOSE =
+ buildConf("spark.connect.session.allowReconnectAfterClose")
+ .doc(
+ "When true, the session id is not tombstoned when this session is
closed, for example " +
Review Comment:
Rewrote the conf doc in f2ebb5a9e5/b4f98b9d37. It now says the id is
released after close, that a client of the old session gets
`INVALID_HANDLE.SESSION_CHANGED`, and that the user must create a new
`SparkSession` with the same id. That session is empty (UDFs, temp views,
cached data and so on must be rebuilt) and must set this conf again.
##########
sql/connect/server/src/main/scala/org/apache/spark/sql/connect/config/Connect.scala:
##########
@@ -150,6 +150,18 @@ object Connect {
.intConf
.createWithDefaultString("1000")
+ val CONNECT_SESSION_ALLOW_RECONNECT_AFTER_CLOSE =
+ buildConf("spark.connect.session.allowReconnectAfterClose")
Review Comment:
Added the conf to the Spark Connect table in `docs/configuration.md` in
f2ebb5a9e5.
##########
sql/connect/server/src/test/scala/org/apache/spark/sql/connect/service/SparkConnectSessionManagerSuite.scala:
##########
@@ -159,6 +164,143 @@ 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. */
+ 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)
+
+ override private[connect] def close(): Unit = {
+ closeCalls.incrementAndGet()
+ closeStarted.countDown()
+ assert(release.await(1, TimeUnit.MINUTES))
+ super.close()
+ }
+ }
+
+ private def startBlockingCloseSession(key: SessionKey):
BlockingCloseSessionHolder = {
+ val holder = new BlockingCloseSessionHolder(key)
+ holder.initializeSession()
+
holder.session.conf.set(Connect.CONNECT_SESSION_ALLOW_RECONNECT_AFTER_CLOSE.key,
"true")
+ SparkConnectService.sessionManager.putSessionHolderForTesting(holder)
+ holder
+ }
+
+ test("session id stays tombstoned until the owning close completes") {
+ val sessionManager = SparkConnectService.sessionManager
+ val key = SessionKey("user", UUID.randomUUID().toString)
+ val holder = startBlockingCloseSession(key)
+
+ val owner = new Thread(() => sessionManager.closeSession(key))
+ owner.start()
+ try {
+ assert(holder.closeStarted.await(1, TimeUnit.MINUTES))
+ // The old session is removed but still closing: it must not be
recreated in this window.
+ assertSessionClosed(key)
+
+ // Callers that lose the removal, such as a concurrent ReleaseSession or
a maintenance run
+ // that observed the holder before it was removed, must not close it
again or touch its
+ // tombstone.
+ assert(!sessionManager.tryRemoveSessionHolder(holder))
+ sessionManager.closeSession(key)
+ sessionManager.periodicMaintenance(defaultInactiveTimeoutMs = 0L)
Review Comment:
Added `maintenance that observed a session but lost its removal does not
close it` in f2ebb5a9e5. It blocks maintenance after it has seen the holder,
lets a concurrent close win, then checks that maintenance does not close or
re-tombstone anything. I checked that the test fails if the removal gate in
`periodicMaintenance` is removed.
##########
sql/connect/server/src/main/scala/org/apache/spark/sql/connect/service/SparkConnectSessionManager.scala:
##########
@@ -224,44 +224,58 @@ class SparkConnectSessionManager extends Logging {
session
}
- // Removes session from sessionStore and returns it.
- private def removeSessionHolder(
- key: SessionKey,
- allowReconnect: Boolean = false): Option[SessionHolder] = {
- var sessionHolder: Option[SessionHolder] = None
-
- // The session holder should remain in the session store until it is added
to the closed session
- // cache, because of a subtle data race: a new session with the same key
can be created if the
- // closed session cache does not contain the key right after the key has
been removed from the
- // session store.
- sessionHolder = Option(sessionStore.get(key))
-
- sessionHolder.foreach { s =>
- if (!allowReconnect) {
- // Put into closedSessionsCache to prevent the same session from being
recreated by
- // getOrCreateIsolatedSession when reconnection isn't allowed.
- closedSessionsCache.put(s.key, s.getSessionHolderInfo)
- }
+ // Removes the session with the given key from sessionStore, if any, and
returns it.
+ private[service] def removeSessionHolder(key: SessionKey):
Option[SessionHolder] = {
+ Option(sessionStore.get(key)).filter(tryRemoveSessionHolder)
+ }
- // Then, remove the session holder from the session store.
- sessionStore.remove(key)
- }
- sessionHolder
+ // Removes the given session from sessionStore if it is still stored there.
Returns true only
+ // for the one caller that removed it; that caller must then call
shutdownSessionHolder.
+ //
+ // The session is tombstoned before it is removed, and the tombstone is kept
until
+ // shutdownSessionHolder has finished closing it. Otherwise a session with
the same key could be
+ // created while the old session is still closing, and the key-based cleanup
in
+ // SessionHolder.close() would tear it down.
+ private[service] def tryRemoveSessionHolder(sessionHolder: SessionHolder):
Boolean = {
+ // Tombstone and remove the session in one atomic step for the key,
matched by identity, so
+ // that two callers cannot both remove the same session, a caller that
loses cannot re-put a
+ // tombstone that the owner already dropped, and a later session with the
same key is never
+ // removed.
+ val removed = new AtomicBoolean(false)
+ sessionStore.computeIfPresent(
+ sessionHolder.key,
+ (_, current) => {
+ if (current eq sessionHolder) {
+ closedSessionsCache.put(current.key, current.getSessionHolderInfo)
+ removed.set(true)
+ null
+ } else {
+ current
+ }
+ })
+ removed.get()
}
// Shuts down the session after removing.
- private def shutdownSessionHolder(
+ private[service] def shutdownSessionHolder(
Review Comment:
Made `removeSessionHolder` and `shutdownSessionHolder` `private` again in
f2ebb5a9e5. `tryRemoveSessionHolder` is `private[service]` because the race
tests call it directly.
--
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]
