NamedObjects: create one instance of JobServerNamedObjects instead of a singleton object

Author: TimMaltGermany

job-server-extras/test/spark.jobserver/NamedObjectsJobSpec.scala

@@ -10,42 +10,50 @@ import collection.JavaConversions._
 
 class NamedObjectsJobSpec extends JobSpecBase(JobManagerSpec.getNewSystem) {
 
-  private val emptyConfig = ConfigFactory.parseString("spark.jobserver.named-object-creation-timeout = 60 s")
+  //private val emptyConfig = ConfigFactory.parseString("spark.jobserver.named-object-creation-timeout = 60 s")
 
   before {
     dao = new InMemoryDAO
     daoActor = system.actorOf(JobDAOActor.props(dao))
     manager = system.actorOf(JobManagerActor.props(JobManagerSpec.getContextConfig(adhoc = false)))
     supervisor = TestProbe().ref
+
+    manager ! JobManagerActor.Initialize(daoActor, None)
+    
+    expectMsgClass(classOf[JobManagerActor.Initialized])
+
+    uploadTestJar()
+
   }
 
   val jobName = "spark.jobserver.NamedObjectsTestJob"
 
+  private def getCreateConfig(createDF: Boolean, createRDD: Boolean) : Config = {
+    ConfigFactory.parseString("spark.jobserver.named-object-creation-timeout = 60 s, " + 
+        NamedObjectsTestJobConfig.CREATE_DF + " = " + createDF + ", " +
+        NamedObjectsTestJobConfig.CREATE_RDD + " = " + createRDD)
+  }
+  
+  private def getDeleteConfig(names: List[String]) : Config = {
+    ConfigFactory.parseString("spark.jobserver.named-object-creation-timeout = 60 s, " + 
+        NamedObjectsTestJobConfig.DELETE+" = [" + names.mkString(", ") + "]")
+  }
+  
   describe("NamedObjects (RDD)") {
     it("should survive from one job to another one") {
 
-      manager ! JobManagerActor.Initialize(daoActor, None)
-      expectMsgClass(classOf[JobManagerActor.Initialized])
-
-      uploadTestJar()
-      manager ! JobManagerActor.StartJob("demo", jobName, emptyConfig.withValue(NamedObjectsTestJobConfig.CREATE_DF, ConfigValueFactory.fromAnyRef(false))
-        .withValue(NamedObjectsTestJobConfig.CREATE_RDD, ConfigValueFactory.fromAnyRef(true)),
-        errorEvents ++ syncEvents)
+      manager ! JobManagerActor.StartJob("demo", jobName, getCreateConfig(false, true), errorEvents ++ syncEvents)
       val JobResult(_, names: Array[String]) = expectMsgClass(classOf[JobResult])
       names should contain("rdd1")
 
-      manager ! JobManagerActor.StartJob("demo", jobName, emptyConfig.withValue(NamedObjectsTestJobConfig.CREATE_DF, ConfigValueFactory.fromAnyRef(false))
-        .withValue(NamedObjectsTestJobConfig.CREATE_RDD, ConfigValueFactory.fromAnyRef(false)),
-        errorEvents ++ syncEvents)
+      manager ! JobManagerActor.StartJob("demo", jobName, getCreateConfig(false, false), errorEvents ++ syncEvents)
       val JobResult(_, names2: Array[String]) = expectMsgClass(classOf[JobResult])
 
       names2 should contain("rdd1")
       names2 should not contain("df1")
 
       //clean-up
-      manager ! JobManagerActor.StartJob("demo", jobName, emptyConfig.withValue(NamedObjectsTestJobConfig.DELETE, ConfigValueFactory.fromIterable(List("rdd1")))
-        .withValue(NamedObjectsTestJobConfig.CREATE_RDD, ConfigValueFactory.fromAnyRef(false)),
-        errorEvents ++ syncEvents)
+      manager ! JobManagerActor.StartJob("demo", jobName, getDeleteConfig(List("rdd1")), errorEvents ++ syncEvents)
       val JobResult(_, names3: Array[String]) = expectMsgClass(classOf[JobResult])
 
       names3 should not contain("rdd1")
@@ -55,57 +63,39 @@ class NamedObjectsJobSpec extends JobSpecBase(JobManagerSpec.getNewSystem) {
 
   describe("NamedObjects (DataFrame)") {
     it("should survive from one job to another one") {
-      manager ! JobManagerActor.Initialize(daoActor, None)
-      expectMsgClass(classOf[JobManagerActor.Initialized])
 
-      uploadTestJar()
-      manager ! JobManagerActor.StartJob("demo", jobName, emptyConfig.withValue(NamedObjectsTestJobConfig.CREATE_DF, ConfigValueFactory.fromAnyRef(true))
-        .withValue(NamedObjectsTestJobConfig.CREATE_RDD, ConfigValueFactory.fromAnyRef(false)),
-        errorEvents ++ syncEvents)
+      manager ! JobManagerActor.StartJob("demo", jobName, getCreateConfig(true, false), errorEvents ++ syncEvents)
       val JobResult(_, names: Array[String]) = expectMsgClass(classOf[JobResult])
 
       names should contain("df1")
 
-      manager ! JobManagerActor.StartJob("demo", jobName, emptyConfig.withValue(NamedObjectsTestJobConfig.CREATE_DF, ConfigValueFactory.fromAnyRef(false))
-        .withValue(NamedObjectsTestJobConfig.CREATE_RDD, ConfigValueFactory.fromAnyRef(false)),
-        errorEvents ++ syncEvents)
+      manager ! JobManagerActor.StartJob("demo", jobName, getCreateConfig(false, false), errorEvents ++ syncEvents)
       val JobResult(_, names2: Array[String]) = expectMsgClass(classOf[JobResult])
 
       names2 should equal(names)
 
       //clean-up
-      manager ! JobManagerActor.StartJob("demo", jobName, emptyConfig.withValue(NamedObjectsTestJobConfig.DELETE, 
-          ConfigValueFactory.fromIterable(List("df1"))),
-        errorEvents ++ syncEvents)
+      manager ! JobManagerActor.StartJob("demo", jobName, getDeleteConfig(List("df1")), errorEvents ++ syncEvents)
       val JobResult(_, names3: Array[String]) = expectMsgClass(classOf[JobResult])
     }
   }
 
   describe("NamedObjects (DataFrame + RDD)") {
     it("should survive from one job to another one") {
-      manager ! JobManagerActor.Initialize(daoActor, None)
-      expectMsgClass(classOf[JobManagerActor.Initialized])
 
-      uploadTestJar()
-      manager ! JobManagerActor.StartJob("demo", jobName, emptyConfig.withValue(NamedObjectsTestJobConfig.CREATE_DF, ConfigValueFactory.fromAnyRef(true))
-        .withValue(NamedObjectsTestJobConfig.CREATE_RDD, ConfigValueFactory.fromAnyRef(true)),
-        errorEvents ++ syncEvents)
+      manager ! JobManagerActor.StartJob("demo", jobName, getCreateConfig(true, true), errorEvents ++ syncEvents)
       val JobResult(_, names: Array[String]) = expectMsgClass(classOf[JobResult])
 
       names should contain("rdd1")
       names should contain("df1")
 
-      manager ! JobManagerActor.StartJob("demo", jobName, emptyConfig.withValue(NamedObjectsTestJobConfig.CREATE_DF, ConfigValueFactory.fromAnyRef(false))
-        .withValue(NamedObjectsTestJobConfig.CREATE_RDD, ConfigValueFactory.fromAnyRef(false)),
-        errorEvents ++ syncEvents)
+      manager ! JobManagerActor.StartJob("demo", jobName, getCreateConfig(false, false), errorEvents ++ syncEvents)
       val JobResult(_, names2: Array[String]) = expectMsgClass(classOf[JobResult])
 
       names2 should equal(names)
 
       //clean-up
-      manager ! JobManagerActor.StartJob("demo", jobName, emptyConfig.withValue(NamedObjectsTestJobConfig.DELETE,
-          ConfigValueFactory.fromIterable(List("rdd1", "df1"))),
-        errorEvents ++ syncEvents)
+      manager ! JobManagerActor.StartJob("demo", jobName, getDeleteConfig(List("rdd1", "df1")), errorEvents ++ syncEvents)
       val JobResult(_, names3: Array[String]) = expectMsgClass(classOf[JobResult])
     }
   }

job-server/src/spark.jobserver/JobManagerActor.scala

@@ -95,6 +95,7 @@ class JobManagerActor(contextConfig: Config) extends InstrumentedActor {
   protected var resultActor: ActorRef = _
   private var daoActor: ActorRef = _
 
+  private val jobServerNamedObjects = new JobServerNamedObjects(context.system)
 
   override def postStop() {
     logger.info("Shutting down SparkContext {}", contextName)
@@ -257,7 +258,7 @@ class JobManagerActor(contextConfig: Config) extends InstrumentedActor {
       if (job.isInstanceOf[NamedObjectSupport]) {
         val namedObjects = job.asInstanceOf[NamedObjectSupport].namedObjectsPrivate
         if (namedObjects.get() == null) {
-          namedObjects.compareAndSet(null, new JobServerNamedObjects(context.system))
+          namedObjects.compareAndSet(null, jobServerNamedObjects)
         }
       }
 

job-server/src/spark.jobserver/JobServerNamedObjects.scala

@@ -18,8 +18,6 @@ import spray.util._
  */
 class JobServerNamedObjects(system: ActorSystem) extends NamedObjects {
 
-  import JobServerNamedObjects._
-
   val logger = LoggerFactory.getLogger(getClass)
 
   implicit val ec: ExecutionContext = system.dispatcher
@@ -32,6 +30,11 @@ class JobServerNamedObjects(system: ActorSystem) extends NamedObjects {
     config.getDuration("spark.jobserver.named-object-creation-timeout",
       SECONDS), SECONDS)
 
+  // we must store a reference to each NamedObject even though only its ID is used here
+  // this reference prevents the object from being GCed and cleaned by sparks ContextCleaner
+  // or some other GC for other types of objects
+  private val namesToObjects: Cache[NamedObject] = LruCache()
+
   override def getOrElseCreate[O <: NamedObject](name: String, objGen: => O)
                                  (implicit timeout: Timeout = defaultTimeout,
                                            persister: NamedObjectPersister[O]): O = {
@@ -93,15 +96,5 @@ class JobServerNamedObjects(system: ActorSystem) extends NamedObjects {
       case answer: Iterable[String] @unchecked => answer
     }
   }
-}
 
-/**
- * companion object that hold reference to cache with named object so that we can reference
- * named objects across jobs
- */
-object JobServerNamedObjects {
-  // we must store a reference to each NamedObject even though only its ID is used here
-  // this reference prevents the object from being GCed and cleaned by sparks ContextCleaner
-  // or some other GC for other types of objects
-  val namesToObjects: Cache[NamedObject] = LruCache()
 }