Skip to content

Commit 507d52c

Browse files
aboitreauddevflow.devflow-routing-intake
aboitreaud
and
devflow.devflow-routing-intake
authored
Mix Databricks repair attempt into Spark parent trace id (#12022)
Mix Databricks repair attempt into Spark parent trace id Test databricks repair attempt extraction from unity.scope.data Harden repair attempt extraction and move its test out of the Spock fixture spotless apply Fix CI: add spark-core test dependency and drop Java 9+ readAllBytes Merge branch 'master' into adrien.boitreaud/databricks-repair-run-spark-trace-correlation Rename indexOf helper to indexOfDatabricksAttemptKey Co-authored-by: devflow.devflow-routing-intake <devflow.devflow-routing-intake@kubernetes.us1.ddbuild.io>
1 parent af2d214 commit 507d52c

7 files changed

Lines changed: 224 additions & 9 deletions

File tree

‎dd-java-agent/instrumentation/spark/spark-common/build.gradle‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -28,6 +28,7 @@ dependencies {
2828
testFixturesCompileOnly(libs.bundles.spock)
2929

3030
testImplementation project(':dd-java-agent:instrumentation-testing')
31+
testImplementation group: 'org.apache.spark', name: 'spark-core_2.12', version: '2.4.0'
3132
testImplementation group: 'org.apache.spark', name: 'spark-launcher_2.12', version: '2.4.0'
3233
testImplementation libs.bundles.junit5
3334
testImplementation libs.bundles.mockito

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

Lines changed: 69 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -23,10 +23,12 @@
2323
import java.lang.invoke.MethodHandle;
2424
import java.lang.invoke.MethodHandles;
2525
import java.lang.invoke.MethodType;
26+
import java.nio.charset.StandardCharsets;
2627
import java.time.OffsetDateTime;
2728
import java.time.format.DateTimeParseException;
2829
import java.util.ArrayList;
2930
import java.util.Arrays;
31+
import java.util.Base64;
3032
import java.util.Collection;
3133
import java.util.HashMap;
3234
import java.util.Iterator;
@@ -451,7 +453,11 @@ private void addDatabricksSpecificTags(
451453
builder.withTag("databricks_task_run_id", databricksTaskRunId);
452454

453455
AgentSpanContext parentContext =
454-
new DatabricksParentContext(databricksJobId, databricksJobRunId, databricksTaskRunId);
456+
new DatabricksParentContext(
457+
databricksJobId,
458+
databricksJobRunId,
459+
databricksTaskRunId,
460+
getDatabricksJobRunAttempt(properties));
455461

456462
if (parentContext.getTraceId() != DDTraceId.ZERO) {
457463
if (withParentContext) {
@@ -1229,6 +1235,68 @@ private static String getDatabricksTaskRunId(Properties properties, String jobRu
12291235
return null;
12301236
}
12311237

1238+
private static final byte[] JOB_RUN_ATTEMPT_NUM_KEY =
1239+
"jobRunAttemptNum".getBytes(StandardCharsets.UTF_8);
1240+
1241+
/**
1242+
* Returns the Databricks repair attempt index (0 for the original run, 1+ for each repair). It is
1243+
* not exposed as a first-class Spark property: the only source is the base64 + java-serialized
1244+
* "unity.scope.data" local property, under the key "jobRunAttemptNum". We extract it best-effort
1245+
* by looking at the raw serialized bytes right after that key, without a full deserialization of
1246+
* the stream (e.g. we don't resolve TC_REFERENCE backreferences); on any failure or unexpected
1247+
* shape we return 0, which reproduces the original (non-repair-aware) trace id so correlation is
1248+
* never worse than before.
1249+
*/
1250+
// Package-private for testing.
1251+
static int getDatabricksJobRunAttempt(Properties properties) {
1252+
String scopeData = properties.getProperty("unity.scope.data");
1253+
if (scopeData == null) {
1254+
return 0;
1255+
}
1256+
try {
1257+
byte[] decoded = Base64.getDecoder().decode(scopeData);
1258+
int keyIdx = indexOfDatabricksAttemptKey(decoded, JOB_RUN_ATTEMPT_NUM_KEY);
1259+
if (keyIdx < 0) {
1260+
return 0;
1261+
}
1262+
// The (key, value) tuple is serialized back-to-back with no gap, so the value's type tag sits
1263+
// immediately after the key's bytes. It's normally TC_STRING (0x74) with a 2-byte big-endian
1264+
// length prefix, but if this exact attempt string already appeared earlier in the stream
1265+
// (e.g.
1266+
// another Databricks tag also has value "1"), Java's serialization writes a TC_REFERENCE
1267+
// (0x71)
1268+
// back-pointer instead of repeating the string. We don't resolve backreferences -- rather
1269+
// than
1270+
// risk scanning forward and matching an unrelated byte, treat anything other than an
1271+
// immediate
1272+
// TC_STRING as unparseable and fall back to attempt 0.
1273+
int valueStart = keyIdx + JOB_RUN_ATTEMPT_NUM_KEY.length;
1274+
if (valueStart + 3 > decoded.length || decoded[valueStart] != 0x74) {
1275+
return 0;
1276+
}
1277+
int len = ((decoded[valueStart + 1] & 0xff) << 8) | (decoded[valueStart + 2] & 0xff);
1278+
if (len > 0 && len <= 9 && valueStart + 3 + len <= decoded.length) {
1279+
return Integer.parseInt(new String(decoded, valueStart + 3, len, StandardCharsets.UTF_8));
1280+
}
1281+
} catch (Exception e) {
1282+
log.debug("Unable to extract databricks job run attempt from unity.scope.data", e);
1283+
}
1284+
return 0;
1285+
}
1286+
1287+
private static int indexOfDatabricksAttemptKey(byte[] haystack, byte[] needle) {
1288+
outer:
1289+
for (int i = 0; i <= haystack.length - needle.length; i++) {
1290+
for (int j = 0; j < needle.length; j++) {
1291+
if (haystack[i + j] != needle[j]) {
1292+
continue outer;
1293+
}
1294+
}
1295+
return i;
1296+
}
1297+
return -1;
1298+
}
1299+
12321300
private String stackTraceToString(Throwable e) {
12331301
StringWriter stringWriter = new StringWriter();
12341302
e.printStackTrace(new PrintWriter(stringWriter));

‎dd-java-agent/instrumentation/spark/spark-common/src/main/java/datadog/trace/instrumentation/spark/DatabricksParentContext.java‎

Lines changed: 14 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -27,7 +27,8 @@ public class DatabricksParentContext implements AgentSpanContext {
2727
private final DDTraceId traceId;
2828
private final long spanId;
2929

30-
public DatabricksParentContext(String jobId, String jobRunId, String taskRunId) {
30+
public DatabricksParentContext(
31+
String jobId, String jobRunId, String taskRunId, int attemptNumber) {
3132
MessageDigest digest = null;
3233
try {
3334
digest = MessageDigest.getInstance("SHA-1");
@@ -36,7 +37,7 @@ public DatabricksParentContext(String jobId, String jobRunId, String taskRunId)
3637
}
3738

3839
if (digest != null && jobId != null && taskRunId != null) {
39-
traceId = computeTraceId(digest, jobId, jobRunId, taskRunId);
40+
traceId = computeTraceId(digest, jobId, jobRunId, taskRunId, attemptNumber);
4041
spanId = computeSpanId(digest, jobId, taskRunId);
4142
} else {
4243
traceId = DDTraceId.ZERO;
@@ -45,11 +46,20 @@ public DatabricksParentContext(String jobId, String jobRunId, String taskRunId)
4546
}
4647

4748
private DDTraceId computeTraceId(
48-
MessageDigest digest, String jobId, String jobRunId, String taskRunId) {
49+
MessageDigest digest, String jobId, String jobRunId, String taskRunId, int attemptNumber) {
4950
byte[] inputBytes;
5051

5152
if (jobRunId != null) {
52-
inputBytes = (jobId + jobRunId).getBytes(StandardCharsets.UTF_8);
53+
// Databricks reuses the same jobRunId when a run is repaired, so the original and repaired
54+
// attempts would otherwise hash to the same trace id. The crawler side (dogweb) mixes the
55+
// repair attempt index into the trace id to give each attempt its own trace; mirror that
56+
// exactly so the spark spans land on the matching trace. attempt 0 stays bare for backward
57+
// compatibility with runs that were never repaired.
58+
String input = jobId + jobRunId;
59+
if (attemptNumber > 0) {
60+
input += "-" + attemptNumber;
61+
}
62+
inputBytes = input.getBytes(StandardCharsets.UTF_8);
5363
} else {
5464
inputBytes = (jobId + taskRunId).getBytes(StandardCharsets.UTF_8);
5565
}

‎dd-java-agent/instrumentation/spark/spark-common/src/test/java/datadog/trace/instrumentation/spark/DatabricksRepairAttemptTest.java‎

Lines changed: 128 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,128 @@
1+
package datadog.trace.instrumentation.spark;
2+
3+
import static org.junit.jupiter.api.Assertions.assertEquals;
4+
import static org.junit.jupiter.api.Assertions.assertNotEquals;
5+
6+
import datadog.trace.api.DDSpanId;
7+
import datadog.trace.api.DDTraceId;
8+
import java.io.ByteArrayOutputStream;
9+
import java.io.IOException;
10+
import java.io.InputStream;
11+
import java.io.UncheckedIOException;
12+
import java.nio.charset.StandardCharsets;
13+
import java.util.Properties;
14+
import org.junit.jupiter.api.Test;
15+
16+
class DatabricksRepairAttemptTest {
17+
18+
private static String loadResource(String path) {
19+
try (InputStream stream = DatabricksRepairAttemptTest.class.getResourceAsStream(path)) {
20+
if (stream == null) {
21+
throw new IllegalStateException("missing test resource " + path);
22+
}
23+
ByteArrayOutputStream bytes = new ByteArrayOutputStream();
24+
byte[] chunk = new byte[4096];
25+
int read;
26+
while ((read = stream.read(chunk)) != -1) {
27+
bytes.write(chunk, 0, read);
28+
}
29+
return new String(bytes.toByteArray(), StandardCharsets.UTF_8).trim();
30+
} catch (IOException e) {
31+
throw new UncheckedIOException(e);
32+
}
33+
}
34+
35+
@Test
36+
void extractsAttemptZeroFromOriginalRunPayload() {
37+
Properties properties = new Properties();
38+
properties.setProperty(
39+
"unity.scope.data", loadResource("/databricks/unity-scope-data-original.txt"));
40+
41+
assertEquals(0, AbstractDatadogSparkListener.getDatabricksJobRunAttempt(properties));
42+
}
43+
44+
@Test
45+
void extractsAttemptOneFromRepairedRunPayload() {
46+
Properties properties = new Properties();
47+
properties.setProperty(
48+
"unity.scope.data", loadResource("/databricks/unity-scope-data-repaired.txt"));
49+
50+
assertEquals(1, AbstractDatadogSparkListener.getDatabricksJobRunAttempt(properties));
51+
}
52+
53+
@Test
54+
void fallsBackToZeroWhenPropertyIsMissing() {
55+
Properties properties = new Properties();
56+
57+
assertEquals(0, AbstractDatadogSparkListener.getDatabricksJobRunAttempt(properties));
58+
}
59+
60+
@Test
61+
void fallsBackToZeroWhenPayloadIsNotValidBase64() {
62+
Properties properties = new Properties();
63+
properties.setProperty("unity.scope.data", "not-valid-base64-@@@");
64+
65+
assertEquals(0, AbstractDatadogSparkListener.getDatabricksJobRunAttempt(properties));
66+
}
67+
68+
@Test
69+
void fallsBackToZeroWhenKeyIsAbsent() {
70+
// Base64 of a serialized map that never mentions jobRunAttemptNum at all.
71+
Properties properties = new Properties();
72+
properties.setProperty(
73+
"unity.scope.data",
74+
java.util.Base64.getEncoder()
75+
.encodeToString("no attempt info in here".getBytes(StandardCharsets.UTF_8)));
76+
77+
assertEquals(0, AbstractDatadogSparkListener.getDatabricksJobRunAttempt(properties));
78+
}
79+
80+
@Test
81+
void fallsBackToZeroWhenValueIsABackreferenceRatherThanAString() {
82+
// If the attempt value string was already written earlier in the serialized stream, Java
83+
// serialization emits a TC_REFERENCE (0x71) back-pointer instead of repeating the string. We
84+
// don't resolve backreferences, so this must safely fall back to attempt 0 rather than risk
85+
// scanning forward and matching an unrelated byte as the value.
86+
byte[] key = "jobRunAttemptNum".getBytes(StandardCharsets.UTF_8);
87+
byte[] payload = new byte[key.length + 5];
88+
System.arraycopy(key, 0, payload, 0, key.length);
89+
payload[key.length] = 0x71; // TC_REFERENCE
90+
payload[key.length + 1] = 0x00;
91+
payload[key.length + 2] = 0x00;
92+
payload[key.length + 3] = 0x00;
93+
payload[key.length + 4] = 0x01;
94+
95+
Properties properties = new Properties();
96+
properties.setProperty(
97+
"unity.scope.data", java.util.Base64.getEncoder().encodeToString(payload));
98+
99+
assertEquals(0, AbstractDatadogSparkListener.getDatabricksJobRunAttempt(properties));
100+
}
101+
102+
@Test
103+
void computesDistinctTraceIdsForOriginalAndRepairedAttemptOfTheSameRun() {
104+
// A repaired run reuses the same jobId/parentRunId as the original (that's the whole reason
105+
// spans collided before this fix); it gets a fresh taskRunId per attempt, but that alone isn't
106+
// enough since the job span (the trace root) is keyed on jobId+parentRunId, not taskRunId.
107+
String jobId = "1234";
108+
String parentRunId = "5678";
109+
110+
DatabricksParentContext original = new DatabricksParentContext(jobId, parentRunId, "9012", 0);
111+
DatabricksParentContext repaired = new DatabricksParentContext(jobId, parentRunId, "3456", 1);
112+
113+
assertEquals(DDTraceId.from("8944764253919609482"), original.getTraceId());
114+
assertNotEquals(original.getTraceId(), repaired.getTraceId());
115+
assertNotEquals(DDSpanId.ZERO, repaired.getSpanId());
116+
}
117+
118+
@Test
119+
void attemptZeroReproducesTheOriginalNonRepairAwareTraceId() {
120+
// Backward compatibility: runs that were never repaired must keep the exact trace id they
121+
// always had, since attempt 0 is the overwhelmingly common case.
122+
DatabricksParentContext withDefaultAttempt =
123+
new DatabricksParentContext("1234", "5678", "9012", 0);
124+
125+
assertEquals(DDTraceId.from("8944764253919609482"), withDefaultAttempt.getTraceId());
126+
assertEquals(DDSpanId.from("15104224823446433673"), withDefaultAttempt.getSpanId());
127+
}
128+
}

‎dd-java-agent/instrumentation/spark/spark-common/src/test/resources/databricks/unity-scope-data-original.txt‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1 @@
1+
rO0ABXNyABdqYXZhLnV0aWwuTGlua2VkSGFzaE1hcDTATlwQbMD7AgABWgALYWNjZXNzT3JkZXJ4cgARamF2YS51dGlsLkhhc2hNYXAFB9rBwxZg0QMAAkYACmxvYWRGYWN0b3JJAAl0aHJlc2hvbGR4cD9AAAAAAAAMdwgAAAAQAAAABXQACHByaW9yS2V5dAACNDJ0AA1qb2JSdW5BdHRlbXB0dAABOXQAEGpvYlJ1bkF0dGVtcHROdW10AAEwdAAVam9iUnVuT3JpZ2luYWxBdHRlbXB0dAABM3QAC3RyYWlsaW5nS2V5dAAFc2V2ZW54AA==

‎dd-java-agent/instrumentation/spark/spark-common/src/test/resources/databricks/unity-scope-data-repaired.txt‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1 @@
1+
rO0ABXNyABdqYXZhLnV0aWwuTGlua2VkSGFzaE1hcDTATlwQbMD7AgABWgALYWNjZXNzT3JkZXJ4cgARamF2YS51dGlsLkhhc2hNYXAFB9rBwxZg0QMAAkYACmxvYWRGYWN0b3JJAAl0aHJlc2hvbGR4cD9AAAAAAAAMdwgAAAAQAAAABXQACHByaW9yS2V5dAACNDJ0AA1qb2JSdW5BdHRlbXB0dAABOXQAEGpvYlJ1bkF0dGVtcHROdW10AAExdAAVam9iUnVuT3JpZ2luYWxBdHRlbXB0dAABM3QAC3RyYWlsaW5nS2V5dAAFc2V2ZW54AA==

‎dd-java-agent/instrumentation/spark/spark-common/src/testFixtures/groovy/datadog/trace/instrumentation/spark/AbstractSparkTest.groovy‎

Lines changed: 10 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -500,10 +500,13 @@ abstract class AbstractSparkTest extends InstrumentationSpecification {
500500

501501
def "compute the databricks parent context"() {
502502
setup:
503-
def contextWithJobRunId = new DatabricksParentContext("1234", "5678", "9012")
504-
def contextWithoutJobRunId = new DatabricksParentContext("1234", null, "9012")
505-
def contextWithoutJobId = new DatabricksParentContext(null, "5678", "9012")
506-
def contextWithoutTaskRunId = new DatabricksParentContext(null, "5678", null)
503+
def contextWithJobRunId = new DatabricksParentContext("1234", "5678", "9012", 0)
504+
def contextWithoutJobRunId = new DatabricksParentContext("1234", null, "9012", 0)
505+
def contextWithoutJobId = new DatabricksParentContext(null, "5678", "9012", 0)
506+
def contextWithoutTaskRunId = new DatabricksParentContext(null, "5678", null, 0)
507+
// A repaired run reuses the same jobRunId but reports a non-zero attempt, which must yield a
508+
// different trace id (its own trace) while a repaired task keeps its own fresh taskRunId span id.
509+
def contextRepairedAttempt = new DatabricksParentContext("1234", "5678", "9012", 1)
507510

508511
expect:
509512
contextWithJobRunId.getTraceId() == DDTraceId.from("8944764253919609482")
@@ -517,6 +520,9 @@ abstract class AbstractSparkTest extends InstrumentationSpecification {
517520

518521
contextWithoutTaskRunId.getTraceId() == DDTraceId.ZERO
519522
contextWithoutTaskRunId.getSpanId() == DDSpanId.ZERO
523+
524+
contextRepairedAttempt.getTraceId() != contextWithJobRunId.getTraceId()
525+
contextRepairedAttempt.getSpanId() == contextWithJobRunId.getSpanId()
520526
}
521527

522528
private Dataset<Row> generateSampleDataframe(SparkSession spark) {

0 commit comments

Comments
 (0)