synchronize setting listener

mobuchowski · mobuchowski · commit e598da6a830e · 2025-03-28T11:06:22.000+01:00
Signed-off-by: Maciej Obuchowski &lt;maciej.obuchowski@datadoghq.com&gt;
diff --git a/dd-java-agent/instrumentation/spark/spark_2.12/src/main/java/datadog/trace/instrumentation/spark/Spark212Instrumentation.java b/dd-java-agent/instrumentation/spark/spark_2.12/src/main/java/datadog/trace/instrumentation/spark/Spark212Instrumentation.java
@@ -44,8 +44,8 @@ public void methodAdvice(MethodTransformer transformer) {
   public static class InjectListener {
     @Advice.OnMethodEnter(suppress = Throwable.class)
     public static void enter(@Advice.This SparkContext sparkContext) {
-      // checking whether OpenLineage integration is available and that it supports tags
-      // Disabling this mechanism for this PR. Will be enabled with provided with Config option.
+      // checking whether OpenLineage integration is enabled, available and that it supports tags
+
       if (Config.get().isDataJobsOpenLineageEnabled()
           && Utils.classIsLoadable("io.openlineage.spark.agent.OpenLineageSparkListener")
           && Utils.classIsLoadable(
diff --git a/dd-java-agent/instrumentation/spark/spark_2.13/src/main/java/datadog/trace/instrumentation/spark/Spark213Instrumentation.java b/dd-java-agent/instrumentation/spark/spark_2.13/src/main/java/datadog/trace/instrumentation/spark/Spark213Instrumentation.java
@@ -44,10 +44,7 @@ public void methodAdvice(MethodTransformer transformer) {
   public static class InjectListener {
     @Advice.OnMethodEnter(suppress = Throwable.class)
     public static void enter(@Advice.This SparkContext sparkContext) {
-      // Checking whether OpenLineage integration is available and that it supports tags
-      // Disabling this mechanism for this PR. Will be enabled with provided with Config
-      // option.
-
+      // checking whether OpenLineage integration is enabled, available and that it supports tags
       if (Config.get().isDataJobsOpenLineageEnabled()
           && Utils.classIsLoadable("io.openlineage.spark.agent.OpenLineageSparkListener")
           && Utils.classIsLoadable(
diff --git a/dd-java-agent/instrumentation/spark/src/main/java/datadog/trace/instrumentation/spark/AbstractDatadogSparkListener.java b/dd-java-agent/instrumentation/spark/src/main/java/datadog/trace/instrumentation/spark/AbstractDatadogSparkListener.java
@@ -36,6 +36,7 @@
 import java.util.Optional;
 import java.util.Properties;
 import java.util.UUID;
+import java.util.concurrent.atomic.AtomicReference;
 import java.util.function.Consumer;
 import org.apache.spark.ExceptionFailure;
 import org.apache.spark.SparkConf;
@@ -69,7 +70,8 @@ public abstract class AbstractDatadogSparkListener extends SparkListener {
   private static final Logger log = LoggerFactory.getLogger(AbstractDatadogSparkListener.class);
   private static final ObjectMapper objectMapper = new ObjectMapper();
   public static volatile AbstractDatadogSparkListener listener = null;
-  public static volatile SparkListenerInterface openLineageSparkListener = null;
+  public static volatile AtomicReference<SparkListenerInterface> openLineageSparkListener =
+      new AtomicReference<>(null);
   public static volatile SparkConf openLineageSparkConf = null;
 
   public static volatile boolean finishTraceOnApplicationEnd = true;
@@ -167,21 +169,45 @@ public AbstractDatadogSparkListener(SparkConf sparkConf, String appId, String sp
 
   public void setupOpenLineage(DDTraceId traceId) {
     log.debug("Setting up OpenLineage tags");
-    if (openLineageSparkListener != null) {
-      openLineageSparkConf.set("spark.openlineage.transport.type", "composite");
-      openLineageSparkConf.set("spark.openlineage.transport.continueOnFailure", "true");
-      openLineageSparkConf.set("spark.openlineage.transport.transports.agent.type", "http");
-      openLineageSparkConf.set(
+    log.info(
+        "Classloader for current variable: ({}) {}",
+        System.identityHashCode(
+            AbstractDatadogSparkListener.openLineageSparkListener
+                .get()
+                .getClass()
+                .getClassLoader()),
+        AbstractDatadogSparkListener.openLineageSparkListener.get().getClass().getClassLoader());
+    log.info(
+        "Current thread class loader: ({}) {}",
+        System.identityHashCode(Thread.currentThread().getContextClassLoader()),
+        Thread.currentThread().getContextClassLoader());
+    if (AbstractDatadogSparkListener.openLineageSparkListener.get() != null) {
+      AbstractDatadogSparkListener.openLineageSparkConf.set(
+          "spark.openlineage.transport.type", "composite");
+      AbstractDatadogSparkListener.openLineageSparkConf.set(
+          "spark.openlineage.transport.continueOnFailure", "true");
+      AbstractDatadogSparkListener.openLineageSparkConf.set(
+          "spark.openlineage.transport.transports.agent.type", "http");
+      AbstractDatadogSparkListener.openLineageSparkConf.set(
           "spark.openlineage.transport.transports.agent.url", getAgentHttpUrl());
-      openLineageSparkConf.set(
+      AbstractDatadogSparkListener.openLineageSparkConf.set(
           "spark.openlineage.transport.transports.agent.endpoint", AGENT_OL_ENDPOINT);
-      openLineageSparkConf.set("spark.openlineage.transport.transports.agent.compression", "gzip");
-      openLineageSparkConf.set(
+      AbstractDatadogSparkListener.openLineageSparkConf.set(
+          "spark.openlineage.transport.transports.agent.compression", "gzip");
+      AbstractDatadogSparkListener.openLineageSparkConf.set(
           "spark.openlineage.run.tags",
-          "_dd.trace_id:" + traceId.toString() + ";_dd.ol_intake.emit_spans:false");
+          "_dd.trace_id:"
+              + traceId.toString()
+              + ";_dd.ol_intake.emit_spans:false;dd.ol_service:"
+              + sparkServiceName);
       return;
     }
-    log.debug("No OpenLineageSparkListener!");
+    log.info(
+        "There is no OpenLineage Spark Listener in the context. Skipping setting tags. {}",
+        AbstractDatadogSparkListener.openLineageSparkListener.get());
+    log.info(
+        "There is no OpenLineage SparkConf in the context. Skipping setting tags. {}",
+        AbstractDatadogSparkListener.openLineageSparkConf);
   }
 
   /** Resource name of the spark job. Provide an implementation based on a specific scala version */
@@ -212,7 +238,7 @@ public synchronized void onApplicationStart(SparkListenerApplicationStart applic
               .map(context -> context.getTraceId())
               .orElse(predeterminedTraceIdContext.getTraceId()));
     }
-    notifyOl(x -> openLineageSparkListener.onApplicationStart(x), applicationStart);
+    notifyOl(x -> openLineageSparkListener.get().onApplicationStart(x), applicationStart);
   }
 
   private void initApplicationSpanIfNotInitialized() {
@@ -273,7 +299,7 @@ public void onApplicationEnd(SparkListenerApplicationEnd applicationEnd) {
     log.info(
         "Received spark application end event, finish trace on this event: {}",
         finishTraceOnApplicationEnd);
-    notifyOl(x -> openLineageSparkListener.onApplicationEnd(x), applicationEnd);
+    notifyOl(x -> openLineageSparkListener.get().onApplicationEnd(x), applicationEnd);
 
     if (finishTraceOnApplicationEnd) {
       finishApplication(applicationEnd.time(), null, 0, null);
@@ -468,7 +494,7 @@ public synchronized void onJobStart(SparkListenerJobStart jobStart) {
       stageToJob.put(stageId, jobStart.jobId());
     }
     jobSpans.put(jobStart.jobId(), jobSpan);
-    notifyOl(x -> openLineageSparkListener.onJobStart(x), jobStart);
+    notifyOl(x -> openLineageSparkListener.get().onJobStart(x), jobStart);
   }
 
   @Override
@@ -499,7 +525,7 @@ public synchronized void onJobEnd(SparkListenerJobEnd jobEnd) {
     if (metrics != null) {
       metrics.setSpanMetrics(jobSpan);
     }
-    notifyOl(x -> openLineageSparkListener.onJobEnd(x), jobEnd);
+    notifyOl(x -> openLineageSparkListener.get().onJobEnd(x), jobEnd);
 
     jobSpan.finish(jobEnd.time() * 1000);
   }
@@ -669,7 +695,7 @@ public void onTaskEnd(SparkListenerTaskEnd taskEnd) {
     Properties props = stageProperties.get(stageSpanKey);
     sendTaskSpan(stageSpan, taskEnd, props);
 
-    notifyOl(x -> openLineageSparkListener.onTaskEnd(x), taskEnd);
+    notifyOl(x -> openLineageSparkListener.get().onTaskEnd(x), taskEnd);
   }
 
   private void sendTaskSpan(
@@ -753,20 +779,20 @@ public void onOtherEvent(SparkListenerEvent event) {
 
   private <T extends SparkListenerEvent> void notifyOl(Consumer<T> ol, T event) {
     if (!Config.get().isDataJobsOpenLineageEnabled()) {
-      log.trace("Ignoring event {} - OpenLineage not enabled", event);
+      log.debug("Ignoring event {} - OpenLineage not enabled", event);
       return;
     }
 
     if (isRunningOnDatabricks || isStreamingJob) {
       log.debug("Not emitting event when running on databricks or on streaming jobs");
       return;
     }
-    if (openLineageSparkListener != null) {
+    if (AbstractDatadogSparkListener.openLineageSparkListener.get() != null) {
       log.debug(
           "Passing event `{}` to OpenLineageSparkListener", event.getClass().getCanonicalName());
       ol.accept(event);
     } else {
-      log.trace("OpenLineageSparkListener is null");
+      log.debug("OpenLineageSparkListener is null");
     }
   }
 
@@ -818,7 +844,7 @@ private synchronized void updateAdaptiveSQLPlan(SparkListenerEvent event) {
   private synchronized void onSQLExecutionStart(SparkListenerSQLExecutionStart sqlStart) {
     sqlPlans.put(sqlStart.executionId(), sqlStart.sparkPlanInfo());
     sqlQueries.put(sqlStart.executionId(), sqlStart);
-    notifyOl(x -> openLineageSparkListener.onOtherEvent(x), sqlStart);
+    notifyOl(x -> openLineageSparkListener.get().onOtherEvent(x), sqlStart);
   }
 
   private synchronized void onSQLExecutionEnd(SparkListenerSQLExecutionEnd sqlEnd) {
@@ -831,7 +857,7 @@ private synchronized void onSQLExecutionEnd(SparkListenerSQLExecutionEnd sqlEnd)
       if (metrics != null) {
         metrics.setSpanMetrics(span);
       }
-      notifyOl(x -> openLineageSparkListener.onOtherEvent(x), sqlEnd);
+      notifyOl(x -> openLineageSparkListener.get().onOtherEvent(x), sqlEnd);
 
       span.finish(sqlEnd.time() * 1000);
     }
diff --git a/dd-java-agent/instrumentation/spark/src/main/java/datadog/trace/instrumentation/spark/AbstractSparkInstrumentation.java b/dd-java-agent/instrumentation/spark/src/main/java/datadog/trace/instrumentation/spark/AbstractSparkInstrumentation.java
@@ -119,17 +119,12 @@ public static class LiveListenerBusAdvice {
     @Advice.OnMethodEnter(suppress = Throwable.class, skipOn = Advice.OnNonDefaultValue.class)
     // If OL is disabled in tracer config but user set it up manually don't interfere
     public static boolean enter(@Advice.Argument(0) SparkListenerInterface listener) {
-      if (!Config.get().isDataJobsOpenLineageEnabled()
-          || 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 adding to ListenerBus");
+      if (Config.get().isDataJobsOpenLineageEnabled()
+          && listener != null
+          && "io.openlineage.spark.agent.OpenLineageSparkListener"
+              .equals(listener.getClass().getCanonicalName())) {
+        LoggerFactory.getLogger("LiveListenerBusAdvice")
+            .debug("Detected OpenLineage listener, skipping adding to ListenerBus");
         return true;
       }
       return false;
diff --git a/dd-java-agent/instrumentation/spark/src/main/java/datadog/trace/instrumentation/spark/OpenLineageInstrumentation.java b/dd-java-agent/instrumentation/spark/src/main/java/datadog/trace/instrumentation/spark/OpenLineageInstrumentation.java
@@ -9,10 +9,10 @@
 import datadog.trace.agent.tooling.Instrumenter;
 import datadog.trace.agent.tooling.InstrumenterModule;
 import datadog.trace.api.Config;
-import java.lang.reflect.Field;
 import net.bytebuddy.asm.Advice;
 import org.apache.spark.SparkConf;
 import org.apache.spark.scheduler.SparkListenerInterface;
+import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
 @AutoService(InstrumenterModule.class)
@@ -46,9 +46,7 @@ public boolean defaultEnabled() {
 
   @Override
   public String[] knownMatchingTypes() {
-    return new String[] {
-      "io.openlineage.spark.agent.OpenLineageSparkListener", "org.apache.spark.util.Utils"
-    };
+    return new String[] {"io.openlineage.spark.agent.OpenLineageSparkListener"};
   }
 
   @Override
@@ -63,18 +61,44 @@ public void methodAdvice(MethodTransformer transformer) {
 
   public static class OpenLineageSparkListenerAdvice {
     @Advice.OnMethodExit(suppress = Throwable.class)
-    public static void exit(@Advice.This Object self) throws IllegalAccessException {
-      LoggerFactory.getLogger(Config.class).debug("Checking for OpenLineageSparkListener");
+    public static void exit(@Advice.This Object self, @Advice.FieldValue("conf") SparkConf conf) {
       try {
-        Field conf = self.getClass().getDeclaredField("conf");
-        conf.setAccessible(true);
-        AbstractDatadogSparkListener.openLineageSparkConf = (SparkConf) conf.get(self);
-        AbstractDatadogSparkListener.openLineageSparkListener = (SparkListenerInterface) self;
-        LoggerFactory.getLogger(Config.class)
-            .debug("Detected OpenLineageSparkListener, passed to DatadogSparkListener");
-      } catch (NoSuchFieldException | IllegalAccessException e) {
+        Logger log = LoggerFactory.getLogger(Config.class);
+        AbstractDatadogSparkListener.openLineageSparkConf = conf;
+
+        log.debug(
+            "ADSL classloader ({}) {}",
+            System.identityHashCode(AbstractDatadogSparkListener.class.getClassLoader()),
+            AbstractDatadogSparkListener.class.getClassLoader());
+        log.debug(
+            "Current classloader ({}) {}",
+            System.identityHashCode(Thread.currentThread().getContextClassLoader()),
+            Thread.currentThread().getContextClassLoader().toString());
+        log.debug("Got OpenLineageSparkListener: {}", self);
+        AbstractDatadogSparkListener.openLineageSparkListener.set((SparkListenerInterface) self);
+        log.debug(
+            "Set OpenLineageSparkListener: {}",
+            AbstractDatadogSparkListener.openLineageSparkListener.get());
+        log.debug("Check for self: {}", self);
+        if (AbstractDatadogSparkListener.openLineageSparkListener.get() != null) {
+          log.debug(
+              "Detected OpenLineageSparkListener, passed to DatadogSparkListener {} on ClassLoader ({}) {}",
+              AbstractDatadogSparkListener.openLineageSparkListener.get().getClass().getName(),
+              System.identityHashCode(
+                  AbstractDatadogSparkListener.openLineageSparkListener
+                      .get()
+                      .getClass()
+                      .getClassLoader()),
+              AbstractDatadogSparkListener.openLineageSparkListener
+                  .get()
+                  .getClass()
+                  .getClassLoader());
+        } else {
+          log.debug("WTF it's null");
+        }
+      } catch (Exception e) {
         LoggerFactory.getLogger(Config.class)
-            .debug("Failed to pass OpenLineageSparkListener to DatadogSparkListener", e);
+            .error("Failed to set OpenLineageSparkListener: {}", e.getMessage());
       }
     }
   }
