From 103f0b88f6c4b2b6a547c66ef508e424647a7fd4 Mon Sep 17 00:00:00 2001
From: Adrien Boitreaud <adrien.boitreaud@datadoghq.com>
Date: Tue, 21 Jul 2026 15:49:26 +0200
Subject: [PATCH 1/6] Mix Databricks repair attempt into Spark parent trace id

---
 .../spark/AbstractDatadogSparkListener.java   | 60 ++++++++++++++++++-
 .../spark/DatabricksParentContext.java        | 18 ++++--
 .../spark/AbstractSparkTest.groovy            | 14 +++--
 3 files changed, 83 insertions(+), 9 deletions(-)

diff --git a/dd-java-agent/instrumentation/spark/spark-common/src/main/java/datadog/trace/instrumentation/spark/AbstractDatadogSparkListener.java b/dd-java-agent/instrumentation/spark/spark-common/src/main/java/datadog/trace/instrumentation/spark/AbstractDatadogSparkListener.java
index 817aa1f0d84..931915ddf48 100644
--- a/dd-java-agent/instrumentation/spark/spark-common/src/main/java/datadog/trace/instrumentation/spark/AbstractDatadogSparkListener.java
+++ b/dd-java-agent/instrumentation/spark/spark-common/src/main/java/datadog/trace/instrumentation/spark/AbstractDatadogSparkListener.java
@@ -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;
@@ -442,7 +444,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) {
@@ -1220,6 +1226,58 @@ 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 scanning the serialized bytes; on any failure we return 0, which reproduces the original
+   * (non-repair-aware) trace id so correlation is never worse than before.
+   */
+  private 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 = indexOf(decoded, JOB_RUN_ATTEMPT_NUM_KEY);
+      if (keyIdx < 0) {
+        return 0;
+      }
+      // In the java-serialized stream the value follows its key as the next string, written as
+      // TC_STRING (0x74) with a 2-byte big-endian length prefix.
+      for (int i = keyIdx + JOB_RUN_ATTEMPT_NUM_KEY.length; i + 3 <= decoded.length; i++) {
+        if (decoded[i] != 0x74) {
+          continue;
+        }
+        int len = ((decoded[i + 1] & 0xff) << 8) | (decoded[i + 2] & 0xff);
+        if (len > 0 && len <= 9 && i + 3 + len <= decoded.length) {
+          return Integer.parseInt(new String(decoded, i + 3, len, StandardCharsets.UTF_8));
+        }
+        break;
+      }
+    } catch (Exception e) {
+      log.debug("Unable to extract databricks job run attempt from unity.scope.data", e);
+    }
+    return 0;
+  }
+
+  private static int indexOf(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));
diff --git a/dd-java-agent/instrumentation/spark/spark-common/src/main/java/datadog/trace/instrumentation/spark/DatabricksParentContext.java b/dd-java-agent/instrumentation/spark/spark-common/src/main/java/datadog/trace/instrumentation/spark/DatabricksParentContext.java
index 19b83f07e4e..b0f5de5e505 100644
--- a/dd-java-agent/instrumentation/spark/spark-common/src/main/java/datadog/trace/instrumentation/spark/DatabricksParentContext.java
+++ b/dd-java-agent/instrumentation/spark/spark-common/src/main/java/datadog/trace/instrumentation/spark/DatabricksParentContext.java
@@ -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);
     }
diff --git a/dd-java-agent/instrumentation/spark/spark-common/src/testFixtures/groovy/datadog/trace/instrumentation/spark/AbstractSparkTest.groovy b/dd-java-agent/instrumentation/spark/spark-common/src/testFixtures/groovy/datadog/trace/instrumentation/spark/AbstractSparkTest.groovy
index 95748c054d3..ec369fa0e58 100644
--- a/dd-java-agent/instrumentation/spark/spark-common/src/testFixtures/groovy/datadog/trace/instrumentation/spark/AbstractSparkTest.groovy
+++ b/dd-java-agent/instrumentation/spark/spark-common/src/testFixtures/groovy/datadog/trace/instrumentation/spark/AbstractSparkTest.groovy
@@ -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) {

From 238340498616152d9eb18528decbe2402f6c4510 Mon Sep 17 00:00:00 2001
From: Adrien Boitreaud <adrien.boitreaud@datadoghq.com>
Date: Tue, 21 Jul 2026 16:48:30 +0200
Subject: [PATCH 2/6] Test databricks repair attempt extraction from
 unity.scope.data

---
 .../spark/AbstractDatadogSparkListener.java   |  3 +-
 .../spark/AbstractSparkTest.groovy            | 54 +++++++++++++++++++
 .../databricks/unity-scope-data-original.txt  |  1 +
 .../databricks/unity-scope-data-repaired.txt  |  1 +
 4 files changed, 58 insertions(+), 1 deletion(-)
 create mode 100644 dd-java-agent/instrumentation/spark/spark-common/src/testFixtures/resources/databricks/unity-scope-data-original.txt
 create mode 100644 dd-java-agent/instrumentation/spark/spark-common/src/testFixtures/resources/databricks/unity-scope-data-repaired.txt

diff --git a/dd-java-agent/instrumentation/spark/spark-common/src/main/java/datadog/trace/instrumentation/spark/AbstractDatadogSparkListener.java b/dd-java-agent/instrumentation/spark/spark-common/src/main/java/datadog/trace/instrumentation/spark/AbstractDatadogSparkListener.java
index 931915ddf48..1549df9f61b 100644
--- a/dd-java-agent/instrumentation/spark/spark-common/src/main/java/datadog/trace/instrumentation/spark/AbstractDatadogSparkListener.java
+++ b/dd-java-agent/instrumentation/spark/spark-common/src/main/java/datadog/trace/instrumentation/spark/AbstractDatadogSparkListener.java
@@ -1236,7 +1236,8 @@ private static String getDatabricksTaskRunId(Properties properties, String jobRu
    * by scanning the serialized bytes; on any failure we return 0, which reproduces the original
    * (non-repair-aware) trace id so correlation is never worse than before.
    */
-  private static int getDatabricksJobRunAttempt(Properties properties) {
+  // Package-private for testing against real unity.scope.data payloads.
+  static int getDatabricksJobRunAttempt(Properties properties) {
     String scopeData = properties.getProperty("unity.scope.data");
     if (scopeData == null) {
       return 0;
diff --git a/dd-java-agent/instrumentation/spark/spark-common/src/testFixtures/groovy/datadog/trace/instrumentation/spark/AbstractSparkTest.groovy b/dd-java-agent/instrumentation/spark/spark-common/src/testFixtures/groovy/datadog/trace/instrumentation/spark/AbstractSparkTest.groovy
index ec369fa0e58..e28f221b088 100644
--- a/dd-java-agent/instrumentation/spark/spark-common/src/testFixtures/groovy/datadog/trace/instrumentation/spark/AbstractSparkTest.groovy
+++ b/dd-java-agent/instrumentation/spark/spark-common/src/testFixtures/groovy/datadog/trace/instrumentation/spark/AbstractSparkTest.groovy
@@ -21,6 +21,8 @@ import org.apache.spark.sql.SparkSession
 import org.apache.spark.sql.types.StructType
 import spock.lang.IgnoreIf
 
+import java.util.Properties
+
 @IgnoreIf(reason="https://issues.apache.org/jira/browse/HADOOP-18174", value = {
   JavaVirtualMachine.isJ9()
 })
@@ -525,6 +527,58 @@ abstract class AbstractSparkTest extends InstrumentationSpecification {
     contextRepairedAttempt.getSpanId() == contextWithJobRunId.getSpanId()
   }
 
+  private static String loadResource(String path) {
+    def stream = AbstractSparkTest.getResourceAsStream(path)
+    assert stream != null : "missing test resource $path"
+    stream.withCloseable { it.getText("UTF-8").trim() }
+  }
+
+  def "extract databricks repair attempt from serialized unity.scope.data payloads"() {
+    // The repair attempt index is only reachable inside the base64 + java-serialized
+    // "unity.scope.data" local property. These fixtures are synthetic, secret-free blobs that
+    // reproduce the exact java-serialization format Databricks emits (a map whose "jobRunAttemptNum"
+    // key is immediately followed by its value as a TC_STRING, alongside decoy keys), so the byte-scan
+    // extraction is exercised without committing real driver captures (which carry tokens/PII) to this
+    // public repo.
+    setup:
+    def originalScopeData = loadResource("/databricks/unity-scope-data-original.txt")
+    def repairedScopeData = loadResource("/databricks/unity-scope-data-repaired.txt")
+
+    def originalProps = new Properties()
+    originalProps.setProperty("unity.scope.data", originalScopeData)
+    def repairedProps = new Properties()
+    repairedProps.setProperty("unity.scope.data", repairedScopeData)
+    def missingProps = new Properties()
+    def garbageProps = new Properties()
+    garbageProps.setProperty("unity.scope.data", "not-valid-base64-@@@")
+
+    when:
+    def originalAttempt = AbstractDatadogSparkListener.getDatabricksJobRunAttempt(originalProps)
+    def repairedAttempt = AbstractDatadogSparkListener.getDatabricksJobRunAttempt(repairedProps)
+
+    then:
+    originalAttempt == 0
+    repairedAttempt == 1
+    // best-effort: absent or unparseable payloads fall back to the original (attempt 0) trace id
+    AbstractDatadogSparkListener.getDatabricksJobRunAttempt(missingProps) == 0
+    AbstractDatadogSparkListener.getDatabricksJobRunAttempt(garbageProps) == 0
+
+    when:
+    // Real ids from the captured run: same job_id and parentRunId across both attempts, so only the
+    // extracted attempt index makes the repaired run land on its own trace.
+    def jobId = "572654733963082"
+    def parentRunId = "969804783110783"
+    def originalContext = new DatabricksParentContext(jobId, parentRunId, "807469286696866", originalAttempt)
+    def repairedContext = new DatabricksParentContext(jobId, parentRunId, "1118998436174647", repairedAttempt)
+
+    then:
+    originalContext.getTraceId() == DDTraceId.from("1778014590613446263")
+    repairedContext.getTraceId() == DDTraceId.from("6773783437690493483")
+    originalContext.getTraceId() != repairedContext.getTraceId()
+    originalContext.getSpanId() == DDSpanId.from("17781921289443961499")
+    repairedContext.getSpanId() == DDSpanId.from("9672347302126591758")
+  }
+
   private Dataset<Row> generateSampleDataframe(SparkSession spark) {
     def structType = new StructType()
     structType = structType.add("col", "String", false)
diff --git a/dd-java-agent/instrumentation/spark/spark-common/src/testFixtures/resources/databricks/unity-scope-data-original.txt b/dd-java-agent/instrumentation/spark/spark-common/src/testFixtures/resources/databricks/unity-scope-data-original.txt
new file mode 100644
index 00000000000..89d0ef43e3a
--- /dev/null
+++ b/dd-java-agent/instrumentation/spark/spark-common/src/testFixtures/resources/databricks/unity-scope-data-original.txt
@@ -0,0 +1 @@
+rO0ABXNyABdqYXZhLnV0aWwuTGlua2VkSGFzaE1hcDTATlwQbMD7AgABWgALYWNjZXNzT3JkZXJ4cgARamF2YS51dGlsLkhhc2hNYXAFB9rBwxZg0QMAAkYACmxvYWRGYWN0b3JJAAl0aHJlc2hvbGR4cD9AAAAAAAAMdwgAAAAQAAAABXQACHByaW9yS2V5dAACNDJ0AA1qb2JSdW5BdHRlbXB0dAABOXQAEGpvYlJ1bkF0dGVtcHROdW10AAEwdAAVam9iUnVuT3JpZ2luYWxBdHRlbXB0dAABM3QAC3RyYWlsaW5nS2V5dAAFc2V2ZW54AA==
diff --git a/dd-java-agent/instrumentation/spark/spark-common/src/testFixtures/resources/databricks/unity-scope-data-repaired.txt b/dd-java-agent/instrumentation/spark/spark-common/src/testFixtures/resources/databricks/unity-scope-data-repaired.txt
new file mode 100644
index 00000000000..1b1650fd5dc
--- /dev/null
+++ b/dd-java-agent/instrumentation/spark/spark-common/src/testFixtures/resources/databricks/unity-scope-data-repaired.txt
@@ -0,0 +1 @@
+rO0ABXNyABdqYXZhLnV0aWwuTGlua2VkSGFzaE1hcDTATlwQbMD7AgABWgALYWNjZXNzT3JkZXJ4cgARamF2YS51dGlsLkhhc2hNYXAFB9rBwxZg0QMAAkYACmxvYWRGYWN0b3JJAAl0aHJlc2hvbGR4cD9AAAAAAAAMdwgAAAAQAAAABXQACHByaW9yS2V5dAACNDJ0AA1qb2JSdW5BdHRlbXB0dAABOXQAEGpvYlJ1bkF0dGVtcHROdW10AAExdAAVam9iUnVuT3JpZ2luYWxBdHRlbXB0dAABM3QAC3RyYWlsaW5nS2V5dAAFc2V2ZW54AA==

From 51d60ebb199bb8d815329f999652743e1c5754ec Mon Sep 17 00:00:00 2001
From: Adrien Boitreaud <adrien.boitreaud@datadoghq.com>
Date: Wed, 22 Jul 2026 11:10:36 +0200
Subject: [PATCH 3/6] Harden repair attempt extraction and move its test out of
 the Spock fixture

---
 .../spark/AbstractDatadogSparkListener.java   |  35 ++---
 .../spark/DatabricksRepairAttemptTest.java    | 124 ++++++++++++++++++
 .../databricks/unity-scope-data-original.txt  |   0
 .../databricks/unity-scope-data-repaired.txt  |   0
 .../spark/AbstractSparkTest.groovy            |  54 --------
 5 files changed, 144 insertions(+), 69 deletions(-)
 create mode 100644 dd-java-agent/instrumentation/spark/spark-common/src/test/java/datadog/trace/instrumentation/spark/DatabricksRepairAttemptTest.java
 rename dd-java-agent/instrumentation/spark/spark-common/src/{testFixtures => test}/resources/databricks/unity-scope-data-original.txt (100%)
 rename dd-java-agent/instrumentation/spark/spark-common/src/{testFixtures => test}/resources/databricks/unity-scope-data-repaired.txt (100%)

diff --git a/dd-java-agent/instrumentation/spark/spark-common/src/main/java/datadog/trace/instrumentation/spark/AbstractDatadogSparkListener.java b/dd-java-agent/instrumentation/spark/spark-common/src/main/java/datadog/trace/instrumentation/spark/AbstractDatadogSparkListener.java
index 1549df9f61b..09a5dfba8cc 100644
--- a/dd-java-agent/instrumentation/spark/spark-common/src/main/java/datadog/trace/instrumentation/spark/AbstractDatadogSparkListener.java
+++ b/dd-java-agent/instrumentation/spark/spark-common/src/main/java/datadog/trace/instrumentation/spark/AbstractDatadogSparkListener.java
@@ -1232,11 +1232,13 @@ private static String getDatabricksTaskRunId(Properties properties, String jobRu
   /**
    * 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 scanning the serialized bytes; on any failure we return 0, which reproduces the original
-   * (non-repair-aware) trace id so correlation is never worse than before.
+   * "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 against real unity.scope.data payloads.
+  // Package-private for testing.
   static int getDatabricksJobRunAttempt(Properties properties) {
     String scopeData = properties.getProperty("unity.scope.data");
     if (scopeData == null) {
@@ -1248,17 +1250,20 @@ static int getDatabricksJobRunAttempt(Properties properties) {
       if (keyIdx < 0) {
         return 0;
       }
-      // In the java-serialized stream the value follows its key as the next string, written as
-      // TC_STRING (0x74) with a 2-byte big-endian length prefix.
-      for (int i = keyIdx + JOB_RUN_ATTEMPT_NUM_KEY.length; i + 3 <= decoded.length; i++) {
-        if (decoded[i] != 0x74) {
-          continue;
-        }
-        int len = ((decoded[i + 1] & 0xff) << 8) | (decoded[i + 2] & 0xff);
-        if (len > 0 && len <= 9 && i + 3 + len <= decoded.length) {
-          return Integer.parseInt(new String(decoded, i + 3, len, StandardCharsets.UTF_8));
-        }
-        break;
+      // 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);
diff --git a/dd-java-agent/instrumentation/spark/spark-common/src/test/java/datadog/trace/instrumentation/spark/DatabricksRepairAttemptTest.java b/dd-java-agent/instrumentation/spark/spark-common/src/test/java/datadog/trace/instrumentation/spark/DatabricksRepairAttemptTest.java
new file mode 100644
index 00000000000..b07f6a129a5
--- /dev/null
+++ b/dd-java-agent/instrumentation/spark/spark-common/src/test/java/datadog/trace/instrumentation/spark/DatabricksRepairAttemptTest.java
@@ -0,0 +1,124 @@
+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.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);
+      }
+      byte[] bytes = stream.readAllBytes();
+      return new String(bytes, 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());
+  }
+}
diff --git a/dd-java-agent/instrumentation/spark/spark-common/src/testFixtures/resources/databricks/unity-scope-data-original.txt b/dd-java-agent/instrumentation/spark/spark-common/src/test/resources/databricks/unity-scope-data-original.txt
similarity index 100%
rename from dd-java-agent/instrumentation/spark/spark-common/src/testFixtures/resources/databricks/unity-scope-data-original.txt
rename to dd-java-agent/instrumentation/spark/spark-common/src/test/resources/databricks/unity-scope-data-original.txt
diff --git a/dd-java-agent/instrumentation/spark/spark-common/src/testFixtures/resources/databricks/unity-scope-data-repaired.txt b/dd-java-agent/instrumentation/spark/spark-common/src/test/resources/databricks/unity-scope-data-repaired.txt
similarity index 100%
rename from dd-java-agent/instrumentation/spark/spark-common/src/testFixtures/resources/databricks/unity-scope-data-repaired.txt
rename to dd-java-agent/instrumentation/spark/spark-common/src/test/resources/databricks/unity-scope-data-repaired.txt
diff --git a/dd-java-agent/instrumentation/spark/spark-common/src/testFixtures/groovy/datadog/trace/instrumentation/spark/AbstractSparkTest.groovy b/dd-java-agent/instrumentation/spark/spark-common/src/testFixtures/groovy/datadog/trace/instrumentation/spark/AbstractSparkTest.groovy
index e28f221b088..ec369fa0e58 100644
--- a/dd-java-agent/instrumentation/spark/spark-common/src/testFixtures/groovy/datadog/trace/instrumentation/spark/AbstractSparkTest.groovy
+++ b/dd-java-agent/instrumentation/spark/spark-common/src/testFixtures/groovy/datadog/trace/instrumentation/spark/AbstractSparkTest.groovy
@@ -21,8 +21,6 @@ import org.apache.spark.sql.SparkSession
 import org.apache.spark.sql.types.StructType
 import spock.lang.IgnoreIf
 
-import java.util.Properties
-
 @IgnoreIf(reason="https://issues.apache.org/jira/browse/HADOOP-18174", value = {
   JavaVirtualMachine.isJ9()
 })
@@ -527,58 +525,6 @@ abstract class AbstractSparkTest extends InstrumentationSpecification {
     contextRepairedAttempt.getSpanId() == contextWithJobRunId.getSpanId()
   }
 
-  private static String loadResource(String path) {
-    def stream = AbstractSparkTest.getResourceAsStream(path)
-    assert stream != null : "missing test resource $path"
-    stream.withCloseable { it.getText("UTF-8").trim() }
-  }
-
-  def "extract databricks repair attempt from serialized unity.scope.data payloads"() {
-    // The repair attempt index is only reachable inside the base64 + java-serialized
-    // "unity.scope.data" local property. These fixtures are synthetic, secret-free blobs that
-    // reproduce the exact java-serialization format Databricks emits (a map whose "jobRunAttemptNum"
-    // key is immediately followed by its value as a TC_STRING, alongside decoy keys), so the byte-scan
-    // extraction is exercised without committing real driver captures (which carry tokens/PII) to this
-    // public repo.
-    setup:
-    def originalScopeData = loadResource("/databricks/unity-scope-data-original.txt")
-    def repairedScopeData = loadResource("/databricks/unity-scope-data-repaired.txt")
-
-    def originalProps = new Properties()
-    originalProps.setProperty("unity.scope.data", originalScopeData)
-    def repairedProps = new Properties()
-    repairedProps.setProperty("unity.scope.data", repairedScopeData)
-    def missingProps = new Properties()
-    def garbageProps = new Properties()
-    garbageProps.setProperty("unity.scope.data", "not-valid-base64-@@@")
-
-    when:
-    def originalAttempt = AbstractDatadogSparkListener.getDatabricksJobRunAttempt(originalProps)
-    def repairedAttempt = AbstractDatadogSparkListener.getDatabricksJobRunAttempt(repairedProps)
-
-    then:
-    originalAttempt == 0
-    repairedAttempt == 1
-    // best-effort: absent or unparseable payloads fall back to the original (attempt 0) trace id
-    AbstractDatadogSparkListener.getDatabricksJobRunAttempt(missingProps) == 0
-    AbstractDatadogSparkListener.getDatabricksJobRunAttempt(garbageProps) == 0
-
-    when:
-    // Real ids from the captured run: same job_id and parentRunId across both attempts, so only the
-    // extracted attempt index makes the repaired run land on its own trace.
-    def jobId = "572654733963082"
-    def parentRunId = "969804783110783"
-    def originalContext = new DatabricksParentContext(jobId, parentRunId, "807469286696866", originalAttempt)
-    def repairedContext = new DatabricksParentContext(jobId, parentRunId, "1118998436174647", repairedAttempt)
-
-    then:
-    originalContext.getTraceId() == DDTraceId.from("1778014590613446263")
-    repairedContext.getTraceId() == DDTraceId.from("6773783437690493483")
-    originalContext.getTraceId() != repairedContext.getTraceId()
-    originalContext.getSpanId() == DDSpanId.from("17781921289443961499")
-    repairedContext.getSpanId() == DDSpanId.from("9672347302126591758")
-  }
-
   private Dataset<Row> generateSampleDataframe(SparkSession spark) {
     def structType = new StructType()
     structType = structType.add("col", "String", false)

From d581603fb7613f97757f1d8dfa1ef6d9d5a93aa3 Mon Sep 17 00:00:00 2001
From: Adrien Boitreaud <adrien.boitreaud@datadoghq.com>
Date: Wed, 22 Jul 2026 14:58:11 +0200
Subject: [PATCH 4/6] spotless apply

---
 .../spark/AbstractDatadogSparkListener.java   | 22 +++++++++++--------
 .../spark/DatabricksRepairAttemptTest.java    |  6 ++---
 2 files changed, 15 insertions(+), 13 deletions(-)

diff --git a/dd-java-agent/instrumentation/spark/spark-common/src/main/java/datadog/trace/instrumentation/spark/AbstractDatadogSparkListener.java b/dd-java-agent/instrumentation/spark/spark-common/src/main/java/datadog/trace/instrumentation/spark/AbstractDatadogSparkListener.java
index 09a5dfba8cc..69317662aa9 100644
--- a/dd-java-agent/instrumentation/spark/spark-common/src/main/java/datadog/trace/instrumentation/spark/AbstractDatadogSparkListener.java
+++ b/dd-java-agent/instrumentation/spark/spark-common/src/main/java/datadog/trace/instrumentation/spark/AbstractDatadogSparkListener.java
@@ -1232,11 +1232,11 @@ private static String getDatabricksTaskRunId(Properties properties, String jobRu
   /**
    * 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.
+   * "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) {
@@ -1252,10 +1252,14 @@ static int getDatabricksJobRunAttempt(Properties properties) {
       }
       // 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
+      // 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) {
diff --git a/dd-java-agent/instrumentation/spark/spark-common/src/test/java/datadog/trace/instrumentation/spark/DatabricksRepairAttemptTest.java b/dd-java-agent/instrumentation/spark/spark-common/src/test/java/datadog/trace/instrumentation/spark/DatabricksRepairAttemptTest.java
index b07f6a129a5..60273a186f5 100644
--- a/dd-java-agent/instrumentation/spark/spark-common/src/test/java/datadog/trace/instrumentation/spark/DatabricksRepairAttemptTest.java
+++ b/dd-java-agent/instrumentation/spark/spark-common/src/test/java/datadog/trace/instrumentation/spark/DatabricksRepairAttemptTest.java
@@ -101,10 +101,8 @@ void computesDistinctTraceIdsForOriginalAndRepairedAttemptOfTheSameRun() {
     String jobId = "1234";
     String parentRunId = "5678";
 
-    DatabricksParentContext original =
-        new DatabricksParentContext(jobId, parentRunId, "9012", 0);
-    DatabricksParentContext repaired =
-        new DatabricksParentContext(jobId, parentRunId, "3456", 1);
+    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());

From d74fb417b068ab3e06f742fb753c2bd3103bab60 Mon Sep 17 00:00:00 2001
From: Adrien Boitreaud <adrien.boitreaud@datadoghq.com>
Date: Wed, 22 Jul 2026 15:16:11 +0200
Subject: [PATCH 5/6] Fix CI: add spark-core test dependency and drop Java 9+
 readAllBytes

---
 .../instrumentation/spark/spark-common/build.gradle    |  1 +
 .../spark/DatabricksRepairAttemptTest.java             | 10 ++++++++--
 2 files changed, 9 insertions(+), 2 deletions(-)

diff --git a/dd-java-agent/instrumentation/spark/spark-common/build.gradle b/dd-java-agent/instrumentation/spark/spark-common/build.gradle
index f19ebd38b4a..0b26627fcea 100644
--- a/dd-java-agent/instrumentation/spark/spark-common/build.gradle
+++ b/dd-java-agent/instrumentation/spark/spark-common/build.gradle
@@ -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'
 }
 
diff --git a/dd-java-agent/instrumentation/spark/spark-common/src/test/java/datadog/trace/instrumentation/spark/DatabricksRepairAttemptTest.java b/dd-java-agent/instrumentation/spark/spark-common/src/test/java/datadog/trace/instrumentation/spark/DatabricksRepairAttemptTest.java
index 60273a186f5..8593853733d 100644
--- a/dd-java-agent/instrumentation/spark/spark-common/src/test/java/datadog/trace/instrumentation/spark/DatabricksRepairAttemptTest.java
+++ b/dd-java-agent/instrumentation/spark/spark-common/src/test/java/datadog/trace/instrumentation/spark/DatabricksRepairAttemptTest.java
@@ -5,6 +5,7 @@
 
 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;
@@ -19,8 +20,13 @@ private static String loadResource(String path) {
       if (stream == null) {
         throw new IllegalStateException("missing test resource " + path);
       }
-      byte[] bytes = stream.readAllBytes();
-      return new String(bytes, StandardCharsets.UTF_8).trim();
+      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);
     }

From 1d6a64fd7b39dc5d49e7ed40ba33513a86e998aa Mon Sep 17 00:00:00 2001
From: Adrien Boitreaud <adrien.boitreaud@datadoghq.com>
Date: Thu, 23 Jul 2026 12:34:38 +0200
Subject: [PATCH 6/6] Rename indexOf helper to indexOfDatabricksAttemptKey

---
 .../instrumentation/spark/AbstractDatadogSparkListener.java   | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/dd-java-agent/instrumentation/spark/spark-common/src/main/java/datadog/trace/instrumentation/spark/AbstractDatadogSparkListener.java b/dd-java-agent/instrumentation/spark/spark-common/src/main/java/datadog/trace/instrumentation/spark/AbstractDatadogSparkListener.java
index b90d2c8a4d6..87fe6096e52 100644
--- a/dd-java-agent/instrumentation/spark/spark-common/src/main/java/datadog/trace/instrumentation/spark/AbstractDatadogSparkListener.java
+++ b/dd-java-agent/instrumentation/spark/spark-common/src/main/java/datadog/trace/instrumentation/spark/AbstractDatadogSparkListener.java
@@ -1255,7 +1255,7 @@ static int getDatabricksJobRunAttempt(Properties properties) {
     }
     try {
       byte[] decoded = Base64.getDecoder().decode(scopeData);
-      int keyIdx = indexOf(decoded, JOB_RUN_ATTEMPT_NUM_KEY);
+      int keyIdx = indexOfDatabricksAttemptKey(decoded, JOB_RUN_ATTEMPT_NUM_KEY);
       if (keyIdx < 0) {
         return 0;
       }
@@ -1284,7 +1284,7 @@ static int getDatabricksJobRunAttempt(Properties properties) {
     return 0;
   }
 
-  private static int indexOf(byte[] haystack, byte[] needle) {
+  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++) {