diff --git a/kyuubi-server/src/main/scala/org/apache/kyuubi/server/KyuubiServer.scala b/kyuubi-server/src/main/scala/org/apache/kyuubi/server/KyuubiServer.scala index df163bd1eb5..a27e18bbf6c 100644 --- a/kyuubi-server/src/main/scala/org/apache/kyuubi/server/KyuubiServer.scala +++ b/kyuubi-server/src/main/scala/org/apache/kyuubi/server/KyuubiServer.scala @@ -131,12 +131,10 @@ object KyuubiServer extends Logging { } private[kyuubi] def refreshUnlimitedUsers(): Unit = synchronized { - val existingUnlimitedUsers = - kyuubiServer.conf.get(KyuubiConf.SERVER_LIMIT_CONNECTIONS_USER_UNLIMITED_LIST).toSet - val refreshedUnlimitedUsers = KyuubiConf().loadFileDefaults().get( - KyuubiConf.SERVER_LIMIT_CONNECTIONS_USER_UNLIMITED_LIST).toSet - kyuubiServer.backendService.sessionManager.asInstanceOf[KyuubiSessionManager] - .refreshUnlimitedUsers(refreshedUnlimitedUsers) + val sessionMgr = kyuubiServer.backendService.sessionManager.asInstanceOf[KyuubiSessionManager] + val existingUnlimitedUsers = sessionMgr.getUnlimitedUsers() + sessionMgr.refreshUnlimitedUsers(KyuubiConf().loadFileDefaults()) + val refreshedUnlimitedUsers = sessionMgr.getUnlimitedUsers() info(s"Refreshed unlimited users from $existingUnlimitedUsers to $refreshedUnlimitedUsers") } } diff --git a/kyuubi-server/src/main/scala/org/apache/kyuubi/session/KyuubiSessionManager.scala b/kyuubi-server/src/main/scala/org/apache/kyuubi/session/KyuubiSessionManager.scala index be7873d3188..f664614f5dd 100644 --- a/kyuubi-server/src/main/scala/org/apache/kyuubi/session/KyuubiSessionManager.scala +++ b/kyuubi-server/src/main/scala/org/apache/kyuubi/session/KyuubiSessionManager.scala @@ -299,9 +299,14 @@ class KyuubiSessionManager private (name: String) extends SessionManager(name) { userUnlimitedList) } - private[kyuubi] def refreshUnlimitedUsers(userUnlimitedList: Set[String]): Unit = { - limiter.foreach(SessionLimiter.resetUnlimitedUsers(_, userUnlimitedList)) - batchLimiter.foreach(SessionLimiter.resetUnlimitedUsers(_, userUnlimitedList)) + private[kyuubi] def getUnlimitedUsers(): Set[String] = { + limiter.orElse(batchLimiter).map(SessionLimiter.getUnlimitedUsers).getOrElse(Set.empty) + } + + private[kyuubi] def refreshUnlimitedUsers(conf: KyuubiConf): Unit = { + val unlimitedUsers = conf.get(SERVER_LIMIT_CONNECTIONS_USER_UNLIMITED_LIST).toSet + limiter.foreach(SessionLimiter.resetUnlimitedUsers(_, unlimitedUsers)) + batchLimiter.foreach(SessionLimiter.resetUnlimitedUsers(_, unlimitedUsers)) } private def applySessionLimiter( diff --git a/kyuubi-server/src/main/scala/org/apache/kyuubi/session/SessionLimiter.scala b/kyuubi-server/src/main/scala/org/apache/kyuubi/session/SessionLimiter.scala index 6cf739c3955..96ca36df176 100644 --- a/kyuubi-server/src/main/scala/org/apache/kyuubi/session/SessionLimiter.scala +++ b/kyuubi-server/src/main/scala/org/apache/kyuubi/session/SessionLimiter.scala @@ -138,10 +138,14 @@ object SessionLimiter { unlimitedUsers) } - def resetUnlimitedUsers(limiter: SessionLimiter, unlimitedUsers: Set[String]): Unit = { + def resetUnlimitedUsers(limiter: SessionLimiter, unlimitedUsers: Set[String]): Unit = limiter match { case l: SessionLimiterWithUnlimitedUsersImpl => l.setUnlimitedUsers(unlimitedUsers) case _ => } + + def getUnlimitedUsers(limiter: SessionLimiter): Set[String] = limiter match { + case l: SessionLimiterWithUnlimitedUsersImpl => l.unlimitedUsers + case _ => Set.empty } }