LIVY-239: Moving the logic to generate session IDs from Session Manager

to SessionStore

Task-url: https://issues.cloudera.org/browse/LIVY-239

Author: meisam

server/src/main/scala/com/cloudera/livy/server/recovery/SessionStore.scala

@@ -47,10 +47,6 @@ class SessionStore(
     store.set(sessionPath(sessionType, m.id), m)
   }
 
-  def saveNextSessionId(sessionType: String, id: Int): Unit = {
-    store.set(sessionManagerPath(sessionType), SessionManagerState(id))
-  }
-
   /**
    * Return all sessions stored in the store with specified session type.
    */
@@ -69,14 +65,17 @@ class SessionStore(
 
   /**
    * Return the next unused session id with specified session type.
-   * If checks the SessionManagerState stored and returns the next free session id.
+   * It checks the SessionManagerState stored and returns the next free session id.
    * If no SessionManagerState is stored, it returns 0.
+    * It saves the new session ID to the session store.
    *
    * @throws Exception If SessionManagerState stored is corrupted, it throws an error.
    */
-  def getNextSessionId(sessionType: String): Int = {
-    store.get[SessionManagerState](sessionManagerPath(sessionType))
+  def getNextSessionId(sessionType: String): Int = synchronized {
+    val nextSessionId = store.get[SessionManagerState](sessionManagerPath(sessionType))
       .map(_.nextSessionId).getOrElse(0)
+    store.set(sessionManagerPath(sessionType), SessionManagerState(nextSessionId + 1))
+    nextSessionId
   }
 
   /**

server/src/main/scala/com/cloudera/livy/sessions/SessionManager.scala

@@ -69,7 +69,6 @@ class SessionManager[S <: Session, R <: RecoveryMetadata : ClassTag](
 
   protected implicit def executor: ExecutionContext = ExecutionContext.global
 
-  protected[this] final val idCounter = new AtomicInteger(0)
   protected[this] final val sessions = mutable.LinkedHashMap[Int, S]()
 
   private[this] final val sessionTimeout =
@@ -78,11 +77,8 @@ class SessionManager[S <: Session, R <: RecoveryMetadata : ClassTag](
   mockSessions.getOrElse(recover()).foreach(register)
   new GarbageCollector().start()
 
-  def nextId(): Int = synchronized {
-    val id = idCounter.getAndIncrement()
-    sessionStore.saveNextSessionId(sessionType, idCounter.get())
-    id
-  }
+  // sessionStore.getNextSessionId is guaranteed to return  atomic and returns unique IDs.
+  def nextId(): Int = sessionStore.getNextSessionId(sessionType)
 
   def register(session: S): S = {
     info(s"Registering new session ${session.id}")
@@ -136,18 +132,13 @@ class SessionManager[S <: Session, R <: RecoveryMetadata : ClassTag](
   }
 
   private def recover(): Seq[S] = {
-    // Recover next session id from state store and create SessionManager.
-    idCounter.set(sessionStore.getNextSessionId(sessionType))
 
     // Retrieve session recovery metadata from state store.
     val sessionMetadata = sessionStore.getAllSessions[R](sessionType)
 
     // Recover session from session recovery metadata.
     val recoveredSessions = sessionMetadata.flatMap(_.toOption).map(sessionRecovery)
 
-    info(s"Recovered ${recoveredSessions.length} $sessionType sessions." +
-      s" Next session id: $idCounter")
-
     // Print recovery error.
     val recoveryFailure = sessionMetadata.filter(_.isFailure).map(_.failed.get)
     recoveryFailure.foreach(ex => error(ex.getMessage, ex.getCause))