Browse Source

[livy] Convert session ids to integers

Erick Tryzelaar 10 years ago
parent
commit
dca1e78c32

+ 3 - 3
apps/spark/java/livy-server/src/main/scala/com/cloudera/hue/livy/server/sessions/ProcessSession.scala

@@ -37,13 +37,13 @@ object ProcessSession extends Logging {
   val CONF_LIVY_REPL_CALLBACK_URL = "livy.repl.callback-url"
   val CONF_LIVY_REPL_CALLBACK_URL = "livy.repl.callback-url"
   val CONF_LIVY_REPL_DRIVER_CLASS_PATH = "livy.repl.driverClassPath"
   val CONF_LIVY_REPL_DRIVER_CLASS_PATH = "livy.repl.driverClassPath"
 
 
-  def create(livyConf: LivyConf, id: String, kind: Kind, proxyUser: Option[String] = None): Session = {
+  def create(livyConf: LivyConf, id: Int, kind: Kind, proxyUser: Option[String] = None): Session = {
     val process = startProcess(livyConf, id, kind, proxyUser)
     val process = startProcess(livyConf, id, kind, proxyUser)
     new ProcessSession(id, kind, proxyUser, process)
     new ProcessSession(id, kind, proxyUser, process)
   }
   }
 
 
   // Loop until we've started a process with a valid port.
   // Loop until we've started a process with a valid port.
-  private def startProcess(livyConf: LivyConf, id: String, kind: Kind, proxyUser: Option[String]): Process = {
+  private def startProcess(livyConf: LivyConf, id: Int, kind: Kind, proxyUser: Option[String]): Process = {
 
 
     val builder = new SparkSubmitProcessBuilder()
     val builder = new SparkSubmitProcessBuilder()
 
 
@@ -72,7 +72,7 @@ object ProcessSession extends Logging {
   }
   }
 }
 }
 
 
-private class ProcessSession(id: String,
+private class ProcessSession(id: Int,
                              kind: Kind,
                              kind: Kind,
                              proxyUser: Option[String],
                              proxyUser: Option[String],
                              process: Process) extends WebSession(id, kind, proxyUser) {
                              process: Process) extends WebSession(id, kind, proxyUser) {

+ 1 - 1
apps/spark/java/livy-server/src/main/scala/com/cloudera/hue/livy/server/sessions/Session.scala

@@ -35,7 +35,7 @@ object Session {
 }
 }
 
 
 trait Session {
 trait Session {
-  def id: String
+  def id: Int
 
 
   def kind: Kind
   def kind: Kind
 
 

+ 4 - 10
apps/spark/java/livy-server/src/main/scala/com/cloudera/hue/livy/server/sessions/SessionFactory.scala

@@ -18,8 +18,6 @@
 
 
 package com.cloudera.hue.livy.server.sessions
 package com.cloudera.hue.livy.server.sessions
 
 
-import java.util.UUID
-
 import com.cloudera.hue.livy.LivyConf
 import com.cloudera.hue.livy.LivyConf
 import com.cloudera.hue.livy.sessions.Kind
 import com.cloudera.hue.livy.sessions.Kind
 import com.cloudera.hue.livy.yarn.Client
 import com.cloudera.hue.livy.yarn.Client
@@ -27,7 +25,7 @@ import com.cloudera.hue.livy.yarn.Client
 import scala.concurrent.{ExecutionContext, Future}
 import scala.concurrent.{ExecutionContext, Future}
 
 
 trait SessionFactory {
 trait SessionFactory {
-  def createSession(kind: Kind, proxyUser: Option[String] = None): Future[Session]
+  def createSession(id: Int, kind: Kind, proxyUser: Option[String] = None): Future[Session]
 
 
   def close(): Unit = {}
   def close(): Unit = {}
 }
 }
@@ -36,9 +34,8 @@ class ThreadSessionFactory(livyConf: LivyConf) extends SessionFactory {
 
 
   implicit def executor: ExecutionContext = ExecutionContext.global
   implicit def executor: ExecutionContext = ExecutionContext.global
 
 
-  override def createSession(kind: Kind, proxyUser: Option[String] = None): Future[Session] = {
+  override def createSession(id: Int, kind: Kind, proxyUser: Option[String] = None): Future[Session] = {
     Future {
     Future {
-      val id = UUID.randomUUID().toString
       ThreadSession.create(id, kind)
       ThreadSession.create(id, kind)
     }
     }
   }
   }
@@ -48,9 +45,8 @@ class ProcessSessionFactory(livyConf: LivyConf) extends SessionFactory {
 
 
   implicit def executor: ExecutionContext = ExecutionContext.global
   implicit def executor: ExecutionContext = ExecutionContext.global
 
 
-  override def createSession(kind: Kind, proxyUser: Option[String] = None): Future[Session] = {
+  override def createSession(id: Int, kind: Kind, proxyUser: Option[String] = None): Future[Session] = {
     Future {
     Future {
-      val id = UUID.randomUUID().toString
       ProcessSession.create(livyConf, id, kind, proxyUser)
       ProcessSession.create(livyConf, id, kind, proxyUser)
     }
     }
   }
   }
@@ -62,9 +58,7 @@ class YarnSessionFactory(livyConf: LivyConf) extends SessionFactory {
 
 
   val client = new Client(livyConf)
   val client = new Client(livyConf)
 
 
-  override def createSession(kind: Kind, proxyUser: Option[String] = None): Future[Session] = {
-    val id = UUID.randomUUID().toString
-
+  override def createSession(id: Int, kind: Kind, proxyUser: Option[String] = None): Future[Session] = {
     Future {
     Future {
       YarnSession.create(livyConf, client, id, kind, proxyUser)
       YarnSession.create(livyConf, client, id, kind, proxyUser)
     }
     }

+ 12 - 7
apps/spark/java/livy-server/src/main/scala/com/cloudera/hue/livy/server/sessions/SessionManager.scala

@@ -18,10 +18,13 @@
 
 
 package com.cloudera.hue.livy.server.sessions
 package com.cloudera.hue.livy.server.sessions
 
 
+import java.util.concurrent.ConcurrentHashMap
+import java.util.concurrent.atomic.AtomicInteger
+
 import com.cloudera.hue.livy.Logging
 import com.cloudera.hue.livy.Logging
 import com.cloudera.hue.livy.sessions.Kind
 import com.cloudera.hue.livy.sessions.Kind
 
 
-import scala.collection.concurrent.TrieMap
+import scala.collection.JavaConversions._
 import scala.concurrent.duration.Duration
 import scala.concurrent.duration.Duration
 import scala.concurrent.{Await, ExecutionContext, ExecutionContextExecutor, Future}
 import scala.concurrent.{Await, ExecutionContext, ExecutionContextExecutor, Future}
 
 
@@ -37,13 +40,14 @@ class SessionManager(factory: SessionFactory) extends Logging {
 
 
   private implicit def executor: ExecutionContextExecutor = ExecutionContext.global
   private implicit def executor: ExecutionContextExecutor = ExecutionContext.global
 
 
-  private val sessions = new TrieMap[String, Session]()
+  private[this] val _idCounter = new AtomicInteger()
+  private[this] val sessions = new ConcurrentHashMap[Int, Session]()
 
 
   private val garbageCollector = new GarbageCollector(this)
   private val garbageCollector = new GarbageCollector(this)
   garbageCollector.start()
   garbageCollector.start()
 
 
-  def get(id: String): Option[Session] = {
-    sessions.get(id)
+  def get(sessionId: Int): Option[Session] = {
+    Option(sessions.get(sessionId))
   }
   }
 
 
   def getSessions = {
   def getSessions = {
@@ -55,7 +59,8 @@ class SessionManager(factory: SessionFactory) extends Logging {
   }
   }
 
 
   def createSession(kind: Kind, proxyUser: Option[String] = None): Future[Session] = {
   def createSession(kind: Kind, proxyUser: Option[String] = None): Future[Session] = {
-    val session = factory.createSession(kind, proxyUser = proxyUser)
+    val id = _idCounter.getAndIncrement
+    val session = factory.createSession(id, kind, proxyUser = proxyUser)
 
 
     session.map({ case(session: Session) =>
     session.map({ case(session: Session) =>
       info("created session %s" format session.id)
       info("created session %s" format session.id)
@@ -69,8 +74,8 @@ class SessionManager(factory: SessionFactory) extends Logging {
     garbageCollector.shutdown()
     garbageCollector.shutdown()
   }
   }
 
 
-  def delete(sessionId: String): Future[Unit] = {
-    sessions.get(sessionId) match {
+  def delete(sessionId: Int): Future[Unit] = {
+    get(sessionId) match {
       case Some(session) => delete(session)
       case Some(session) => delete(session)
       case None => Future.successful(Unit)
       case None => Future.successful(Unit)
     }
     }

+ 22 - 9
apps/spark/java/livy-server/src/main/scala/com/cloudera/hue/livy/server/sessions/SessionServlet.scala

@@ -57,7 +57,9 @@ class SessionServlet(sessionManager: SessionManager)
   }
   }
 
 
   val getSession = get("/:sessionId") {
   val getSession = get("/:sessionId") {
-    sessionManager.get(params("sessionId")) match {
+    val sessionId = params("sessionId").toInt
+
+    sessionManager.get(sessionId) match {
       case Some(session) => session
       case Some(session) => session
       case None => NotFound("Session not found")
       case None => NotFound("Session not found")
     }
     }
@@ -79,9 +81,10 @@ class SessionServlet(sessionManager: SessionManager)
   }
   }
 
 
   post("/:sessionId/callback") {
   post("/:sessionId/callback") {
+    val sessionId = params("sessionId").toInt
     val callback = parsedBody.extract[CallbackRequest]
     val callback = parsedBody.extract[CallbackRequest]
 
 
-    sessionManager.get(params("sessionId")) match {
+    sessionManager.get(sessionId) match {
       case Some(session) =>
       case Some(session) =>
         if (session.state == Starting()) {
         if (session.state == Starting()) {
           session.url = new URL(callback.url)
           session.url = new URL(callback.url)
@@ -94,7 +97,8 @@ class SessionServlet(sessionManager: SessionManager)
   }
   }
 
 
   post("/:sessionId/stop") {
   post("/:sessionId/stop") {
-    sessionManager.get(params("sessionId")) match {
+    val sessionId = params("sessionId").toInt
+    sessionManager.get(sessionId) match {
       case Some(session) =>
       case Some(session) =>
         val future = session.stop()
         val future = session.stop()
 
 
@@ -104,7 +108,8 @@ class SessionServlet(sessionManager: SessionManager)
   }
   }
 
 
   post("/:sessionId/interrupt") {
   post("/:sessionId/interrupt") {
-    sessionManager.get(params("sessionId")) match {
+    val sessionId = params("sessionId").toInt
+    sessionManager.get(sessionId) match {
       case Some(session) =>
       case Some(session) =>
         val future = for {
         val future = for {
           _ <- session.interrupt()
           _ <- session.interrupt()
@@ -117,15 +122,19 @@ class SessionServlet(sessionManager: SessionManager)
   }
   }
 
 
   delete("/:sessionId") {
   delete("/:sessionId") {
+    val sessionId = params("sessionId").toInt
+
     val future = for {
     val future = for {
-      _ <- sessionManager.delete(params("sessionId"))
+      _ <- sessionManager.delete(sessionId)
     } yield Accepted()
     } yield Accepted()
 
 
     new AsyncResult() { val is = for { _ <- future } yield NoContent() }
     new AsyncResult() { val is = for { _ <- future } yield NoContent() }
   }
   }
 
 
   get("/:sessionId/statements") {
   get("/:sessionId/statements") {
-    sessionManager.get(params("sessionId")) match {
+    val sessionId = params("sessionId").toInt
+
+    sessionManager.get(sessionId) match {
       case Some(session: Session) =>
       case Some(session: Session) =>
         Map(
         Map(
           "statements" -> session.statements()
           "statements" -> session.statements()
@@ -135,9 +144,12 @@ class SessionServlet(sessionManager: SessionManager)
   }
   }
 
 
   val getStatement = get("/:sessionId/statements/:statementId") {
   val getStatement = get("/:sessionId/statements/:statementId") {
-    sessionManager.get(params("sessionId")) match {
+    val sessionId = params("sessionId").toInt
+    val statementId = params("statementId").toInt
+
+    sessionManager.get(sessionId) match {
       case Some(session) =>
       case Some(session) =>
-        session.statement(params("statementId").toInt) match {
+        session.statement(statementId) match {
           case Some(statement) => statement
           case Some(statement) => statement
           case None => NotFound("Statement not found")
           case None => NotFound("Statement not found")
         }
         }
@@ -146,9 +158,10 @@ class SessionServlet(sessionManager: SessionManager)
   }
   }
 
 
   post("/:sessionId/statements") {
   post("/:sessionId/statements") {
+    val sessionId = params("sessionId").toInt
     val req = parsedBody.extract[ExecuteRequest]
     val req = parsedBody.extract[ExecuteRequest]
 
 
-    sessionManager.get(params("sessionId")) match {
+    sessionManager.get(sessionId) match {
       case Some(session) =>
       case Some(session) =>
         val statement = session.executeStatement(req)
         val statement = session.executeStatement(req)
 
 

+ 2 - 2
apps/spark/java/livy-server/src/main/scala/com/cloudera/hue/livy/server/sessions/ThreadSession.scala

@@ -32,7 +32,7 @@ object ThreadSession {
   val LIVY_HOME = System.getenv("LIVY_HOME")
   val LIVY_HOME = System.getenv("LIVY_HOME")
   val LIVY_REPL = LIVY_HOME + "/bin/livy-repl"
   val LIVY_REPL = LIVY_HOME + "/bin/livy-repl"
 
 
-  def create(id: String, kind: Kind): Session = {
+  def create(id: Int, kind: Kind): Session = {
     val session = kind match {
     val session = kind match {
       case Spark() =>
       case Spark() =>
         SparkSession.create()
         SparkSession.create()
@@ -43,7 +43,7 @@ object ThreadSession {
   }
   }
 }
 }
 
 
-private class ThreadSession(val id: String,
+private class ThreadSession(val id: Int,
                             val kind: Kind,
                             val kind: Kind,
                             session: com.cloudera.hue.livy.repl.Session) extends Session {
                             session: com.cloudera.hue.livy.repl.Session) extends Session {
 
 

+ 1 - 1
apps/spark/java/livy-server/src/main/scala/com/cloudera/hue/livy/server/sessions/WebSession.scala

@@ -32,7 +32,7 @@ import scala.collection.mutable.ArrayBuffer
 import scala.concurrent.duration.Duration
 import scala.concurrent.duration.Duration
 import scala.concurrent.{Future, _}
 import scala.concurrent.{Future, _}
 
 
-class WebSession(val id: String,
+class WebSession(val id: Int,
                  val kind: Kind,
                  val kind: Kind,
                  val proxyUser: Option[String]) extends Session with Logging {
                  val proxyUser: Option[String]) extends Session with Logging {
 
 

+ 2 - 2
apps/spark/java/livy-server/src/main/scala/com/cloudera/hue/livy/server/sessions/YarnSession.scala

@@ -35,7 +35,7 @@ object YarnSession {
   private val CONF_LIVY_JAR = "livy.yarn.jar"
   private val CONF_LIVY_JAR = "livy.yarn.jar"
   private lazy val regex = """Application report for (\w+)""".r.unanchored
   private lazy val regex = """Application report for (\w+)""".r.unanchored
 
 
-  def create(livyConf: LivyConf, client: Client, id: String, kind: Kind, proxyUser: Option[String] = None): Session = {
+  def create(livyConf: LivyConf, client: Client, id: Int, kind: Kind, proxyUser: Option[String] = None): Session = {
     val callbackUrl = System.getProperty("livy.server.callback-url")
     val callbackUrl = System.getProperty("livy.server.callback-url")
     val url = f"$callbackUrl/sessions/$id/callback"
     val url = f"$callbackUrl/sessions/$id/callback"
 
 
@@ -73,7 +73,7 @@ object YarnSession {
   }
   }
 }
 }
 
 
-private class YarnSession(id: String,
+private class YarnSession(id: Int,
                           kind: Kind,
                           kind: Kind,
                           proxyUser: Option[String],
                           proxyUser: Option[String],
                           job: Future[Job]) extends WebSession(id, kind, proxyUser) {
                           job: Future[Job]) extends WebSession(id, kind, proxyUser) {

+ 1 - 1
apps/spark/java/livy-server/src/test/scala/com/cloudera/hue/livy/server/ProcessSessionSpec.scala

@@ -28,5 +28,5 @@ class ProcessSessionSpec extends BaseSessionSpec with FunSpecLike with Matchers
   val livyConf = new LivyConf()
   val livyConf = new LivyConf()
   livyConf.set("livy.repl.driverClassPath", sys.props("java.class.path"))
   livyConf.set("livy.repl.driverClassPath", sys.props("java.class.path"))
 
 
-  def createSession() = ProcessSession.create(livyConf, "0", Spark())
+  def createSession() = ProcessSession.create(livyConf, 0, Spark())
 }
 }

+ 3 - 5
apps/spark/java/livy-server/src/test/scala/com/cloudera/hue/livy/server/SessionServletSpec.scala

@@ -35,7 +35,7 @@ import scala.concurrent.Future
 
 
 class SessionServletSpec extends ScalatraSuite with FunSpecLike {
 class SessionServletSpec extends ScalatraSuite with FunSpecLike {
 
 
-  class MockSession extends Session {
+  class MockSession(val id: Int) extends Session {
     var _state: State = Idle()
     var _state: State = Idle()
 
 
     var _idCounter = new AtomicInteger()
     var _idCounter = new AtomicInteger()
@@ -45,8 +45,6 @@ class SessionServletSpec extends ScalatraSuite with FunSpecLike {
 
 
     override def state = _state
     override def state = _state
 
 
-    override def id: String = ???
-
     override def stop(): Future[Unit] = ???
     override def stop(): Future[Unit] = ???
 
 
     override def url_=(url: URL): Unit = ???
     override def url_=(url: URL): Unit = ???
@@ -79,8 +77,8 @@ class SessionServletSpec extends ScalatraSuite with FunSpecLike {
   }
   }
 
 
   class MockSessionFactory() extends SessionFactory {
   class MockSessionFactory() extends SessionFactory {
-    override def createSession(kind: Kind, proxyUser: Option[String]): Future[Session] = {
-      Future.successful(new MockSession())
+    override def createSession(id: Int, kind: Kind, proxyUser: Option[String]): Future[Session] = {
+      Future.successful(new MockSession(id))
     }
     }
   }
   }
 
 

+ 1 - 1
apps/spark/java/livy-server/src/test/scala/com/cloudera/hue/livy/server/ThreadSessionSpec.scala

@@ -24,5 +24,5 @@ import org.scalatest.{BeforeAndAfter, FunSpecLike, Matchers}
 
 
 class ThreadSessionSpec extends BaseSessionSpec with FunSpecLike with Matchers with BeforeAndAfter {
 class ThreadSessionSpec extends BaseSessionSpec with FunSpecLike with Matchers with BeforeAndAfter {
 
 
-  def createSession() = ThreadSession.create("0", Spark())
+  def createSession() = ThreadSession.create(0, Spark())
 }
 }