stop injecting OL if it's not available

Signed-off-by: Maciej Obuchowski <[email protected]>

Author: mobuchowski

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

@@ -7,6 +7,7 @@
 import datadog.trace.agent.tooling.InstrumenterModule;
 import net.bytebuddy.asm.Advice;
 import org.apache.spark.SparkContext;
+import org.apache.spark.util.Utils;
 
 @AutoService(InstrumenterModule.class)
 public class Spark212Instrumentation extends AbstractSparkInstrumentation {
@@ -17,6 +18,7 @@ public String[] helperClassNames() {
       packageName + ".DatabricksParentContext",
       packageName + ".OpenlineageParentContext",
       packageName + ".DatadogSpark212Listener",
+      packageName + ".PredeterminedTraceIdContext",
       packageName + ".RemoveEldestHashMap",
       packageName + ".SparkAggregatedTaskMetrics",
       packageName + ".SparkConfAllowList",
@@ -41,18 +43,23 @@ public void methodAdvice(MethodTransformer transformer) {
   public static class InjectListener {
     @Advice.OnMethodEnter(suppress = Throwable.class)
     public static void enter(@Advice.This SparkContext sparkContext) {
-      if (!sparkContext.conf().contains("spark.extraListeners")) {
-        sparkContext
-            .conf()
-            .set("spark.extraListeners", "io.openlineage.spark.agent.OpenLineageSparkListener");
-      } else {
-        String extraListeners = sparkContext.conf().get("spark.extraListeners");
-        if (!extraListeners.contains("io.openlineage.spark.agent.OpenLineageSparkListener")) {
+      // checking whether OpenLineage integration is available and that it supports tags
+      if (Utils.classIsLoadable("io.openlineage.spark.agent.OpenLineageSparkListener")
+          && Utils.classIsLoadable(
+              "io.openlineage.spark.agent.facets.builder.TagsRunFacetBuilder")) {
+        if (!sparkContext.conf().contains("spark.extraListeners")) {
           sparkContext
               .conf()
-              .set(
-                  "spark.extraListeners",
-                  extraListeners + ",io.openlineage.spark.agent.OpenLineageSparkListener");
+              .set("spark.extraListeners", "io.openlineage.spark.agent.OpenLineageSparkListener");
+        } else {
+          String extraListeners = sparkContext.conf().get("spark.extraListeners");
+          if (!extraListeners.contains("io.openlineage.spark.agent.OpenLineageSparkListener")) {
+            sparkContext
+                .conf()
+                .set(
+                    "spark.extraListeners",
+                    extraListeners + ",io.openlineage.spark.agent.OpenLineageSparkListener");
+          }
         }
       }
 

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

@@ -7,6 +7,7 @@
 import datadog.trace.agent.tooling.InstrumenterModule;
 import net.bytebuddy.asm.Advice;
 import org.apache.spark.SparkContext;
+import org.apache.spark.util.Utils;
 
 @AutoService(InstrumenterModule.class)
 public class Spark213Instrumentation extends AbstractSparkInstrumentation {
@@ -17,6 +18,7 @@ public String[] helperClassNames() {
       packageName + ".DatabricksParentContext",
       packageName + ".OpenlineageParentContext",
       packageName + ".DatadogSpark213Listener",
+      packageName + ".PredeterminedTraceIdContext",
       packageName + ".RemoveEldestHashMap",
       packageName + ".SparkAggregatedTaskMetrics",
       packageName + ".SparkConfAllowList",
@@ -41,18 +43,23 @@ public void methodAdvice(MethodTransformer transformer) {
   public static class InjectListener {
     @Advice.OnMethodEnter(suppress = Throwable.class)
     public static void enter(@Advice.This SparkContext sparkContext) {
-      if (!sparkContext.conf().contains("spark.extraListeners")) {
-        sparkContext
-            .conf()
-            .set("spark.extraListeners", "io.openlineage.spark.agent.OpenLineageSparkListener");
-      } else {
-        String extraListeners = sparkContext.conf().get("spark.extraListeners");
-        if (!extraListeners.contains("io.openlineage.spark.agent.OpenLineageSparkListener")) {
+      // checking whether OpenLineage integration is available and that it supports tags
+      if (Utils.classIsLoadable("io.openlineage.spark.agent.OpenLineageSparkListener")
+          && Utils.classIsLoadable(
+              "io.openlineage.spark.agent.facets.builder.TagsRunFacetBuilder")) {
+        if (!sparkContext.conf().contains("spark.extraListeners")) {
           sparkContext
               .conf()
-              .set(
-                  "spark.extraListeners",
-                  extraListeners + ",io.openlineage.spark.agent.OpenLineageSparkListener");
+              .set("spark.extraListeners", "io.openlineage.spark.agent.OpenLineageSparkListener");
+        } else {
+          String extraListeners = sparkContext.conf().get("spark.extraListeners");
+          if (!extraListeners.contains("io.openlineage.spark.agent.OpenLineageSparkListener")) {
+            sparkContext
+                .conf()
+                .set(
+                    "spark.extraListeners",
+                    extraListeners + ",io.openlineage.spark.agent.OpenLineageSparkListener");
+          }
         }
       }
 

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

@@ -88,7 +88,7 @@ public abstract class AbstractDatadogSparkListener extends SparkListener {
 
   // This is created by constructor, and used if we're not in other known
   // parent context like Databricks, OpenLineage
-  private final PredeterminedTraceIdParentContext predeterminedTraceIdParentContext;
+  private final PredeterminedTraceIdContext predeterminedTraceIdContext;
 
   private AgentSpan applicationSpan;
   private SparkListenerApplicationStart applicationStart;
@@ -145,9 +145,8 @@ public AbstractDatadogSparkListener(SparkConf sparkConf, String appId, String sp
     databricksClusterName = sparkConf.get("spark.databricks.clusterUsageTags.clusterName", null);
     databricksServiceName = getDatabricksServiceName(sparkConf, databricksClusterName);
     sparkServiceName = getSparkServiceName(sparkConf, isRunningOnDatabricks);
-    predeterminedTraceIdParentContext =
-        new PredeterminedTraceIdParentContext(
-            Config.get().getIdGenerationStrategy().generateTraceId());
+    predeterminedTraceIdContext =
+        new PredeterminedTraceIdContext(Config.get().getIdGenerationStrategy().generateTraceId());
 
     // If JVM exiting with System.exit(code), it bypass the code closing the application span
     //
@@ -167,7 +166,7 @@ public AbstractDatadogSparkListener(SparkConf sparkConf, String appId, String sp
   }
 
   public void setupOpenLineage(DDTraceId traceId) {
-    log.error("Setting up OpenLineage tags");
+    log.debug("Setting up OpenLineage tags");
     if (openLineageSparkListener != null) {
       openLineageSparkConf.set("spark.openlineage.transport.type", "composite");
       openLineageSparkConf.set("spark.openlineage.transport.continueOnFailure", "true");
@@ -182,7 +181,7 @@ public void setupOpenLineage(DDTraceId traceId) {
           "_dd.trace_id:" + traceId.toString() + ";_dd.ol_intake.emit_spans:false");
       return;
     }
-    log.error("No OpenLineageSparkListener!");
+    log.debug("No OpenLineageSparkListener!");
   }
 
   /** Resource name of the spark job. Provide an implementation based on a specific scala version */
@@ -210,7 +209,7 @@ public synchronized void onApplicationStart(SparkListenerApplicationStart applic
     setupOpenLineage(
         OpenlineageParentContext.from(sparkConf)
             .map(context -> context.getTraceId())
-            .orElse(predeterminedTraceIdParentContext.getTraceId()));
+            .orElse(predeterminedTraceIdContext.getTraceId()));
     notifyOl(x -> openLineageSparkListener.onApplicationStart(x), applicationStart);
   }
 
@@ -219,7 +218,7 @@ private void initApplicationSpanIfNotInitialized() {
       return;
     }
 
-    log.error("Starting tracer application span.");
+    log.debug("Starting tracer application span.");
 
     AgentTracer.SpanBuilder builder = buildSparkSpan("spark.application", null);
 
@@ -242,7 +241,7 @@ private void initApplicationSpanIfNotInitialized() {
     if (openlineageParentContext.isPresent()) {
       captureOpenlineageContextIfPresent(builder, openlineageParentContext.get());
     } else {
-      builder.asChildOf(predeterminedTraceIdParentContext);
+      builder.asChildOf(predeterminedTraceIdContext);
     }
 
     applicationSpan = builder.start();

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

@@ -32,6 +32,7 @@ public String[] knownMatchingTypes() {
       "org.apache.spark.SparkContext",
       "org.apache.spark.deploy.SparkSubmit",
       "org.apache.spark.deploy.yarn.ApplicationMaster",
+      "org.apache.spark.util.Utils",
       "org.apache.spark.util.SparkClassUtils",
       "org.apache.spark.scheduler.LiveListenerBus"
     };

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

@@ -29,6 +29,7 @@ public String[] helperClassNames() {
       packageName + ".AbstractDatadogSparkListener",
       packageName + ".DatabricksParentContext",
       packageName + ".OpenlineageParentContext",
+      packageName + ".PredeterminedTraceIdContext",
       packageName + ".RemoveEldestHashMap",
       packageName + ".SparkAggregatedTaskMetrics",
       packageName + ".SparkConfAllowList",

dd-java-agent/instrumentation/spark/src/main/java/datadog/trace/instrumentation/spark/PredeterminedTraceIdParentContext.java renamed to dd-java-agent/instrumentation/spark/src/main/java/datadog/trace/instrumentation/spark/PredeterminedTraceIdContext.java

@@ -5,12 +5,13 @@
 import datadog.trace.api.sampling.PrioritySampling;
 import datadog.trace.bootstrap.instrumentation.api.AgentSpanContext;
 import datadog.trace.bootstrap.instrumentation.api.AgentTraceCollector;
+import datadog.trace.bootstrap.instrumentation.api.AgentTracer;
 import java.util.Map;
 
-public class PredeterminedTraceIdParentContext implements AgentSpanContext {
+public class PredeterminedTraceIdContext implements AgentSpanContext {
   private final DDTraceId traceId;
 
-  public PredeterminedTraceIdParentContext(DDTraceId traceId) {
+  public PredeterminedTraceIdContext(DDTraceId traceId) {
     this.traceId = traceId;
   }
 
@@ -26,7 +27,7 @@ public long getSpanId() {
 
   @Override
   public AgentTraceCollector getTraceCollector() {
-    return null;
+    return AgentTracer.NoopAgentTraceCollector.INSTANCE;
   }
 
   @Override

dd-java-agent/instrumentation/spark/src/testFixtures/groovy/datadog/trace/instrumentation/spark/AbstractSparkTest.groovy

@@ -53,6 +53,7 @@ abstract class AbstractSparkTest extends AgentTestRunner {
           resourceName "spark.application"
           spanType "spark"
           errored false
+          assert span.context().getTraceId() != DDTraceId.ZERO
           assert span.context().getSamplingPriority() == PrioritySampling.USER_KEEP
           assert span.context().getPropagationTags().createTagMap()["_dd.p.dm"] == (-SamplingMechanism.DATA_JOBS).toString()
           parent()