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 68a3880a9..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,7 +39,7 @@ 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} import org.apache.livy.sessions.SessionManager.SESSION_RECOVERY_MODE_OFF @@ -176,6 +176,11 @@ 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 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 // Returns 404 if the file does not exist @@ -256,6 +261,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/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/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/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) + } + } + } +} 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 + } } }