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

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)

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