Sfoglia il codice sorgente

[livy] Create a SparkManager

Erick Tryzelaar 10 anni fa
parent
commit
e091dfb

+ 8 - 40
apps/spark/java/livy-server/src/main/scala/com/cloudera/hue/livy/server/Main.scala

@@ -24,13 +24,7 @@ import javax.servlet.ServletContext
 import com.cloudera.hue.livy._
 import com.cloudera.hue.livy.server.batch.BatchSessionServlet
 import com.cloudera.hue.livy.server.interactive.InteractiveSessionServlet
-import com.cloudera.hue.livy.sessions.SessionManager
-import com.cloudera.hue.livy.sessions.batch.BatchSession
-import com.cloudera.hue.livy.sessions.interactive.InteractiveSession
-import com.cloudera.hue.livy.spark.SparkProcessBuilderFactory
-import com.cloudera.hue.livy.spark.batch.{BatchSessionProcessFactory, BatchSessionYarnFactory}
-import com.cloudera.hue.livy.spark.interactive.{InteractiveSessionYarnFactory, InteractiveSessionProcessFactory}
-import com.cloudera.hue.livy.yarn.Client
+import com.cloudera.hue.livy.spark.SparkManager
 import org.scalatra._
 import org.scalatra.metrics.MetricsBootstrap
 import org.scalatra.metrics.MetricsSupportExtensions._
@@ -155,38 +149,15 @@ class ScalatraBootstrap
   with Logging
   with MetricsBootstrap {
 
-  var sessionManager: SessionManager[InteractiveSession] = null
-  var batchManager: SessionManager[BatchSession] = null
+  var sparkManager: SparkManager = null
 
   override def init(context: ServletContext): Unit = {
     try {
       val livyConf = new LivyConf()
-      val sessionFactoryKind = livyConf.sessionKind()
+      sparkManager = SparkManager(livyConf)
 
-      info(f"Using $sessionFactoryKind sessions")
-
-      val processFactory = new SparkProcessBuilderFactory(livyConf)
-
-      val (sessionFactory, batchFactory) = sessionFactoryKind match {
-        case LivyConf.Process() =>
-          val interactiveFactory = new InteractiveSessionProcessFactory(processFactory)
-          val batchFactory = new BatchSessionProcessFactory(processFactory)
-
-          (interactiveFactory, batchFactory)
-
-        case LivyConf.Yarn() =>
-          val client = new Client(livyConf)
-          val interactiveFactory = new InteractiveSessionYarnFactory(client, processFactory)
-          val batchFactory = new BatchSessionYarnFactory(client, processFactory)
-
-          (interactiveFactory, batchFactory)
-      }
-
-      sessionManager = new SessionManager(livyConf, sessionFactory)
-      batchManager = new SessionManager(livyConf, batchFactory)
-
-      context.mount(new InteractiveSessionServlet(sessionManager), "/sessions/*")
-      context.mount(new BatchSessionServlet(batchManager), "/batches/*")
+      context.mount(new InteractiveSessionServlet(sparkManager.interactiveManager), "/sessions/*")
+      context.mount(new BatchSessionServlet(sparkManager.batchManager), "/batches/*")
       context.mountMetricsAdminServlet("/")
 
       context.initParameters(org.scalatra.EnvironmentKey) = livyConf.get("livy.environment", "development")
@@ -198,12 +169,9 @@ class ScalatraBootstrap
   }
 
   override def destroy(context: ServletContext): Unit = {
-    if (sessionManager != null) {
-      sessionManager.shutdown()
-    }
-
-    if (batchManager != null) {
-      batchManager.shutdown()
+    if (sparkManager != null) {
+      sparkManager.shutdown()
+      sparkManager = null
     }
   }
 }

+ 59 - 0
apps/spark/java/livy-spark/src/main/scala/com/cloudera/hue/livy/spark/SparkManager.scala

@@ -0,0 +1,59 @@
+package com.cloudera.hue.livy.spark
+
+import com.cloudera.hue.livy.LivyConf
+import com.cloudera.hue.livy.LivyConf.{Process, Yarn}
+import com.cloudera.hue.livy.sessions.SessionManager
+import com.cloudera.hue.livy.sessions.batch.BatchSession
+import com.cloudera.hue.livy.sessions.interactive.InteractiveSession
+import com.cloudera.hue.livy.spark.batch.{BatchSessionProcessFactory, BatchSessionYarnFactory}
+import com.cloudera.hue.livy.spark.interactive.{InteractiveSessionProcessFactory, InteractiveSessionYarnFactory}
+import com.cloudera.hue.livy.yarn.Client
+
+object SparkManager {
+  def apply(livyConf: LivyConf): SparkManager = {
+    val processFactory = SparkProcessBuilderFactory(livyConf)
+
+    livyConf.sessionKind() match {
+      case Process() => new SparkProcessManager(processFactory)
+      case Yarn() => new SparkYarnManager(processFactory)
+    }
+  }
+}
+
+trait SparkManager {
+  def batchManager: SessionManager[BatchSession]
+
+  def interactiveManager: SessionManager[InteractiveSession]
+
+  def shutdown()
+}
+
+private class SparkProcessManager(processFactory: SparkProcessBuilderFactory) extends SparkManager {
+  private[this] val batchFactory = new BatchSessionProcessFactory(processFactory)
+  private[this] val interactiveFactory = new InteractiveSessionProcessFactory(processFactory)
+
+  val batchManager = new SessionManager(processFactory.livyConf, batchFactory)
+
+  val interactiveManager = new SessionManager(processFactory.livyConf, interactiveFactory)
+
+  override def shutdown(): Unit = {
+    batchManager.shutdown()
+    interactiveManager.shutdown()
+  }
+}
+
+private class SparkYarnManager(processFactory: SparkProcessBuilderFactory) extends SparkManager {
+  private[this] val client = new Client(processFactory.livyConf)
+  private[this] val batchFactory = new BatchSessionYarnFactory(client, processFactory)
+  private[this] val interactiveFactory = new InteractiveSessionYarnFactory(client, processFactory)
+
+  val batchManager = new SessionManager(processFactory.livyConf, batchFactory)
+
+  val interactiveManager = new SessionManager(processFactory.livyConf, interactiveFactory)
+
+  override def shutdown(): Unit = {
+    batchManager.shutdown()
+    interactiveManager.shutdown()
+    client.close()
+  }
+}

+ 0 - 1
apps/spark/java/livy-spark/src/main/scala/com/cloudera/hue/livy/spark/batch/BatchSessionYarnFactory.scala

@@ -18,7 +18,6 @@
 
 package com.cloudera.hue.livy.spark.batch
 
-import com.cloudera.hue.livy.LivyConf
 import com.cloudera.hue.livy.spark.{SparkProcess, SparkProcessBuilderFactory}
 import com.cloudera.hue.livy.yarn.Client
 

+ 3 - 7
apps/spark/java/livy-spark/src/main/scala/com/cloudera/hue/livy/spark/interactive/InteractiveSessionYarnFactory.scala

@@ -29,13 +29,9 @@ class InteractiveSessionYarnFactory(client: Client, processFactory: SparkProcess
 
    implicit def executor: ExecutionContext = ExecutionContext.global
 
-   protected  override def create(id: Int,
-                                  process: SparkProcess,
-                                  request: CreateInteractiveRequest): InteractiveSession = {
+   protected override def create(id: Int,
+                                 process: SparkProcess,
+                                 request: CreateInteractiveRequest): InteractiveSession = {
      InteractiveSessionYarn(client, id, process, request)
    }
-
-   override def close(): Unit = {
-     client.close()
-   }
  }