Spark instrumentation: capture emr step id when spark runs as EMR step (#10670)

make span attr for emr step id captured from workdir

capture step id for launcher spans

extract helper in shared location

step id extraction error resilience

add class as helper

debug log

rename to EmrUtils

Add unit tests

Gate emr_step_id attribute to emr-detected runtimes

spotless apply

switch to keys contain .emr. check

remove isRunningOnEMR

regex is specific enough to avoid collision

early returns to avoid deep nesting

Merge branch 'master' into adrien.boitreaud/capture-emr-step-id

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

Author: aboitreaud

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

@@ -276,6 +276,7 @@ private void initApplicationSpanIfNotInitialized() {
     }
 
     captureApplicationParameters(builder);
+    captureEmrStepId(builder);
 
     Optional<OpenlineageParentContext> openlineageParentContext =
         OpenlineageParentContext.from(sparkConf);
@@ -1210,6 +1211,13 @@ private void captureApplicationParameters(AgentTracer.SpanBuilder builder) {
     builder.withTag("config.spark_version", sparkVersion);
   }
 
+  private static void captureEmrStepId(AgentTracer.SpanBuilder builder) {
+    String stepId = EmrUtils.getEmrStepId();
+    if (stepId != null) {
+      builder.withTag("emr_step_id", stepId);
+    }
+  }
+
   private void captureJobParameters(AgentTracer.SpanBuilder builder, Properties properties) {
     for (Tuple2<String, String> conf : sparkConf.getAll()) {
       if (SparkConfAllowList.canCaptureJobParameter(conf._1)) {

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

@@ -0,0 +1,42 @@
+package datadog.trace.instrumentation.spark;
+
+import java.nio.file.Path;
+import java.nio.file.Paths;
+import java.util.regex.Matcher;
+import java.util.regex.Pattern;
+import javax.annotation.Nullable;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+/** Extracts the AWS EMR Step ID from the working directory name (e.g. s-07767992IY7VC5NVV854). */
+class EmrUtils {
+
+  private static final Logger log = LoggerFactory.getLogger(EmrUtils.class);
+
+  /** EMR step ID is a 20 character string with numbers and uppercase letters only */
+  private static final Pattern EMR_STEP_ID_PATTERN = Pattern.compile("^(s-[0-9A-Z]{20})$");
+
+  @Nullable
+  static String getEmrStepId() {
+    try {
+      String userDir = System.getProperty("user.dir");
+      if (userDir == null) {
+        return null;
+      }
+      Path workDir = Paths.get(userDir).getFileName();
+      if (workDir == null) {
+        return null;
+      }
+      Matcher matcher = EMR_STEP_ID_PATTERN.matcher(workDir.toString());
+      if (matcher.matches()) {
+        log.debug("EMR step ID extracted: {}", matcher.group(1));
+        return matcher.group(1);
+      }
+    } catch (Throwable t) {
+      log.debug("Unable to extract EMR step ID from working directory", t);
+    }
+    return null;
+  }
+
+  private EmrUtils() {}
+}

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

@@ -32,7 +32,9 @@ public String instrumentedType() {
   @Override
   public String[] helperClassNames() {
     return new String[] {
-      packageName + ".SparkConfAllowList", packageName + ".SparkLauncherListener",
+      packageName + ".EmrUtils",
+      packageName + ".SparkConfAllowList",
+      packageName + ".SparkLauncherListener",
     };
   }
 

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

@@ -39,6 +39,7 @@ public static synchronized void createLauncherSpan(Object launcher) {
             .start();
     span.setSamplingPriority(PrioritySampling.USER_KEEP, SamplingMechanism.DATA_JOBS);
     setLauncherConfigTags(span, launcher);
+    captureEmrStepId(span);
     launcherSpan = span;
 
     if (!shutdownHookRegistered) {
@@ -131,6 +132,13 @@ public void infoChanged(SparkAppHandle handle) {
     }
   }
 
+  private static void captureEmrStepId(AgentSpan span) {
+    String stepId = EmrUtils.getEmrStepId();
+    if (stepId != null) {
+      span.setTag("emr_step_id", stepId);
+    }
+  }
+
   /**
    * 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

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

@@ -0,0 +1,41 @@
+package datadog.trace.instrumentation.spark;
+
+import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertNull;
+
+import org.junit.jupiter.api.AfterEach;
+import org.junit.jupiter.api.BeforeEach;
+import org.junit.jupiter.api.Test;
+
+class EmrUtilsTest {
+
+  private String originalUserDir;
+
+  @BeforeEach
+  void saveUserDir() {
+    originalUserDir = System.getProperty("user.dir");
+  }
+
+  @AfterEach
+  void restoreUserDir() {
+    System.setProperty("user.dir", originalUserDir);
+  }
+
+  @Test
+  void returnsStepIdWhenWorkdirMatchesEmrPattern() {
+    System.setProperty("user.dir", "/mnt/var/lib/hadoop/steps/s-07767992IY7VC5NVV854");
+    assertEquals("s-07767992IY7VC5NVV854", EmrUtils.getEmrStepId());
+  }
+
+  @Test
+  void returnsNullWhenWorkdirDoesNotMatchEmrPattern() {
+    System.setProperty("user.dir", "/home/hadoop");
+    assertNull(EmrUtils.getEmrStepId());
+  }
+
+  @Test
+  void returnsNullForApplicationIdWorkdir() {
+    System.setProperty("user.dir", "/home/hadoop/application_1234567890_0001");
+    assertNull(EmrUtils.getEmrStepId());
+  }
+}

dd-java-agent/instrumentation/spark/spark_2.12/src/main/java/datadog/trace/instrumentation/spark/Spark212Instrumentation.java

@@ -26,6 +26,7 @@ public String[] helperClassNames() {
       packageName + ".AbstractSparkPlanSerializer",
       packageName + ".AbstractSparkPlanUtils",
       packageName + ".DatabricksParentContext",
+      packageName + ".EmrUtils",
       packageName + ".OpenlineageParentContext",
       packageName + ".DatadogSpark212Listener",
       packageName + ".PredeterminedTraceIdContext",

dd-java-agent/instrumentation/spark/spark_2.13/src/main/java/datadog/trace/instrumentation/spark/Spark213Instrumentation.java

@@ -26,6 +26,7 @@ public String[] helperClassNames() {
       packageName + ".AbstractSparkPlanSerializer",
       packageName + ".AbstractSparkPlanUtils",
       packageName + ".DatabricksParentContext",
+      packageName + ".EmrUtils",
       packageName + ".OpenlineageParentContext",
       packageName + ".DatadogSpark213Listener",
       packageName + ".PredeterminedTraceIdContext",