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/AbstractSparkInstrumentation.java

@@ -106,6 +106,17 @@ public static void exit(@Advice.Thrown Throwable throwable) {
       if (AbstractDatadogSparkListener.listener != null) {
         AbstractDatadogSparkListener.listener.finishApplication(
             System.currentTimeMillis(), throwable, 0, null);
+      } else {
+        try {
+          Class<?> adviceClass =
+              Thread.currentThread()
+                  .getContextClassLoader()
+                  .loadClass("datadog.trace.instrumentation.spark.SparkLauncherAdvice");
+          java.lang.reflect.Method finishMethod =
+              adviceClass.getMethod("finishLauncherSpan", Throwable.class);
+          finishMethod.invoke(null, throwable);
+        } catch (Exception ignored) {
+        }
       }
     }
   }

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,232 @@
+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 java.lang.reflect.Field;
+import java.util.Map;
+import java.util.regex.Pattern;
+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);
+
+  // Same default pattern as spark.redaction.regex in Spark source
+  private static final Pattern CONF_REDACTION_PATTERN =
+      Pattern.compile("(?i)secret|password|token|access.key|api.key");
+
+  /** The launcher span, accessible from SparkExitAdvice via reflection. */
+  public static volatile AgentSpan launcherSpan;
+
+  private static volatile boolean shutdownHookRegistered = false;
+
+  /** Extract SparkLauncher configuration via reflection and set as span tags. */
+  private static void setLauncherConfigTags(AgentSpan span, Object launcher) {
+    try {
+      // SparkLauncher extends AbstractLauncher which has a 'builder' field
+      Field builderField = launcher.getClass().getSuperclass().getDeclaredField("builder");
+      builderField.setAccessible(true);
+      Object builder = builderField.get(launcher);
+      if (builder == null) {
+        return;
+      }
+
+      Class<?> builderClass = builder.getClass();
+      // Fields are on AbstractCommandBuilder (parent of SparkSubmitCommandBuilder)
+      Class<?> abstractBuilderClass = builderClass.getSuperclass();
+
+      setStringFieldAsTag(span, builder, abstractBuilderClass, "master", "master");
+      setStringFieldAsTag(span, builder, abstractBuilderClass, "deployMode", "deploy_mode");
+      setStringFieldAsTag(span, builder, abstractBuilderClass, "appName", "application_name");
+      setStringFieldAsTag(span, builder, abstractBuilderClass, "mainClass", "main_class");
+      setStringFieldAsTag(span, builder, abstractBuilderClass, "appResource", "app_resource");
+
+      // Extract spark conf entries and redact sensitive values
+      try {
+        Field confField = abstractBuilderClass.getDeclaredField("conf");
+        confField.setAccessible(true);
+        @SuppressWarnings("unchecked")
+        Map<String, String> conf = (Map<String, String>) confField.get(builder);
+        if (conf != null) {
+          for (Map.Entry<String, String> entry : conf.entrySet()) {
+            if (SparkConfAllowList.canCaptureJobParameter(entry.getKey())) {
+              String value = entry.getValue();
+              if (CONF_REDACTION_PATTERN.matcher(entry.getKey()).find()
+                  || CONF_REDACTION_PATTERN.matcher(value).find()) {
+                value = "[redacted]";
+              }
+              span.setTag("config." + entry.getKey().replace('.', '_'), value);
+            }
+          }
+        }
+      } catch (NoSuchFieldException e) {
+        log.debug("Could not find conf field on builder", e);
+      }
+    } catch (Exception e) {
+      log.debug("Failed to extract SparkLauncher configuration", e);
+    }
+  }
+
+  private static void setStringFieldAsTag(
+      AgentSpan span, Object obj, Class<?> clazz, String fieldName, String tagName) {
+    try {
+      Field field = clazz.getDeclaredField(fieldName);
+      field.setAccessible(true);
+      Object value = field.get(obj);
+      if (value != null) {
+        span.setTag(tagName, value.toString());
+      }
+    } catch (Exception e) {
+      log.debug("Could not read field {} from builder", fieldName, e);
+    }
+  }
+
+  public static synchronized void createLauncherSpan(String resource, Object launcher) {
+    if (launcherSpan != null) {
+      return;
+    }
+
+    AgentTracer.TracerAPI tracer = AgentTracer.get();
+    AgentSpan span =
+        tracer
+            .buildSpan("spark.launcher.launch")
+            .withSpanType("spark")
+            .withResourceName(resource)
+            .start();
+    span.setSamplingPriority(PrioritySampling.USER_KEEP, SamplingMechanism.DATA_JOBS);
+
+    if (launcher != null) {
+      setLauncherConfigTags(span, launcher);
+    }
+
+    launcherSpan = span;
+
+    if (!shutdownHookRegistered) {
+      shutdownHookRegistered = true;
+      Runtime.getRuntime()
+          .addShutdownHook(
+              new Thread(
+                  () -> {
+                    synchronized (SparkLauncherAdvice.class) {
+                      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 synchronized void finishLauncherSpan(Throwable throwable) {
+    AgentSpan span = launcherSpan;
+    if (span == null) {
+      return;
+    }
+    if (throwable != null) {
+      span.addThrowable(throwable);
+    }
+    span.finish();
+    launcherSpan = null;
+  }
+
+  public static class StartApplicationAdvice {
+    @Advice.OnMethodExit(suppress = Throwable.class, onThrowable = Throwable.class)
+    public static void exit(
+        @Advice.This Object launcher,
+        @Advice.Return SparkAppHandle handle,
+        @Advice.Thrown Throwable throwable) {
+      createLauncherSpan("SparkLauncher.startApplication", launcher);
+
+      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.This Object launcher, @Advice.Thrown Throwable throwable) {
+      createLauncherSpan("SparkLauncher.launch", launcher);
+
+      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");
+  }
+}