[LIVY-239] Generate session IDs in SessionStore, not in SessionManager

server/src/main/scala/org/apache/livy/server/recovery/BlackholeStateStore.scala

@@ -17,6 +17,8 @@
 
 package org.apache.livy.server.recovery
 
+import java.util.concurrent.atomic.AtomicLong
+
 import scala.reflect.ClassTag
 
 import org.apache.livy.LivyConf
@@ -26,11 +28,16 @@ import org.apache.livy.LivyConf
  * Livy will use this when session recovery is disabled.
  */
 class BlackholeStateStore(livyConf: LivyConf) extends StateStore(livyConf) {
+
+  private val atomicLong: AtomicLong = new AtomicLong(-1L)
+
   def set(key: String, value: Object): Unit = {}
 
   def get[T: ClassTag](key: String): Option[T] = None
 
   def getChildren(key: String): Seq[String] = List.empty[String]
 
   def remove(key: String): Unit = {}
+
+  override def nextValue(key: String): Long = atomicLong.incrementAndGet()
 }

server/src/main/scala/org/apache/livy/server/recovery/FileSystemStateStore.scala

@@ -129,4 +129,10 @@ class FileSystemStateStore(
   }
 
   private def absPath(key: String): Path = new Path(fsUri.getPath(), key)
+
+  override def nextValue(key: String): Long = synchronized {
+    val incrementedValue = 1 + get[Long](key).getOrElse(-1L)
+    set(key, incrementedValue.asInstanceOf[Object])
+    incrementedValue
+  }
 }

server/src/main/scala/org/apache/livy/server/recovery/SessionStore.scala

@@ -26,8 +26,6 @@ import scala.util.control.NonFatal
 import org.apache.livy.{LivyConf, Logging}
 import org.apache.livy.sessions.Session.RecoveryMetadata
 
-private[recovery] case class SessionManagerState(nextSessionId: Int)
-
 /**
  * SessionStore provides high level functions to get/save session state from/to StateStore.
  */
@@ -46,10 +44,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.
    */
@@ -67,15 +61,14 @@ class SessionStore(
   }
 
   /**
-   * Return the next unused session id with specified session type.
-   * If checks the SessionManagerState stored and returns the next free session id.
-   * If no SessionManagerState is stored, it returns 0.
+   * Return the next unused session ID from the state store for the specified session type.
+   * Return 0 if no value is stored in the state store.
+   * Save the next unused session ID to the session store before returning the current value.
    *
-   * @throws Exception If SessionManagerState stored is corrupted, it throws an error.
+   * @throws Exception If session store is corrupted or unreachable, it throws an error.
    */
   def getNextSessionId(sessionType: String): Int = {
-    store.get[SessionManagerState](sessionManagerPath(sessionType))
-      .map(_.nextSessionId).getOrElse(0)
+    store.nextValue(nextSessionIdPath(sessionType)).toInt
   }
 
   /**
@@ -85,8 +78,8 @@ class SessionStore(
     store.remove(sessionPath(sessionType, id))
   }
 
-  private def sessionManagerPath(sessionType: String): String =
-    s"$STORE_VERSION/$sessionType/state"
+  private def nextSessionIdPath(sessionType: String): String =
+    s"$STORE_VERSION/$sessionType/nextSessionId"
 
   private def sessionPath(sessionType: String): String =
     s"$STORE_VERSION/$sessionType"

server/src/main/scala/org/apache/livy/server/recovery/StateStore.scala

@@ -70,6 +70,14 @@ abstract class StateStore(livyConf: LivyConf) extends JsonMapper {
    * @throws Exception Throw when persisting the state store fails.
    */
   def remove(key: String): Unit
+
+  /**
+    * Gets the Long value for the given key, increments the value
+    * and stores the new value before returning the value.
+    *
+    * @return incremented value
+    */
+  def nextValue(key: String): Long
 }
 
 /**

server/src/main/scala/org/apache/livy/server/recovery/ZooKeeperStateStore.scala

@@ -18,10 +18,10 @@ package org.apache.livy.server.recovery
 
 import scala.collection.JavaConverters._
 import scala.reflect.ClassTag
-import scala.util.Try
 
 import org.apache.curator.framework.{CuratorFramework, CuratorFrameworkFactory}
 import org.apache.curator.framework.api.UnhandledErrorListener
+import org.apache.curator.framework.recipes.atomic.{DistributedAtomicLong => DistributedLong}
 import org.apache.curator.retry.RetryNTimes
 import org.apache.zookeeper.KeeperException.NoNodeException
 
@@ -114,5 +114,15 @@ class ZooKeeperStateStore(
     }
   }
 
+  override def nextValue(key: String): Long = {
+    val distributedSessionId = new DistributedLong(curatorClient, key, retryPolicy)
+    distributedSessionId.increment() match {
+      case atomicValue if atomicValue.succeeded() =>
+        atomicValue.postValue()
+      case _ =>
+        throw new java.io.IOException(s"Failed to atomically nextValueget the next value for $key.")
+    }
+  }
+
   private def prefixKey(key: String) = s"/$zkKeyPrefix/$key"
 }

server/src/main/scala/org/apache/livy/sessions/SessionManager.scala

@@ -71,7 +71,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 sessionsByName = mutable.HashMap[String, S]()
 
@@ -85,11 +84,7 @@ 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
-  }
+  def nextId(): Int = sessionStore.getNextSessionId(sessionType)
 
   def register(session: S): S = {
     info(s"Registering new session ${session.id}")
@@ -169,18 +164,12 @@ 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))

server/src/test/scala/org/apache/livy/server/batch/BatchServletSpec.scala

@@ -25,7 +25,10 @@ import javax.servlet.http.HttpServletResponse._
 
 import scala.concurrent.duration.Duration
 
+import org.mockito.Matchers.anyString
 import org.mockito.Mockito._
+import org.mockito.invocation.InvocationOnMock
+import org.mockito.stubbing.Answer
 import org.scalatest.mock.MockitoSugar.mock
 
 import org.apache.livy.{LivyConf, Utils}
@@ -54,6 +57,17 @@ class BatchServletSpec extends BaseSessionServletSpec[BatchSession, BatchRecover
   override def createServlet(): BatchSessionServlet = {
     val livyConf = createConf()
     val sessionStore = mock[SessionStore]
+    val nextSessionIdAnswer = new Answer[Int] {
+      private var lastUsedId = -1
+
+      def answer(invokation: InvocationOnMock): Int = {
+        lastUsedId += 1
+        lastUsedId
+      }
+    }
+
+    when(sessionStore.getNextSessionId(anyString)).thenAnswer(nextSessionIdAnswer)
+
     val accessManager = new AccessManager(livyConf)
     new BatchSessionServlet(
       new BatchSessionManager(livyConf, sessionStore, Some(Seq.empty)),

server/src/test/scala/org/apache/livy/server/recovery/BlackholeStateStoreSpec.scala

@@ -45,6 +45,12 @@ class BlackholeStateStoreSpec extends FunSpec with LivyBaseUnitTestSuite {
       stateStore.remove("")
     }
 
+    it("nextValue should return consecutive numbers starting from 0") {
+      stateStore.nextValue("") shouldBe 0
+      stateStore.nextValue("") shouldBe 1
+      stateStore.nextValue("") shouldBe 2
+    }
+
     it("should deserialize sessions without name") {
       val jsonbytes =
         """

server/src/test/scala/org/apache/livy/server/recovery/SessionStoreSpec.scala

@@ -33,7 +33,7 @@ class SessionStoreSpec extends FunSpec with LivyBaseUnitTestSuite {
 
     val sessionType = "test"
     val sessionPath = s"v1/$sessionType"
-    val sessionManagerPath = s"v1/$sessionType/state"
+    val sessionManagerPath = s"v1/$sessionType/nextSessionId"
 
     val conf = new LivyConf()
     it("should set session state and session counter when saving a session.") {
@@ -88,12 +88,11 @@ class SessionStoreSpec extends FunSpec with LivyBaseUnitTestSuite {
       val stateStore = mock[StateStore]
       val sessionStore = new SessionStore(conf, stateStore)
 
-      when(stateStore.get[SessionManagerState](sessionManagerPath)).thenReturn(None)
+      when(stateStore.nextValue(sessionManagerPath)).thenReturn(0L)
       sessionStore.getNextSessionId(sessionType) shouldBe 0
 
-      val sms = SessionManagerState(100)
-      when(stateStore.get[SessionManagerState](sessionManagerPath)).thenReturn(Some(sms))
-      sessionStore.getNextSessionId(sessionType) shouldBe sms.nextSessionId
+      when(stateStore.nextValue(sessionManagerPath)).thenReturn(100)
+      sessionStore.getNextSessionId(sessionType) shouldBe 100
     }
 
     it("should remove session") {

server/src/test/scala/org/apache/livy/sessions/SessionManagerSpec.scala

@@ -17,12 +17,17 @@
 
 package org.apache.livy.sessions
 
+import java.util.concurrent.atomic.AtomicInteger
+
 import scala.concurrent.{Await, ExecutionContext, Future}
 import scala.concurrent.duration._
 import scala.language.postfixOps
 import scala.util.{Failure, Try}
 
+import org.mockito.Matchers.anyString
 import org.mockito.Mockito.{doReturn, never, verify, when}
+import org.mockito.invocation.InvocationOnMock
+import org.mockito.stubbing.Answer
 import org.scalatest.{FunSpec, Matchers}
 import org.scalatest.concurrent.Eventually._
 import org.scalatest.mock.MockitoSugar.mock
@@ -39,10 +44,19 @@ class SessionManagerSpec extends FunSpec with Matchers with LivyBaseUnitTestSuit
   private def createSessionManager(): (LivyConf, SessionManager[MockSession, RecoveryMetadata]) = {
     val livyConf = new LivyConf()
     livyConf.set(LivyConf.SESSION_TIMEOUT, "100ms")
+    val sessionStore = mock[SessionStore]
+    when(sessionStore.getNextSessionId(anyString())).thenAnswer( new Answer[Int] {
+
+      private val nextSessionId = new AtomicInteger(0)
+      override def answer(invocationOnMock: InvocationOnMock): Int = {
+        nextSessionId.getAndDecrement()
+      }
+    })
+
     val manager = new SessionManager[MockSession, RecoveryMetadata](
       livyConf,
       { _ => assert(false).asInstanceOf[MockSession] },
-      mock[SessionStore],
+      sessionStore,
       "test",
       Some(Seq.empty))
     (livyConf, manager)
@@ -72,6 +86,8 @@ class SessionManagerSpec extends FunSpec with Matchers with LivyBaseUnitTestSuit
       val name = "Mock-session"
       val session1 = new MockSession(manager.nextId(), null, livyConf, Some(name))
       val session2 = new MockSession(manager.nextId(), null, livyConf, Some(name))
+      session1.id shouldNot be(session2.id)
+
       manager.register(session1)
       an[IllegalArgumentException] should be thrownBy manager.register(session2)
       manager.get(session1.id).isDefined should be(true)