juliuszsompolski commented on code in PR #43985:
URL: https://github.com/apache/spark/pull/43985#discussion_r1409826101
##########
connector/connect/server/src/main/scala/org/apache/spark/sql/connect/service/SparkConnectSessionManager.scala:
##########
@@ -95,47 +92,134 @@ class SparkConnectSessionManager extends Logging {
Option(getSession(key, None))
}
- private def getSession(
- key: SessionKey,
- default: Option[Callable[SessionHolder]]): SessionHolder = {
- val session = default match {
- case Some(callable) => sessionStore.get(key, callable)
- case None => sessionStore.getIfPresent(key)
- }
- // record access time before returning
- session match {
- case null =>
- null
- case s: SessionHolder =>
- s.updateAccessTime()
- s
+ private def getSession(key: SessionKey, default: Option[() =>
SessionHolder]): SessionHolder = {
+ sessionsLock.synchronized {
+ // try to get existing session from store
+ val sessionOpt = sessionStore.get(key)
+ // create using default if missing
+ val session = sessionOpt match {
+ case Some(s) => s
+ case None =>
+ default match {
+ case Some(callable) =>
+ val session = callable()
+ sessionStore.put(key, session)
+ session
+ case None =>
+ null
+ }
+ }
+ // record access time before returning
+ session match {
+ case null =>
+ null
+ case s: SessionHolder =>
+ s.updateAccessTime()
+ s
+ }
}
}
def closeSession(key: SessionKey): Unit = {
- // Invalidate will trigger RemoveSessionListener
- sessionStore.invalidate(key)
- }
-
- private class RemoveSessionListener extends RemovalListener[SessionKey,
SessionHolder] {
- override def onRemoval(notification: RemovalNotification[SessionKey,
SessionHolder]): Unit = {
- val sessionHolder = notification.getValue
- sessionsLock.synchronized {
+ var sessionHolder: Option[SessionHolder] = None
+ sessionsLock.synchronized {
+ sessionHolder = sessionStore.remove(key)
+ sessionHolder.foreach { s =>
// First put into closedSessionsCache, so that it cannot get
accidentally recreated by
// getOrCreateIsolatedSession.
- closedSessionsCache.put(sessionHolder.key,
sessionHolder.getSessionHolderInfo)
+ closedSessionsCache.put(s.key, s.getSessionHolderInfo)
}
- // Rest of the cleanup outside sessionLock - the session cannot be
accessed anymore by
- // getOrCreateIsolatedSession.
- sessionHolder.close()
+ }
+ // Rest of the cleanup outside sessionLock - the session cannot be
accessed anymore by
+ // getOrCreateIsolatedSession.
+ sessionHolder.foreach { s =>
+ s.close()
+ // Update in closedSessionsCache: above it wasn't updated with
closedTime etc. yet.
+ closedSessionsCache.put(s.key, s.getSessionHolderInfo)
}
}
- def shutdown(): Unit = {
+ private[connect] def shutdown(): Unit = sessionsLock.synchronized {
+ scheduledExecutor.foreach { executor =>
+ executor.shutdown()
+ executor.awaitTermination(1, TimeUnit.MINUTES)
+ }
+ scheduledExecutor = None
+ // note: this does not cleanly shut down the sessions, but the server is
shutting down.
+ sessionStore.clear()
+ closedSessionsCache.invalidateAll()
+ }
+
+ def listActiveSessions: Seq[SessionHolderInfo] = sessionsLock.synchronized {
+ sessionStore.values.map(_.getSessionHolderInfo).toSeq
+ }
+
+ def listClosedSessions: Seq[SessionHolderInfo] = sessionsLock.synchronized {
+ closedSessionsCache.asMap.asScala.values.toSeq
+ }
+
+ /**
+ * Schedules periodic maintenance checks if it is not already scheduled.
+ *
+ * The checks are looking to remove sessions that expired.
+ */
+ private def schedulePeriodicChecks(): Unit = sessionsLock.synchronized {
+ scheduledExecutor match {
+ case Some(_) => // Already running.
+ case None =>
+ val interval =
SparkEnv.get.conf.get(CONNECT_SESSION_MANAGER_MAINTENANCE_INTERVAL).toLong
+ // Default for sessions that don't have expiration time set.
+ val defaultInactiveTimeout =
+
SparkEnv.get.conf.get(CONNECT_SESSION_MANAGER_DEFAULT_SESSION_TIMEOUT).toLong
+ logInfo(s"Starting thread for cleanup of expired sessions every
$interval ms")
+ scheduledExecutor = Some(Executors.newSingleThreadScheduledExecutor())
+ scheduledExecutor.get.scheduleAtFixedRate(
+ () => {
+ try periodicMaintenance(defaultInactiveTimeout)
+ catch {
+ case NonFatal(ex) => logWarning("Unexpected exception in
periodic task", ex)
+ }
+ },
+ interval,
+ interval,
+ TimeUnit.MILLISECONDS)
+ }
+ }
+
+ // Visible for testing.
+ private[connect] def periodicMaintenance(
+ defaultInactiveTimeout: Long,
+ ignoreExpirationTime: Boolean = false): Unit = {
+ logInfo("Started periodic run of SparkConnectSessionManager maintenance.")
+ // Find any sessions that expired and should be removed.
+ val toRemove = new mutable.ArrayBuffer[SessionHolder]()
sessionsLock.synchronized {
- sessionStore.invalidateAll()
- closedSessionsCache.invalidateAll()
+ val nowMs = System.currentTimeMillis()
+
+ sessionStore.values.foreach { sessionHolder =>
+ val info = sessionHolder.getSessionHolderInfo
+ if (info.customExpirationTime.isDefined && !ignoreExpirationTime) {
Review Comment:
> Why is this expiry time rather than custom inactivity time?
Good question. In the previous PR https://github.com/apache/spark/pull/43913
where I was adding an ExtendSession RPC, I found it a less ambiguous interface
to pass expiry time rather by extending by X. This way the client and server
could know what the deadline is precisely.
Here I kept it this way.
What you're saying is that we could also set custom inactivity time, and
then could do it only once. I.e. make this override the server side default
config.
Yeah, I guess this could be more consistent with the other config being an
inactivity time, and also allow it to be set only once instead of having to
keep resetting the expiry time. I'll change it.
> Who sets this in PR?
At this point, as we discussed offline to not introduce the ExtendSession
RPC at this point, nobody. But I am leaving it here for use by extensions. If
someone wanted to implement custom session expiry rules, they could still do it
e.g. using a custom Interceptor.
--
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: reviews-unsubscr...@spark.apache.org
For queries about this service, please contact Infrastructure at:
us...@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscr...@spark.apache.org
For additional commands, e-mail: reviews-h...@spark.apache.org
