Merge pull request #25 from data-integrations/feature_release/CDAP-12941-fix-dynamic-spark-cl

(CDAP-12941) Fix dynamic spark plugin classloading

Author: chtyim

pom.xml

@@ -22,7 +22,7 @@
 
   <groupId>io.cdap.plugin</groupId>
   <artifactId>dynamic-spark</artifactId>
-  <version>2.2.0</version>
+  <version>2.2.1</version>
 
   <properties>
     <!-- properties for script build step that creates the config files for the artifacts -->

src/main/java/io/cdap/plugin/spark/dynamic/ScalaSparkCodeExecutor.java

@@ -170,31 +170,37 @@ public void onEvent(SparkListenerEvent event) {
    */
   public Object execute(SparkExecutionPluginContext context,
                         JavaRDD<StructuredRecord> javaRDD) throws InvocationTargetException, IllegalAccessException {
-    // RDD case
-    if (!isDataFrame) {
-      if (takeContext) {
-        //noinspection unchecked
-        return ((RDD<StructuredRecord>) method.invoke(null, javaRDD.rdd(), context)).toJavaRDD();
-      } else {
-        //noinspection unchecked
-        return ((RDD<StructuredRecord>) method.invoke(null, javaRDD.rdd())).toJavaRDD();
+    ClassLoader oldCL = Thread.currentThread().getContextClassLoader();
+    Thread.currentThread().setContextClassLoader(interpreter.getClassLoader());
+    try {
+      // RDD case
+      if (!isDataFrame) {
+        if (takeContext) {
+          //noinspection unchecked
+          return ((RDD<StructuredRecord>) method.invoke(null, javaRDD.rdd(), context)).toJavaRDD();
+        } else {
+          //noinspection unchecked
+          return ((RDD<StructuredRecord>) method.invoke(null, javaRDD.rdd())).toJavaRDD();
+        }
       }
-    }
 
-    // DataFrame case
-    Schema inputSchema = context.getInputSchema();
-    if (inputSchema == null) {
-      // Should already been checked in initialize. This is to safeguard in case the call sequence changed in future.
-      throw new IllegalArgumentException("Input schema must be provided for using DataFrame in Spark Compute");
-    }
+      // DataFrame case
+      Schema inputSchema = context.getInputSchema();
+      if (inputSchema == null) {
+        // Should already been checked in initialize. This is to safeguard in case the call sequence changed in future.
+        throw new IllegalArgumentException("Input schema must be provided for using DataFrame in Spark Compute");
+      }
 
-    SQLContext sqlContext = getSQLContext(context.getSparkContext().sc());
+      SQLContext sqlContext = getSQLContext(context.getSparkContext().sc());
 
-    StructType rowType = DataFrames.toDataType(inputSchema);
-    JavaRDD<Row> rowRDD = javaRDD.map(new RecordToRow(rowType));
+      StructType rowType = DataFrames.toDataType(inputSchema);
+      JavaRDD<Row> rowRDD = javaRDD.map(new RecordToRow(rowType));
 
-    Object dataFrame = createDataFrame(sqlContext, rowRDD, rowType);
-    return takeContext ? method.invoke(null, dataFrame, context) : method.invoke(null, dataFrame);
+      Object dataFrame = createDataFrame(sqlContext, rowRDD, rowType);
+      return takeContext ? method.invoke(null, dataFrame, context) : method.invoke(null, dataFrame);
+    } finally {
+      Thread.currentThread().setContextClassLoader(oldCL);
+    }
   }
 
   private String generateSourceClass(String className) {

src/main/java/io/cdap/plugin/spark/dynamic/ScalaSparkProgram.java

@@ -28,8 +28,6 @@
 import io.cdap.cdap.api.spark.SparkMain;
 import io.cdap.cdap.api.spark.dynamic.CompilationFailureException;
 import io.cdap.cdap.api.spark.dynamic.SparkInterpreter;
-import org.slf4j.Logger;
-import org.slf4j.LoggerFactory;
 
 import java.io.File;
 import java.io.IOException;
@@ -47,8 +45,6 @@
 @Description("Executes user-provided Spark program")
 public class ScalaSparkProgram implements JavaSparkMain {
 
-  private static final Logger LOG = LoggerFactory.getLogger(ScalaSparkProgram.class);
-
   private final Config config;
 
   public ScalaSparkProgram(Config config) throws CompilationFailureException, IOException {
@@ -144,15 +140,18 @@ private Callable<Void> getMethodCallable(ClassLoader classLoader, String mainCla
         arg = sec == null ? null : RuntimeArguments.toPosixArray(sec.getRuntimeArguments());
       }
 
-      return new Callable<Void>() {
-        @Override
-        public Void call() throws Exception {
+      return () -> {
+        ClassLoader oldCl = Thread.currentThread().getContextClassLoader();
+        Thread.currentThread().setContextClassLoader(classLoader);
+        try {
           Object instance = null;
           if (!Modifier.isStatic(method.getModifiers())) {
             instance = cls.newInstance();
           }
           method.invoke(instance, arg);
           return null;
+        } finally {
+          Thread.currentThread().setContextClassLoader(oldCl);
         }
       };
 

src/test/java/io/cdap/plugin/spark/dynamic/ScalaSparkTest.java

@@ -214,6 +214,60 @@ public void testScalaProgramDependency() throws Exception {
     workflowManager.waitForRun(ProgramRunStatus.COMPLETED, 5, TimeUnit.MINUTES);
   }
 
+  @Test
+  public void testScalaSparkProgramClosure() throws Exception {
+    StringWriter codeWriter = new StringWriter();
+    try (PrintWriter printer = new PrintWriter(codeWriter, true)) {
+      printer.println("import io.cdap.cdap.api.spark._");
+      printer.println("import org.apache.spark._");
+      printer.println("import org.apache.spark.rdd.RDD");
+      printer.println("import org.slf4j._");
+
+      printer.println("class SparkProgram extends SparkMain {");
+      printer.println("  import SparkProgram._");
+
+      printer.println("  override def run(implicit sec: SparkExecutionContext): Unit = {");
+      printer.println("    LOG.info(\"Spark Program Started\")");
+
+      printer.println("    val sc = new SparkContext");
+      printer.println("    val points = sc.parallelize(Seq((\"a\", Array(1, 2)), (\"a\", Array(3, 4))))");
+
+      printer.println("    val sq = points.mapValues(t => Array(t.apply(0) * t.apply(0), t.apply(1) * t.apply(1)))");
+      printer.println("    LOG.info(\"squared = {}\", sq.collect)");
+
+      printer.println("    val squaredNested = points.mapValues(t => t.map(x => x * x))");
+      printer.println("    LOG.info(\"squaredNested = {}\", squaredNested.collect)");
+
+      printer.println("    LOG.info(\"Spark Program Completed\")");
+      printer.println("  }");
+      printer.println("}");
+
+      printer.println("object SparkProgram {");
+      printer.println("  val LOG = LoggerFactory.getLogger(getClass())");
+      printer.println("}");
+    }
+
+    // Pipeline configuration
+    ETLBatchConfig etlConfig = ETLBatchConfig.builder()
+      .addStage(new ETLStage("action", new ETLPlugin("ScalaSparkProgram", "sparkprogram", ImmutableMap.of(
+        "scalaCode", codeWriter.toString(),
+        "mainClass", "SparkProgram"
+      ))))
+      .build();
+
+    // Deploy the pipeline
+    ArtifactSummary artifactSummary = new ArtifactSummary(DATAPIPELINE_ARTIFACT_ID.getArtifact(),
+                                                          DATAPIPELINE_ARTIFACT_ID.getVersion());
+    AppRequest<ETLBatchConfig> appRequest = new AppRequest<>(artifactSummary, etlConfig);
+    ApplicationId appId = NamespaceId.DEFAULT.app("ScalaSparkProgramApp");
+    ApplicationManager appManager = deployApplication(appId, appRequest);
+
+    // Run the pipeline
+    WorkflowManager workflowManager = appManager.getWorkflowManager(SmartWorkflow.NAME);
+    workflowManager.start();
+    workflowManager.waitForRun(ProgramRunStatus.COMPLETED, 5, TimeUnit.MINUTES);
+  }
+
   @Test
   public void testScalaSparkCompute() throws Exception {
     Schema inputSchema = Schema.recordOf(