From fe7bdc73f35eb6ab7b8b48c48fa2502292689cc0 Mon Sep 17 00:00:00 2001 From: Venugopal Allenki Date: Tue, 12 May 2026 20:38:32 +0530 Subject: [PATCH 1/2] Add SessionManager refresh and /recovery operator endpoint (cherry picked from commit d30c64365a4155c83e9e639dedc4c6eda33d37b2) Co-authored-by: Soumyadeep Roy --- .../org/apache/livy/server/LivyServer.scala | 37 ++++++- .../apache/livy/sessions/SessionManager.scala | 43 +++++++- .../livy/sessions/SessionManagerSpec.scala | 99 +++++++++++++++++++ 3 files changed, 177 insertions(+), 2 deletions(-) diff --git a/server/src/main/scala/org/apache/livy/server/LivyServer.scala b/server/src/main/scala/org/apache/livy/server/LivyServer.scala index 68a3880a9..daf0abbb2 100644 --- a/server/src/main/scala/org/apache/livy/server/LivyServer.scala +++ b/server/src/main/scala/org/apache/livy/server/LivyServer.scala @@ -41,7 +41,7 @@ import org.apache.livy.server.batch.BatchSessionServlet import org.apache.livy.server.interactive.InteractiveSessionServlet import org.apache.livy.server.recovery.{SessionStore, StateStore, ZooKeeperManager} import org.apache.livy.server.ui.UIServlet -import org.apache.livy.sessions.{BatchSessionManager, InteractiveSessionManager} +import org.apache.livy.sessions.{BatchSessionManager, InteractiveSessionManager, SessionManager} import org.apache.livy.sessions.SessionManager.SESSION_RECOVERY_MODE_OFF import org.apache.livy.utils.{SparkKubernetesApp, SparkYarnApp} import org.apache.livy.utils.LivySparkUtils._ @@ -176,6 +176,39 @@ class LivyServer extends Logging { } } + // Operator endpoint: re-scan the recovery state store and import any sessions + // written by another Livy server. Guarded by livy.superusers. + val recoveryServlet = new JsonServlet { + before() { + contentType = "application/json" + val user = request.getRemoteUser() + if (!accessManager.checkSuperUser(user)) { + halt(403, Map("msg" -> s"User '$user' not authorized for recovery endpoints.")) + } + } + + private def resultMap(r: SessionManager.RefreshResult): Map[String, Int] = + Map("added" -> r.added, "total" -> r.total, "failed" -> r.failed) + + post("/sessions/refresh") { + info(s"Interactive session refresh triggered by user='${request.getRemoteUser()}'") + resultMap(interactiveSessionManager.refresh()) + } + + post("/batches/refresh") { + info(s"Batch session refresh triggered by user='${request.getRemoteUser()}'") + resultMap(batchSessionManager.refresh()) + } + + post("/refresh") { + info(s"Full session refresh triggered by user='${request.getRemoteUser()}'") + Map( + "batches" -> resultMap(batchSessionManager.refresh()), + "sessions" -> resultMap(interactiveSessionManager.refresh()) + ) + } + } + // Servlet for hosting static files such as html, css, and js // Necessary since Jetty cannot set it's resource base inside a jar // Returns 404 if the file does not exist @@ -256,6 +289,8 @@ class LivyServer extends Logging { metricRegistry, interactiveSessionManager, batchSessionManager) mount(context, livyVersionServlet, "/version/*") + + mount(context, recoveryServlet, "/recovery/*") } catch { case e: Throwable => error("Exception thrown when initializing server", e) diff --git a/server/src/main/scala/org/apache/livy/sessions/SessionManager.scala b/server/src/main/scala/org/apache/livy/sessions/SessionManager.scala index 1dc1d820e..cb41e8c85 100644 --- a/server/src/main/scala/org/apache/livy/sessions/SessionManager.scala +++ b/server/src/main/scala/org/apache/livy/sessions/SessionManager.scala @@ -36,6 +36,9 @@ import org.apache.livy.sessions.Session.RecoveryMetadata object SessionManager { val SESSION_RECOVERY_MODE_OFF = "off" val SESSION_RECOVERY_MODE_RECOVERY = "recovery" + + /** Outcome of a [[SessionManager.refresh]] call. */ + case class RefreshResult(added: Int, total: Int, failed: Int) } class BatchSessionManager( @@ -100,8 +103,14 @@ class SessionManager[S <: Session, R <: RecoveryMetadata : ClassTag]( } def register(session: S): S = { - info(s"Registering new session ${session.id}") synchronized { + sessions.get(session.id) match { + case Some(existing) => + debug(s"Session ${session.id} already registered; skipping duplicate registration.") + return existing + case None => + } + info(s"Registering new session ${session.id}") session.name.foreach { sessionName => if (sessionsByName.contains(sessionName)) { val errMsg = s"Duplicate session name: ${session.name} for session ${session.id}" @@ -229,6 +238,38 @@ class SessionManager[S <: Session, R <: RecoveryMetadata : ClassTag]( recoveredSessions } + /** + * Re-scan the state store and import sessions written by another Livy server. + * Add-only: in-memory sessions are not removed even if their state-store entry is + * gone. The id counter is advanced forward only. + */ + def refresh(): RefreshResult = { + // Read the state store outside the SessionManager monitor so we don't block the + // garbage collector and heartbeat watchdog while doing N small EFS / ZK reads. + val storeNextId = sessionStore.getNextSessionId(sessionType) + val sessionMetadata = sessionStore.getAllSessions[R](sessionType) + + val recoveryFailure = sessionMetadata.filter(_.isFailure).map(_.failed.get) + recoveryFailure.foreach(ex => warn(s"Refresh failure for $sessionType: ${ex.getMessage}", ex)) + + synchronized { + if (storeNextId > idCounter.get) { + idCounter.set(storeNextId) + } + + val before = sessions.size + sessionMetadata.flatMap(_.toOption) + .filterNot(m => sessions.contains(m.id)) + .map(sessionRecovery) + .foreach(register) + + val added = sessions.size - before + info(s"Refreshed $sessionType sessions: added=$added, total=${sessions.size}," + + s" failed=${recoveryFailure.size}, next session id=$idCounter") + RefreshResult(added, sessions.size, recoveryFailure.size) + } + } + private class GarbageCollector extends Thread("session gc thread") { setDaemon(true) diff --git a/server/src/test/scala/org/apache/livy/sessions/SessionManagerSpec.scala b/server/src/test/scala/org/apache/livy/sessions/SessionManagerSpec.scala index 8e5557018..3cde221a4 100644 --- a/server/src/test/scala/org/apache/livy/sessions/SessionManagerSpec.scala +++ b/server/src/test/scala/org/apache/livy/sessions/SessionManagerSpec.scala @@ -210,6 +210,16 @@ class SessionManagerSpec extends AnyFunSpec with Matchers with LivyBaseUnitTestS when(session.lastActivity).thenReturn(System.nanoTime()) when(session.state).thenReturn(state) } + + it("should be idempotent on register() for a session id already tracked") { + val (livyConf, manager) = createSessionManager() + val session = new MockSession(manager.nextId(), null, livyConf, Some("foo")) + val first = manager.register(session) + val again = manager.register(session) + again should be theSameInstanceAs first + manager.size() should be (1) + session.stopped should be (false) + } } describe("BatchSessionManager") { @@ -306,5 +316,94 @@ class SessionManagerSpec extends AnyFunSpec with Matchers with LivyBaseUnitTestS verify(session, never).stop() } + + it("refresh should be a no-op when state store has no new sessions") { + val conf = new LivyConf() + val sessionStore = mock[SessionStore] + when(sessionStore.getNextSessionId("batch")).thenReturn(0) + when(sessionStore.getAllSessions[BatchRecoveryMetadata]("batch")) + .thenReturn(Seq.empty) + + val sm = new BatchSessionManager(conf, sessionStore) + val result = sm.refresh() + result.added shouldBe 0 + result.total shouldBe 0 + result.failed shouldBe 0 + } + + it("refresh should import only previously-unseen sessions and skip already-known ids") { + val conf = new LivyConf() + conf.set(LivyConf.LIVY_SPARK_MASTER.key, "yarn-cluster") + + val sessionStore = mock[SessionStore] + // Initial recovery: only id=0. + when(sessionStore.getNextSessionId("batch")).thenReturn(1) + when(sessionStore.getAllSessions[BatchRecoveryMetadata]("batch")) + .thenReturn(Seq(Try(makeMetadata(0, "t0")))) + + val sm = new BatchSessionManager(conf, sessionStore) + sm.size() shouldBe 1 + sm.get(0) shouldBe defined + + // After a peer writes id=1 to the state store, refresh should add it without + // re-registering id=0. + when(sessionStore.getNextSessionId("batch")).thenReturn(2) + when(sessionStore.getAllSessions[BatchRecoveryMetadata]("batch")) + .thenReturn(Seq(Try(makeMetadata(0, "t0")), Try(makeMetadata(1, "t1")))) + + val r1 = sm.refresh() + r1.added shouldBe 1 + r1.total shouldBe 2 + sm.get(1) shouldBe defined + + // A second refresh with no new entries should add nothing. + val r2 = sm.refresh() + r2.added shouldBe 0 + r2.total shouldBe 2 + } + + it("refresh should advance the id counter forward but never backward") { + val conf = new LivyConf() + val sessionStore = mock[SessionStore] + when(sessionStore.getNextSessionId("batch")).thenReturn(0) + when(sessionStore.getAllSessions[BatchRecoveryMetadata]("batch")) + .thenReturn(Seq.empty) + + val sm = new BatchSessionManager(conf, sessionStore) + + // Local counter ahead of state store. + sm.nextId() shouldBe 0 + sm.nextId() shouldBe 1 + sm.nextId() shouldBe 2 + + // State store reports a smaller next id; refresh must NOT regress local counter. + when(sessionStore.getNextSessionId("batch")).thenReturn(1) + sm.refresh() + sm.nextId() shouldBe 3 + + // State store reports a larger next id; refresh SHOULD jump local counter forward. + when(sessionStore.getNextSessionId("batch")).thenReturn(99) + sm.refresh() + sm.nextId() shouldBe 99 + } + + it("refresh should report deserialization failures via failed count") { + val conf = new LivyConf() + conf.set(LivyConf.LIVY_SPARK_MASTER.key, "yarn-cluster") + + val sessionStore = mock[SessionStore] + when(sessionStore.getNextSessionId("batch")).thenReturn(1) + when(sessionStore.getAllSessions[BatchRecoveryMetadata]("batch")) + .thenReturn(Seq( + Try(makeMetadata(0, "t0")), + Failure(new java.io.IOException("corrupted entry id=42")), + Failure(new java.io.IOException("corrupted entry id=43")))) + + val sm = new BatchSessionManager(conf, sessionStore) + val r = sm.refresh() + r.added shouldBe 0 + r.total shouldBe 1 + r.failed shouldBe 2 + } } } From 842db2446208fa9ebbae44bb60fff2c205a6e2ce Mon Sep 17 00:00:00 2001 From: soumyadeep-roy Date: Wed, 26 Aug 2026 17:34:39 +0530 Subject: [PATCH 2/2] [LIVY-1073] Address review feedback on /recovery endpoints Extracts the previously-inline anonymous recovery servlet in LivyServer into a standalone RecoveryServlet class, following the existing BatchSessionServlet/InteractiveSessionServlet convention. This makes it unit-testable and adds RecoveryServletSpec, covering the 403 (non-superuser) vs 200 (superuser) authorization path that had no test coverage. Also has the servlet return 409 instead of a silently-empty added=0/failed=0 result when livy.server.recovery.mode=off, since in that case StateStore resolves to a no-op BlackholeStateStore and the previous behavior was misleading. Documents the three /recovery/* endpoints in docs/rest-api.md, which had none. Co-authored-by: Soumyadeep Roy --- docs/rest-api.md | 43 ++++++ .../org/apache/livy/server/LivyServer.scala | 36 +---- .../server/recovery/RecoveryServlet.scala | 74 +++++++++ .../server/recovery/RecoveryServletSpec.scala | 140 ++++++++++++++++++ 4 files changed, 261 insertions(+), 32 deletions(-) create mode 100644 server/src/main/scala/org/apache/livy/server/recovery/RecoveryServlet.scala create mode 100644 server/src/test/scala/org/apache/livy/server/recovery/RecoveryServletSpec.scala diff --git a/docs/rest-api.md b/docs/rest-api.md index b4d2042d8..a2fbdede9 100644 --- a/docs/rest-api.md +++ b/docs/rest-api.md @@ -694,6 +694,49 @@ instance. +### POST /recovery/sessions/refresh + +Operator endpoint. Re-scans the recovery state store and imports any interactive +sessions written by another Livy server that this server doesn't yet have in memory. +Add-only: existing in-memory sessions are never removed by this call. + +**Authorization:** restricted to users listed in `livy.superusers`; all other callers +get a `403`. Also requires `livy.server.recovery.mode` to be set to something other +than `off` (i.e. a real state store must be configured); otherwise this returns `409`, +since there would be nothing to refresh from. + +#### Response Body + + + + + + + + + + + + + + + + + + +
NameDescriptionType
addedNumber of sessions imported from the state store that weren't already in memoryint
totalTotal number of interactive sessions in memory after the refreshint
failedNumber of state store entries that failed to deserializeint
+ +### POST /recovery/batches/refresh + +Same as `POST /recovery/sessions/refresh`, but for batch sessions. Same authorization +and recovery-mode requirements, and the same response body shape. + +### POST /recovery/refresh + +Runs both of the above in a single call. Returns a JSON object with `sessions` and +`batches` keys, each holding a response body of the same shape as the individual +endpoints above. + ## REST Objects ### Session diff --git a/server/src/main/scala/org/apache/livy/server/LivyServer.scala b/server/src/main/scala/org/apache/livy/server/LivyServer.scala index daf0abbb2..50107ebd0 100644 --- a/server/src/main/scala/org/apache/livy/server/LivyServer.scala +++ b/server/src/main/scala/org/apache/livy/server/LivyServer.scala @@ -39,9 +39,9 @@ import org.apache.livy._ import org.apache.livy.server.auth.LdapAuthenticationHandlerImpl import org.apache.livy.server.batch.BatchSessionServlet import org.apache.livy.server.interactive.InteractiveSessionServlet -import org.apache.livy.server.recovery.{SessionStore, StateStore, ZooKeeperManager} +import org.apache.livy.server.recovery.{RecoveryServlet, SessionStore, StateStore, ZooKeeperManager} import org.apache.livy.server.ui.UIServlet -import org.apache.livy.sessions.{BatchSessionManager, InteractiveSessionManager, SessionManager} +import org.apache.livy.sessions.{BatchSessionManager, InteractiveSessionManager} import org.apache.livy.sessions.SessionManager.SESSION_RECOVERY_MODE_OFF import org.apache.livy.utils.{SparkKubernetesApp, SparkYarnApp} import org.apache.livy.utils.LivySparkUtils._ @@ -178,36 +178,8 @@ class LivyServer extends Logging { // Operator endpoint: re-scan the recovery state store and import any sessions // written by another Livy server. Guarded by livy.superusers. - val recoveryServlet = new JsonServlet { - before() { - contentType = "application/json" - val user = request.getRemoteUser() - if (!accessManager.checkSuperUser(user)) { - halt(403, Map("msg" -> s"User '$user' not authorized for recovery endpoints.")) - } - } - - private def resultMap(r: SessionManager.RefreshResult): Map[String, Int] = - Map("added" -> r.added, "total" -> r.total, "failed" -> r.failed) - - post("/sessions/refresh") { - info(s"Interactive session refresh triggered by user='${request.getRemoteUser()}'") - resultMap(interactiveSessionManager.refresh()) - } - - post("/batches/refresh") { - info(s"Batch session refresh triggered by user='${request.getRemoteUser()}'") - resultMap(batchSessionManager.refresh()) - } - - post("/refresh") { - info(s"Full session refresh triggered by user='${request.getRemoteUser()}'") - Map( - "batches" -> resultMap(batchSessionManager.refresh()), - "sessions" -> resultMap(interactiveSessionManager.refresh()) - ) - } - } + val recoveryServlet = new RecoveryServlet( + livyConf, accessManager, interactiveSessionManager, batchSessionManager) // Servlet for hosting static files such as html, css, and js // Necessary since Jetty cannot set it's resource base inside a jar diff --git a/server/src/main/scala/org/apache/livy/server/recovery/RecoveryServlet.scala b/server/src/main/scala/org/apache/livy/server/recovery/RecoveryServlet.scala new file mode 100644 index 000000000..9b4366758 --- /dev/null +++ b/server/src/main/scala/org/apache/livy/server/recovery/RecoveryServlet.scala @@ -0,0 +1,74 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.livy.server.recovery + +import javax.servlet.http.HttpServletRequest + +import org.apache.livy.{LivyConf, Logging} +import org.apache.livy.server.{AccessManager, JsonServlet} +import org.apache.livy.sessions.{BatchSessionManager, InteractiveSessionManager} +import org.apache.livy.sessions.SessionManager.{RefreshResult, SESSION_RECOVERY_MODE_OFF} + +/** + * Operator endpoint that re-scans the recovery state store and imports any sessions + * written by another Livy server. Guarded by `livy.superusers`. + */ +class RecoveryServlet( + livyConf: LivyConf, + accessManager: AccessManager, + interactiveSessionManager: InteractiveSessionManager, + batchSessionManager: BatchSessionManager) + extends JsonServlet + with Logging { + + protected def remoteUser(req: HttpServletRequest): String = req.getRemoteUser() + + before() { + contentType = "application/json" + val user = remoteUser(request) + if (!accessManager.checkSuperUser(user)) { + halt(403, Map("msg" -> s"User '$user' not authorized for recovery endpoints.")) + } + if (livyConf.get(LivyConf.RECOVERY_MODE) == SESSION_RECOVERY_MODE_OFF) { + halt(409, Map("msg" -> + "Recovery is disabled (livy.server.recovery.mode=off); there is no state to refresh.")) + } + } + + private def resultMap(r: RefreshResult): Map[String, Int] = + Map("added" -> r.added, "total" -> r.total, "failed" -> r.failed) + + post("/sessions/refresh") { + info(s"Interactive session refresh triggered by user='${remoteUser(request)}'") + resultMap(interactiveSessionManager.refresh()) + } + + post("/batches/refresh") { + info(s"Batch session refresh triggered by user='${remoteUser(request)}'") + resultMap(batchSessionManager.refresh()) + } + + post("/refresh") { + info(s"Full session refresh triggered by user='${remoteUser(request)}'") + Map( + "batches" -> resultMap(batchSessionManager.refresh()), + "sessions" -> resultMap(interactiveSessionManager.refresh()) + ) + } + +} diff --git a/server/src/test/scala/org/apache/livy/server/recovery/RecoveryServletSpec.scala b/server/src/test/scala/org/apache/livy/server/recovery/RecoveryServletSpec.scala new file mode 100644 index 000000000..61c70d5b6 --- /dev/null +++ b/server/src/test/scala/org/apache/livy/server/recovery/RecoveryServletSpec.scala @@ -0,0 +1,140 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.livy.server.recovery + +import javax.servlet.http.HttpServletRequest + +import org.mockito.Mockito.when +import org.scalatestplus.mockito.MockitoSugar.mock + +import org.apache.livy.LivyConf +import org.apache.livy.server.{AccessManager, BaseJsonServletSpec} +import org.apache.livy.server.batch.BatchRecoveryMetadata +import org.apache.livy.server.interactive.InteractiveRecoveryMetadata +import org.apache.livy.sessions.{BatchSessionManager, InteractiveSessionManager} +import org.apache.livy.sessions.SessionManager.{SESSION_RECOVERY_MODE_OFF, + SESSION_RECOVERY_MODE_RECOVERY} + +object RecoveryServletSpec { + val REMOTE_USER_HEADER = "X-Livy-RecoveryServlet-User" + val ADMIN = "__admin__" + val REGULAR_USER = "__user__" +} + +/** Reads the test-only remote-user header instead of the (unavailable in tests) container user. */ +class TestRecoveryServlet( + livyConf: LivyConf, + accessManager: AccessManager, + interactiveSessionManager: InteractiveSessionManager, + batchSessionManager: BatchSessionManager) + extends RecoveryServlet(livyConf, accessManager, interactiveSessionManager, batchSessionManager) { + + override protected def remoteUser(req: HttpServletRequest): String = { + req.getHeader(RecoveryServletSpec.REMOTE_USER_HEADER) + } +} + +trait RecoveryServletSpecBase extends BaseJsonServletSpec { + + import RecoveryServletSpec._ + + protected def recoveryMode: String + + protected def headersFor(user: String): Map[String, String] = + defaultHeaders ++ Map(REMOTE_USER_HEADER -> user) + + private def mockSessionStore(): SessionStore = { + val sessionStore = mock[SessionStore] + when(sessionStore.getAllSessions[BatchRecoveryMetadata]("batch")).thenReturn(Seq.empty) + when(sessionStore.getAllSessions[InteractiveRecoveryMetadata]("interactive")) + .thenReturn(Seq.empty) + when(sessionStore.getNextSessionId("batch")).thenReturn(0) + when(sessionStore.getNextSessionId("interactive")).thenReturn(0) + sessionStore + } + + private val livyConf = new LivyConf() + .set(LivyConf.SUPERUSERS, ADMIN) + .set(LivyConf.RECOVERY_MODE, recoveryMode) + private val accessManager = new AccessManager(livyConf) + private val sessionStore = mockSessionStore() + private val batchSessionManager = new BatchSessionManager(livyConf, sessionStore) + private val interactiveSessionManager = new InteractiveSessionManager(livyConf, sessionStore) + + addServlet( + new TestRecoveryServlet(livyConf, accessManager, interactiveSessionManager, + batchSessionManager), + "/*") +} + +class RecoveryServletEnabledSpec extends RecoveryServletSpecBase { + + import RecoveryServletSpec._ + + override protected def recoveryMode: String = SESSION_RECOVERY_MODE_RECOVERY + + describe("RecoveryServlet with recovery enabled") { + + it("rejects non-superusers with 403") { + post("/refresh", headers = headersFor(REGULAR_USER)) { + status should be (403) + } + } + + it("allows superusers and returns refresh counts with 200") { + post("/refresh", headers = headersFor(ADMIN)) { + status should be (200) + body should include ("\"sessions\"") + body should include ("\"batches\"") + } + } + + it("allows superusers on the single-manager endpoints with 200") { + post("/sessions/refresh", headers = headersFor(ADMIN)) { + status should be (200) + body should include ("\"added\"") + } + post("/batches/refresh", headers = headersFor(ADMIN)) { + status should be (200) + body should include ("\"added\"") + } + } + } +} + +class RecoveryServletDisabledSpec extends RecoveryServletSpecBase { + + import RecoveryServletSpec._ + + override protected def recoveryMode: String = SESSION_RECOVERY_MODE_OFF + + describe("RecoveryServlet with recovery disabled") { + + it("returns 409 when recovery is disabled, even for superusers") { + post("/refresh", headers = headersFor(ADMIN)) { + status should be (409) + } + } + + it("checks authorization before the recovery-mode check") { + post("/refresh", headers = headersFor(REGULAR_USER)) { + status should be (403) + } + } + } +}