From 2c8c4bb6542d6d8a9b117758cec8f5416ebf023d Mon Sep 17 00:00:00 2001
From: KevinSmile <kevinwang013@hotmail.com>
Date: Sat, 5 Sep 2020 18:25:45 +0800
Subject: [PATCH] [SPARK-32804][Launcher] fix run-example command builder bug

---
 .../launcher/SparkSubmitCommandBuilder.java    | 15 +++++++++++++--
 .../SparkSubmitCommandBuilderSuite.java        | 18 ++++++++++++++++++
 2 files changed, 31 insertions(+), 2 deletions(-)

diff --git a/launcher/src/main/java/org/apache/spark/launcher/SparkSubmitCommandBuilder.java b/launcher/src/main/java/org/apache/spark/launcher/SparkSubmitCommandBuilder.java
index 383c3f60a595..43e7f8debe17 100644
--- a/launcher/src/main/java/org/apache/spark/launcher/SparkSubmitCommandBuilder.java
+++ b/launcher/src/main/java/org/apache/spark/launcher/SparkSubmitCommandBuilder.java
@@ -139,7 +139,7 @@ class SparkSubmitCommandBuilder extends AbstractCommandBuilder {
 
         case RUN_EXAMPLE:
           isExample = true;
-          appResource = SparkLauncher.NO_RESOURCE;
+          appResource = findExamplesAppJar();
           submitArgs = args.subList(1, args.size());
       }
 
@@ -241,9 +241,11 @@ List<String> buildSparkSubmitArgs() {
     }
 
     args.addAll(parsedArgs);
+
     if (appResource != null) {
       args.add(appResource);
     }
+
     args.addAll(appArgs);
 
     return args;
@@ -401,6 +403,15 @@ private boolean isThriftServer(String mainClass) {
       mainClass.equals("org.apache.spark.sql.hive.thriftserver.HiveThriftServer2"));
   }
 
+  private String findExamplesAppJar() {
+    for (String exampleJar : findExamplesJars()) {
+      if (new File(exampleJar).getName().startsWith("spark-examples")) {
+        return exampleJar;
+      }
+    }
+    throw new IllegalStateException("Failed to find examples' main app jar.");
+  }
+
   private List<String> findExamplesJars() {
     boolean isTesting = "1".equals(getenv("SPARK_TESTING"));
     List<String> examplesJars = new ArrayList<>();
@@ -513,7 +524,7 @@ protected boolean handleUnknown(String opt) {
           className = EXAMPLE_CLASS_PREFIX + className;
         }
         mainClass = className;
-        appResource = SparkLauncher.NO_RESOURCE;
+        appResource = findExamplesAppJar();
         return false;
       } else if (errorOnUnknownArgs) {
         checkArgument(!opt.startsWith("-"), "Unrecognized option: %s", opt);
diff --git a/launcher/src/test/java/org/apache/spark/launcher/SparkSubmitCommandBuilderSuite.java b/launcher/src/test/java/org/apache/spark/launcher/SparkSubmitCommandBuilderSuite.java
index 752e8d4c23f8..6cd089e256b9 100644
--- a/launcher/src/test/java/org/apache/spark/launcher/SparkSubmitCommandBuilderSuite.java
+++ b/launcher/src/test/java/org/apache/spark/launcher/SparkSubmitCommandBuilderSuite.java
@@ -245,6 +245,24 @@ public void testExamplesRunner() throws Exception {
     assertEquals("42", cmd.get(cmd.size() - 1));
   }
 
+  @Test
+  public void testExamplesRunnerPrimaryResource() throws Exception {
+    List<String> sparkSubmitArgs = Arrays.asList(
+            SparkSubmitCommandBuilder.RUN_EXAMPLE,
+            parser.MASTER + "=foo",
+            parser.DEPLOY_MODE + "=cluster",
+            "SparkPi",
+            "100");
+
+    List<String> cmd = newCommandBuilder(sparkSubmitArgs).buildSparkSubmitArgs();
+    assertEquals(SparkSubmitCommandBuilder.EXAMPLE_CLASS_PREFIX + "SparkPi",
+            findArgValue(cmd, parser.CLASS));
+    assertEquals("cluster", findArgValue(cmd, parser.DEPLOY_MODE));
+    String primaryResource = cmd.get(cmd.size() - 2);
+    assertTrue(new File(primaryResource).getName().startsWith("spark-examples"));
+    assertFalse(cmd.contains(SparkLauncher.NO_RESOURCE));
+  }
+
   @Test(expected = IllegalArgumentException.class)
   public void testMissingAppResource() {
     new SparkSubmitCommandBuilder().buildSparkSubmitArgs();