SparkContainerLauncher.java

/*
 * Copyright © 2016 Cask Data, Inc.
 *
 * Licensed 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 co.cask.cdap.app.runtime.spark.distributed;

import co.cask.cdap.app.runtime.spark.SparkRunnerClassLoader;
import co.cask.cdap.common.lang.ClassLoaders;
import com.google.common.base.Throwables;

import java.net.MalformedURLException;
import java.net.URI;
import java.net.URL;
import java.util.ArrayList;
import java.util.List;

/**
 * This class launches Spark YARN containers with classes loaded through the {@link SparkRunnerClassLoader}.
 */
public final class SparkContainerLauncher {

  /**
   * Launches the given main class. The main class will be loaded through the {@link SparkRunnerClassLoader}.
   *
   * @param mainClassName the main class to launch
   * @param args arguments for the main class
   */
  @SuppressWarnings("unused")
  public static void launch(String mainClassName, String[] args) {
    List<URL> urls = ClassLoaders.getClassLoaderURLs(ClassLoader.getSystemClassLoader(), new ArrayList<URL>());

    // Remove the URL that contains the given main classname to avoid infinite recursion.
    // This is needed because we generate a class with the same main classname in order to intercept the main()
    // method call from the container launch script.
    URL resource = ClassLoader.getSystemClassLoader().getResource(mainClassName.replace('.', '/') + ".class");
    if (resource == null) {
      throw new IllegalStateException("Failed to find resource for main class " + mainClassName);
    }

    if (!urls.remove(getClassPathURL(mainClassName, resource))) {
      throw new IllegalStateException("Failed to remove main class resource " + resource);
    }

    // Creates the SparkRunnerClassLoader for class rewriting and it will be used for the rest of the execution.
    // Use the extension classloader as the parent instead of the system classloader because
    // Spark classes are in the system classloader which we want to rewrite.
    ClassLoader classLoader = new SparkRunnerClassLoader(urls.toArray(new URL[urls.size()]),
                                                         ClassLoader.getSystemClassLoader().getParent(),
                                                         false, false);
    // Sets the context classloader and launch the actual Spark main class.
    Thread.currentThread().setContextClassLoader(classLoader);
    try {
      classLoader.loadClass(mainClassName).getMethod("main", String[].class).invoke(null, new Object[]{args});
    } catch (Exception e) {
      throw new RuntimeException("Failed to call " + mainClassName + ".main(String[])", e);
    }
  }

  private static URL getClassPathURL(String className, URL classUrl) {
    try {
      if ("file".equals(classUrl.getProtocol())) {
        String path = classUrl.getFile();
        // Compute the directory container the class.
        int endIdx = path.length() - className.length() - ".class".length();
        if (endIdx > 1) {
          // If it is not the root directory, return the end index to remove the trailing '/'.
          endIdx--;
        }
        return new URL("file", "", -1, path.substring(0, endIdx));
      }
      if ("jar".equals(classUrl.getProtocol())) {
        String path = classUrl.getFile();
        return URI.create(path.substring(0, path.indexOf("!/"))).toURL();
      }
    } catch (MalformedURLException e) {
      throw Throwables.propagate(e);
    }
    throw new IllegalStateException("Unsupported class URL: " + classUrl);
  }
}