diff --git a/dd-java-agent/instrumentation/spark/spark-common/build.gradle b/dd-java-agent/instrumentation/spark/spark-common/build.gradle index 224484243e8..71c8803f978 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' testImplementation libs.bundles.junit5 testImplementation libs.bundles.mockito 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 2f033394a53..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 @@ -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)); 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/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..8593853733d --- /dev/null +++ b/dd-java-agent/instrumentation/spark/spark-common/src/test/java/datadog/trace/instrumentation/spark/DatabricksRepairAttemptTest.java @@ -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()); + } +} diff --git a/dd-java-agent/instrumentation/spark/spark-common/src/test/resources/databricks/unity-scope-data-original.txt b/dd-java-agent/instrumentation/spark/spark-common/src/test/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/test/resources/databricks/unity-scope-data-original.txt @@ -0,0 +1 @@ +rO0ABXNyABdqYXZhLnV0aWwuTGlua2VkSGFzaE1hcDTATlwQbMD7AgABWgALYWNjZXNzT3JkZXJ4cgARamF2YS51dGlsLkhhc2hNYXAFB9rBwxZg0QMAAkYACmxvYWRGYWN0b3JJAAl0aHJlc2hvbGR4cD9AAAAAAAAMdwgAAAAQAAAABXQACHByaW9yS2V5dAACNDJ0AA1qb2JSdW5BdHRlbXB0dAABOXQAEGpvYlJ1bkF0dGVtcHROdW10AAEwdAAVam9iUnVuT3JpZ2luYWxBdHRlbXB0dAABM3QAC3RyYWlsaW5nS2V5dAAFc2V2ZW54AA== diff --git a/dd-java-agent/instrumentation/spark/spark-common/src/test/resources/databricks/unity-scope-data-repaired.txt b/dd-java-agent/instrumentation/spark/spark-common/src/test/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/test/resources/databricks/unity-scope-data-repaired.txt @@ -0,0 +1 @@ +rO0ABXNyABdqYXZhLnV0aWwuTGlua2VkSGFzaE1hcDTATlwQbMD7AgABWgALYWNjZXNzT3JkZXJ4cgARamF2YS51dGlsLkhhc2hNYXAFB9rBwxZg0QMAAkYACmxvYWRGYWN0b3JJAAl0aHJlc2hvbGR4cD9AAAAAAAAMdwgAAAAQAAAABXQACHByaW9yS2V5dAACNDJ0AA1qb2JSdW5BdHRlbXB0dAABOXQAEGpvYlJ1bkF0dGVtcHROdW10AAExdAAVam9iUnVuT3JpZ2luYWxBdHRlbXB0dAABM3QAC3RyYWlsaW5nS2V5dAAFc2V2ZW54AA== 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 f311d437a84..8e98656a45f 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 generateSampleDataframe(SparkSession spark) {