diff --git a/dd-java-agent/instrumentation-testing/src/main/groovy/datadog/trace/agent/test/ValidatingRequestContextDecorator.groovy b/dd-java-agent/instrumentation-testing/src/main/groovy/datadog/trace/agent/test/ValidatingRequestContextDecorator.groovy index 9e4d1436182..66ee4cca217 100644 --- a/dd-java-agent/instrumentation-testing/src/main/groovy/datadog/trace/agent/test/ValidatingRequestContextDecorator.groovy +++ b/dd-java-agent/instrumentation-testing/src/main/groovy/datadog/trace/agent/test/ValidatingRequestContextDecorator.groovy @@ -87,24 +87,12 @@ class ValidatingRequestContextDecorator implements RequestContext { delegate.setTagTop(key, value, sanitize) } - @Override - void setTagCurrent(String key, Object value, boolean sanitize) { - check() - delegate.setTagCurrent(key, value, sanitize) - } - @Override Object getTagTop(String key, boolean sanitize) { check() return delegate.getTagTop(key, sanitize) } - @Override - Object getTagCurrent(String key, boolean sanitize) { - check() - return delegate.getTagCurrent(key, sanitize) - } - @Override void setDataTop(String key, Object value) { check() @@ -123,28 +111,10 @@ class ValidatingRequestContextDecorator implements RequestContext { delegate.setMetaStructTop(key, value) } - @Override - void setMetaStructCurrent(String key, Object value) { - check() - delegate.setMetaStructCurrent(key, value) - } - @Override void effectivelyBlocked() { check() delegate.effectivelyBlocked() } - - @Override - void setDataCurrent(String key, Object value) { - check() - delegate.setDataCurrent(key, value) - } - - @Override - Object getDataCurrent(String key) { - check() - return delegate.getDataCurrent(key) - } } } diff --git a/dd-trace-api/src/main/java/datadog/trace/api/internal/TraceSegment.java b/dd-trace-api/src/main/java/datadog/trace/api/internal/TraceSegment.java index d6e262cfc78..437fc85024c 100644 --- a/dd-trace-api/src/main/java/datadog/trace/api/internal/TraceSegment.java +++ b/dd-trace-api/src/main/java/datadog/trace/api/internal/TraceSegment.java @@ -2,9 +2,8 @@ /** * A {@code TraceSegment} represents the local, i.e. in the scope of this {@code Tracer}, part of a - * a {@code Trace}. It can consist of multiple spans, and the {@code TraceSegment} instance can only - * be used in relation to the {@code Span} from which it was acquired. This means that the {@code - * TraceSegment} should not be stored and used at a later time, or passed around to async functions. + * {@code Trace}. It can consist of multiple spans. Operations act on the segment as a whole, and + * "top" refers to the segment's local root span. */ public interface TraceSegment { @@ -44,43 +43,6 @@ default Object getTagTop(String key) { */ Object getTagTop(String key, boolean sanitize); - /** - * Add a tag to the current span in this {@code TraceSegment}. - * - * @param key key of the tag - * @param value value of the tag - */ - default void setTagCurrent(String key, Object value) { - setTagCurrent(key, value, false); - } - - /** - * Add a tag to the current span in this {@code TraceSegment}. with optional key sanitization. - * - * @param key key of the tag - * @param value value of the tag - * @param sanitize indicates is key need to be sanitized - */ - void setTagCurrent(String key, Object value, boolean sanitize); - - /** - * Get the tag value from the current span in this {@code TraceSegment}. - * - * @param key key of the tag - */ - default Object getTagCurrent(String key) { - return getTagCurrent(key, false); - } - - /** - * Get the tag value from the current span in this {@code TraceSegment}. with optional key - * sanitization. - * - * @param key key of the tag - * @param sanitize indicates is key need to be sanitized - */ - Object getTagCurrent(String key, boolean sanitize); - /** * Add data to the top of this {@code TraceSegment}. The {@code toString} representation of the * {@code value} must be valid top level JSON, i.e. an {@code Object} or an {@code Array}. @@ -101,41 +63,15 @@ default Object getTagCurrent(String key) { /** Mark the request as effectively blocked, by setting the tag appsec.blocked */ void effectivelyBlocked(); - /** - * Add data to the current span in this {@code TraceSegment}. The {@code toString} representation - * of the {@code value} must be valid top level JSON, i.e. an {@code Object} or an {@code Array}. - * - * @param key key of the data - * @param value value of the data - */ - void setDataCurrent(String key, Object value); - - /** - * Gets the value of the tag from the current span in this {@code TraceSegment}. - * - * @param key key of the data - * @return value of the data - */ - Object getDataCurrent(String key); - /** * Add a field to the meta_struct of the top of this {@code TraceSegment}. * * @param field field name * @param value value of the data - * @see #setMetaStructCurrent(String, Object) (String, Object) - */ - void setMetaStructTop(String field, Object value); - - /** - * Add a field to the current span meta_struct in this {@code TraceSegment}. - * - * @param field field name - * @param value value of the data * @see "datadog.trace.common.writer.ddagent.TraceMapperV0_4.MetaStructWriter" * @see "datadog.trace.core.CoreSpan#setMetaStruct(String, Object)" */ - void setMetaStructCurrent(String field, Object value); + void setMetaStructTop(String field, Object value); class NoOp implements TraceSegment { public static final TraceSegment INSTANCE = new NoOp(); @@ -145,19 +81,11 @@ private NoOp() {} @Override public void setTagTop(String key, Object value, boolean sanitize) {} - @Override - public void setTagCurrent(String key, Object value, boolean sanitize) {} - @Override public Object getTagTop(String key, boolean sanitize) { return null; } - @Override - public Object getTagCurrent(String key, boolean sanitize) { - return null; - } - @Override public void setDataTop(String key, Object value) {} @@ -169,18 +97,7 @@ public Object getDataTop(String key) { @Override public void effectivelyBlocked() {} - @Override - public void setDataCurrent(String key, Object value) {} - - @Override - public Object getDataCurrent(String key) { - return null; - } - @Override public void setMetaStructTop(String key, Object value) {} - - @Override - public void setMetaStructCurrent(String key, Object value) {} } } 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 2f10cceba4b..9c691aa367e 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 @@ -2166,17 +2166,10 @@ protected static final DDSpanContext buildSpanContext( if (serviceName == null) { serviceName = parentServiceName; } - RequestContext requestContext = - ((DDSpanContext) resolvedParentSpanContext).getRequestContext(); - if (requestContext != null) { - requestContextDataAppSec = requestContext.getData(RequestContextSlot.APPSEC); - requestContextDataIast = requestContext.getData(RequestContextSlot.IAST); - ciVisibilityContextData = requestContext.getData(RequestContextSlot.CI_VISIBILITY); - } else { - requestContextDataAppSec = null; - requestContextDataIast = null; - ciVisibilityContextData = null; - } + // Request-scoped product data lives on the shared trace collector + requestContextDataAppSec = null; + requestContextDataIast = null; + ciVisibilityContextData = null; propagationTags = tracer.propagationTagsFactory.empty(); } else { long endToEndStartTime; 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 1e236d3a120..d0f71b4ab01 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 @@ -23,9 +23,7 @@ import datadog.trace.api.cache.DDCaches; import datadog.trace.api.config.TracerConfig; import datadog.trace.api.datastreams.PathwayContext; -import datadog.trace.api.gateway.BlockResponseFunction; import datadog.trace.api.gateway.RequestContext; -import datadog.trace.api.gateway.RequestContextSlot; import datadog.trace.api.internal.TraceSegment; import datadog.trace.api.llmobs.LLMObsPropagationValues; import datadog.trace.api.sampling.PrioritySampling; @@ -33,7 +31,6 @@ import datadog.trace.bootstrap.instrumentation.api.AgentSpanContext; import datadog.trace.bootstrap.instrumentation.api.AppendableSpanLinks; import datadog.trace.bootstrap.instrumentation.api.Baggage; -import datadog.trace.bootstrap.instrumentation.api.ClientIpAddressData; import datadog.trace.bootstrap.instrumentation.api.ProfilerContext; import datadog.trace.bootstrap.instrumentation.api.ProfilingContextIntegration; import datadog.trace.bootstrap.instrumentation.api.ResourceNamePriorities; @@ -43,9 +40,6 @@ import datadog.trace.core.propagation.PropagationTags; import datadog.trace.core.taginterceptor.TagInterceptor; import datadog.trace.core.tagprocessor.TagsPostProcessorFactory; -import datadog.trace.util.TagsHelper; -import java.io.Closeable; -import java.io.IOException; import java.util.Collections; import java.util.HashMap; import java.util.List; @@ -67,8 +61,7 @@ * across Span boundaries and (2) any Datadog fields that are needed to identify or contextualize * the associated Span instance */ -public class DDSpanContext - implements AgentSpanContext, RequestContext, TraceSegment, ProfilerContext { +public class DDSpanContext implements AgentSpanContext, ProfilerContext { private static final Logger log = LoggerFactory.getLogger(DDSpanContext.class); public static final String PRIORITY_SAMPLING_KEY = "_sampling_priority_v1"; @@ -174,22 +167,12 @@ public class DDSpanContext /** The origin of the trace. (eg. Synthetics, CI App) */ private volatile CharSequence origin; - /** RequestContext data for the InstrumentationGateway */ - private final Object requestContextDataAppSec; - - private final Object requestContextDataIast; - private final Object ciVisibilityContextData; - private final boolean disableSamplingMechanismValidation; private final PropagationTags propagationTags; private volatile PathwayContext pathwayContext; - private volatile BlockResponseFunction blockResponseFunction; - - private volatile ClientIpAddressData clientIpAddressData; - private final ProfilingContextIntegration profilingContextIntegration; private final boolean injectBaggageAsTags; private final boolean injectLinksAsTags; @@ -404,9 +387,12 @@ public DDSpanContext( } this.w3cBaggage = w3cBaggage; - this.requestContextDataAppSec = requestContextDataAppSec; - this.requestContextDataIast = requestContextDataIast; - this.ciVisibilityContextData = CiVisibilityContextData; + if (requestContextDataAppSec != null + || requestContextDataIast != null + || CiVisibilityContextData != null) { + traceCollector.initRequestContextData( + requestContextDataAppSec, requestContextDataIast, CiVisibilityContextData); + } assert pathwayContext != null; this.pathwayContext = pathwayContext; @@ -877,7 +863,7 @@ public TraceCollector getTraceCollector() { } public RequestContext getRequestContext() { - return this; + return traceCollector; } @Override @@ -1740,68 +1726,8 @@ public String toString() { return s.toString(); } - /** RequestContext Implementation */ - @Override - public Object getData(RequestContextSlot slot) { - if (slot == RequestContextSlot.APPSEC) { - return this.requestContextDataAppSec; - } else if (slot == RequestContextSlot.CI_VISIBILITY) { - return this.ciVisibilityContextData; - } else if (slot == RequestContextSlot.IAST) { - return this.requestContextDataIast; - } - return null; - } - - @Override - public void close() throws IOException { - Exception exc = null; - if (this.requestContextDataAppSec instanceof Closeable) { - try { - ((Closeable) this.requestContextDataAppSec).close(); - } catch (IOException | RuntimeException e) { - exc = e; - } - } - if (this.requestContextDataIast instanceof Closeable) { - try { - ((Closeable) this.requestContextDataIast).close(); - } catch (IOException | RuntimeException e) { - exc = e; - } - } - if (exc != null) { - if (exc instanceof RuntimeException) { - throw (RuntimeException) exc; - } else { - throw (IOException) exc; - } - } - } - - @Override public TraceSegment getTraceSegment() { - return this; - } - - @Override - public void setBlockResponseFunction(BlockResponseFunction blockResponseFunction) { - getRootSpanContextOrThis().blockResponseFunction = blockResponseFunction; - } - - @Override - public BlockResponseFunction getBlockResponseFunction() { - return getRootSpanContextOrThis().blockResponseFunction; - } - - @Override - public void setClientIpAddressData(final ClientIpAddressData clientIpAddressData) { - getRootSpanContextOrThis().clientIpAddressData = clientIpAddressData; - } - - @Override - public ClientIpAddressData getClientIpAddressData() { - return getRootSpanContextOrThis().clientIpAddressData; + return traceCollector; } public PropagationTags getPropagationTags() { @@ -1813,82 +1739,14 @@ public LLMObsPropagationValues getExtractedLLMObsValues() { return getPropagationTags().getExtractedLLMObsValues(); } - /** TraceSegment Implementation */ - @Override - public void setTagTop(String key, Object value, boolean sanitize) { - getRootSpanContextOrThis().setTagCurrent(key, value, sanitize); - } - - @Override - public Object getTagTop(String key, boolean sanitize) { - return getRootSpanContextOrThis().getTagCurrent(key, sanitize); - } - - @Override - public void setTagCurrent(String key, Object value, boolean sanitize) { - if (sanitize) { - key = TagsHelper.sanitize(key); - } - this.setTag(key, value); - } - - @Override - public Object getTagCurrent(String key, boolean sanitize) { - if (sanitize) { - key = TagsHelper.sanitize(key); - } - return this.getTag(key); - } - - @Override - public void setDataTop(String key, Object value) { - getRootSpanContextOrThis().setDataCurrent(key, value); - } - - @Override - public Object getDataTop(String key) { - return getRootSpanContextOrThis().getDataCurrent(key); - } - - @Override - public void effectivelyBlocked() { - setTagTop("appsec.blocked", "true"); - } - - @Override - public void setDataCurrent(String key, Object value) { - this.setTag(getTagName(key), value); - } - - @Override - public Object getDataCurrent(String key) { - return this.getTag(getTagName(key)); - } - - private String getTagName(String key) { - // TODO is this decided? - return "_dd." + key + ".json"; - } - - @Override - public void setMetaStructTop(String field, Object value) { - getRootSpanContextOrThis().setMetaStructCurrent(field, value); - } - + /** Returns the meta_struct field for {@code key}, creating it if absent. */ @SuppressWarnings("unchecked") - @Override - public T getOrCreateMetaStructTop(final String key, final Function defaultValue) { - final DDSpanContext top = getRootSpanContextOrThis(); - if (top.metaStruct == EMPTY_META_STRUCT) { + T getOrCreateMetaStruct(final String key, final Function defaultValue) { + if (metaStruct == EMPTY_META_STRUCT) { // this will safely create the metastruct if the field has not been initialized already - top.setMetaStruct(key, defaultValue.apply(key)); + setMetaStruct(key, defaultValue.apply(key)); } - return (T) top.metaStruct.computeIfAbsent(key, defaultValue); - } - - @Override - public void setMetaStructCurrent(String field, Object value) { - setMetaStruct(field, value); + return (T) metaStruct.computeIfAbsent(key, defaultValue); } @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 4cb9a3edb42..713f7769247 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 @@ -3,14 +3,28 @@ import datadog.trace.api.Config; import datadog.trace.api.DDTraceId; import datadog.trace.api.ProductTraceSource; +import datadog.trace.api.gateway.BlockResponseFunction; +import datadog.trace.api.gateway.RequestContext; +import datadog.trace.api.gateway.RequestContextSlot; +import datadog.trace.api.internal.TraceSegment; import datadog.trace.api.sampling.PrioritySampling; import datadog.trace.api.time.TimeSource; import datadog.trace.bootstrap.instrumentation.api.AgentTraceCollector; +import datadog.trace.bootstrap.instrumentation.api.ClientIpAddressData; import datadog.trace.common.sampling.PrioritySampler; +import datadog.trace.util.TagsHelper; +import java.io.Closeable; +import java.io.IOException; import java.util.concurrent.atomic.AtomicLongFieldUpdater; +import java.util.concurrent.atomic.AtomicReferenceFieldUpdater; +import java.util.function.Function; import javax.annotation.Nonnull; -public abstract class TraceCollector implements AgentTraceCollector { +/** + * The local part of a trace: the spans of one trace within this tracer. Holds trace-level state + * shared by those spans, and is their {@link TraceSegment} and {@link RequestContext}. + */ +public abstract class TraceCollector implements AgentTraceCollector, TraceSegment, RequestContext { interface Factory { /** Used by tests and benchmarks. */ @@ -35,6 +49,23 @@ enum PublishState { private static final AtomicLongFieldUpdater END_TO_END_START_TIME = AtomicLongFieldUpdater.newUpdater(TraceCollector.class, "endToEndStartTime"); + private volatile Object requestContextDataAppSec; + private static final AtomicReferenceFieldUpdater REQUEST_DATA_APPSEC = + AtomicReferenceFieldUpdater.newUpdater( + TraceCollector.class, Object.class, "requestContextDataAppSec"); + + private volatile Object requestContextDataIast; + private static final AtomicReferenceFieldUpdater REQUEST_DATA_IAST = + AtomicReferenceFieldUpdater.newUpdater( + TraceCollector.class, Object.class, "requestContextDataIast"); + + private volatile Object ciVisibilityContextData; + private static final AtomicReferenceFieldUpdater CI_VISIBILITY_DATA = + AtomicReferenceFieldUpdater.newUpdater( + TraceCollector.class, Object.class, "ciVisibilityContextData"); + private volatile BlockResponseFunction blockResponseFunction; + private volatile ClientIpAddressData clientIpAddressData; + protected TraceCollector( CoreTracer tracer, CoreTracer.ConfigSnapshot traceConfig, TimeSource timeSource) { this.tracer = tracer; @@ -104,4 +135,147 @@ long getEndToEndStartTime() { abstract DDSpan getRootSpan(); abstract PublishState onPublish(final DDSpan span); + + /** + * Records the request-scoped product data the trace was started with. Only the first non-null + * value for each slot is kept. + */ + void initRequestContextData( + Object requestContextDataAppSec, + Object requestContextDataIast, + Object ciVisibilityContextData) { + if (requestContextDataAppSec != null) { + REQUEST_DATA_APPSEC.compareAndSet(this, null, requestContextDataAppSec); + } + if (requestContextDataIast != null) { + REQUEST_DATA_IAST.compareAndSet(this, null, requestContextDataIast); + } + if (ciVisibilityContextData != null) { + CI_VISIBILITY_DATA.compareAndSet(this, null, ciVisibilityContextData); + } + } + + /** RequestContext Implementation */ + @SuppressWarnings("unchecked") + @Override + public T getData(RequestContextSlot slot) { + if (slot == RequestContextSlot.APPSEC) { + return (T) requestContextDataAppSec; + } else if (slot == RequestContextSlot.CI_VISIBILITY) { + return (T) ciVisibilityContextData; + } else if (slot == RequestContextSlot.IAST) { + return (T) requestContextDataIast; + } + return null; + } + + @Override + public TraceSegment getTraceSegment() { + return this; + } + + @Override + public void setBlockResponseFunction(BlockResponseFunction blockResponseFunction) { + this.blockResponseFunction = blockResponseFunction; + } + + @Override + public BlockResponseFunction getBlockResponseFunction() { + return blockResponseFunction; + } + + @Override + public void setClientIpAddressData(ClientIpAddressData clientIpAddressData) { + this.clientIpAddressData = clientIpAddressData; + } + + @Override + public ClientIpAddressData getClientIpAddressData() { + return clientIpAddressData; + } + + @Override + public T getOrCreateMetaStructTop(String key, Function defaultValue) { + DDSpanContext top = rootSpanContext(); + return top != null ? top.getOrCreateMetaStruct(key, defaultValue) : null; + } + + @Override + public void close() throws IOException { + Exception exc = null; + if (requestContextDataAppSec instanceof Closeable) { + try { + ((Closeable) requestContextDataAppSec).close(); + } catch (IOException | RuntimeException e) { + exc = e; + } + } + if (requestContextDataIast instanceof Closeable) { + try { + ((Closeable) requestContextDataIast).close(); + } catch (IOException | RuntimeException e) { + exc = e; + } + } + if (exc != null) { + if (exc instanceof RuntimeException) { + throw (RuntimeException) exc; + } else { + throw (IOException) exc; + } + } + } + + /** TraceSegment Implementation */ + @Override + public void setTagTop(String key, Object value, boolean sanitize) { + DDSpanContext top = rootSpanContext(); + if (top != null) { + top.setTag(sanitize ? TagsHelper.sanitize(key) : key, value); + } + } + + @Override + public Object getTagTop(String key, boolean sanitize) { + DDSpanContext top = rootSpanContext(); + return top != null ? top.getTag(sanitize ? TagsHelper.sanitize(key) : key) : null; + } + + @Override + public void setDataTop(String key, Object value) { + DDSpanContext top = rootSpanContext(); + if (top != null) { + top.setTag(dataTagName(key), value); + } + } + + @Override + public Object getDataTop(String key) { + DDSpanContext top = rootSpanContext(); + return top != null ? top.getTag(dataTagName(key)) : null; + } + + @Override + public void effectivelyBlocked() { + setTagTop("appsec.blocked", "true"); + } + + @Override + public void setMetaStructTop(String field, Object value) { + DDSpanContext top = rootSpanContext(); + if (top != null) { + top.setMetaStruct(field, value); + } + } + + private static String dataTagName(String key) { + // TODO is this decided? + return "_dd." + key + ".json"; + } + + /** The local root span's context, or null before the first span registers. */ + private DDSpanContext rootSpanContext() { + DDSpan rootSpan = getRootSpan(); + return rootSpan != null ? rootSpan.spanContext() : null; + } } diff --git a/dd-trace-core/src/test/java/datadog/trace/common/writer/TraceMapperTest.java b/dd-trace-core/src/test/java/datadog/trace/common/writer/TraceMapperTest.java index a88d7d9000f..d1b669edc7a 100644 --- a/dd-trace-core/src/test/java/datadog/trace/common/writer/TraceMapperTest.java +++ b/dd-trace-core/src/test/java/datadog/trace/common/writer/TraceMapperTest.java @@ -35,7 +35,7 @@ void testTraceMapperV05() throws Exception { .withTag("elasticsearch.version", "7.0") .start(); span.setBaggageItem("baggage", "item"); - span.spanContext().setDataTop("mydata", "[1,2,3]"); + span.spanContext().getTraceSegment().setDataTop("mydata", "[1,2,3]"); List trace = Collections.singletonList(span); TraceMapperV0_5 traceMapper = new TraceMapperV0_5(); diff --git a/dd-trace-core/src/test/java/datadog/trace/core/CoreSpanBuilderTest.java b/dd-trace-core/src/test/java/datadog/trace/core/CoreSpanBuilderTest.java index 425c3616be9..aa010e62df9 100644 --- a/dd-trace-core/src/test/java/datadog/trace/core/CoreSpanBuilderTest.java +++ b/dd-trace-core/src/test/java/datadog/trace/core/CoreSpanBuilderTest.java @@ -22,6 +22,7 @@ import static org.junit.jupiter.api.Assertions.assertFalse; import static org.junit.jupiter.api.Assertions.assertNotEquals; import static org.junit.jupiter.api.Assertions.assertNull; +import static org.junit.jupiter.api.Assertions.assertSame; import static org.junit.jupiter.api.Assertions.assertTrue; import static org.junit.jupiter.params.provider.Arguments.arguments; import static org.mockito.Mockito.mock; @@ -558,7 +559,7 @@ void globalSpanTagsPopulatedOnEachSpan(String tagString, Map tag } @Test - void canOverwriteRequestContextDataWithBuilderFromEmpty() { + void requestContextDataFromBuilderFillsEmptyTrace() { AgentSpan span1 = tracer.startSpan("test", "span1"); assertNull(span1.getRequestContext().getData(RequestContextSlot.APPSEC)); @@ -574,16 +575,17 @@ void canOverwriteRequestContextDataWithBuilderFromEmpty() { .withRequestContextData(RequestContextSlot.IAST, "override") .start(); - assertEquals("override", span2.getRequestContext().getData(RequestContextSlot.APPSEC)); - assertEquals("override", span2.getRequestContext().getData(RequestContextSlot.CI_VISIBILITY)); - assertEquals("override", span2.getRequestContext().getData(RequestContextSlot.IAST)); + assertSame(span1.getRequestContext(), span2.getRequestContext()); + assertEquals("override", span1.getRequestContext().getData(RequestContextSlot.APPSEC)); + assertEquals("override", span1.getRequestContext().getData(RequestContextSlot.CI_VISIBILITY)); + assertEquals("override", span1.getRequestContext().getData(RequestContextSlot.IAST)); span2.finish(); span1.finish(); } @Test - void canOverwriteRequestContextDataWithBuilder() { + void requestContextDataIsSharedByLocalTrace() { TagContext context = new TagContext() .withCiVisibilityContextData("value") @@ -606,15 +608,33 @@ void canOverwriteRequestContextDataWithBuilder() { .withRequestContextData(RequestContextSlot.IAST, "override") .start(); - assertEquals("override", span3.getRequestContext().getData(RequestContextSlot.APPSEC)); - assertEquals("override", span3.getRequestContext().getData(RequestContextSlot.CI_VISIBILITY)); - assertEquals("override", span3.getRequestContext().getData(RequestContextSlot.IAST)); + assertEquals("value", span3.getRequestContext().getData(RequestContextSlot.APPSEC)); + assertEquals("value", span3.getRequestContext().getData(RequestContextSlot.CI_VISIBILITY)); + assertEquals("value", span3.getRequestContext().getData(RequestContextSlot.IAST)); span3.finish(); span2.finish(); span1.finish(); } + @Test + void requestContextDataFromBuilderStartsNewTrace() { + AgentSpan span = + tracer + .buildSpan("test", "span") + .asChildOf(new TagContext()) + .withRequestContextData(RequestContextSlot.APPSEC, "appsec") + .withRequestContextData(RequestContextSlot.CI_VISIBILITY, "civis") + .withRequestContextData(RequestContextSlot.IAST, "iast") + .start(); + + assertEquals("appsec", span.getRequestContext().getData(RequestContextSlot.APPSEC)); + assertEquals("civis", span.getRequestContext().getData(RequestContextSlot.CI_VISIBILITY)); + assertEquals("iast", span.getRequestContext().getData(RequestContextSlot.IAST)); + + span.finish(); + } + private Map productTags() { Map productTags = new HashMap<>(); productTags.put(PROFILING_ENABLED, Config.get().isProfilingEnabled() ? 1 : 0); diff --git a/dd-trace-core/src/test/java/datadog/trace/core/DDSpanContextTest.java b/dd-trace-core/src/test/java/datadog/trace/core/DDSpanContextTest.java index 221d350f032..e5e92124333 100644 --- a/dd-trace-core/src/test/java/datadog/trace/core/DDSpanContextTest.java +++ b/dd-trace-core/src/test/java/datadog/trace/core/DDSpanContextTest.java @@ -25,6 +25,7 @@ import static org.junit.jupiter.api.Assertions.assertFalse; import static org.junit.jupiter.api.Assertions.assertInstanceOf; import static org.junit.jupiter.api.Assertions.assertNull; +import static org.junit.jupiter.api.Assertions.assertSame; import static org.junit.jupiter.api.Assertions.assertTrue; import static org.mockito.Mockito.mock; import static org.mockito.Mockito.times; @@ -242,7 +243,7 @@ void forceKeepReallyKeepsTrace() { } @Test - void setTraceSegmentTagsAndDataOnCorrectSpan() { + void setTraceSegmentTagsAndDataOnTopSpan() { ExtractedContext extracted = (ExtractedContext) new ExtractedContext( @@ -265,26 +266,19 @@ void setTraceSegmentTagsAndDataOnCorrectSpan() { currentTS.setDataTop("ctd", "[1]"); currentTS.setTagTop("ctt", "t1"); - currentTS.setDataCurrent("ccd", "[2]"); - currentTS.setTagCurrent("cct", "t2"); topTS.setDataTop("ttd", "[3]"); topTS.setTagTop("ttt", "t3"); - topTS.setDataCurrent("tcd", "[4]"); - topTS.setTagCurrent("tct", "t4"); + + assertSame(topTS, currentTS); Map expectedTopTags = new HashMap<>(); expectedTopTags.put(dataTag("ctd"), "[1]"); expectedTopTags.put("ctt", "t1"); expectedTopTags.put(dataTag("ttd"), "[3]"); expectedTopTags.put("ttt", "t3"); - expectedTopTags.put(dataTag("tcd"), "[4]"); - expectedTopTags.put("tct", "t4"); assertTagmap(topC.getTags(), expectedTopTags, true); - Map expectedCurrentTags = new HashMap<>(); - expectedCurrentTags.put(dataTag("ccd"), "[2]"); - expectedCurrentTags.put("cct", "t2"); - assertTagmap(currentC.getTags(), expectedCurrentTags, true); + assertTagmap(currentC.getTags(), new HashMap<>(), true); current.finish(); top.finish(); diff --git a/dd-trace-core/src/test/java/datadog/trace/core/PendingTraceBufferTest.java b/dd-trace-core/src/test/java/datadog/trace/core/PendingTraceBufferTest.java index e72e5fa3a7c..0c5baa1e967 100644 --- a/dd-trace-core/src/test/java/datadog/trace/core/PendingTraceBufferTest.java +++ b/dd-trace-core/src/test/java/datadog/trace/core/PendingTraceBufferTest.java @@ -32,6 +32,7 @@ import datadog.trace.api.DDTraceId; import datadog.trace.api.datastreams.NoopPathwayContext; import datadog.trace.api.flare.TracerFlare; +import datadog.trace.api.time.ControllableTimeSource; import datadog.trace.api.time.SystemTimeSource; import datadog.trace.core.monitor.HealthMetrics; import datadog.trace.core.propagation.PropagationTags; @@ -426,82 +427,91 @@ void testingTracerFlareDumpWithMultipleTraces() throws IOException, InterruptedE TracerFlare.Reporter dumpReporter = mock(TracerFlare.Reporter.class); TracerFlare.addReporter(dumpReporter); - PendingTrace trace1 = factory.create(DDTraceId.ONE); - DDSpan parent1 = newSpanOf(trace1, UNSET, System.currentTimeMillis() * 1000); - DDSpan child1 = newSpanOf(parent1); - PendingTrace trace2 = factory.create(DDTraceId.from(2)); - DDSpan parent2 = newSpanOf(trace2, UNSET, System.currentTimeMillis() * 2000); - DDSpan child2 = newSpanOf(parent2); - // first flare dump with two traces - parent1.finish(); - parent2.finish(); - buffer.start(); - Map entries1 = buildAndExtractZip(); - - verify(dumpReporter).prepareForFlare(); - verify(dumpReporter).addReportToFlare(any()); - verify(dumpReporter).cleanupAfterFlare(); - - assertEquals(1, entries1.size()); - String pendingTraceText1 = (String) entries1.get("pending_traces.txt"); - assertTrue( - pendingTraceText1.startsWith( - "[{\"service\":\"fakeService\",\"name\":\"fakeOperation\",\"resource\":\"fakeResource\",\"trace_id\":1,\"span_id\":1,\"parent_id\":0")); - - List> parsedTraces1 = parseTraceLines(pendingTraceText1); - assertEquals(2, parsedTraces1.size()); - assertEquals(1L, ((Number) parsedTraces1.get(0).get("trace_id")).longValue()); - assertEquals(2L, ((Number) parsedTraces1.get(1).get("trace_id")).longValue()); - assertTrue( - ((Number) parsedTraces1.get(0).get("start")).longValue() - < ((Number) parsedTraces1.get(1).get("start")).longValue()); - - clearInvocations(dumpReporter); - - // New pending traces are needed here because generating the first flare takes long enough that - // the - // earlier pending traces are flushed (within 500ms). - // second flare dump with new pending traces - // Finish the first set of traces - child1.finish(); - child2.finish(); - // Create new pending traces - PendingTrace trace3 = factory.create(DDTraceId.from(3)); - DDSpan parent3 = newSpanOf(trace3, UNSET, System.currentTimeMillis() * 3000); - DDSpan child3 = newSpanOf(parent3); - PendingTrace trace4 = factory.create(DDTraceId.from(4)); - DDSpan parent4 = newSpanOf(trace4, UNSET, System.currentTimeMillis() * 4000); - DDSpan child4 = newSpanOf(parent4); - parent3.finish(); - parent4.finish(); - Map entries2 = buildAndExtractZip(); - - verify(dumpReporter).prepareForFlare(); - verify(dumpReporter).addReportToFlare(any()); - verify(dumpReporter).cleanupAfterFlare(); - - assertEquals(1, entries2.size()); - String pendingTraceText2 = (String) entries2.get("pending_traces.txt"); - List> parsedTraces2 = parseTraceLines(pendingTraceText2); - assertEquals(2, parsedTraces2.size()); - - child3.finish(); - child4.finish(); + // The buffer writes out a trace untouched for 500ms, so with a real clock a dump only sees the + // traces if it runs sooner, which every other registered reporter can delay. A clock held still + // keeps them pending until each dump is done. + ControllableTimeSource timeSource = new ControllableTimeSource(); + PendingTraceBuffer flareBuffer = + PendingTraceBuffer.delaying(timeSource, mock(Config.class), null, null); + PendingTrace.Factory flareFactory = + new PendingTrace.Factory(tracer, flareBuffer, timeSource, false, HealthMetrics.NO_OP); + try { + PendingTrace trace1 = flareFactory.create(DDTraceId.ONE); + DDSpan parent1 = newSpanOf(trace1, UNSET, System.currentTimeMillis() * 1000); + DDSpan child1 = newSpanOf(parent1); + PendingTrace trace2 = flareFactory.create(DDTraceId.from(2)); + DDSpan parent2 = newSpanOf(trace2, UNSET, System.currentTimeMillis() * 2000); + DDSpan child2 = newSpanOf(parent2); + // first flare dump with two traces + parent1.finish(); + parent2.finish(); + flareBuffer.start(); + Map entries1 = buildAndExtractZip(); + + verify(dumpReporter).prepareForFlare(); + verify(dumpReporter).addReportToFlare(any()); + verify(dumpReporter).cleanupAfterFlare(); + + assertEquals(1, entries1.size()); + String pendingTraceText1 = (String) entries1.get("pending_traces.txt"); + assertTrue( + pendingTraceText1.startsWith( + "[{\"service\":\"fakeService\",\"name\":\"fakeOperation\",\"resource\":\"fakeResource\",\"trace_id\":1,\"span_id\":1,\"parent_id\":0")); + + List> parsedTraces1 = parseTraceLines(pendingTraceText1); + assertEquals(2, parsedTraces1.size()); + assertEquals(1L, ((Number) parsedTraces1.get(0).get("trace_id")).longValue()); + assertEquals(2L, ((Number) parsedTraces1.get(1).get("trace_id")).longValue()); + assertTrue( + ((Number) parsedTraces1.get(0).get("start")).longValue() + < ((Number) parsedTraces1.get(1).get("start")).longValue()); + + clearInvocations(dumpReporter); + + // second flare dump with new pending traces + // Finish the first set of traces + child1.finish(); + child2.finish(); + // Create new pending traces + PendingTrace trace3 = flareFactory.create(DDTraceId.from(3)); + DDSpan parent3 = newSpanOf(trace3, UNSET, System.currentTimeMillis() * 3000); + DDSpan child3 = newSpanOf(parent3); + PendingTrace trace4 = flareFactory.create(DDTraceId.from(4)); + DDSpan parent4 = newSpanOf(trace4, UNSET, System.currentTimeMillis() * 4000); + DDSpan child4 = newSpanOf(parent4); + parent3.finish(); + parent4.finish(); + Map entries2 = buildAndExtractZip(); + + verify(dumpReporter).prepareForFlare(); + verify(dumpReporter).addReportToFlare(any()); + verify(dumpReporter).cleanupAfterFlare(); + + assertEquals(1, entries2.size()); + String pendingTraceText2 = (String) entries2.get("pending_traces.txt"); + List> parsedTraces2 = parseTraceLines(pendingTraceText2); + assertEquals(2, parsedTraces2.size()); + + child3.finish(); + child4.finish(); + + long deadline = System.currentTimeMillis() + 3000; + while ((trace1.size() != 0 || trace2.size() != 0 || trace3.size() != 0 || trace4.size() != 0) + && System.currentTimeMillis() < deadline) { + Thread.sleep(50); + } - long deadline = System.currentTimeMillis() + 3000; - while ((trace1.size() != 0 || trace2.size() != 0 || trace3.size() != 0 || trace4.size() != 0) - && System.currentTimeMillis() < deadline) { - Thread.sleep(50); + assertEquals(0, trace1.size()); + assertEquals(0, trace1.getPendingReferenceCount()); + assertEquals(0, trace2.size()); + assertEquals(0, trace2.getPendingReferenceCount()); + assertEquals(0, trace3.size()); + assertEquals(0, trace3.getPendingReferenceCount()); + assertEquals(0, trace4.size()); + assertEquals(0, trace4.getPendingReferenceCount()); + } finally { + flareBuffer.close(); } - - assertEquals(0, trace1.size()); - assertEquals(0, trace1.getPendingReferenceCount()); - assertEquals(0, trace2.size()); - assertEquals(0, trace2.getPendingReferenceCount()); - assertEquals(0, trace3.size()); - assertEquals(0, trace3.getPendingReferenceCount()); - assertEquals(0, trace4.size()); - assertEquals(0, trace4.getPendingReferenceCount()); } private DDSpan addContinuation(DDSpan span) { diff --git a/internal-api/src/main/java/datadog/trace/bootstrap/instrumentation/api/AgentTracer.java b/internal-api/src/main/java/datadog/trace/bootstrap/instrumentation/api/AgentTracer.java index 71767ea7fae..12bacc58cb0 100644 --- a/internal-api/src/main/java/datadog/trace/bootstrap/instrumentation/api/AgentTracer.java +++ b/internal-api/src/main/java/datadog/trace/bootstrap/instrumentation/api/AgentTracer.java @@ -441,6 +441,11 @@ public interface SpanBuilder { SpanBuilder withSpanType(CharSequence spanType); + /** + * Sets request-scoped data for the local trace this span belongs to. Intended for spans that + * start a new local trace; on a span joining an existing local trace, it only fills a slot that + * trace has not set. + */ SpanBuilder withRequestContextData(RequestContextSlot slot, T data); SpanBuilder withLink(AgentSpanLink link); diff --git a/internal-api/src/test/java/datadog/trace/api/gateway/BlockResponseFunctionTest.java b/internal-api/src/test/java/datadog/trace/api/gateway/BlockResponseFunctionTest.java index 22b0744bd48..fe55af41114 100644 --- a/internal-api/src/test/java/datadog/trace/api/gateway/BlockResponseFunctionTest.java +++ b/internal-api/src/test/java/datadog/trace/api/gateway/BlockResponseFunctionTest.java @@ -172,14 +172,6 @@ public Object getTagTop(String key, boolean sanitize) { return null; } - @Override - public void setTagCurrent(String key, Object value, boolean sanitize) {} - - @Override - public Object getTagCurrent(String key, boolean sanitize) { - return null; - } - @Override public void setDataTop(String key, Object value) {} @@ -193,19 +185,8 @@ public void effectivelyBlocked() { effectivelyBlockedCalls++; } - @Override - public void setDataCurrent(String key, Object value) {} - - @Override - public Object getDataCurrent(String key) { - return null; - } - @Override public void setMetaStructTop(String field, Object value) {} - - @Override - public void setMetaStructCurrent(String field, Object value) {} } private static final class CountingAppSecContext implements AppSecContext {