From 31f3bc4a197c2da786ff0072e1c475f88da10fb1 Mon Sep 17 00:00:00 2001 From: "Fathi Salmi, Meisam" Date: Mon, 16 Jul 2018 17:23:48 -0700 Subject: [PATCH 1/6] LIVY-239: Generate session IDs in SessionStore, not in SessionManager --- .../server/recovery/BlackholeStateStore.scala | 7 +++++++ .../server/recovery/FileSystemStateStore.scala | 6 ++++++ .../livy/server/recovery/SessionStore.scala | 17 +++++------------ .../livy/server/recovery/StateStore.scala | 8 ++++++++ .../server/recovery/ZooKeeperStateStore.scala | 13 +++++++++++++ .../apache/livy/sessions/SessionManager.scala | 13 +------------ .../livy/server/batch/BatchServletSpec.scala | 14 ++++++++++++++ .../livy/server/recovery/SessionStoreSpec.scala | 7 +++---- 8 files changed, 57 insertions(+), 28 deletions(-) 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..cfda2ad5d 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 increment(ket: 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..270167f38 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 increment(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..5183b1b97 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.increment(sessionManagerPath(sessionType)).toInt } /** 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..412e08d9c 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 increment(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..11b980481 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 @@ -19,9 +19,12 @@ package org.apache.livy.server.recovery import scala.collection.JavaConverters._ import scala.reflect.ClassTag import scala.util.Try +import scala.util.matching.Regex +import org.apache.curator.RetryPolicy 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 +117,15 @@ class ZooKeeperStateStore( } } + override def increment(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 increment the 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 d482c335d..cb496e40f 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 sessionTimeoutCheck = livyConf.getBoolean(LivyConf.SESSION_TIMEOUT_CHECK) @@ -83,11 +82,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}") @@ -153,18 +148,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 2c37c19d0..aec1abf33 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/SessionStoreSpec.scala b/server/src/test/scala/org/apache/livy/server/recovery/SessionStoreSpec.scala index 5eeb2cfc3..5d8c8e22a 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 @@ -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.increment(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.increment(sessionManagerPath)).thenReturn(100) + sessionStore.getNextSessionId(sessionType) shouldBe 100 } it("should remove session") { From 6bce0079ec34240c316ce512a6b21bcbdafe6cf9 Mon Sep 17 00:00:00 2001 From: "Fathi Salmi, Meisam(mfathisalmi)" Date: Tue, 17 Jul 2018 13:17:39 -0700 Subject: [PATCH 2/6] [LIVY-239] Fix the typo in BlackholeStateStore change ket --> key --- .../org/apache/livy/server/recovery/BlackholeStateStore.scala | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) 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 cfda2ad5d..8caa79d09 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 @@ -39,5 +39,5 @@ class BlackholeStateStore(livyConf: LivyConf) extends StateStore(livyConf) { def remove(key: String): Unit = {} - override def increment(ket: String): Long = atomicLong.incrementAndGet() + override def increment(key: String): Long = atomicLong.incrementAndGet() } From 38880d89653a8fa344bbb162348b60542a9c3c19 Mon Sep 17 00:00:00 2001 From: "Fathi Salmi, Meisam(mfathisalmi)" Date: Tue, 17 Jul 2018 13:18:35 -0700 Subject: [PATCH 3/6] [LIVY-239] Test increment in blackhole state store --- .../livy/server/recovery/BlackholeStateStoreSpec.scala | 6 ++++++ 1 file changed, 6 insertions(+) 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 e40bb1c03..dfff21945 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 @@ -43,5 +43,11 @@ class BlackholeStateStoreSpec extends FunSpec with LivyBaseUnitTestSuite { it("remove should not throw") { stateStore.remove("") } + + it("increment should return consecutive numbers starting from 0") { + stateStore.increment("") shouldBe 0 + stateStore.increment("") shouldBe 1 + stateStore.increment("") shouldBe 2 + } } } From 5e038e54501670303becd102bade8c03dabc897b Mon Sep 17 00:00:00 2001 From: Meisam Fathi Date: Thu, 7 Feb 2019 13:13:15 -0800 Subject: [PATCH 4/6] [LIVY-555] Do not accept new sessions on Livy shutdown. - Rename `StateStore.increment()` to `StateStore.nextValue()` - Change the path for storing nextSessionId to - `$STORE_VERSION/$sessionType/nextSessionId`. Task-url: https://issues.apache.org/jira/projects/LIVY/issues/LIVY-555 --- .../apache/livy/server/recovery/BlackholeStateStore.scala | 2 +- .../livy/server/recovery/FileSystemStateStore.scala | 2 +- .../org/apache/livy/server/recovery/SessionStore.scala | 6 +++--- .../org/apache/livy/server/recovery/StateStore.scala | 2 +- .../apache/livy/server/recovery/ZooKeeperStateStore.scala | 4 ++-- .../livy/server/recovery/BlackholeStateStoreSpec.scala | 8 ++++---- .../apache/livy/server/recovery/SessionStoreSpec.scala | 6 +++--- 7 files changed, 15 insertions(+), 15 deletions(-) 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 8caa79d09..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 @@ -39,5 +39,5 @@ class BlackholeStateStore(livyConf: LivyConf) extends StateStore(livyConf) { def remove(key: String): Unit = {} - override def increment(key: String): Long = atomicLong.incrementAndGet() + 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 270167f38..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 @@ -130,7 +130,7 @@ class FileSystemStateStore( private def absPath(key: String): Path = new Path(fsUri.getPath(), key) - override def increment(key: String): Long = synchronized { + 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 5183b1b97..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 @@ -68,7 +68,7 @@ class SessionStore( * @throws Exception If session store is corrupted or unreachable, it throws an error. */ def getNextSessionId(sessionType: String): Int = { - store.increment(sessionManagerPath(sessionType)).toInt + store.nextValue(nextSessionIdPath(sessionType)).toInt } /** @@ -78,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 412e08d9c..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 @@ -77,7 +77,7 @@ abstract class StateStore(livyConf: LivyConf) extends JsonMapper { * * @return incremented value */ - def increment(key: String): Long + 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 11b980481..33a4068aa 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 @@ -117,13 +117,13 @@ class ZooKeeperStateStore( } } - override def increment(key: String): Long = { + 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 increment the value for $key") + throw new java.io.IOException(s"Failed to atomically nextValueget the next value for $key.") } } 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 dfff21945..909fea1a2 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 @@ -44,10 +44,10 @@ class BlackholeStateStoreSpec extends FunSpec with LivyBaseUnitTestSuite { stateStore.remove("") } - it("increment should return consecutive numbers starting from 0") { - stateStore.increment("") shouldBe 0 - stateStore.increment("") shouldBe 1 - stateStore.increment("") shouldBe 2 + it("nextValue should return consecutive numbers starting from 0") { + stateStore.nextValue("") shouldBe 0 + stateStore.nextValue("") shouldBe 1 + stateStore.nextValue("") shouldBe 2 } } } 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 5d8c8e22a..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,10 +88,10 @@ class SessionStoreSpec extends FunSpec with LivyBaseUnitTestSuite { val stateStore = mock[StateStore] val sessionStore = new SessionStore(conf, stateStore) - when(stateStore.increment(sessionManagerPath)).thenReturn(0L) + when(stateStore.nextValue(sessionManagerPath)).thenReturn(0L) sessionStore.getNextSessionId(sessionType) shouldBe 0 - when(stateStore.increment(sessionManagerPath)).thenReturn(100) + when(stateStore.nextValue(sessionManagerPath)).thenReturn(100) sessionStore.getNextSessionId(sessionType) shouldBe 100 } From 24900504ab8f0fcfe0982d6d9856e1f7d1b01a2b Mon Sep 17 00:00:00 2001 From: Meisam Fathi Date: Thu, 7 Feb 2019 14:02:42 -0800 Subject: [PATCH 5/6] [LIVY-555] Merge with apache/master branch and resolve conflicts - Remove unused and unnecessary imports. Task-url: https://issues.apache.org/jira/projects/LIVY/issues/LIVY-555 --- .../org/apache/livy/server/recovery/ZooKeeperStateStore.scala | 3 --- 1 file changed, 3 deletions(-) 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 33a4068aa..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,7 @@ package org.apache.livy.server.recovery import scala.collection.JavaConverters._ import scala.reflect.ClassTag -import scala.util.Try -import scala.util.matching.Regex -import org.apache.curator.RetryPolicy import org.apache.curator.framework.{CuratorFramework, CuratorFrameworkFactory} import org.apache.curator.framework.api.UnhandledErrorListener import org.apache.curator.framework.recipes.atomic.{DistributedAtomicLong => DistributedLong} From 3baf2646a3cc107e0b6d232dc6a69e6fea8894b4 Mon Sep 17 00:00:00 2001 From: Meisam Fathi Date: Thu, 7 Feb 2019 15:30:46 -0800 Subject: [PATCH 6/6] [LIVY-555] Generated unique sessionIds in test cases - Properly mock SessionStore in test cases. Task-url: https://issues.apache.org/jira/projects/LIVY/issues/LIVY-555 --- .../livy/sessions/SessionManagerSpec.scala | 18 +++++++++++++++++- 1 file changed, 17 insertions(+), 1 deletion(-) 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)