DataDog · aboitreaud · Jul 21, 2026 · Jul 21, 2026 · Jul 22, 2026 · Jul 22, 2026
@@ -28,6 +28,7 @@ dependencies {
   testFixturesCompileOnly(libs.bundles.spock)
 
   testImplementation project(':dd-java-agent:instrumentation-testing')
+  testImplementation group: 'org.apache.spark', name: 'spark-core_2.12', version: '2.4.0'
   testImplementation group: 'org.apache.spark', name: 'spark-launcher_2.12', version: '2.4.0'
   testImplementation libs.bundles.junit5
   testImplementation libs.bundles.mockito

@@ -23,10 +23,12 @@
 import java.lang.invoke.MethodHandle;
 import java.lang.invoke.MethodHandles;
 import java.lang.invoke.MethodType;
+import java.nio.charset.StandardCharsets;
 import java.time.OffsetDateTime;
 import java.time.format.DateTimeParseException;
 import java.util.ArrayList;
 import java.util.Arrays;
+import java.util.Base64;
 import java.util.Collection;
 import java.util.HashMap;
 import java.util.Iterator;
@@ -451,7 +453,11 @@ private void addDatabricksSpecificTags(
       builder.withTag("databricks_task_run_id", databricksTaskRunId);
 
       AgentSpanContext parentContext =
-          new DatabricksParentContext(databricksJobId, databricksJobRunId, databricksTaskRunId);
+          new DatabricksParentContext(
+              databricksJobId,
+              databricksJobRunId,
+              databricksTaskRunId,
+              getDatabricksJobRunAttempt(properties));
 
       if (parentContext.getTraceId() != DDTraceId.ZERO) {
         if (withParentContext) {
@@ -1229,6 +1235,68 @@ private static String getDatabricksTaskRunId(Properties properties, String jobRu
     return null;
   }
 
+  private static final byte[] JOB_RUN_ATTEMPT_NUM_KEY =
+      "jobRunAttemptNum".getBytes(StandardCharsets.UTF_8);
+
+  /**
+   * Returns the Databricks repair attempt index (0 for the original run, 1+ for each repair). It is
+   * not exposed as a first-class Spark property: the only source is the base64 + java-serialized
+   * "unity.scope.data" local property, under the key "jobRunAttemptNum". We extract it best-effort
+   * by looking at the raw serialized bytes right after that key, without a full deserialization of
+   * the stream (e.g. we don't resolve TC_REFERENCE backreferences); on any failure or unexpected
+   * shape we return 0, which reproduces the original (non-repair-aware) trace id so correlation is
+   * never worse than before.
+   */
+  // Package-private for testing.
+  static int getDatabricksJobRunAttempt(Properties properties) {
+    String scopeData = properties.getProperty("unity.scope.data");
+    if (scopeData == null) {
+      return 0;
+    }
+    try {
+      byte[] decoded = Base64.getDecoder().decode(scopeData);
+      int keyIdx = indexOfDatabricksAttemptKey(decoded, JOB_RUN_ATTEMPT_NUM_KEY);
+      if (keyIdx < 0) {
+        return 0;
+      }
+      // The (key, value) tuple is serialized back-to-back with no gap, so the value's type tag sits
+      // immediately after the key's bytes. It's normally TC_STRING (0x74) with a 2-byte big-endian
+      // length prefix, but if this exact attempt string already appeared earlier in the stream
+      // (e.g.
+      // another Databricks tag also has value "1"), Java's serialization writes a TC_REFERENCE
+      // (0x71)
+      // back-pointer instead of repeating the string. We don't resolve backreferences -- rather
+      // than
+      // risk scanning forward and matching an unrelated byte, treat anything other than an
+      // immediate
+      // TC_STRING as unparseable and fall back to attempt 0.
+      int valueStart = keyIdx + JOB_RUN_ATTEMPT_NUM_KEY.length;
+      if (valueStart + 3 > decoded.length || decoded[valueStart] != 0x74) {
+        return 0;
+      }
+      int len = ((decoded[valueStart + 1] & 0xff) << 8) | (decoded[valueStart + 2] & 0xff);
+      if (len > 0 && len <= 9 && valueStart + 3 + len <= decoded.length) {
+        return Integer.parseInt(new String(decoded, valueStart + 3, len, StandardCharsets.UTF_8));
+      }
+    } catch (Exception e) {
+      log.debug("Unable to extract databricks job run attempt from unity.scope.data", e);
+    }
+    return 0;
+  }
+
+  private static int indexOfDatabricksAttemptKey(byte[] haystack, byte[] needle) {
+    outer:
+    for (int i = 0; i <= haystack.length - needle.length; i++) {
+      for (int j = 0; j < needle.length; j++) {
+        if (haystack[i + j] != needle[j]) {
+          continue outer;
+        }
+      }
+      return i;
+    }
+    return -1;
+  }
+
   private String stackTraceToString(Throwable e) {
     StringWriter stringWriter = new StringWriter();
     e.printStackTrace(new PrintWriter(stringWriter));

@@ -27,7 +27,8 @@ public class DatabricksParentContext implements AgentSpanContext {
   private final DDTraceId traceId;
   private final long spanId;
 
-  public DatabricksParentContext(String jobId, String jobRunId, String taskRunId) {
+  public DatabricksParentContext(
+      String jobId, String jobRunId, String taskRunId, int attemptNumber) {
     MessageDigest digest = null;
     try {
       digest = MessageDigest.getInstance("SHA-1");
@@ -36,7 +37,7 @@ public DatabricksParentContext(String jobId, String jobRunId, String taskRunId)
     }
 
     if (digest != null && jobId != null && taskRunId != null) {
-      traceId = computeTraceId(digest, jobId, jobRunId, taskRunId);
+      traceId = computeTraceId(digest, jobId, jobRunId, taskRunId, attemptNumber);
       spanId = computeSpanId(digest, jobId, taskRunId);
     } else {
       traceId = DDTraceId.ZERO;
@@ -45,11 +46,20 @@ public DatabricksParentContext(String jobId, String jobRunId, String taskRunId)
   }
 
   private DDTraceId computeTraceId(
-      MessageDigest digest, String jobId, String jobRunId, String taskRunId) {
+      MessageDigest digest, String jobId, String jobRunId, String taskRunId, int attemptNumber) {
     byte[] inputBytes;
 
     if (jobRunId != null) {
-      inputBytes = (jobId + jobRunId).getBytes(StandardCharsets.UTF_8);
+      // Databricks reuses the same jobRunId when a run is repaired, so the original and repaired
+      // attempts would otherwise hash to the same trace id. The crawler side (dogweb) mixes the
+      // repair attempt index into the trace id to give each attempt its own trace; mirror that
+      // exactly so the spark spans land on the matching trace. attempt 0 stays bare for backward
+      // compatibility with runs that were never repaired.
+      String input = jobId + jobRunId;
+      if (attemptNumber > 0) {
+        input += "-" + attemptNumber;
+      }
+      inputBytes = input.getBytes(StandardCharsets.UTF_8);
     } else {
       inputBytes = (jobId + taskRunId).getBytes(StandardCharsets.UTF_8);
     }

@@ -0,0 +1,128 @@
+package datadog.trace.instrumentation.spark;
+
+import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertNotEquals;
+
+import datadog.trace.api.DDSpanId;
+import datadog.trace.api.DDTraceId;
+import java.io.ByteArrayOutputStream;
+import java.io.IOException;
+import java.io.InputStream;
+import java.io.UncheckedIOException;
+import java.nio.charset.StandardCharsets;
+import java.util.Properties;
+import org.junit.jupiter.api.Test;
+
+class DatabricksRepairAttemptTest {
+
+  private static String loadResource(String path) {
+    try (InputStream stream = DatabricksRepairAttemptTest.class.getResourceAsStream(path)) {
+      if (stream == null) {
+        throw new IllegalStateException("missing test resource " + path);
+      }
+      ByteArrayOutputStream bytes = new ByteArrayOutputStream();
+      byte[] chunk = new byte[4096];
+      int read;
+      while ((read = stream.read(chunk)) != -1) {
+        bytes.write(chunk, 0, read);
+      }
+      return new String(bytes.toByteArray(), StandardCharsets.UTF_8).trim();
+    } catch (IOException e) {
+      throw new UncheckedIOException(e);
+    }
+  }
+
+  @Test
+  void extractsAttemptZeroFromOriginalRunPayload() {
+    Properties properties = new Properties();
+    properties.setProperty(
+        "unity.scope.data", loadResource("/databricks/unity-scope-data-original.txt"));
+
+    assertEquals(0, AbstractDatadogSparkListener.getDatabricksJobRunAttempt(properties));
+  }
+
+  @Test
+  void extractsAttemptOneFromRepairedRunPayload() {
+    Properties properties = new Properties();
+    properties.setProperty(
+        "unity.scope.data", loadResource("/databricks/unity-scope-data-repaired.txt"));
+
+    assertEquals(1, AbstractDatadogSparkListener.getDatabricksJobRunAttempt(properties));
+  }
+
+  @Test
+  void fallsBackToZeroWhenPropertyIsMissing() {
+    Properties properties = new Properties();
+
+    assertEquals(0, AbstractDatadogSparkListener.getDatabricksJobRunAttempt(properties));
+  }
+
+  @Test
+  void fallsBackToZeroWhenPayloadIsNotValidBase64() {
+    Properties properties = new Properties();
+    properties.setProperty("unity.scope.data", "not-valid-base64-@@@");
+
+    assertEquals(0, AbstractDatadogSparkListener.getDatabricksJobRunAttempt(properties));
+  }
+
+  @Test
+  void fallsBackToZeroWhenKeyIsAbsent() {
+    // Base64 of a serialized map that never mentions jobRunAttemptNum at all.
+    Properties properties = new Properties();
+    properties.setProperty(
+        "unity.scope.data",
+        java.util.Base64.getEncoder()
+            .encodeToString("no attempt info in here".getBytes(StandardCharsets.UTF_8)));
+
+    assertEquals(0, AbstractDatadogSparkListener.getDatabricksJobRunAttempt(properties));
+  }
+
+  @Test
+  void fallsBackToZeroWhenValueIsABackreferenceRatherThanAString() {
+    // If the attempt value string was already written earlier in the serialized stream, Java
+    // serialization emits a TC_REFERENCE (0x71) back-pointer instead of repeating the string. We
+    // don't resolve backreferences, so this must safely fall back to attempt 0 rather than risk
+    // scanning forward and matching an unrelated byte as the value.
+    byte[] key = "jobRunAttemptNum".getBytes(StandardCharsets.UTF_8);
+    byte[] payload = new byte[key.length + 5];
+    System.arraycopy(key, 0, payload, 0, key.length);
+    payload[key.length] = 0x71; // TC_REFERENCE
+    payload[key.length + 1] = 0x00;
+    payload[key.length + 2] = 0x00;
+    payload[key.length + 3] = 0x00;
+    payload[key.length + 4] = 0x01;
+
+    Properties properties = new Properties();
+    properties.setProperty(
+        "unity.scope.data", java.util.Base64.getEncoder().encodeToString(payload));
+
+    assertEquals(0, AbstractDatadogSparkListener.getDatabricksJobRunAttempt(properties));
+  }
+
+  @Test
+  void computesDistinctTraceIdsForOriginalAndRepairedAttemptOfTheSameRun() {
+    // A repaired run reuses the same jobId/parentRunId as the original (that's the whole reason
+    // spans collided before this fix); it gets a fresh taskRunId per attempt, but that alone isn't
+    // enough since the job span (the trace root) is keyed on jobId+parentRunId, not taskRunId.
+    String jobId = "1234";
+    String parentRunId = "5678";
+
+    DatabricksParentContext original = new DatabricksParentContext(jobId, parentRunId, "9012", 0);
+    DatabricksParentContext repaired = new DatabricksParentContext(jobId, parentRunId, "3456", 1);
+
+    assertEquals(DDTraceId.from("8944764253919609482"), original.getTraceId());
+    assertNotEquals(original.getTraceId(), repaired.getTraceId());
+    assertNotEquals(DDSpanId.ZERO, repaired.getSpanId());
+  }
+
+  @Test
+  void attemptZeroReproducesTheOriginalNonRepairAwareTraceId() {
+    // Backward compatibility: runs that were never repaired must keep the exact trace id they
+    // always had, since attempt 0 is the overwhelmingly common case.
+    DatabricksParentContext withDefaultAttempt =
+        new DatabricksParentContext("1234", "5678", "9012", 0);
+
+    assertEquals(DDTraceId.from("8944764253919609482"), withDefaultAttempt.getTraceId());
+    assertEquals(DDSpanId.from("15104224823446433673"), withDefaultAttempt.getSpanId());
+  }
+}
@@ -0,0 +1 @@
+rO0ABXNyABdqYXZhLnV0aWwuTGlua2VkSGFzaE1hcDTATlwQbMD7AgABWgALYWNjZXNzT3JkZXJ4cgARamF2YS51dGlsLkhhc2hNYXAFB9rBwxZg0QMAAkYACmxvYWRGYWN0b3JJAAl0aHJlc2hvbGR4cD9AAAAAAAAMdwgAAAAQAAAABXQACHByaW9yS2V5dAACNDJ0AA1qb2JSdW5BdHRlbXB0dAABOXQAEGpvYlJ1bkF0dGVtcHROdW10AAEwdAAVam9iUnVuT3JpZ2luYWxBdHRlbXB0dAABM3QAC3RyYWlsaW5nS2V5dAAFc2V2ZW54AA==
@@ -0,0 +1 @@
+rO0ABXNyABdqYXZhLnV0aWwuTGlua2VkSGFzaE1hcDTATlwQbMD7AgABWgALYWNjZXNzT3JkZXJ4cgARamF2YS51dGlsLkhhc2hNYXAFB9rBwxZg0QMAAkYACmxvYWRGYWN0b3JJAAl0aHJlc2hvbGR4cD9AAAAAAAAMdwgAAAAQAAAABXQACHByaW9yS2V5dAACNDJ0AA1qb2JSdW5BdHRlbXB0dAABOXQAEGpvYlJ1bkF0dGVtcHROdW10AAExdAAVam9iUnVuT3JpZ2luYWxBdHRlbXB0dAABM3QAC3RyYWlsaW5nS2V5dAAFc2V2ZW54AA==
@@ -500,10 +500,13 @@ abstract class AbstractSparkTest extends InstrumentationSpecification {
 
   def "compute the databricks parent context"() {
     setup:
-    def contextWithJobRunId = new DatabricksParentContext("1234", "5678", "9012")
-    def contextWithoutJobRunId = new DatabricksParentContext("1234", null, "9012")
-    def contextWithoutJobId = new DatabricksParentContext(null, "5678", "9012")
-    def contextWithoutTaskRunId = new DatabricksParentContext(null, "5678", null)
+    def contextWithJobRunId = new DatabricksParentContext("1234", "5678", "9012", 0)
+    def contextWithoutJobRunId = new DatabricksParentContext("1234", null, "9012", 0)
+    def contextWithoutJobId = new DatabricksParentContext(null, "5678", "9012", 0)
+    def contextWithoutTaskRunId = new DatabricksParentContext(null, "5678", null, 0)
+    // A repaired run reuses the same jobRunId but reports a non-zero attempt, which must yield a
+    // different trace id (its own trace) while a repaired task keeps its own fresh taskRunId span id.
+    def contextRepairedAttempt = new DatabricksParentContext("1234", "5678", "9012", 1)
 
     expect:
     contextWithJobRunId.getTraceId() == DDTraceId.from("8944764253919609482")
@@ -517,6 +520,9 @@ abstract class AbstractSparkTest extends InstrumentationSpecification {
 
     contextWithoutTaskRunId.getTraceId() == DDTraceId.ZERO
     contextWithoutTaskRunId.getSpanId() == DDSpanId.ZERO
+
+    contextRepairedAttempt.getTraceId() != contextWithJobRunId.getTraceId()
+    contextRepairedAttempt.getSpanId() == contextWithJobRunId.getSpanId()
   }
 
   private Dataset<Row> generateSampleDataframe(SparkSession spark) {
Original file line number	Diff line number	Diff line change
		@@ -0,0 +1 @@
		rO0ABXNyABdqYXZhLnV0aWwuTGlua2VkSGFzaE1hcDTATlwQbMD7AgABWgALYWNjZXNzT3JkZXJ4cgARamF2YS51dGlsLkhhc2hNYXAFB9rBwxZg0QMAAkYACmxvYWRGYWN0b3JJAAl0aHJlc2hvbGR4cD9AAAAAAAAMdwgAAAAQAAAABXQACHByaW9yS2V5dAACNDJ0AA1qb2JSdW5BdHRlbXB0dAABOXQAEGpvYlJ1bkF0dGVtcHROdW10AAEwdAAVam9iUnVuT3JpZ2luYWxBdHRlbXB0dAABM3QAC3RyYWlsaW5nS2V5dAAFc2V2ZW54AA==