Add a timeout for initialization

Also move sparkRBackend.stop into a finally block

Author: shivaram

pkg/src/src/main/scala/edu/berkeley/cs/amplab/sparkr/SparkRRunner.scala

@@ -3,6 +3,7 @@ package edu.berkeley.cs.amplab.sparkr
 import java.io._
 import java.net.URI
 import java.util.concurrent.Semaphore
+import java.util.concurrent.TimeUnit
 
 import scala.collection.mutable.ArrayBuffer
 import scala.collection.JavaConversions._
@@ -19,6 +20,8 @@ object SparkRRunner {
 
     val otherArgs = args.slice(1, args.length)
 
+    // Time to wait for SparkR backend to initialize in seconds
+    val backendTimeout = sys.env.getOrElse("SPARKR_BACKEND_TIMEOUT", "120").toInt
     // TODO: Can we get this from SparkConf ?
     val sparkRBackendPort = sys.env.getOrElse("SPARKR_BACKEND_PORT", "12345").toInt
     val rCommand = "Rscript"
@@ -52,20 +55,26 @@ object SparkRRunner {
 
     sparkRBackendThread.start()
     // Wait for SparkRBackend initialization to finish
-    sparkRBackendThread.finishedInit.acquire()
-
-    // Launch R
-    val builder = new ProcessBuilder(Seq(rCommand, rFileNormalized) ++ otherArgs)
-    val env = builder.environment()
-    env.put("EXISTING_SPARKR_BACKEND_PORT", sparkRBackendPort.toString)
-    builder.redirectErrorStream(true) // Ugly but needed for stdout and stderr to synchronize
-    val process = builder.start()
-
-    new RedirectThread(process.getInputStream, System.out, "redirect output").start()
-
-    val returnCode = process.waitFor()
-    sparkRBackendThread.stopBackend()
-    System.exit(returnCode)
+    if (sparkRBackendThread.finishedInit.tryAcquire(backendTimeout, TimeUnit.SECONDS)) {
+      // Launch R
+      val returnCode = try {
+        val builder = new ProcessBuilder(Seq(rCommand, rFileNormalized) ++ otherArgs)
+        val env = builder.environment()
+        env.put("EXISTING_SPARKR_BACKEND_PORT", sparkRBackendPort.toString)
+        builder.redirectErrorStream(true) // Ugly but needed for stdout and stderr to synchronize
+        val process = builder.start()
+
+        new RedirectThread(process.getInputStream, System.out, "redirect output").start()
+
+        process.waitFor()
+      } finally {
+        sparkRBackendThread.stopBackend()
+      }
+      System.exit(returnCode)
+    } else {
+      System.err.println("SparkR backend did not initialize in " + backendTimeout + " seconds")
+      System.exit(-1)
+    }
   }
 
   private class RedirectThread(