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 f14a3d9c2ef..f154198f320 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 @@ -166,14 +166,11 @@ public class DDSpanContext implements AgentSpanContext, ProfilerContext { private volatile boolean topLevel; - /** The origin of the trace. (eg. Synthetics, CI App) */ - private volatile CharSequence origin; - private final boolean disableSamplingMechanismValidation; /** - * This span's own propagation tags. Local child spans leave this null until something writes to - * it, since trace-level state is read from the local root. Access via {@link + * This span's own propagation tags, when they differ from the trace's. Only a non-root span that + * is the target of a write (e.g. {@code _dd.p.ts}) gets one. Access via {@link * #ownPropagationTags()}. */ private volatile PropagationTags propagationTags; @@ -439,7 +436,10 @@ public DDSpanContext( this.disableSamplingMechanismValidation = disableSamplingMechanismValidation; if (propagationTags != null) { propagationTags.updateTraceIdHighOrderBits(this.traceId.toHighOrderLong()); - this.propagationTags = propagationTags; + if (traceCollector.initPropagationTags(propagationTags) != propagationTags) { + // the trace already has propagation tags, so these are this span's own + this.propagationTags = propagationTags; + } } this.injectBaggageAsTags = injectBaggageAsTags; this.injectLinksAsTags = injectLinksAsTags; @@ -454,9 +454,18 @@ public DDSpanContext( } } - /** Returns this span's own propagation tags, creating an empty instance on first use. */ + /** Returns the trace's propagation tags, creating an empty instance on first use. */ + private PropagationTags tracePropagationTags() { + PropagationTags tags = traceCollector.getPropagationTagsIfPresent(); + return tags != null ? tags : traceCollector.initPropagationTags(newEmptyPropagationTags()); + } + + /** + * Returns this span's own propagation tags: the trace's for the local root, otherwise an instance + * created on first use. + */ private PropagationTags ownPropagationTags() { - PropagationTags tags = propagationTags; + PropagationTags tags = ownPropagationTagsIfPresent(); if (tags != null) { return tags; } @@ -464,6 +473,15 @@ private PropagationTags ownPropagationTags() { return PROPAGATION_TAGS.compareAndSet(this, null, created) ? created : propagationTags; } + /** This span's own propagation tags, or null if a non-root span has none yet. */ + private PropagationTags ownPropagationTagsIfPresent() { + PropagationTags tags = propagationTags; + if (tags != null) { + return tags; + } + return isLocalRoot() ? tracePropagationTags() : null; + } + @VisibleForTesting boolean hasOwnPropagationTags() { return propagationTags != null; @@ -708,12 +726,7 @@ public void forceKeep() { } public void forceKeep(byte samplingMechanism) { - // set trace level sampling priority - getRootSpanContextOrThis().forceKeepThisSpan(samplingMechanism); - } - - private void forceKeepThisSpan(byte samplingMechanism) { - ownPropagationTags().forceKeep(samplingMechanism); + tracePropagationTags().forceKeep(samplingMechanism); } public void addPropagatedTraceSource(final int value) { @@ -728,31 +741,10 @@ public void updateDebugPropagation(String value) { * @return if sampling priority was set by this method invocation */ public boolean setSamplingPriority(final int newPriority, final int newMechanism) { - DDSpanContext spanContext = getRootSpanContextOrThis(); - // set trace level sampling priority - return spanContext.setThisSpanSamplingPriority(newPriority, newMechanism); - } - - private DDSpanContext getRootSpanContextOrThis() { - DDSpanContext rootSpanContext = getRootSpanContextIfDifferent(); - return rootSpanContext != null ? rootSpanContext : this; - } - - private DDSpanContext getRootSpanContextIfDifferent() { - if (traceCollector != null) { - final DDSpan rootSpan = traceCollector.getRootSpan(); - if (null != rootSpan && rootSpan.spanContext() != this) { - return rootSpan.spanContext(); - } - } - return null; - } - - private boolean setThisSpanSamplingPriority(final int newPriority, final int newMechanism) { if (!validateSamplingPriority(newPriority, newMechanism)) { return false; } - PropagationTags tags = ownPropagationTags(); + PropagationTags tags = tracePropagationTags(); boolean updated = tags.tryUpdateTraceSamplingPriority( newPriority, @@ -771,18 +763,36 @@ private boolean setThisSpanSamplingPriority(final int newPriority, final int new return true; } + private DDSpanContext getRootSpanContextOrThis() { + DDSpanContext rootSpanContext = getRootSpanContextIfDifferent(); + return rootSpanContext != null ? rootSpanContext : this; + } + + /** Whether this is the local root span's context, or no root span has registered yet. */ + private boolean isLocalRoot() { + return getRootSpanContextIfDifferent() == null; + } + + private DDSpanContext getRootSpanContextIfDifferent() { + if (traceCollector != null) { + final DDSpan rootSpan = traceCollector.getRootSpan(); + if (null != rootSpan && rootSpan.spanContext() != this) { + return rootSpan.spanContext(); + } + } + return null; + } + public boolean setSamplingPriority( final int newPriority, final int newMechanism, final double sampleRate, final long traceIdLowOrderBits, final boolean rateLimiterRejected) { - DDSpanContext spanContext = getRootSpanContextOrThis(); - if (!spanContext.validateSamplingPriority(newPriority, newMechanism)) { + if (!validateSamplingPriority(newPriority, newMechanism)) { return false; } - return spanContext - .ownPropagationTags() + return tracePropagationTags() .tryUpdateProbabilitySamplingDecision( newPriority, newMechanism, @@ -820,7 +830,7 @@ private boolean validateSamplingPriority(final int newPriority, final int newMec @Override public int getSamplingPriority() { - return getRootSpanContextOrThis().ownPropagationTags().getSamplingPriority(); + return tracePropagationTags().getSamplingPriority(); } public void setSpanSamplingPriority(double rate, int limit) { @@ -847,16 +857,11 @@ public boolean lockSamplingPriority() { // this is now effectively a no-op - there is no locking. // the priority is just CAS'd against UNSET/UNKNOWN, unless it's forced to USER_KEEP/MANUAL // but is maintained for backwards compatibility, and returns false when it used to - final DDSpan rootSpan = traceCollector.getRootSpan(); - if (null != rootSpan && rootSpan.spanContext() != this) { - return rootSpan.spanContext().lockSamplingPriority(); - } - - return ownPropagationTags().getSamplingPriority() != PrioritySampling.UNSET; + return tracePropagationTags().getSamplingPriority() != PrioritySampling.UNSET; } public CharSequence getOrigin() { - return getRootSpanContextOrThis().origin; + return traceCollector.getOrigin(); } public void beginEndToEnd() { @@ -990,9 +995,8 @@ public String getSpanKindString() { } public void setOrigin(final CharSequence origin) { - DDSpanContext context = getRootSpanContextOrThis(); - context.origin = origin; - context.ownPropagationTags().updateTraceOrigin(origin); + traceCollector.setOrigin(origin); + tracePropagationTags().updateTraceOrigin(origin); } public void setMetric(final CharSequence key, final Number value) { @@ -1478,8 +1482,7 @@ public TagMap getTags() { tags.put(DDTags.THREAD_NAME, threadName.toString()); int currentSamplingPriority = getSamplingPriority(); // add _sample_rate tag only on the root/local span owning the decision - if (getRootSpanContextIfDifferent() == null - && currentSamplingPriority != PrioritySampling.UNSET) { + if (isLocalRoot() && currentSamplingPriority != PrioritySampling.UNSET) { tags.put(SAMPLE_RATE_KEY, currentSamplingPriority); } if (httpStatusCode != 0) { @@ -1544,7 +1547,7 @@ void processTagsAndBaggage( if ((options & ROOT_PROPAGATION_TAGS) != 0) { serializedPropagationTags = getPropagationTags(); } else if ((options & OWN_PROPAGATION_TAGS) != 0) { - PropagationTags own = propagationTags; + PropagationTags own = ownPropagationTagsIfPresent(); // Don't retain an instance just to serialize it serializedPropagationTags = own != null ? own : newEmptyPropagationTags(); } else { @@ -1771,7 +1774,7 @@ public TraceSegment getTraceSegment() { } public PropagationTags getPropagationTags() { - return getRootSpanContextOrThis().ownPropagationTags(); + return tracePropagationTags(); } @Override diff --git a/dd-trace-core/src/main/java/datadog/trace/core/TraceCollector.java b/dd-trace-core/src/main/java/datadog/trace/core/TraceCollector.java index 713f7769247..87ca4b9a956 100644 --- a/dd-trace-core/src/main/java/datadog/trace/core/TraceCollector.java +++ b/dd-trace-core/src/main/java/datadog/trace/core/TraceCollector.java @@ -12,6 +12,7 @@ import datadog.trace.bootstrap.instrumentation.api.AgentTraceCollector; import datadog.trace.bootstrap.instrumentation.api.ClientIpAddressData; import datadog.trace.common.sampling.PrioritySampler; +import datadog.trace.core.propagation.PropagationTags; import datadog.trace.util.TagsHelper; import java.io.Closeable; import java.io.IOException; @@ -63,6 +64,15 @@ enum PublishState { private static final AtomicReferenceFieldUpdater CI_VISIBILITY_DATA = AtomicReferenceFieldUpdater.newUpdater( TraceCollector.class, Object.class, "ciVisibilityContextData"); + private volatile PropagationTags propagationTags; + private static final AtomicReferenceFieldUpdater + PROPAGATION_TAGS = + AtomicReferenceFieldUpdater.newUpdater( + TraceCollector.class, PropagationTags.class, "propagationTags"); + + /** The origin of the trace. (eg. Synthetics, CI App) */ + private volatile CharSequence origin; + private volatile BlockResponseFunction blockResponseFunction; private volatile ClientIpAddressData clientIpAddressData; @@ -136,6 +146,30 @@ long getEndToEndStartTime() { abstract PublishState onPublish(final DDSpan span); + /** The trace's propagation tags, or null if no span has supplied them yet. */ + PropagationTags getPropagationTagsIfPresent() { + return propagationTags; + } + + /** + * Sets the trace's propagation tags unless already set. + * + * @return the trace's propagation tags, which may not be {@code propagationTags} + */ + PropagationTags initPropagationTags(PropagationTags propagationTags) { + return PROPAGATION_TAGS.compareAndSet(this, null, propagationTags) + ? propagationTags + : this.propagationTags; + } + + CharSequence getOrigin() { + return origin; + } + + void setOrigin(CharSequence origin) { + this.origin = origin; + } + /** * Records the request-scoped product data the trace was started with. Only the first non-null * value for each slot is kept. diff --git a/dd-trace-core/src/test/java/datadog/trace/common/writer/PayloadDispatcherImplTest.java b/dd-trace-core/src/test/java/datadog/trace/common/writer/PayloadDispatcherImplTest.java index cf756003e7c..574cea13611 100644 --- a/dd-trace-core/src/test/java/datadog/trace/common/writer/PayloadDispatcherImplTest.java +++ b/dd-trace-core/src/test/java/datadog/trace/common/writer/PayloadDispatcherImplTest.java @@ -29,8 +29,11 @@ import datadog.trace.test.util.DDJavaSpecification; import java.lang.reflect.Constructor; import java.nio.ByteBuffer; +import java.util.Arrays; import java.util.Collections; +import java.util.HashSet; import java.util.List; +import java.util.Set; import java.util.concurrent.TimeUnit; import java.util.concurrent.atomic.AtomicBoolean; import org.junit.jupiter.api.Test; @@ -39,6 +42,11 @@ class PayloadDispatcherImplTest extends DDJavaSpecification { + private static final Set TRACE_STATE_METHODS = + new HashSet<>( + Arrays.asList( + "getPropagationTagsIfPresent", "initPropagationTags", "getOrigin", "setOrigin")); + static final MonitoringImpl monitoring = new MonitoringImpl(StatsDClient.NO_OP, 1, TimeUnit.SECONDS); @@ -187,6 +195,10 @@ DDSpan realSpan() throws Exception { mock( PendingTrace.class, invocation -> { + if (TRACE_STATE_METHODS.contains(invocation.getMethod().getName())) { + // trace-level state the span context reads back from its collector + return invocation.callRealMethod(); + } Class returnType = invocation.getMethod().getReturnType(); if (CoreTracer.class.isAssignableFrom(returnType)) { // Use RETURNS_DEFAULTS so getTagInterceptor() returns null (matching Groovy Stub diff --git a/dd-trace-core/src/test/java/datadog/trace/core/DDSpanTest.java b/dd-trace-core/src/test/java/datadog/trace/core/DDSpanTest.java index 1d252859739..c82eb0d07d2 100644 --- a/dd-trace-core/src/test/java/datadog/trace/core/DDSpanTest.java +++ b/dd-trace-core/src/test/java/datadog/trace/core/DDSpanTest.java @@ -39,7 +39,6 @@ import datadog.trace.core.util.TestThrowables; import java.io.Closeable; import java.io.IOException; -import java.lang.reflect.Field; import java.util.Collection; import java.util.Collections; import java.util.HashMap; @@ -276,7 +275,7 @@ void prioritySamplingMetricSetOnlyOnRootSpan() { assertFalse(child2.hasSamplingPriority()); } - static Stream originSetOnlyOnRootSpanArguments() { + static Stream originIsSharedByLocalTraceArguments() { return Stream.of( arguments("TagContext", new TagContext("some-origin", TagMap.fromMap(new HashMap<>()))), arguments( @@ -291,9 +290,8 @@ static Stream originSetOnlyOnRootSpanArguments() { } @ParameterizedTest - @MethodSource("originSetOnlyOnRootSpanArguments") - void originSetOnlyOnRootSpan(String scenario, AgentSpanContext extractedContext) - throws Exception { + @MethodSource("originIsSharedByLocalTraceArguments") + void originIsSharedByLocalTrace(String scenario, AgentSpanContext extractedContext) { DDSpanContext parent = (DDSpanContext) tracer @@ -305,14 +303,10 @@ void originSetOnlyOnRootSpan(String scenario, AgentSpanContext extractedContext) (DDSpanContext) tracer.buildSpan("datadog", "testChild1").asChildOf(parent).start().spanContext(); + assertSame(parent.getTraceCollector(), child.getTraceCollector()); + assertEquals("some-origin", parent.getTraceCollector().getOrigin().toString()); assertEquals("some-origin", parent.getOrigin().toString()); - // Access field directly instead of getter. - Field originField = DDSpanContext.class.getDeclaredField("origin"); - originField.setAccessible(true); - assertEquals("some-origin", originField.get(parent).toString()); assertEquals("some-origin", child.getOrigin().toString()); - // Access field directly instead of getter. - assertNull(originField.get(child)); } static Stream isRootSpanArguments() {