Merge pull request apache#139 from shivaram/fix-backend-exit

[SPARKR-185] Fix SparkR backend to exit in more cases

Author: shivaram

pkg/R/sparkR.R

@@ -11,21 +11,42 @@ sparkR.onLoad <- function(libname, pkgname) {
   .sparkREnv$assemblyJarPath <- assemblyJarPath
 }
 
-#' Stop this Spark context
-#' Also terminates the backend this R session is connected to
-sparkR.stop <- function(sparkREnv) {
-  if (exists(".sparkRjsc", envir = .sparkREnv)) {
-    sc <- get(".sparkRjsc", envir = .sparkREnv)
-    callJMethod(sc, "stop")
+# Utility function that returns TRUE if we have an active connection to the
+# backend and FALSE otherwise
+connExists <- function(env) {
+  tryCatch({
+    exists(".sparkRCon", envir = env) && isOpen(env[[".sparkRCon"]])
+  }, error = function(err) {
+    return(FALSE)
+  })
+}
+
+# Stop the Spark context.
+# Also terminates the backend this R session is connected to
+sparkR.stop <- function(env) {
+  cat("Stopping SparkR\n")
+
+  if (!connExists(env)) {
+    # When the workspace is saved in R, the connections are closed
+    # *before* the finalizer is run. In these cases, we reconnect
+    # to the backend, so we can shut it down.
+    connectBackend("localhost", .sparkREnv$sparkRBackendPort)
   }
 
-  if (exists(".sparkRCon", envir = .sparkREnv)) {
-    callJStatic("SparkRHandler", "stopBackend")
-    # Also close the connection and remove it from our env
-    conn <- get(".sparkRCon", .sparkREnv)
-    close(conn)
-    rm(".sparkRCon", envir = .sparkREnv)
+  if (exists(".sparkRjsc", envir = env)) {
+    sc <- get(".sparkRjsc", envir = env)
+    callJMethod(sc, "stop")
   }
+
+  callJStatic("SparkRHandler", "stopBackend")
+  # Also close the connection and remove it from our env
+  conn <- get(".sparkRCon", env)
+  close(conn)
+  rm(".sparkRCon", envir = env)
+
+  # Finally, sleep for 1 sec to let backend finish exiting.
+  # Without this we get port conflicts in RStudio when we try to 'Restart R'.
+  Sys.sleep(1)
 }
 
 #' Initialize a new Spark Context.
@@ -79,7 +100,8 @@ sparkR.init <- function(
                 args = as.character(sparkRBackendPort),
                 javaOpts = paste("-Xmx", sparkMem, sep = ""))
   Sys.sleep(2) # Wait for backend to come up
-  connectBackend("localhost", 12345) # Connect to it
+  .sparkREnv$sparkRBackendPort <- sparkRBackendPort
+  connectBackend("localhost", sparkRBackendPort) # Connect to it
 
   if (nchar(sparkHome) != 0) {
     sparkHome <- normalizePath(sparkHome)

pkg/R/sparkRBackend.R

@@ -47,7 +47,7 @@ isRemoveMethod <- function(isStatic, objId, methodName) {
 # methodName - name of method to be invoked
 invokeJava <- function(isStatic, objId, methodName, ...) {
   if (!exists(".sparkRCon", .sparkREnv)) {
-    stop("No connection to backend found")
+    stop("No connection to backend found. Please re-run sparkR.init")
   }
 
   # If this isn't a removeJObject call
@@ -62,6 +62,7 @@ invokeJava <- function(isStatic, objId, methodName, ...) {
     }
   }
 
+
   rc <- rawConnection(raw(0), "r+")
 
   writeBoolean(rc, isStatic)
@@ -78,12 +79,12 @@ invokeJava <- function(isStatic, objId, methodName, ...) {
   writeInt(conn, length(bytesToSend))
   writeBin(bytesToSend, conn)
 
+  close(rc) # TODO: Can we close this before ?
+
   # TODO: check the status code to output error information
   returnStatus <- readInt(conn)
   stopifnot(returnStatus == 0)
   ret <- readObject(conn)
 
-  close(rc) # TODO: Can we close this before ?
-
   ret
 }