classloader and code review fixes

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

Author: mobuchowski

dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/Agent.java

@@ -252,7 +252,7 @@ public static void start(
 
       if (Config.get().isDataJobsOpenLineageEnabled()) {
         setSystemPropertyDefault(
-            propertyNameToSystemPropertyName("integration.openlineage-spark.enabled"), "true");
+            propertyNameToSystemPropertyName("integration.spark-openlineage.enabled"), "true");
       }
 
       String javaCommand = System.getProperty("sun.java.command");

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

@@ -46,11 +46,6 @@ public static class InjectListener {
     @Advice.OnMethodEnter(suppress = Throwable.class)
     public static void enter(@Advice.This SparkContext sparkContext) {
       Logger log = LoggerFactory.getLogger("Spark212InjectListener");
-      log.debug(
-          "AbstractDatadogSparkListener classloader is: ({}) {}",
-          System.identityHashCode(AbstractDatadogSparkListener.class.getClassLoader()),
-          AbstractDatadogSparkListener.class.getClassLoader());
-
       if (Config.get().isDataJobsOpenLineageEnabled()
           && AbstractDatadogSparkListener.classIsLoadable(
               "io.openlineage.spark.agent.OpenLineageSparkListener")

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

@@ -47,11 +47,6 @@ public static class InjectListener {
     public static void enter(@Advice.This SparkContext sparkContext) {
       // checking whether OpenLineage integration is enabled, available and that it supports tags
       Logger log = LoggerFactory.getLogger("Spark212InjectListener");
-      log.debug(
-          "AbstractDatadogSparkListener classloader is: ({}) {}",
-          System.identityHashCode(AbstractDatadogSparkListener.class.getClassLoader()),
-          AbstractDatadogSparkListener.class.getClassLoader());
-
       if (Config.get().isDataJobsOpenLineageEnabled()
           && AbstractDatadogSparkListener.classIsLoadable(
               "io.openlineage.spark.agent.OpenLineageSparkListener")

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

@@ -12,6 +12,7 @@
 import datadog.trace.api.DDTraceId;
 import datadog.trace.api.sampling.PrioritySampling;
 import datadog.trace.api.sampling.SamplingMechanism;
+import datadog.trace.bootstrap.InstanceStore;
 import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
 import datadog.trace.bootstrap.instrumentation.api.AgentSpanContext;
 import datadog.trace.bootstrap.instrumentation.api.AgentTracer;
@@ -69,8 +70,6 @@ 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 SparkConf openLineageSparkConf = null;
 
   public static volatile boolean finishTraceOnApplicationEnd = true;
   public static volatile boolean isPysparkShell = false;
@@ -80,6 +79,9 @@ public abstract class AbstractDatadogSparkListener extends SparkListener {
   private final String RUNTIME_TAGS_PREFIX = "spark.datadog.tags.";
   private static final String AGENT_OL_ENDPOINT = "openlineage/api/v1/lineage";
 
+  public volatile SparkListenerInterface openLineageSparkListener = null;
+  public volatile SparkConf openLineageSparkConf = null;
+
   private final SparkConf sparkConf;
   private final String sparkVersion;
   private final String appId;
@@ -180,8 +182,10 @@ public void setupOpenLineage(DDTraceId traceId) {
           "spark.openlineage.run.tags",
           "_dd.trace_id:"
               + traceId.toString()
-              + ";_dd.ol_intake.emit_spans:false;dd.ol_service:"
-              + sparkServiceName);
+              + ";_dd.ol_intake.emit_spans:false;_dd.ol_service:"
+              + sparkServiceName
+              + ";_dd.ol_app_id:"
+              + appId);
       return;
     }
     log.debug(
@@ -214,6 +218,12 @@ public void setupOpenLineage(DDTraceId traceId) {
   public synchronized void onApplicationStart(SparkListenerApplicationStart applicationStart) {
     this.applicationStart = applicationStart;
 
+    if (openLineageSparkListener == null) {
+      openLineageSparkListener =
+          InstanceStore.of(SparkListenerInterface.class).get("openLineageListener");
+      openLineageSparkConf = InstanceStore.of(SparkConf.class).get("openLineageSparkConf");
+    }
+
     if (openLineageSparkListener != null) {
       setupOpenLineage(
           OpenlineageParentContext.from(sparkConf)

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

@@ -9,6 +9,7 @@
 import datadog.trace.agent.tooling.Instrumenter;
 import datadog.trace.agent.tooling.InstrumenterModule;
 import datadog.trace.api.Config;
+import datadog.trace.bootstrap.InstanceStore;
 import net.bytebuddy.asm.Advice;
 import org.apache.spark.deploy.SparkSubmitArguments;
 import org.apache.spark.scheduler.SparkListenerInterface;
@@ -121,15 +122,12 @@ public static class LiveListenerBusAdvice {
     // 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) {
       Logger log = LoggerFactory.getLogger("LiveListenerBusAdvice");
-      log.debug(
-          "AbstractDatadogSparkListener classloader for LiveListenerBusAdvice is: ({}) {}",
-          System.identityHashCode(AbstractDatadogSparkListener.class.getClassLoader()),
-          AbstractDatadogSparkListener.class.getClassLoader());
       if (Config.get().isDataJobsOpenLineageEnabled()
           && listener != null
           && "io.openlineage.spark.agent.OpenLineageSparkListener"
               .equals(listener.getClass().getCanonicalName())) {
         log.debug("Detected OpenLineage listener, skipping adding it to ListenerBus");
+        InstanceStore.of(SparkListenerInterface.class).put("openLineageListener", listener);
         return true;
       }
       return false;

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

@@ -9,10 +9,9 @@
 import datadog.trace.agent.tooling.Instrumenter;
 import datadog.trace.agent.tooling.InstrumenterModule;
 import datadog.trace.api.Config;
-import java.lang.reflect.Field;
+import datadog.trace.bootstrap.InstanceStore;
 import net.bytebuddy.asm.Advice;
 import org.apache.spark.SparkConf;
-import org.apache.spark.scheduler.SparkListenerInterface;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
@@ -21,7 +20,7 @@ public class OpenLineageInstrumentation extends InstrumenterModule.Tracing
     implements Instrumenter.ForKnownTypes, Instrumenter.HasMethodAdvice {
 
   public OpenLineageInstrumentation() {
-    super("openlineage-spark");
+    super("spark-openlineage");
   }
 
   @Override
@@ -70,35 +69,7 @@ public static void exit(@Advice.This Object self, @Advice.FieldValue("conf") Spa
             "OpenLineage - Data Jobs integration disabled. Not manipulating OpenLineageSparkListener");
         return;
       }
-
-      log.debug("Checking for OpenLineageSparkListener");
-      ClassLoader cl = Thread.currentThread().getContextClassLoader();
-      if (cl.getClass().getName().contains("MutableURLClassLoader")
-          || cl.getClass().getName().contains("ChildFirstURLClassLoader")) {
-        log.debug(
-            "Detected MutableURLClassLoader. Setting OpenLineage on AbstractDatadogSparkListener.class of parent classloader");
-        try {
-          log.debug(
-              "Parent classloader: ({}) {}",
-              System.identityHashCode(cl.getParent()),
-              cl.getParent());
-          Class clazz = cl.getParent().loadClass(AbstractDatadogSparkListener.class.getName());
-          Field openLineageSparkListener = clazz.getDeclaredField("openLineageSparkListener");
-          openLineageSparkListener.setAccessible(true);
-          openLineageSparkListener.set(null, self);
-
-          Field openLineageSparkConf = clazz.getDeclaredField("openLineageSparkConf");
-          openLineageSparkConf.setAccessible(true);
-          openLineageSparkConf.set(null, conf);
-        } catch (Throwable e) {
-          log.info("Failed to setup OpenLineage", e);
-        }
-      } else {
-        log.debug(
-            "Detected other classloader than MutableURLClassLoader. Setting OpenLineage on AbstractDatadogSparkListener.class");
-        AbstractDatadogSparkListener.openLineageSparkListener = (SparkListenerInterface) self;
-        AbstractDatadogSparkListener.openLineageSparkConf = conf;
-      }
+      InstanceStore.of(SparkConf.class).put("openLineageSparkConf", conf);
     }
   }
 }

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

@@ -10,7 +10,7 @@ abstract class AbstractSpark32SqlTest extends AgentTestRunner {
   void configurePreAgent() {
     super.configurePreAgent()
     injectSysConfig("dd.integration.spark.enabled", "true")
-    injectSysConfig("dd.integration.openlineage-spark.enabled", "true")
+    injectSysConfig("dd.integration.spark-openlineage.enabled", "true")
   }
 
   def "compute a GROUP BY sql query plan"() {

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

@@ -25,7 +25,7 @@ class AbstractSparkStructuredStreamingTest extends AgentTestRunner {
   void configurePreAgent() {
     super.configurePreAgent()
     injectSysConfig("dd.integration.spark.enabled", "true")
-    injectSysConfig("dd.integration.openlineage-spark.enabled", "true")
+    injectSysConfig("dd.integration.spark-openlineage.enabled", "true")
   }
 
   private SparkSession createSparkSession(String appName) {

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

@@ -34,7 +34,7 @@ abstract class AbstractSparkTest extends AgentTestRunner {
   void configurePreAgent() {
     super.configurePreAgent()
     injectSysConfig("dd.integration.spark.enabled", "true")
-    injectSysConfig("dd.integration.openlineage-spark.enabled", "true")
+    injectSysConfig("dd.integration.spark-openlineage.enabled", "true")
   }
 
   def "generate application span with child job and stages"() {