take care of sparkconf/ol lifecycle

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

Author: mobuchowski

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

@@ -172,11 +172,15 @@ static void setupSparkConf(SparkConf sparkConf) {
         "_dd.trace_id:"
             + listener.applicationSpan.context().getTraceId().toString()
             + ";_dd.ol_intake.emit_spans:false");
+    for (Tuple2<String, String> tuple : sparkConf.getAll()) {
+      log.error("Set Spark conf: Key: " + tuple._1 + ", Value: " + tuple._2);
+    }
   }
 
   public void setupOpenLineage() {
-    log.debug("Setting up OpenLineage-Datadog integration");
+    log.error("Setting up OpenLineage-Datadog integration");
     if (openLineageSparkListener != null) {
+      log.error("No init needed");
       setupSparkConf(openLineageSparkConf);
       return;
     }
@@ -240,6 +244,8 @@ private void initApplicationSpanIfNotInitialized() {
       return;
     }
 
+    log.error("Starting tracer application span.");
+
     AgentTracer.SpanBuilder builder = buildSparkSpan("spark.application", null);
 
     if (applicationStart != null) {
@@ -253,8 +259,6 @@ private void initApplicationSpanIfNotInitialized() {
       }
     }
 
-    notifyOl(x -> this.openLineageSparkListener.onApplicationStart(x), applicationStart);
-
     captureApplicationParameters(builder);
     captureOpenlineageContextIfPresent(builder);
 
@@ -263,6 +267,7 @@ private void initApplicationSpanIfNotInitialized() {
     applicationSpan.setMeasured(true);
     // We need to set it up after we create application span to have correlation.
     setupOpenLineage();
+    notifyOl(x -> openLineageSparkListener.onApplicationStart(x), applicationStart);
   }
 
   private void captureOpenlineageContextIfPresent(AgentTracer.SpanBuilder builder) {
@@ -774,14 +779,15 @@ public void onOtherEvent(SparkListenerEvent event) {
 
   private <T extends SparkListenerEvent> void notifyOl(Consumer<T> ol, T event) {
     if (isRunningOnDatabricks || isStreamingJob) {
-      log.debug("Not emitting event when running on databricks or on streaming jobs");
+      log.error("Not emitting event when running on databricks or on streaming jobs");
       return;
     }
+    initApplicationSpanIfNotInitialized();
     if (openLineageSparkListener != null) {
-      log.debug("Notifying with event `{}`", event.getClass().getCanonicalName());
+      log.error("Notifying with event `{}`", event.getClass().getCanonicalName());
       ol.accept(event);
     } else {
-      log.debug("OpenLineageSparkListener is null");
+      log.error("OpenLineageSparkListener is null");
     }
   }
 

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

@@ -8,8 +8,11 @@
 
 import datadog.trace.agent.tooling.Instrumenter;
 import datadog.trace.agent.tooling.InstrumenterModule;
+import datadog.trace.api.Config;
 import net.bytebuddy.asm.Advice;
 import org.apache.spark.deploy.SparkSubmitArguments;
+import org.apache.spark.scheduler.SparkListenerInterface;
+import org.slf4j.LoggerFactory;
 
 public abstract class AbstractSparkInstrumentation extends InstrumenterModule.Tracing
     implements Instrumenter.ForKnownTypes, Instrumenter.HasMethodAdvice {
@@ -29,7 +32,8 @@ public String[] knownMatchingTypes() {
       "org.apache.spark.SparkContext",
       "org.apache.spark.deploy.SparkSubmit",
       "org.apache.spark.deploy.yarn.ApplicationMaster",
-      "org.apache.spark.util.SparkClassUtils"
+      "org.apache.spark.util.SparkClassUtils",
+      "org.apache.spark.scheduler.LiveListenerBus"
     };
   }
 
@@ -56,6 +60,15 @@ public void methodAdvice(MethodTransformer transformer) {
             .and(named("finish"))
             .and(isDeclaredBy(named("org.apache.spark.deploy.yarn.ApplicationMaster"))),
         AbstractSparkInstrumentation.class.getName() + "$YarnFinishAdvice");
+
+    // LiveListenerBus class is used when running in a YARN cluster
+    transformer.applyAdvice(
+        isMethod()
+            .and(named("addToSharedQueue"))
+            //            .and(takesArgument(0,
+            // named("org.apache.spark.scheduler.SparkListenerInterface")))
+            .and(isDeclaredBy(named("org.apache.spark.scheduler.LiveListenerBus"))),
+        AbstractSparkInstrumentation.class.getName() + "$LiveListenerBusAdvice");
   }
 
   public static class PrepareSubmitEnvAdvice {
@@ -101,4 +114,22 @@ public static void enter(@Advice.Argument(1) int exitCode, @Advice.Argument(2) S
       }
     }
   }
+
+  public static class LiveListenerBusAdvice {
+    @Advice.OnMethodEnter(suppress = Throwable.class, skipOn = Advice.OnNonDefaultValue.class)
+    public static boolean enter(@Advice.Argument(0) SparkListenerInterface listener) {
+      if (listener == null || listener.getClass().getCanonicalName() == null) {
+        return false;
+      }
+      if (listener
+          .getClass()
+          .getCanonicalName()
+          .equals("io.openlineage.spark.agent.OpenLineageSparkListener")) {
+        LoggerFactory.getLogger(Config.class)
+            .debug("Detected OL listener, skipping initialization");
+        return true;
+      }
+      return false;
+    }
+  }
 }