[LIVY-1073] Add SessionManager refresh and /recovery operator endpoint

docs/rest-api.md

@@ -694,6 +694,49 @@ instance.
   </tr>
 </table>
 
+### 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
+
+<table class="table">
+  <tr><th>Name</th><th>Description</th><th>Type</th></tr>
+  <tr>
+    <td>added</td>
+    <td>Number of sessions imported from the state store that weren't already in memory</td>
+    <td>int</td>
+  </tr>
+  <tr>
+    <td>total</td>
+    <td>Total number of interactive sessions in memory after the refresh</td>
+    <td>int</td>
+  </tr>
+  <tr>
+    <td>failed</td>
+    <td>Number of state store entries that failed to deserialize</td>
+    <td>int</td>
+  </tr>
+</table>
+
+### 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

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)

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())
+    )
+  }
+
+}

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)

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)
+      }
+    }
+  }
+}