diff --git a/server/src/main/scala/org/apache/livy/server/recovery/BlackholeStateStore.scala b/server/src/main/scala/org/apache/livy/server/recovery/BlackholeStateStore.scala index df9a712a8..143cae889 100644 --- a/server/src/main/scala/org/apache/livy/server/recovery/BlackholeStateStore.scala +++ b/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,6 +28,9 @@ 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 @@ -33,4 +38,6 @@ class BlackholeStateStore(livyConf: LivyConf) extends StateStore(livyConf) { def getChildren(key: String): Seq[String] = List.empty[String] def remove(key: String): Unit = {} + + override def nextValue(key: String): Long = atomicLong.incrementAndGet() } diff --git a/server/src/main/scala/org/apache/livy/server/recovery/FileSystemStateStore.scala b/server/src/main/scala/org/apache/livy/server/recovery/FileSystemStateStore.scala index d5f8f3dc6..345bf6c9f 100644 --- a/server/src/main/scala/org/apache/livy/server/recovery/FileSystemStateStore.scala +++ b/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 + } } diff --git a/server/src/main/scala/org/apache/livy/server/recovery/SessionStore.scala b/server/src/main/scala/org/apache/livy/server/recovery/SessionStore.scala index 04292957c..b345c1340 100644 --- a/server/src/main/scala/org/apache/livy/server/recovery/SessionStore.scala +++ b/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" diff --git a/server/src/main/scala/org/apache/livy/server/recovery/StateStore.scala b/server/src/main/scala/org/apache/livy/server/recovery/StateStore.scala index a6c3275d7..14edd4d21 100644 --- a/server/src/main/scala/org/apache/livy/server/recovery/StateStore.scala +++ b/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 } /** diff --git a/server/src/main/scala/org/apache/livy/server/recovery/ZooKeeperStateStore.scala b/server/src/main/scala/org/apache/livy/server/recovery/ZooKeeperStateStore.scala index ec6b9df18..9b4c46ba1 100644 --- a/server/src/main/scala/org/apache/livy/server/recovery/ZooKeeperStateStore.scala +++ b/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" } 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 a63cab3e4..1e1371d79 100644 --- a/server/src/main/scala/org/apache/livy/sessions/SessionManager.scala +++ b/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)) diff --git a/server/src/test/scala/org/apache/livy/server/batch/BatchServletSpec.scala b/server/src/test/scala/org/apache/livy/server/batch/BatchServletSpec.scala index ed298000f..4a22bb7dc 100644 --- a/server/src/test/scala/org/apache/livy/server/batch/BatchServletSpec.scala +++ b/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)), diff --git a/server/src/test/scala/org/apache/livy/server/recovery/BlackholeStateStoreSpec.scala b/server/src/test/scala/org/apache/livy/server/recovery/BlackholeStateStoreSpec.scala index 8ee448f5e..073f77c6a 100644 --- a/server/src/test/scala/org/apache/livy/server/recovery/BlackholeStateStoreSpec.scala +++ b/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 = """ diff --git a/server/src/test/scala/org/apache/livy/server/recovery/SessionStoreSpec.scala b/server/src/test/scala/org/apache/livy/server/recovery/SessionStoreSpec.scala index 5eeb2cfc3..b544789d0 100644 --- a/server/src/test/scala/org/apache/livy/server/recovery/SessionStoreSpec.scala +++ b/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") { 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 523e1d7b2..7342e0955 100644 --- a/server/src/test/scala/org/apache/livy/sessions/SessionManagerSpec.scala +++ b/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)