[SPARK-17443][SPARK-11035] Stop Spark Application if launcher goes down and use reflection

core/src/main/scala/org/apache/spark/SparkApp.scala

@@ -0,0 +1,40 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.spark
+
+/**
+ * A client that SparkSubmit uses to launch spark Application.
+ * This is currently supported only in YARN mode.
+ */
+private[spark] trait SparkApp {
+  this: Singleton =>
+
+  /**
+   * The Client should implement this as entry method to provide application,
+   * spark conf and system configuration.
+   *
+   * @param args    - all arguments for SparkApp.
+   * @param conf    - Spark Configuration.
+   * @param envvars - system environment Variables.
+   */
+  def sparkMain(
+    args: Array[String],
+    conf: scala.collection.immutable.Map[String, String],
+    envvars: scala.collection.immutable.Map[String, String]): Unit
+
+}

core/src/main/scala/org/apache/spark/deploy/SparkSubmit.scala

@@ -24,6 +24,7 @@ import java.security.PrivilegedExceptionAction
 import java.text.ParseException
 
 import scala.annotation.tailrec
+import scala.collection.JavaConverters._
 import scala.collection.mutable.{ArrayBuffer, HashMap, Map}
 import scala.util.Properties
 
@@ -685,9 +686,8 @@ object SparkSubmit extends CommandLineUtils {
       addJarToClasspath(jar, loader)
     }
 
-    for ((key, value) <- sysProps) {
-      System.setProperty(key, value)
-    }
+    val threadEnabled = sysProps.getOrElse(SparkLauncher.LAUNCHER_INTERNAL_THREAD_ENABLED,
+      "false").toBoolean
 
     var mainClass: Class[_] = null
 
@@ -719,7 +719,23 @@ object SparkSubmit extends CommandLineUtils {
       printWarning("Subclasses of scala.App may not work correctly. Use a main() method instead.")
     }
 
-    val mainMethod = mainClass.getMethod("main", new Array[String](0).getClass)
+    val sparkAppMainMethodArr = mainClass.getMethods().filter{_.getName() == "sparkMain"}
+    val isSparkApp = sparkAppMainMethodArr.length > 0
+
+    val childSparkConf = sysProps.filter{ p => p._1.startsWith("spark.")}.toMap
+    // If running sparkApp or in thread mode, the System properties should not be cluttered.
+    // This helps keep clean isolation between multiple Spark Apps launched in different threads.
+    if (!isSparkApp || !threadEnabled) {
+      sysProps.foreach { case (key, value) =>
+        System.setProperty(key, value)
+      }
+    }
+
+    val mainMethod = if (isSparkApp) {
+      sparkAppMainMethodArr(0)
+    } else {
+      mainClass.getMethod("main", new Array[String](0).getClass)
+    }
     if (!Modifier.isStatic(mainMethod.getModifiers)) {
       throw new IllegalStateException("The main method in the given main class must be static")
     }
@@ -735,7 +751,11 @@ object SparkSubmit extends CommandLineUtils {
     }
 
     try {
-      mainMethod.invoke(null, childArgs.toArray)
+      if (isSparkApp) {
+        mainMethod.invoke(null, childArgs.toArray, childSparkConf, sys.env)
+      } else {
+        mainMethod.invoke(null, childArgs.toArray)
+      }
     } catch {
       case t: Throwable =>
         findCause(t) match {

core/src/main/scala/org/apache/spark/launcher/LauncherBackend.scala

@@ -17,38 +17,67 @@
 
 package org.apache.spark.launcher
 
+import java.io.IOException
 import java.net.{InetAddress, Socket}
 
 import org.apache.spark.SPARK_VERSION
+import org.apache.spark.internal.Logging
 import org.apache.spark.launcher.LauncherProtocol._
-import org.apache.spark.util.{ThreadUtils, Utils}
+import org.apache.spark.util.{ShutdownHookManager, ThreadUtils, Utils}
 
 /**
  * A class that can be used to talk to a launcher server. Users should extend this class to
  * provide implementation for the abstract methods.
  *
  * See `LauncherServer` for an explanation of how launcher communication works.
  */
-private[spark] abstract class LauncherBackend {
+private[spark] abstract class LauncherBackend extends Logging {
 
   private var clientThread: Thread = _
   private var connection: BackendConnection = _
   private var lastState: SparkAppHandle.State = _
+  private var stopOnShutdown: Boolean = false
   @volatile private var _isConnected = false
 
   def connect(): Unit = {
     val port = sys.env.get(LauncherProtocol.ENV_LAUNCHER_PORT).map(_.toInt)
     val secret = sys.env.get(LauncherProtocol.ENV_LAUNCHER_SECRET)
+    val stopFlag = sys.env.get(LauncherProtocol.ENV_LAUNCHER_STOP_IF_SHUTDOWN).map(_.toBoolean)
     if (port != None && secret != None) {
-      val s = new Socket(InetAddress.getLoopbackAddress(), port.get)
-      connection = new BackendConnection(s)
-      connection.send(new Hello(secret.get, SPARK_VERSION))
-      clientThread = LauncherBackend.threadFactory.newThread(connection)
-      clientThread.start()
-      _isConnected = true
+      connect(port.get, secret.get, stopFlag.getOrElse(false))
     }
   }
 
+  def connect(port: Int, secret: String): Unit = {
+    val s = new Socket(InetAddress.getLoopbackAddress(), port)
+    connection = new BackendConnection(s)
+    connection.send(new Hello(secret, SPARK_VERSION))
+    clientThread = LauncherBackend.threadFactory.newThread(connection)
+    clientThread.start()
+    _isConnected = true
+    if (stopOnShutdown) {
+      logDebug("Adding shutdown hook") // force eager creation of logger
+      var _shutdownHookRef = ShutdownHookManager.addShutdownHook(
+        ShutdownHookManager.SPARK_CONTEXT_SHUTDOWN_PRIORITY) { () =>
+        logInfo("Invoking onStopRequest() from shutdown hook")
+        try {
+          if (_isConnected) {
+            onStopRequest()
+          }
+        }
+        catch {
+          case anotherIOE: IOException =>
+            logError("Error while running LauncherBackend shutdownHook...", anotherIOE)
+        }
+      }
+    }
+  }
+
+  def connect(port: Int, secret: String, stopFlag: Boolean): Unit = {
+    this.stopOnShutdown = stopFlag
+    connect(port, secret)
+  }
+
   def close(): Unit = {
     if (connection != null) {
       try {
@@ -71,6 +100,9 @@ private[spark] abstract class LauncherBackend {
     if (connection != null && lastState != state) {
       connection.send(new SetState(state))
       lastState = state
+      if (!_isConnected && stopOnShutdown) {
+        fireStopRequest()
+      }
     }
   }
 
@@ -110,12 +142,14 @@ private[spark] abstract class LauncherBackend {
     override def close(): Unit = {
       try {
         super.close()
+        if (stopOnShutdown) {
+          fireStopRequest()
+        }
       } finally {
         onDisconnected()
         _isConnected = false
       }
     }
-
   }
 
 }

launcher/src/main/java/org/apache/spark/launcher/AbstractSparkAppHandle.java

@@ -0,0 +1,132 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.spark.launcher;
+
+import java.io.IOException;
+import java.util.ArrayList;
+import java.util.List;
+import java.util.logging.Level;
+import java.util.logging.Logger;
+
+public abstract class AbstractSparkAppHandle implements SparkAppHandle {
+  private static final Logger LOG = Logger.getLogger(AbstractSparkAppHandle.class.getName());
+
+  protected final String secret;
+  protected final LauncherServer server;
+  protected boolean disposed;
+  protected List<Listener> listeners;
+  protected State state;
+  private LauncherConnection connection;
+  private String appId;
+  OutputRedirector redirector;
+
+  public AbstractSparkAppHandle(LauncherServer server, String secret) {
+    this.server = server;
+    this.secret = secret;
+    this.state = State.UNKNOWN;
+  }
+
+  @Override
+  public synchronized void addListener(Listener l) {
+    if (listeners == null) {
+      listeners = new ArrayList<>();
+    }
+    listeners.add(l);
+  }
+
+  @Override
+  public State getState() {
+    return state;
+  }
+
+  @Override
+  public String getAppId() {
+    return appId;
+  }
+
+  @Override
+  public void stop() {
+    CommandBuilderUtils.checkState(connection != null, "Application is still not connected.");
+    try {
+      connection.send(new LauncherProtocol.Stop());
+    } catch (IOException ioe) {
+      throw new RuntimeException(ioe);
+    }
+  }
+
+  @Override
+  public synchronized void disconnect() {
+    if (!disposed) {
+      disposed = true;
+      if (connection != null) {
+        try {
+          connection.close();
+        } catch (IOException ioe) {
+          // no-op.
+        }
+      }
+      server.unregister(this);
+      if (redirector != null) {
+        redirector.stop();
+      }
+    }
+  }
+
+  String getSecret() {
+    return secret;
+  }
+
+  void setConnection(LauncherConnection connection) {
+    this.connection = connection;
+  }
+
+  LauncherServer getServer() {
+    return server;
+  }
+
+  LauncherConnection getConnection() {
+    return connection;
+  }
+
+  void setState(State s) {
+    if (!state.isFinal()) {
+      state = s;
+      fireEvent(false);
+    } else {
+      LOG.log(Level.WARNING, "Backend requested transition from final state {0} to {1}.",
+        new Object[]{state, s});
+    }
+  }
+
+  void setAppId(String appId) {
+    this.appId = appId;
+    fireEvent(true);
+  }
+
+  private synchronized void fireEvent(boolean isInfoChanged) {
+    if (listeners != null) {
+      for (Listener l : listeners) {
+        if (isInfoChanged) {
+          l.infoChanged(this);
+        } else {
+          l.stateChanged(this);
+        }
+      }
+    }
+  }
+}