initial spark launcher instrumentation

dd-java-agent/instrumentation/spark/spark-common/build.gradle

@@ -10,6 +10,7 @@ configurations.configureEach {
 dependencies {
   compileOnly group: 'org.apache.spark', name: 'spark-core_2.12', version: '2.4.0'
   compileOnly group: 'org.apache.spark', name: 'spark-sql_2.12', version: '2.4.0'
+  compileOnly group: 'org.apache.spark', name: 'spark-launcher_2.12', version: '2.4.0'
 
   testFixturesImplementation group: 'com.datadoghq', name: 'sketches-java', version: '0.8.2'
   testFixturesImplementation group: 'com.google.protobuf', name: 'protobuf-java', version: '3.14.0'
@@ -21,7 +22,12 @@ dependencies {
   testFixturesCompileOnly group: 'org.apache.spark', name: 'spark-core_2.12', version: '2.4.0'
   testFixturesCompileOnly group: 'org.apache.spark', name: 'spark-sql_2.12', version: '2.4.0'
   testFixturesCompileOnly group: 'org.apache.spark', name: 'spark-yarn_2.12', version: '2.4.0'
+  testFixturesCompileOnly group: 'org.apache.spark', name: 'spark-launcher_2.12', version: '2.4.0'
 
   testFixturesCompileOnly(libs.bundles.groovy)
   testFixturesCompileOnly(libs.bundles.spock)
+
+  testImplementation project(':dd-java-agent:instrumentation-testing')
+  testImplementation group: 'org.apache.spark', name: 'spark-launcher_2.12', version: '2.4.0'
 }
+

dd-java-agent/instrumentation/spark/spark-common/src/main/java/datadog/trace/instrumentation/spark/SparkExitAdvice.java

@@ -21,6 +21,14 @@ public static void enter(@Advice.Argument(0) int exitCode) {
                 .getMethod(
                     "finishApplication", long.class, Throwable.class, int.class, String.class);
         method.invoke(datadogListener, System.currentTimeMillis(), null, exitCode, null);
+      } else {
+        // No Spark listener means we may be in a launcher-only process (e.g. SparkLauncher on EMR)
+        Class<?> adviceClass =
+            Thread.currentThread()
+                .getContextClassLoader()
+                .loadClass("datadog.trace.instrumentation.spark.SparkLauncherAdvice");
+        Method finishMethod = adviceClass.getDeclaredMethod("finishLauncherSpan", int.class);
+        finishMethod.invoke(null, exitCode);
       }
     } catch (Exception ignored) {
     }

dd-java-agent/instrumentation/spark/spark-common/src/main/java/datadog/trace/instrumentation/spark/SparkLauncherAdvice.java

@@ -0,0 +1,139 @@
+package datadog.trace.instrumentation.spark;
+
+import datadog.trace.api.DDTags;
+import datadog.trace.api.sampling.PrioritySampling;
+import datadog.trace.api.sampling.SamplingMechanism;
+import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
+import datadog.trace.bootstrap.instrumentation.api.AgentTracer;
+import net.bytebuddy.asm.Advice;
+import org.apache.spark.launcher.SparkAppHandle;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+public class SparkLauncherAdvice {
+
+  private static final Logger log = LoggerFactory.getLogger(SparkLauncherAdvice.class);
+
+  /** The launcher span, accessible from SparkExitAdvice via reflection. */
+  public static volatile AgentSpan launcherSpan;
+
+  private static volatile boolean shutdownHookRegistered = false;
+
+  public static synchronized void createLauncherSpan(String resource) {
+    if (launcherSpan != null) {
+      return;
+    }
+
+    AgentTracer.TracerAPI tracer = AgentTracer.get();
+    AgentSpan span =
+        tracer.buildSpan("spark.launcher").withSpanType("spark").withResourceName(resource).start();
+    span.setSamplingPriority(PrioritySampling.USER_KEEP, SamplingMechanism.DATA_JOBS);
+    launcherSpan = span;
+
+    if (!shutdownHookRegistered) {
+      shutdownHookRegistered = true;
+      Runtime.getRuntime()
+          .addShutdownHook(
+              new Thread(
+                  () -> {
+                    AgentSpan s = launcherSpan;
+                    if (s != null) {
+                      log.info("Finishing spark.launcher span from shutdown hook");
+                      s.finish();
+                      launcherSpan = null;
+                    }
+                  }));
+    }
+  }
+
+  public static synchronized void finishLauncherSpan(int exitCode) {
+    AgentSpan span = launcherSpan;
+    if (span == null) {
+      return;
+    }
+    if (exitCode != 0) {
+      span.setError(true);
+      span.setTag(DDTags.ERROR_TYPE, "Spark Launcher Failed with exit code " + exitCode);
+    }
+    span.finish();
+    launcherSpan = null;
+  }
+
+  public static class StartApplicationAdvice {
+    @Advice.OnMethodExit(suppress = Throwable.class, onThrowable = Throwable.class)
+    public static void exit(
+        @Advice.Return SparkAppHandle handle, @Advice.Thrown Throwable throwable) {
+      createLauncherSpan("SparkLauncher.startApplication");
+
+      if (throwable != null) {
+        AgentSpan span = launcherSpan;
+        if (span != null) {
+          span.addThrowable(throwable);
+          span.finish();
+          launcherSpan = null;
+        }
+        return;
+      }
+
+      if (handle != null) {
+        try {
+          handle.addListener(new AppHandleListener());
+        } catch (Exception e) {
+          log.debug("Failed to register SparkAppHandle listener", e);
+        }
+      }
+    }
+  }
+
+  public static class LaunchAdvice {
+    @Advice.OnMethodExit(suppress = Throwable.class, onThrowable = Throwable.class)
+    public static void exit(@Advice.Thrown Throwable throwable) {
+      createLauncherSpan("SparkLauncher.launch");
+
+      if (throwable != null) {
+        AgentSpan span = launcherSpan;
+        if (span != null) {
+          span.addThrowable(throwable);
+          span.finish();
+          launcherSpan = null;
+        }
+      }
+    }
+  }
+
+  static class AppHandleListener implements SparkAppHandle.Listener {
+    @Override
+    public void stateChanged(SparkAppHandle handle) {
+      SparkAppHandle.State state = handle.getState();
+      AgentSpan span = launcherSpan;
+      if (span != null) {
+        span.setTag("spark.launcher.app_state", state.toString());
+
+        String appId = handle.getAppId();
+        if (appId != null) {
+          span.setTag("spark.app_id", appId);
+        }
+
+        if (state.isFinal()) {
+          if (state == SparkAppHandle.State.FAILED
+              || state == SparkAppHandle.State.KILLED
+              || state == SparkAppHandle.State.LOST) {
+            span.setError(true);
+            span.setTag(DDTags.ERROR_TYPE, "Spark Application " + state);
+          }
+        }
+      }
+    }
+
+    @Override
+    public void infoChanged(SparkAppHandle handle) {
+      AgentSpan span = launcherSpan;
+      if (span != null) {
+        String appId = handle.getAppId();
+        if (appId != null) {
+          span.setTag("spark.app_id", appId);
+        }
+      }
+    }
+  }
+}

dd-java-agent/instrumentation/spark/spark-common/src/main/java/datadog/trace/instrumentation/spark/SparkLauncherInstrumentation.java

@@ -0,0 +1,51 @@
+package datadog.trace.instrumentation.spark;
+
+import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.named;
+import static net.bytebuddy.matcher.ElementMatchers.isDeclaredBy;
+import static net.bytebuddy.matcher.ElementMatchers.isMethod;
+
+import com.google.auto.service.AutoService;
+import datadog.trace.agent.tooling.Instrumenter;
+import datadog.trace.agent.tooling.InstrumenterModule;
+import datadog.trace.api.InstrumenterConfig;
+
+@AutoService(InstrumenterModule.class)
+public class SparkLauncherInstrumentation extends InstrumenterModule.Tracing
+    implements Instrumenter.ForSingleType, Instrumenter.HasMethodAdvice {
+
+  public SparkLauncherInstrumentation() {
+    super("spark-launcher");
+  }
+
+  @Override
+  protected boolean defaultEnabled() {
+    return InstrumenterConfig.get().isDataJobsEnabled();
+  }
+
+  @Override
+  public String instrumentedType() {
+    return "org.apache.spark.launcher.SparkLauncher";
+  }
+
+  @Override
+  public String[] helperClassNames() {
+    return new String[] {
+      packageName + ".SparkLauncherAdvice", packageName + ".SparkLauncherAdvice$AppHandleListener",
+    };
+  }
+
+  @Override
+  public void methodAdvice(MethodTransformer transformer) {
+    transformer.applyAdvice(
+        isMethod()
+            .and(named("startApplication"))
+            .and(isDeclaredBy(named("org.apache.spark.launcher.SparkLauncher"))),
+        packageName + ".SparkLauncherAdvice$StartApplicationAdvice");
+
+    transformer.applyAdvice(
+        isMethod()
+            .and(named("launch"))
+            .and(isDeclaredBy(named("org.apache.spark.launcher.SparkLauncher"))),
+        packageName + ".SparkLauncherAdvice$LaunchAdvice");
+  }
+}

dd-java-agent/instrumentation/spark/spark-common/src/test/groovy/datadog/trace/instrumentation/spark/SparkLauncherTest.groovy

@@ -0,0 +1,72 @@
+package datadog.trace.instrumentation.spark
+
+import datadog.trace.agent.test.InstrumentationSpecification
+import datadog.trace.api.sampling.PrioritySampling
+import datadog.trace.api.sampling.SamplingMechanism
+import datadog.trace.bootstrap.instrumentation.api.AgentTracer
+
+class SparkLauncherTest extends InstrumentationSpecification {
+
+  def "generate spark.launcher span from startApplication"() {
+    setup:
+    SparkLauncherAdvice.launcherSpan = null
+
+    when:
+    def tracer = AgentTracer.get()
+    def launcherSpan = tracer
+      .buildSpan("spark.launcher")
+      .withSpanType("spark")
+      .withResourceName("SparkLauncher.startApplication")
+      .start()
+    launcherSpan.setSamplingPriority(
+      PrioritySampling.USER_KEEP,
+      SamplingMechanism.DATA_JOBS)
+    SparkLauncherAdvice.launcherSpan = launcherSpan
+
+    SparkLauncherAdvice.finishLauncherSpan(1)
+
+    then:
+    assertTraces(1) {
+      trace(1) {
+        span {
+          operationName "spark.launcher"
+          spanType "spark"
+          resourceName "SparkLauncher.startApplication"
+          errored true
+          assert span.tags["error.type"] == "Spark Launcher Failed with exit code 1"
+        }
+      }
+    }
+  }
+
+  def "generate spark.launcher span with successful exit"() {
+    setup:
+    SparkLauncherAdvice.launcherSpan = null
+
+    when:
+    def tracer = AgentTracer.get()
+    def launcherSpan = tracer
+      .buildSpan("spark.launcher")
+      .withSpanType("spark")
+      .withResourceName("SparkLauncher.launch")
+      .start()
+    launcherSpan.setSamplingPriority(
+      PrioritySampling.USER_KEEP,
+      SamplingMechanism.DATA_JOBS)
+    SparkLauncherAdvice.launcherSpan = launcherSpan
+
+    SparkLauncherAdvice.finishLauncherSpan(0)
+
+    then:
+    assertTraces(1) {
+      trace(1) {
+        span {
+          operationName "spark.launcher"
+          spanType "spark"
+          resourceName "SparkLauncher.launch"
+          errored false
+        }
+      }
+    }
+  }
+}