classloader and code review fixes

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

@@ -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;
@@ -214,6 +215,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;
 
@@ -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);
     }
   }
 }