Gate emr_step_id attribute to emr-detected runtimes

Author: aboitreaud

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

@@ -28,6 +28,7 @@ dependencies {
   testFixturesCompileOnly(libs.bundles.spock)
 
   testImplementation project(':dd-java-agent:instrumentation-testing')
+  testImplementation group: 'org.apache.spark', name: 'spark-core_2.12', version: '2.4.0'
   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/AbstractDatadogSparkListener.java

@@ -134,6 +134,7 @@ public abstract class AbstractDatadogSparkListener extends SparkListener {
 
   private volatile boolean isStreamingJob = false;
   private final boolean isRunningOnDatabricks;
+  private final boolean isRunningOnEmr;
   private final String databricksClusterName;
   private final String databricksServiceName;
   private final String sparkServiceName;
@@ -156,6 +157,7 @@ public AbstractDatadogSparkListener(SparkConf sparkConf, String appId, String sp
     this.sparkVersion = sparkVersion;
 
     isRunningOnDatabricks = sparkConf.contains("spark.databricks.sparkContextId");
+    isRunningOnEmr = EmrUtils.isRunningOnEmr(sparkConf);
     databricksClusterName = sparkConf.get("spark.databricks.clusterUsageTags.clusterName", null);
     databricksServiceName = getDatabricksServiceName(sparkConf, databricksClusterName);
     sparkServiceName = getSparkServiceName(sparkConf, isRunningOnDatabricks);
@@ -276,7 +278,7 @@ private void initApplicationSpanIfNotInitialized() {
     }
 
     captureApplicationParameters(builder);
-    captureEmrStepId(builder);
+    captureEmrStepId(builder, isRunningOnEmr);
 
     Optional<OpenlineageParentContext> openlineageParentContext =
         OpenlineageParentContext.from(sparkConf);
@@ -1211,7 +1213,10 @@ private void captureApplicationParameters(AgentTracer.SpanBuilder builder) {
     builder.withTag("config.spark_version", sparkVersion);
   }
 
-  private static void captureEmrStepId(AgentTracer.SpanBuilder builder) {
+  private static void captureEmrStepId(AgentTracer.SpanBuilder builder, boolean isRunningOnEmr) {
+    if (!isRunningOnEmr) {
+      return;
+    }
     String stepId = EmrUtils.getEmrStepId();
     if (stepId != null) {
       builder.withTag("emr_step_id", stepId);

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

@@ -5,14 +5,31 @@
 import java.util.regex.Matcher;
 import java.util.regex.Pattern;
 import javax.annotation.Nullable;
+import org.apache.spark.SparkConf;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
-/** Extracts the AWS EMR Step ID from the working directory name (e.g. s-07767992IY7VC5NVV854). */
+/** Utilities for detecting AWS EMR and extracting EMR-specific metadata. */
 class EmrUtils {
 
   private static final Logger log = LoggerFactory.getLogger(EmrUtils.class);
-  private static final Pattern EMR_STEP_ID_PATTERN = Pattern.compile("^(s-[0-9A-Za-z]+)$");
+  /** 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})$");
+
+  /**
+   * Returns true if the Spark job is running on AWS EMR. Detection uses two EMR-exclusive keys:
+   *
+   * <ul>
+   *   <li>{@code spark.sql.emr.internal.extensions}: registers EMR's Spark session extensions,
+   *       present across all known EMR releases (observed in community EMR configs)
+   *   <li>{@code spark.emr.default.executor.cores}: newer EMR default, added as a fallback for
+   *       future releases where the extensions key might change
+   * </ul>
+   */
+  static boolean isRunningOnEmr(SparkConf conf) {
+    return conf.contains("spark.sql.emr.internal.extensions")
+        || conf.contains("spark.emr.default.executor.cores");
+  }
 
   @Nullable
   static String getEmrStepId() {

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

@@ -2,7 +2,9 @@
 
 import static org.junit.jupiter.api.Assertions.assertEquals;
 import static org.junit.jupiter.api.Assertions.assertNull;
+import static org.junit.jupiter.api.Assertions.assertTrue;
 
+import org.apache.spark.SparkConf;
 import org.junit.jupiter.api.AfterEach;
 import org.junit.jupiter.api.BeforeEach;
 import org.junit.jupiter.api.Test;
@@ -38,4 +40,19 @@ void returnsNullForApplicationIdWorkdir() {
     System.setProperty("user.dir", "/home/hadoop/application_1234567890_0001");
     assertNull(EmrUtils.getEmrStepId());
   }
+
+  @Test
+  void isRunningOnEmrReturnsTrueWhenEmrInternalExtensionsPresent() {
+    SparkConf conf = new SparkConf();
+    conf.set("spark.sql.emr.internal.extensions", "com.amazonaws.emr.spark.EmrSparkSessionExtensions");
+    assertTrue(EmrUtils.isRunningOnEmr(conf));
+  }
+
+  @Test
+  void isRunningOnEmrReturnsTrueWhenEmrDefaultExecutorCoresPresent() {
+    SparkConf conf = new SparkConf();
+    conf.set("spark.emr.default.executor.cores", "1");
+    assertTrue(EmrUtils.isRunningOnEmr(conf));
+  }
+
 }