diff --git a/dd-trace-core/src/main/java/datadog/trace/core/CoreTracer.java b/dd-trace-core/src/main/java/datadog/trace/core/CoreTracer.java index 289647c5c74..178042f8669 100644 --- a/dd-trace-core/src/main/java/datadog/trace/core/CoreTracer.java +++ b/dd-trace-core/src/main/java/datadog/trace/core/CoreTracer.java @@ -1310,6 +1310,9 @@ void write(final SpanList trace) { spanToSample.forceKeep(forceKeep); boolean published = forceKeep || traceCollector.sample(spanToSample); if (published) { + if (rootSpan != null && writtenTrace.get(0) != rootSpan) { + writtenTrace.get(0).spanContext().copyDecisionMakerFrom(rootSpan.spanContext()); + } if (!apmTracingEnabled) { // Stamp the billing marker on every span of each exported chunk so the intake does not bill // APM host usage. diff --git a/dd-trace-core/src/main/java/datadog/trace/core/DDSpanContext.java b/dd-trace-core/src/main/java/datadog/trace/core/DDSpanContext.java index 6120502ec09..e30647699e6 100644 --- a/dd-trace-core/src/main/java/datadog/trace/core/DDSpanContext.java +++ b/dd-trace-core/src/main/java/datadog/trace/core/DDSpanContext.java @@ -1393,6 +1393,10 @@ public PropagationTags getPropagationTags() { return getRootSpanContextOrThis().propagationTags; } + void copyDecisionMakerFrom(DDSpanContext source) { + propagationTags.updateAndLockDecisionMaker(source.getPropagationTags()); + } + /** TraceSegment Implementation */ @Override public void setTagTop(String key, Object value, boolean sanitize) { diff --git a/dd-trace-core/src/test/java/datadog/trace/core/LateSpanSamplingMechanismSerializationTest.java b/dd-trace-core/src/test/java/datadog/trace/core/LateSpanSamplingMechanismSerializationTest.java new file mode 100644 index 00000000000..f68ed798494 --- /dev/null +++ b/dd-trace-core/src/test/java/datadog/trace/core/LateSpanSamplingMechanismSerializationTest.java @@ -0,0 +1,130 @@ +package datadog.trace.core; + +import static datadog.trace.api.sampling.PrioritySampling.USER_KEEP; +import static datadog.trace.api.sampling.SamplingMechanism.REMOTE_ADAPTIVE_RULE; +import static datadog.trace.common.sampling.RuleBasedTraceSampler.SAMPLING_RULE_RATE; +import static java.util.Collections.singletonList; +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertTrue; + +import datadog.communication.serialization.GrowableBuffer; +import datadog.communication.serialization.msgpack.MsgPackWriter; +import datadog.trace.common.writer.ListWriter; +import datadog.trace.common.writer.ddagent.TraceMapperV0_4; +import java.io.IOException; +import java.util.ArrayList; +import java.util.List; +import java.util.concurrent.TimeoutException; +import org.junit.jupiter.api.Test; +import org.msgpack.core.MessagePack; +import org.msgpack.core.MessageUnpacker; +import org.msgpack.value.MapValue; +import org.msgpack.value.Value; +import org.msgpack.value.ValueFactory; + +/** + * Verifies that trace-level sampling metadata is preserved when a root and a later child are + * exported in separate chunks (e.g. during partial sampling). + */ +class LateSpanSamplingMechanismSerializationTest extends DDCoreJavaSpecification { + + @Override + protected boolean useStrictTraceWrites() { + return false; + } + + @Test + void rootlessLateChunkRetainsAdaptiveSamplingMechanism() + throws InterruptedException, TimeoutException { + SerializingWriter writer = new SerializingWriter(); + CoreTracer tracer = + tracerBuilder() + .writer(writer) + // Keep the partial-flush threshold above the trace size so the pending-trace idle + // timeout controls the first write. + .partialFlushMinSpans(1000) + .build(); + + try { + DDSpan root = (DDSpan) tracer.buildSpan("test", "root").start(); + DDSpan lateChild = + (DDSpan) tracer.buildSpan("test", "late-child").asChildOf(root.spanContext()).start(); + PendingTrace trace = (PendingTrace) root.spanContext().getTraceCollector(); + + root.setSamplingPriority(USER_KEEP, REMOTE_ADAPTIVE_RULE); + root.setMetric(SAMPLING_RULE_RATE, 0.25); + + // Root finishes while the child is still running -> root is buffered, not yet written. + root.finish(); + assertTrue(writer.isEmpty()); + + // Write the buffered root as the PendingTraceBuffer worker does after its idle timeout. + trace.write(); + writer.waitForTraces(1); + + // The child finishes after the root was written, so it is exported in a root-less chunk. + lateChild.finish(); + trace.write(); + writer.waitForTraces(2); + + assertEquals(singletonList(root), writer.get(0)); + assertEquals(singletonList(lateChild), writer.get(1)); + + MapValue rootChunk = writer.serializedSpans.get(0); + MapValue lateChunk = writer.serializedSpans.get(1); + + assertEquals( + (int) USER_KEEP, + numberValue(rootChunk, "metrics", DDSpanContext.PRIORITY_SAMPLING_KEY).intValue()); + assertEquals("-12", stringValue(rootChunk, "meta", "_dd.p.dm")); + + assertEquals( + (int) USER_KEEP, + numberValue(lateChunk, "metrics", DDSpanContext.PRIORITY_SAMPLING_KEY).intValue()); + // The late chunk's decision maker must match its propagated trace-level sampling priority. + assertEquals("-12", stringValue(lateChunk, "meta", "_dd.p.dm")); + } finally { + tracer.close(); + } + } + + private static Number numberValue(MapValue span, String mapName, String key) { + Value value = value(span, mapName, key); + return value == null ? null : value.asNumberValue().toInt(); + } + + private static String stringValue(MapValue span, String mapName, String key) { + Value value = value(span, mapName, key); + return value == null ? null : value.asStringValue().asString(); + } + + private static Value value(MapValue span, String mapName, String key) { + Value map = span.map().get(ValueFactory.newString(mapName)); + return map.asMapValue().map().get(ValueFactory.newString(key)); + } + + private static final class SerializingWriter extends ListWriter { + private final List serializedSpans = new ArrayList<>(); + + @Override + public void write(List trace) { + serializedSpans.add(serialize(trace)); + super.write(trace); + } + + private static MapValue serialize(List trace) { + TraceMapperV0_4 mapper = new TraceMapperV0_4(); + GrowableBuffer buffer = new GrowableBuffer(16 * 1024); + MsgPackWriter packer = new MsgPackWriter(buffer); + assertTrue(packer.format(trace, mapper)); + + try (MessageUnpacker unpacker = MessagePack.newDefaultUnpacker(buffer.slice())) { + return unpacker.unpackValue().asArrayValue().get(0).asMapValue(); + } catch (IOException e) { + throw new IllegalStateException("Unable to decode serialized trace chunk", e); + } finally { + mapper.reset(); + } + } + } +}