set OL if detected

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

add OL instrumentation class

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

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

@@ -36,6 +36,7 @@
 import java.util.Optional;
 import java.util.Properties;
 import java.util.UUID;
+import java.util.function.Consumer;
 import org.apache.spark.ExceptionFailure;
 import org.apache.spark.SparkConf;
 import org.apache.spark.TaskFailedReason;
@@ -51,6 +52,7 @@
 import org.apache.spark.sql.streaming.StateOperatorProgress;
 import org.apache.spark.sql.streaming.StreamingQueryListener;
 import org.apache.spark.sql.streaming.StreamingQueryProgress;
+import org.apache.spark.util.Utils;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 import scala.Tuple2;
@@ -68,12 +70,16 @@ 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;
 
   private final int MAX_COLLECTION_SIZE = 5000;
   private final int MAX_ACCUMULATOR_SIZE = 50000;
   private final String RUNTIME_TAGS_PREFIX = "spark.datadog.tags.";
+  private static final String AGENT_OL_ENDPOINT = "openlineage/api/v1/lineage";
 
   private final SparkConf sparkConf;
   private final String sparkVersion;
@@ -109,6 +115,7 @@ public abstract class AbstractDatadogSparkListener extends SparkListener {
   private final Map<Long, SparkSQLUtils.AccumulatorWithStage> accumulators =
       new RemoveEldestHashMap<>(MAX_ACCUMULATOR_SIZE);
 
+  private volatile boolean isStreamingJob = false;
   private final boolean isRunningOnDatabricks;
   private final String databricksClusterName;
   private final String databricksServiceName;
@@ -151,8 +158,62 @@ public AbstractDatadogSparkListener(SparkConf sparkConf, String appId, String sp
                     finishApplication(System.currentTimeMillis(), null, 0, null);
                   }
                 }));
+  }
+
+  static void setupSparkConf(SparkConf sparkConf) {
+    sparkConf.set("spark.openlineage.transport.type", "composite");
+    sparkConf.set("spark.openlineage.transport.continueOnFailure", "true");
+    sparkConf.set("spark.openlineage.transport.transports.agent.type", "http");
+    sparkConf.set("spark.openlineage.transport.transports.agent.url", getAgentHttpUrl());
+    sparkConf.set("spark.openlineage.transport.transports.agent.endpoint", AGENT_OL_ENDPOINT);
+    sparkConf.set("spark.openlineage.transport.transports.agent.compression", "gzip");
+    sparkConf.set(
+        "spark.openlineage.run.tags",
+        "_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.error("Setting up OpenLineage-Datadog integration");
+    if (openLineageSparkListener != null) {
+      log.error("No init needed");
+      setupSparkConf(openLineageSparkConf);
+      return;
+    }
+
+    String className = "io.openlineage.spark.agent.OpenLineageSparkListener";
+    Class clazz;
+    try {
+      try {
+        clazz = Class.forName(className, true, Thread.currentThread().getContextClassLoader());
+      } catch (ClassNotFoundException e) {
+        clazz = Class.forName(className, true, Utils.class.getClassLoader());
+      }
+    } catch (ClassNotFoundException e) {
+      log.info("OpenLineage integration is not present on the classpath");
+      return;
+    }
+
+    openLineageSparkConf = sparkConf;
+    if (clazz == null) {
+      log.info("OpenLineage integration is not present on the classpath: class is null");
+      return;
+    }
+    try {
+      setupSparkConf(openLineageSparkConf);
+      openLineageSparkListener =
+          (SparkListenerInterface)
+              clazz.getConstructor(SparkConf.class).newInstance(openLineageSparkConf);
 
-    log.info("Created datadog spark listener: {}", this.getClass().getSimpleName());
+      log.info(
+          "Created OL spark listener: {}", openLineageSparkListener.getClass().getSimpleName());
+    } catch (Exception e) {
+      log.warn("Failed to instantiate OL Spark Listener: {}", e.toString());
+    }
   }
 
   /** Resource name of the spark job. Provide an implementation based on a specific scala version */
@@ -183,6 +244,8 @@ private void initApplicationSpanIfNotInitialized() {
       return;
     }
 
+    log.error("Starting tracer application span.");
+
     AgentTracer.SpanBuilder builder = buildSparkSpan("spark.application", null);
 
     if (applicationStart != null) {
@@ -202,6 +265,9 @@ private void initApplicationSpanIfNotInitialized() {
     applicationSpan = builder.start();
     setDataJobsSamplingPriority(applicationSpan);
     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) {
@@ -233,6 +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);
 
     if (finishTraceOnApplicationEnd) {
       finishApplication(applicationEnd.time(), null, 0, null);
@@ -405,6 +472,7 @@ public synchronized void onJobStart(SparkListenerJobStart jobStart) {
     if (sqlSpan != null) {
       jobSpanBuilder.asChildOf(sqlSpan.context());
     } else if (batchKey != null) {
+      isStreamingJob = true;
       AgentSpan batchSpan =
           getOrCreateStreamingBatchSpan(batchKey, jobStart.time(), jobStart.properties());
       jobSpanBuilder.asChildOf(batchSpan.context());
@@ -426,6 +494,7 @@ public synchronized void onJobStart(SparkListenerJobStart jobStart) {
       stageToJob.put(stageId, jobStart.jobId());
     }
     jobSpans.put(jobStart.jobId(), jobSpan);
+    notifyOl(x -> openLineageSparkListener.onJobStart(x), jobStart);
   }
 
   @Override
@@ -456,6 +525,7 @@ public synchronized void onJobEnd(SparkListenerJobEnd jobEnd) {
     if (metrics != null) {
       metrics.setSpanMetrics(jobSpan);
     }
+    notifyOl(x -> openLineageSparkListener.onJobEnd(x), jobEnd);
 
     jobSpan.finish(jobEnd.time() * 1000);
   }
@@ -624,6 +694,8 @@ public void onTaskEnd(SparkListenerTaskEnd taskEnd) {
 
     Properties props = stageProperties.get(stageSpanKey);
     sendTaskSpan(stageSpan, taskEnd, props);
+
+    notifyOl(x -> openLineageSparkListener.onTaskEnd(x), taskEnd);
   }
 
   private void sendTaskSpan(
@@ -705,6 +777,20 @@ public void onOtherEvent(SparkListenerEvent event) {
     updateAdaptiveSQLPlan(event);
   }
 
+  private <T extends SparkListenerEvent> void notifyOl(Consumer<T> ol, T event) {
+    if (isRunningOnDatabricks || isStreamingJob) {
+      log.error("Not emitting event when running on databricks or on streaming jobs");
+      return;
+    }
+    initApplicationSpanIfNotInitialized();
+    if (openLineageSparkListener != null) {
+      log.error("Notifying with event `{}`", event.getClass().getCanonicalName());
+      ol.accept(event);
+    } else {
+      log.error("OpenLineageSparkListener is null");
+    }
+  }
+
   private static final Class<?> adaptiveExecutionUpdateClass;
   private static final MethodHandle adaptiveExecutionIdMethod;
   private static final MethodHandle adaptiveSparkPlanMethod;
@@ -753,6 +839,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);
   }
 
   private synchronized void onSQLExecutionEnd(SparkListenerSQLExecutionEnd sqlEnd) {
@@ -765,6 +852,7 @@ private synchronized void onSQLExecutionEnd(SparkListenerSQLExecutionEnd sqlEnd)
       if (metrics != null) {
         metrics.setSpanMetrics(span);
       }
+      notifyOl(x -> openLineageSparkListener.onOtherEvent(x), sqlEnd);
 
       span.finish(sqlEnd.time() * 1000);
     }
@@ -1260,6 +1348,15 @@ private static String getDatabricksRunName(SparkConf conf) {
     return null;
   }
 
+  private static String getAgentHttpUrl() {
+    StringBuilder sb =
+        new StringBuilder("http://")
+            .append(Config.get().getAgentHost())
+            .append(":")
+            .append(Config.get().getAgentPort());
+    return sb.toString();
+  }
+
   @SuppressForbidden // called at most once per spark application
   private static String removeUuidFromEndOfString(String input) {
     return input.replaceAll(

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 {
@@ -28,7 +31,9 @@ public String[] knownMatchingTypes() {
     return new String[] {
       "org.apache.spark.SparkContext",
       "org.apache.spark.deploy.SparkSubmit",
-      "org.apache.spark.deploy.yarn.ApplicationMaster"
+      "org.apache.spark.deploy.yarn.ApplicationMaster",
+      "org.apache.spark.util.SparkClassUtils",
+      "org.apache.spark.scheduler.LiveListenerBus"
     };
   }
 
@@ -55,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 {
@@ -100,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;
+    }
+  }
 }

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

@@ -0,0 +1,80 @@
+package datadog.trace.instrumentation.spark;
+
+import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.named;
+import static net.bytebuddy.matcher.ElementMatchers.isConstructor;
+import static net.bytebuddy.matcher.ElementMatchers.isDeclaredBy;
+import static net.bytebuddy.matcher.ElementMatchers.takesArgument;
+
+import com.google.auto.service.AutoService;
+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.LoggerFactory;
+
+@AutoService(InstrumenterModule.class)
+public class OpenLineageInstrumentation extends InstrumenterModule.Tracing
+    implements Instrumenter.ForKnownTypes, Instrumenter.HasMethodAdvice {
+
+  public OpenLineageInstrumentation() {
+    super("openlineage-spark");
+  }
+
+  @Override
+  public String[] helperClassNames() {
+    return new String[] {
+      packageName + ".AbstractDatadogSparkListener",
+      packageName + ".DatabricksParentContext",
+      packageName + ".OpenlineageParentContext",
+      packageName + ".RemoveEldestHashMap",
+      packageName + ".SparkAggregatedTaskMetrics",
+      packageName + ".SparkConfAllowList",
+      packageName + ".SparkSQLUtils",
+      packageName + ".SparkSQLUtils$SparkPlanInfoForStage",
+      packageName + ".SparkSQLUtils$AccumulatorWithStage",
+    };
+  }
+
+  @Override
+  public boolean defaultEnabled() {
+    return true;
+  }
+
+  @Override
+  public String[] knownMatchingTypes() {
+    return new String[] {
+      "io.openlineage.spark.agent.OpenLineageSparkListener", "org.apache.spark.util.Utils"
+    };
+  }
+
+  @Override
+  public void methodAdvice(MethodTransformer transformer) {
+    // LiveListenerBus class is used when running in a YARN cluster
+    transformer.applyAdvice(
+        isConstructor()
+            .and(isDeclaredBy(named("io.openlineage.spark.agent.OpenLineageSparkListener")))
+            .and(takesArgument(0, named("org.apache.spark.SparkConf"))),
+        OpenLineageInstrumentation.class.getName() + "$OpenLineageSparkListenerAdvice");
+  }
+
+  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");
+      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) {
+        LoggerFactory.getLogger(Config.class)
+            .debug("Failed to pass OpenLineageSparkListener to DatadogSparkListener", e);
+      }
+    }
+  }
+}