initial spark launcher instrumentation (#10629)

initial spark launcher instrumentation

use ddtags

Fix tess

move test to the right /test dir

advice should be public

finish launcher span with error via RunMainAdvice

sportLess

synchronize shutdown hook

Capture more spark relevant attrs

Update tests with new attrs

fix sportBugsMain and muzzle

remove SparkLauncher.launch() instrumentation

share common config key redaction method

make public to avoid IllegalAccessError

error type and error message

Add appId and stack trace

extract span building in SparkLaunchListener

wait for throwable and let the span be finished by shutdown hook

spotless apply

fix span lifecycle in the launcher listener

Merge branch 'master' into adrien.boitreaud/spark-launcher-instrumentation

Co-authored-by: devflow.devflow-routing-intake <devflow.devflow-routing-intake@kubernetes.us1.ddbuild.io>

Author: aboitreaud

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,8 @@ public static void exit(@Advice.Thrown Throwable throwable) {
       if (AbstractDatadogSparkListener.listener != null) {
         AbstractDatadogSparkListener.listener.finishApplication(
             System.currentTimeMillis(), throwable, 0, null);
+      } else {
+        SparkLauncherListener.finishSpanWithThrowable(throwable);
       }
     }
   }

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

@@ -17,6 +17,12 @@
  * @see <a href="https://spark.apache.org/docs/latest/configuration.html">Spark Configuration</a>
  */
 class SparkConfAllowList {
+  // Using values from
+  // https://github.com/apache/spark/blob/v3.5.1/core/src/main/scala/org/apache/spark/internal/config/package.scala#L1150-L1158
+  static final String DEFAULT_REDACTION_REGEX = "(?i)secret|password|token|access.key|api.key";
+
+  private static final Pattern DEFAULT_REDACTION_PATTERN = Pattern.compile(DEFAULT_REDACTION_REGEX);
+
   /**
    * Job-specific parameters that can be used to control job execution or provide metadata about the
    * job being executed
@@ -80,11 +86,17 @@ public static boolean canCaptureJobParameter(String parameterName) {
     return allowedJobParams.contains(parameterName);
   }
 
+  /** Redact a value if the key or value matches the default redaction pattern. */
+  public static String redactValue(String key, String value) {
+    if (DEFAULT_REDACTION_PATTERN.matcher(key).find()
+        || DEFAULT_REDACTION_PATTERN.matcher(value).find()) {
+      return "[redacted]";
+    }
+    return value;
+  }
+
   public static List<Map.Entry<String, String>> getRedactedSparkConf(SparkConf conf) {
-    // Using values from
-    // https://github.com/apache/spark/blob/v3.5.1/core/src/main/scala/org/apache/spark/internal/config/package.scala#L1150-L1158
-    String redactionPattern =
-        conf.get("spark.redaction.regex", "(?i)secret|password|token|access.key|api.key");
+    String redactionPattern = conf.get("spark.redaction.regex", DEFAULT_REDACTION_REGEX);
     List<Map.Entry<String, String>> redacted = new ArrayList<>();
     Pattern pattern = Pattern.compile(redactionPattern);
 

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

@@ -0,0 +1,66 @@
+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;
+import net.bytebuddy.asm.Advice;
+import org.apache.spark.launcher.SparkAppHandle;
+
+@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 + ".SparkConfAllowList", packageName + ".SparkLauncherListener",
+    };
+  }
+
+  @Override
+  public void methodAdvice(MethodTransformer transformer) {
+    transformer.applyAdvice(
+        isMethod()
+            .and(named("startApplication"))
+            .and(isDeclaredBy(named("org.apache.spark.launcher.SparkLauncher"))),
+        SparkLauncherInstrumentation.class.getName() + "$StartApplicationAdvice");
+  }
+
+  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) {
+      SparkLauncherListener.createLauncherSpan(launcher);
+
+      if (throwable != null) {
+        SparkLauncherListener.finishSpanWithThrowable(throwable);
+        return;
+      }
+
+      if (handle != null) {
+        handle.addListener(new SparkLauncherListener());
+      }
+    }
+  }
+}

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

@@ -0,0 +1,192 @@
+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 org.apache.spark.launcher.SparkAppHandle;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+/**
+ * Listener for SparkLauncher spans. Tracks the lifecycle of a Spark application submitted via
+ * SparkLauncher.startApplication(). Only a single launcher span can be active at a time. Subsequent
+ * calls to startApplication() from the same or different launcher instances will not create spans;
+ * only the first launch in the JVM is traced
+ */
+public class SparkLauncherListener implements SparkAppHandle.Listener {
+
+  private static final Logger log = LoggerFactory.getLogger(SparkLauncherListener.class);
+
+  static volatile AgentSpan launcherSpan;
+
+  private static volatile boolean shutdownHookRegistered = false;
+
+  public static synchronized void createLauncherSpan(Object launcher) {
+    if (launcherSpan != null) {
+      return;
+    }
+
+    AgentTracer.TracerAPI tracer = AgentTracer.get();
+    AgentSpan span =
+        tracer
+            .buildSpan("spark.launcher.launch")
+            .withSpanType("spark")
+            .withResourceName("SparkLauncher.startApplication")
+            .start();
+    span.setSamplingPriority(PrioritySampling.USER_KEEP, SamplingMechanism.DATA_JOBS);
+    setLauncherConfigTags(span, launcher);
+    launcherSpan = span;
+
+    if (!shutdownHookRegistered) {
+      shutdownHookRegistered = true;
+      Runtime.getRuntime()
+          .addShutdownHook(
+              new Thread(
+                  () -> {
+                    synchronized (SparkLauncherListener.class) {
+                      AgentSpan s = launcherSpan;
+                      if (s != null) {
+                        log.info("Finishing spark.launcher span from shutdown hook");
+                        s.finish();
+                        launcherSpan = null;
+                      }
+                    }
+                  }));
+    }
+  }
+
+  public static synchronized void finishSpan(boolean isError, String errorMessage) {
+    AgentSpan span = launcherSpan;
+    if (span == null) {
+      return;
+    }
+    if (isError) {
+      span.setError(true);
+      span.setTag(DDTags.ERROR_TYPE, "Spark Launcher Failed");
+      span.setTag(DDTags.ERROR_MSG, errorMessage);
+    }
+    span.finish();
+    launcherSpan = null;
+  }
+
+  public static synchronized void finishSpanWithThrowable(Throwable throwable) {
+    AgentSpan span = launcherSpan;
+    if (span == null) {
+      return;
+    }
+    if (throwable != null) {
+      span.addThrowable(throwable);
+    }
+    span.finish();
+    launcherSpan = null;
+  }
+
+  @Override
+  public void stateChanged(SparkAppHandle handle) {
+    synchronized (SparkLauncherListener.class) {
+      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);
+          span.setTag("app_id", appId);
+        }
+
+        if (state.isFinal()) {
+          if (state == SparkAppHandle.State.FAILED
+              || state == SparkAppHandle.State.KILLED
+              || state == SparkAppHandle.State.LOST) {
+            // Set error tags but don't finish yet — RunMainAdvice may add the throwable
+            // with the full stack trace. The span will be finished by RunMainAdvice or
+            // the shutdown hook.
+            span.setError(true);
+            span.setTag(DDTags.ERROR_TYPE, "Spark Launcher Failed");
+            span.setTag(DDTags.ERROR_MSG, "Application " + state);
+          } else {
+            finishSpan(false, null);
+          }
+        }
+      }
+    }
+  }
+
+  @Override
+  public void infoChanged(SparkAppHandle handle) {
+    synchronized (SparkLauncherListener.class) {
+      AgentSpan span = launcherSpan;
+      if (span != null) {
+        String appId = handle.getAppId();
+        if (appId != null) {
+          span.setTag("spark.app_id", appId);
+          span.setTag("app_id", appId);
+        }
+      }
+    }
+  }
+
+  /**
+   * Extract launcher configuration via reflection and set as span tags. Secret redaction uses the
+   * default pattern only (not spark.redaction.regex) because the SparkLauncher conf map is a plain
+   * Map, not a SparkConf, so there is no way to read the user's custom redaction regex at this
+   * point.
+   */
+  private static void setLauncherConfigTags(AgentSpan span, Object launcher) {
+    try {
+      Field builderField = launcher.getClass().getSuperclass().getDeclaredField("builder");
+      builderField.setAccessible(true);
+      Object builder = builderField.get(launcher);
+      if (builder == null) {
+        return;
+      }
+
+      Class<?> builderClass = builder.getClass();
+      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");
+
+      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 = SparkConfAllowList.redactValue(entry.getKey(), entry.getValue());
+              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);
+    }
+  }
+}