spark instrumentation use openlineage trace id if can generate it out of root parent information

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

Author: mobuchowski

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

@@ -168,7 +168,7 @@ public AbstractDatadogSparkListener(SparkConf sparkConf, String appId, String sp
   }
 
   public void setupOpenLineage(DDTraceId traceId) {
-    log.debug("Setting up OpenLineage configuration");
+    log.error("Setting up OpenLineage configuration with trace id {}", traceId);
     if (openLineageSparkListener != null) {
       openLineageSparkConf.set("spark.openlineage.transport.type", "composite");
       openLineageSparkConf.set("spark.openlineage.transport.continueOnFailure", "true");
@@ -284,6 +284,9 @@ private void captureOpenlineageContextIfPresent(
     builder.withTag("openlineage_parent_job_namespace", context.getParentJobNamespace());
     builder.withTag("openlineage_parent_job_name", context.getParentJobName());
     builder.withTag("openlineage_parent_run_id", context.getParentRunId());
+    builder.withTag("openlineage_root_parent_job_namespace", context.getRootParentJobNamespace());
+    builder.withTag("openlineage_root_parent_job_name", context.getRootParentJobName());
+    builder.withTag("openlineage_root_parent_run_id", context.getRootParentRunId());
   }
 
   @Override

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

@@ -7,10 +7,7 @@
 import datadog.trace.bootstrap.instrumentation.api.AgentSpanContext;
 import datadog.trace.bootstrap.instrumentation.api.AgentTraceCollector;
 import datadog.trace.bootstrap.instrumentation.api.AgentTracer;
-import java.nio.ByteBuffer;
-import java.nio.charset.StandardCharsets;
-import java.security.MessageDigest;
-import java.security.NoSuchAlgorithmException;
+import datadog.trace.util.FNV64Hash;
 import java.util.Collections;
 import java.util.Map;
 import java.util.Optional;
@@ -32,11 +29,19 @@ public class OpenlineageParentContext implements AgentSpanContext {
   private final String parentJobNamespace;
   private final String parentJobName;
   private final String parentRunId;
+  private final String rootParentJobNamespace;
+  private final String rootParentJobName;
+  private final String rootParentRunId;
 
   public static final String OPENLINEAGE_PARENT_JOB_NAMESPACE =
       "spark.openlineage.parentJobNamespace";
   public static final String OPENLINEAGE_PARENT_JOB_NAME = "spark.openlineage.parentJobName";
   public static final String OPENLINEAGE_PARENT_RUN_ID = "spark.openlineage.parentRunId";
+  public static final String OPENLINEAGE_ROOT_PARENT_JOB_NAMESPACE =
+      "spark.openlineage.rootParentJobNamespace";
+  public static final String OPENLINEAGE_ROOT_PARENT_JOB_NAME =
+      "spark.openlineage.rootParentJobName";
+  public static final String OPENLINEAGE_ROOT_PARENT_RUN_ID = "spark.openlineage.rootParentRunId";
 
   public static Optional<OpenlineageParentContext> from(SparkConf sparkConf) {
     if (!sparkConf.contains(OPENLINEAGE_PARENT_JOB_NAMESPACE)
@@ -53,60 +58,73 @@ public static Optional<OpenlineageParentContext> from(SparkConf sparkConf) {
       return Optional.empty();
     }
 
+    if (!sparkConf.contains(OPENLINEAGE_ROOT_PARENT_RUN_ID)) {
+      log.error("Have parent info, but not root parent info. Can't construct valid trace id.");
+      return Optional.empty();
+    }
+
+    String rootParentJobNamespace = sparkConf.get(OPENLINEAGE_ROOT_PARENT_JOB_NAMESPACE, "");
+    String rootParentJobName = sparkConf.get(OPENLINEAGE_ROOT_PARENT_JOB_NAME, "");
+    String rootParentRunId = sparkConf.get(OPENLINEAGE_ROOT_PARENT_RUN_ID, "");
+
     return Optional.of(
-        new OpenlineageParentContext(parentJobNamespace, parentJobName, parentRunId));
+        new OpenlineageParentContext(
+            parentJobNamespace,
+            parentJobName,
+            parentRunId,
+            rootParentJobNamespace,
+            rootParentJobName,
+            rootParentRunId));
   }
 
-  OpenlineageParentContext(String parentJobNamespace, String parentJobName, String parentRunId) {
+  OpenlineageParentContext(
+      String parentJobNamespace,
+      String parentJobName,
+      String parentRunId,
+      String rootParentJobNamespace,
+      String rootParentJobName,
+      String rootParentRunId) {
     log.debug(
-        "Creating OpenlineageParentContext with parentJobNamespace: {}, parentJobName: {}, parentRunId: {}",
+        "Creating OpenlineageParentContext with parentJobNamespace: {}, parentJobName: {}, parentRunId: {}, rootParentJobNamespace: {}, rootParentJobName: {}, rootParentRunId: {}",
         parentJobNamespace,
         parentJobName,
-        parentRunId);
+        parentRunId,
+        rootParentJobNamespace,
+        rootParentJobName,
+        rootParentRunId);
 
     this.parentJobNamespace = parentJobNamespace;
     this.parentJobName = parentJobName;
     this.parentRunId = parentRunId;
 
-    MessageDigest digest = null;
-    try {
-      digest = MessageDigest.getInstance("SHA-256");
-    } catch (NoSuchAlgorithmException e) {
-      log.debug("Unable to find SHA-256 algorithm", e);
-    }
+    this.rootParentJobNamespace = rootParentJobNamespace;
+    this.rootParentJobName = rootParentJobName;
+    this.rootParentRunId = rootParentRunId;
 
-    if (digest != null && parentJobNamespace != null && parentRunId != null) {
-      traceId = computeTraceId(digest, parentJobNamespace, parentJobName, parentRunId);
+    if (this.parentRunId != null) {
+      traceId = computeTraceId(this.parentRunId);
       spanId = DDSpanId.ZERO;
 
-      childRootSpanId =
-          computeChildRootSpanId(digest, parentJobNamespace, parentJobName, parentRunId);
+      if (this.rootParentRunId != null) {
+        childRootSpanId = computeSpanId(this.rootParentRunId);
+      } else {
+        childRootSpanId = DDSpanId.ZERO;
+      }
     } else {
       traceId = DDTraceId.ZERO;
       spanId = DDSpanId.ZERO;
-
       childRootSpanId = DDSpanId.ZERO;
     }
 
     log.debug("Created OpenlineageParentContext with traceId: {}, spanId: {}", traceId, spanId);
   }
 
-  private long computeChildRootSpanId(
-      MessageDigest digest, String parentJobNamespace, String parentJobName, String parentRunId) {
-    byte[] inputBytes =
-        (parentJobNamespace + parentJobName + parentRunId).getBytes(StandardCharsets.UTF_8);
-    byte[] hash = digest.digest(inputBytes);
-
-    return ByteBuffer.wrap(hash).getLong();
+  private long computeSpanId(String runId) {
+    return FNV64Hash.generateHash(runId, FNV64Hash.Version.v1A);
   }
 
-  private DDTraceId computeTraceId(
-      MessageDigest digest, String parentJobNamespace, String parentJobName, String parentRunId) {
-    byte[] inputBytes =
-        (parentJobNamespace + parentJobName + parentRunId).getBytes(StandardCharsets.UTF_8);
-    byte[] hash = digest.digest(inputBytes);
-
-    return DDTraceId.from(ByteBuffer.wrap(hash).getLong());
+  private DDTraceId computeTraceId(String runId) {
+    return DDTraceId.from(FNV64Hash.generateHash(runId, FNV64Hash.Version.v1A));
   }
 
   @Override
@@ -159,4 +177,16 @@ public String getParentJobName() {
   public String getParentRunId() {
     return parentRunId;
   }
+
+  public String getRootParentJobNamespace() {
+    return rootParentJobNamespace;
+  }
+
+  public String getRootParentJobName() {
+    return rootParentJobName;
+  }
+
+  public String getRootParentRunId() {
+    return rootParentRunId;
+  }
 }