diff --git a/dd-java-agent/agent-llmobs/src/main/java/datadog/trace/llmobs/LLMObsSystem.java b/dd-java-agent/agent-llmobs/src/main/java/datadog/trace/llmobs/LLMObsSystem.java index d249372c14f..4244544bb4f 100644 --- a/dd-java-agent/agent-llmobs/src/main/java/datadog/trace/llmobs/LLMObsSystem.java +++ b/dd-java-agent/agent-llmobs/src/main/java/datadog/trace/llmobs/LLMObsSystem.java @@ -42,6 +42,14 @@ public static void start(Instrumentation inst, SharedCommunicationObjects sco) { return; } + // LLMObs spans need a real tracer (buildSpan() returns null on the no-op one), so keep the + // no-op SDK when none is installed. Mirrors TracerInstaller.installGlobalTracer, which still + // installs one for CI Visibility. + if (!config.isTraceEnabled() && !config.isCiVisibilityEnabled()) { + LOGGER.debug("LLM Observability is disabled: no tracer is installed (dd.trace.enabled)"); + return; + } + sco.createRemaining(config); String mlApp = config.getLlmObsMlApp(); diff --git a/dd-java-agent/agent-llmobs/src/test/java/datadog/trace/llmobs/LLMObsSystemTest.java b/dd-java-agent/agent-llmobs/src/test/java/datadog/trace/llmobs/LLMObsSystemTest.java new file mode 100644 index 00000000000..34f6050ca2f --- /dev/null +++ b/dd-java-agent/agent-llmobs/src/test/java/datadog/trace/llmobs/LLMObsSystemTest.java @@ -0,0 +1,48 @@ +package datadog.trace.llmobs; + +import static datadog.trace.api.config.CiVisibilityConfig.CIVISIBILITY_ENABLED; +import static datadog.trace.api.config.LlmObsConfig.LLMOBS_ENABLED; +import static datadog.trace.api.config.TraceInstrumentationConfig.TRACE_ENABLED; +import static org.junit.jupiter.api.Assertions.assertSame; +import static org.junit.jupiter.api.Assertions.assertThrows; + +import datadog.trace.api.llmobs.LLMObs; +import datadog.trace.api.llmobs.noop.NoOpLLMObsSpan; +import datadog.trace.test.junit.utils.config.WithConfig; +import datadog.trace.test.util.DDJavaSpecification; +import org.junit.jupiter.api.Test; + +class LLMObsSystemTest extends DDJavaSpecification { + + /** + * LLM Observability spans are backed by tracer spans, so the subsystem must stay off when tracing + * is disabled. Starting it would install the real SDK implementation on top of the no-op tracer, + * whose {@code buildSpan()} returns null, and every span-starting call would throw. + */ + @WithConfig(key = LLMOBS_ENABLED, value = "true") + @WithConfig(key = TRACE_ENABLED, value = "false") + @Test + void staysNoOpWhenTracingIsDisabled() { + // The null SharedCommunicationObjects doubles as an assertion that start() returns before + // touching it: any use would throw. + LLMObsSystem.start(null, null); + + assertSame( + NoOpLLMObsSpan.INSTANCE, LLMObs.startLLMSpan("span", "model", "provider", null, null)); + } + + /** + * CI Visibility installs a {@code CoreTracer} even with tracing disabled (see {@code + * TracerInstaller#installGlobalTracer}), so the guard above must not fire in that combination — + * the tracer LLM Observability needs does exist. + */ + @WithConfig(key = LLMOBS_ENABLED, value = "true") + @WithConfig(key = TRACE_ENABLED, value = "false") + @WithConfig(key = CIVISIBILITY_ENABLED, value = "true") + @Test + void startsWhenCiVisibilityInstallsTheTracer() { + // Reaching the null SharedCommunicationObjects is the assertion: start() only dereferences it + // after the guard, so the NPE proves the subsystem was not short-circuited. + assertThrows(NullPointerException.class, () -> LLMObsSystem.start(null, null)); + } +} diff --git a/dd-trace-core/src/main/java/datadog/trace/common/sampling/ApmTracingDisabledSampler.java b/dd-trace-core/src/main/java/datadog/trace/common/sampling/ApmTracingDisabledSampler.java new file mode 100644 index 00000000000..907072259bb --- /dev/null +++ b/dd-trace-core/src/main/java/datadog/trace/common/sampling/ApmTracingDisabledSampler.java @@ -0,0 +1,43 @@ +package datadog.trace.common.sampling; + +import datadog.trace.api.sampling.PrioritySampling; +import datadog.trace.api.sampling.SamplingMechanism; +import datadog.trace.bootstrap.ActiveSubsystems; +import datadog.trace.core.CoreSpan; +import java.time.Clock; + +/** + * Drops APM traces for {@code dd.apm.tracing.enabled=false}. If Remote Configuration activates + * AppSec at runtime, falls back to the standalone ASM trickle of 1 APM trace per minute, so it + * checks {@link ActiveSubsystems#APPSEC_ACTIVE} per trace. + */ +public class ApmTracingDisabledSampler implements Sampler, PrioritySampler { + + private final AsmStandaloneSampler asmStandaloneSampler; + private final ForcePrioritySampler dropSampler; + + public ApmTracingDisabledSampler(final Clock clock) { + this.asmStandaloneSampler = new AsmStandaloneSampler(clock); + this.dropSampler = + new ForcePrioritySampler(PrioritySampling.SAMPLER_DROP, SamplingMechanism.DEFAULT); + } + + @Override + public > boolean sample(final T span) { + // Both delegates keep sending the trace to the agent so it can collect stats on dropped traces. + return isAsmActive() ? asmStandaloneSampler.sample(span) : dropSampler.sample(span); + } + + @Override + public > void setSamplingPriority(final T span) { + if (isAsmActive()) { + asmStandaloneSampler.setSamplingPriority(span); + } else { + dropSampler.setSamplingPriority(span); + } + } + + private static boolean isAsmActive() { + return ActiveSubsystems.APPSEC_ACTIVE; + } +} diff --git a/dd-trace-core/src/main/java/datadog/trace/common/sampling/AsmStandaloneSampler.java b/dd-trace-core/src/main/java/datadog/trace/common/sampling/AsmStandaloneSampler.java index 54a2d750700..4772b5dc806 100644 --- a/dd-trace-core/src/main/java/datadog/trace/common/sampling/AsmStandaloneSampler.java +++ b/dd-trace-core/src/main/java/datadog/trace/common/sampling/AsmStandaloneSampler.java @@ -51,8 +51,12 @@ public > void setSamplingPriority(final T span) { private boolean shouldSample() { long now = clock.millis(); - return lastSampleTime.updateAndGet( - lastTime -> now - lastTime >= RATE_IN_MILLISECONDS ? now : lastTime) - == now; + // Compare against the time before the update, not after: two traces landing in the same + // millisecond as the last sample would both see the updated value equal to now and both be + // kept, sending 2 APM traces a minute instead of 1. + long lastSample = + lastSampleTime.getAndUpdate( + lastTime -> now - lastTime >= RATE_IN_MILLISECONDS ? now : lastTime); + return now - lastSample >= RATE_IN_MILLISECONDS; } } diff --git a/dd-trace-core/src/main/java/datadog/trace/common/sampling/Sampler.java b/dd-trace-core/src/main/java/datadog/trace/common/sampling/Sampler.java index 070f5f5cad3..8e47348359c 100644 --- a/dd-trace-core/src/main/java/datadog/trace/common/sampling/Sampler.java +++ b/dd-trace-core/src/main/java/datadog/trace/common/sampling/Sampler.java @@ -57,9 +57,17 @@ public static Sampler forConfig( @Nullable final RateByServiceTraceSampler agentSampler) { Sampler sampler; if (config != null) { - if (!config.isApmTracingEnabled() && isAsmEnabled(config)) { - log.debug("APM is disabled. Only 1 trace per minute will be sent."); - return new AsmStandaloneSampler(Clock.systemUTC()); + if (!config.isApmTracingEnabled()) { + if (isAsmEnabled(config)) { + log.debug( + "APM tracing is disabled, but ASM is enabled. Only 1 APM trace per minute will be sent."); + return new AsmStandaloneSampler(Clock.systemUTC()); + } + // Only ASM needs a trickle of APM traces, so drop them all unless AppSec is activated + // later by Remote Configuration. + log.debug( + "APM tracing is disabled. APM traces will be dropped, unless AppSec is activated at runtime."); + return new ApmTracingDisabledSampler(Clock.systemUTC()); } final Map serviceRules = config.getTraceSamplingServiceRules(); final Map operationRules = config.getTraceSamplingOperationRules(); @@ -133,6 +141,7 @@ public static Sampler forConfig( return sampler; } + /** Whether any Application Security product is on at startup. */ private static boolean isAsmEnabled(Config config) { return config.getAppSecActivation() == ProductActivation.FULLY_ENABLED || config.getIastActivation() == ProductActivation.FULLY_ENABLED diff --git a/dd-trace-core/src/main/java/datadog/trace/common/writer/TraceProcessingWorker.java b/dd-trace-core/src/main/java/datadog/trace/common/writer/TraceProcessingWorker.java index 9b70b09bc9e..0fae63585e2 100644 --- a/dd-trace-core/src/main/java/datadog/trace/common/writer/TraceProcessingWorker.java +++ b/dd-trace-core/src/main/java/datadog/trace/common/writer/TraceProcessingWorker.java @@ -128,6 +128,11 @@ MessagePassingBlockingQueue getPrimaryQueue() { return primaryQueue; } + @VisibleForTesting + MessagePassingBlockingQueue getSecondaryQueue() { + return secondaryQueue; + } + private static MessagePassingBlockingQueue createQueue(int capacity) { return Queues.mpscBlockingConsumerArrayQueue(capacity); } @@ -195,6 +200,9 @@ public void onEvent(Object event) { // TODO populate `_sample_rate` metric in a way that accounts for lost/dropped traces payloadDispatcher.addTrace(trace); } else if (event instanceof FlushEvent) { + // The flush marker only goes to the primary queue, so drain the secondary queue too: + // sampled-out traces can carry LLM Observability spans that must still be sent. + consumeFromSecondaryQueue(); payloadDispatcher.flush(); ((FlushEvent) event).sync(); } diff --git a/dd-trace-core/src/main/java/datadog/trace/core/propagation/ContextInterpreter.java b/dd-trace-core/src/main/java/datadog/trace/core/propagation/ContextInterpreter.java index 14a6459ca0f..7e91658ccaf 100644 --- a/dd-trace-core/src/main/java/datadog/trace/core/propagation/ContextInterpreter.java +++ b/dd-trace-core/src/main/java/datadog/trace/core/propagation/ContextInterpreter.java @@ -19,6 +19,7 @@ import datadog.trace.api.DDSpanId; import datadog.trace.api.DDTraceId; import datadog.trace.api.Functions; +import datadog.trace.api.ProductTraceSource; import datadog.trace.api.TagMap; import datadog.trace.api.TraceConfig; import datadog.trace.api.TracePropagationStyle; @@ -72,6 +73,7 @@ public abstract class ContextInterpreter implements AgentPropagation.KeyClassifi private final boolean requestHeaderTagsCommaAllowed; private final int baggageMaxItems; private final int baggageMaxBytes; + private final boolean apmTracingEnabled; protected static final boolean LOG_EXTRACT_HEADER_NAMES = Config.get().isLogExtractHeaderNames(); private static final DDCache CACHE = DDCaches.newFixedSizeCache(64); @@ -89,6 +91,7 @@ protected ContextInterpreter(Config config) { this.requestHeaderTagsCommaAllowed = config.isRequestHeaderTagsCommaAllowed(); this.baggageMaxItems = config.getTraceBaggageMaxItems(); this.baggageMaxBytes = config.getTraceBaggageMaxBytes(); + this.apmTracingEnabled = config.isApmTracingEnabled(); } final TagMap.Ledger tagLedger() { @@ -345,6 +348,17 @@ private TagContext.HttpHeaders getHeaders() { } private int samplingPriorityOrDefault(DDTraceId traceId, int samplingPriority) { + if (!apmTracingEnabled + && (propagationTags == null + || !ProductTraceSource.isProductMarked( + propagationTags.getTraceSource(), + ProductTraceSource.ASM, + ProductTraceSource.AI_GUARD))) { + // With APM tracing disabled, don't inherit an upstream APM decision: let the local sampler + // set the priority. Same products as TraceCollector.setSamplingPriorityIfNecessary. Needed + // when Data Jobs is enabled, where the drop can't override an already-set priority. + return PrioritySampling.UNSET; + } return samplingPriority == PrioritySampling.UNSET || DDTraceId.ZERO.equals(traceId) ? defaultSamplingPriority() : samplingPriority; diff --git a/dd-trace-core/src/test/java/datadog/trace/common/sampling/SamplerTest.java b/dd-trace-core/src/test/java/datadog/trace/common/sampling/SamplerTest.java index e4658e7046b..3f7c9d3739c 100644 --- a/dd-trace-core/src/test/java/datadog/trace/common/sampling/SamplerTest.java +++ b/dd-trace-core/src/test/java/datadog/trace/common/sampling/SamplerTest.java @@ -3,23 +3,38 @@ import static datadog.trace.api.config.AppSecConfig.APPSEC_ENABLED; import static datadog.trace.api.config.AppSecConfig.APPSEC_SCA_ENABLED; import static datadog.trace.api.config.GeneralConfig.APM_TRACING_ENABLED; +import static datadog.trace.api.config.GeneralConfig.DATA_JOBS_ENABLED; import static datadog.trace.api.config.IastConfig.IAST_ENABLED; import static datadog.trace.api.config.OtlpConfig.TRACE_OTEL_EXPORTER; import static datadog.trace.api.config.TracerConfig.PRIORITY_SAMPLING; import static datadog.trace.api.config.TracerConfig.PRIORITY_SAMPLING_FORCE; import static datadog.trace.api.config.TracerConfig.TRACE_SAMPLE_RATE; +import static datadog.trace.api.sampling.PrioritySampling.SAMPLER_DROP; import static datadog.trace.api.sampling.PrioritySampling.SAMPLER_KEEP; +import static datadog.trace.api.sampling.PrioritySampling.USER_KEEP; +import static java.util.Arrays.asList; +import static java.util.Collections.singletonList; import static org.junit.jupiter.api.Assertions.assertEquals; import static org.junit.jupiter.api.Assertions.assertFalse; import static org.junit.jupiter.api.Assertions.assertInstanceOf; import static org.junit.jupiter.api.Assertions.assertNotNull; +import static org.junit.jupiter.api.Assertions.assertTrue; import datadog.trace.api.Config; +import datadog.trace.api.DDTags; +import datadog.trace.api.ProductTraceSource; +import datadog.trace.api.sampling.SamplingMechanism; +import datadog.trace.bootstrap.ActiveSubsystems; +import datadog.trace.bootstrap.instrumentation.api.Tags; import datadog.trace.common.writer.ListWriter; import datadog.trace.core.CoreTracer; import datadog.trace.core.DDSpan; import datadog.trace.test.junit.utils.config.WithConfig; import datadog.trace.test.util.DDJavaSpecification; +import java.util.ArrayList; +import java.util.List; +import java.util.concurrent.TimeoutException; +import java.util.function.Consumer; import org.junit.jupiter.api.Test; class SamplerTest extends DDJavaSpecification { @@ -67,6 +82,98 @@ void asmStandaloneSamplerNotSelectedWhenApmTracingAndAsmNotEnabled() { assertFalse(sampler instanceof AsmStandaloneSampler); } + @WithConfig(key = APM_TRACING_ENABLED, value = "false") + @Test + void apmTracesDroppedWhenApmTracingDisabledAndNoOtherProductEnabled() { + assertApmTracesDropped(); + } + + /** + * {@code manual.keep} force-keeps the span long before the sampler votes, locking the sampling + * priority. With APM tracing disabled the drop still has to win, or the trace would be indexed + * and billed as APM anyway — exactly what the setting opts out of. + */ + @WithConfig(key = APM_TRACING_ENABLED, value = "false") + @Test + void apmTracesDroppedWhenApmTracingDisabledAndTraceManuallyKept() { + Sampler sampler = Sampler.Builder.forConfig(Config.get(), null); + + assertEquals( + SAMPLER_DROP, + (int) samplingPriorityOfTrace(sampler, span -> span.setTag(DDTags.MANUAL_KEEP, true))); + } + + /** + * A trace a product asked for is not an APM trace the setting opts out of, so {@code manual.keep} + * still wins once {@code _dd.p.ts} is marked — here for ASM. + */ + @WithConfig(key = APM_TRACING_ENABLED, value = "false") + @Test + void manuallyKeptTracesKeptWhenApmTracingDisabledAndMarkedForAsm() { + Sampler sampler = Sampler.Builder.forConfig(Config.get(), null); + + assertEquals( + USER_KEEP, + (int) + samplingPriorityOfTrace( + sampler, + span -> { + span.setTag(Tags.PROPAGATED_TRACE_SOURCE, ProductTraceSource.ASM); + span.setTag(DDTags.MANUAL_KEEP, true); + })); + } + + /** + * Data Jobs keeps its traces with a {@link SamplingMechanism#DATA_JOBS} priority rather than a + * {@code _dd.p.ts} mark, so the APM traces drop has to leave it alone. + */ + @WithConfig(key = APM_TRACING_ENABLED, value = "false") + @WithConfig(key = DATA_JOBS_ENABLED, value = "true") + @Test + void dataJobsTracesKeptWhenApmTracingDisabled() { + Sampler sampler = Sampler.Builder.forConfig(Config.get(), null); + + assertEquals( + USER_KEEP, + (int) + samplingPriorityOfTrace( + sampler, span -> span.setSamplingPriority(USER_KEEP, SamplingMechanism.DATA_JOBS))); + } + + /** + * AppSec can be activated by Remote Configuration after the sampler has been built, and that only + * flips {@link ActiveSubsystems#APPSEC_ACTIVE} — the immutable config the sampler was chosen from + * still says AppSec is off. Standalone ASM then needs its trickle of 1 APM trace per minute to + * keep the service in the service catalog, so the first trace has to be kept. + */ + @WithConfig(key = APM_TRACING_ENABLED, value = "false") + @Test + void apmTracesKeptOncePerMinuteWhenAppSecIsActivatedAtRuntime() { + Sampler sampler = Sampler.Builder.forConfig(Config.get(), null); + assertInstanceOf(ApmTracingDisabledSampler.class, sampler); + + // the trickle is one per minute, so only the first of the two traces is kept + withAppSecActive( + () -> + assertEquals( + asList((int) SAMPLER_KEEP, (int) SAMPLER_DROP), + samplingPrioritiesOfTraces(sampler, 2))); + } + + /** Deactivating AppSec at runtime puts the drop back. */ + @WithConfig(key = APM_TRACING_ENABLED, value = "false") + @Test + void apmTracesDroppedWhenAppSecIsDeactivatedAtRuntime() { + Sampler sampler = Sampler.Builder.forConfig(Config.get(), null); + + withAppSecActive( + () -> + assertEquals( + singletonList((int) SAMPLER_KEEP), samplingPrioritiesOfTraces(sampler, 1))); + + assertEquals(singletonList((int) SAMPLER_DROP), samplingPrioritiesOfTraces(sampler, 1)); + } + @Test void asmStandaloneSamplerNotSelectedWhenApmTracingEnabledAndAsmNotEnabled() { Config config = Config.get(); @@ -165,4 +272,100 @@ void spansBuiltWithOtlpEnabledAndPrioritySamplingDisabledHaveNonUnsetSamplingPri tracer.close(); } } + + /** + * Runs {@code assertions} with AppSec active, as Remote Configuration would leave it, always + * restoring the flag afterwards — it is global mutable state shared with every other test. + */ + private static void withAppSecActive(Runnable assertions) { + boolean wasActive = ActiveSubsystems.APPSEC_ACTIVE; + ActiveSubsystems.APPSEC_ACTIVE = true; + try { + assertions.run(); + } finally { + ActiveSubsystems.APPSEC_ACTIVE = wasActive; + } + } + + /** + * Runs {@code traceCount} traces through a tracer using {@code sampler} and returns the priority + * each one was written with. + * + *

The sampler is left to vote from the publish path ({@code + * TraceCollector.setSamplingPriorityIfNecessary}) rather than being called directly: with APM + * tracing disabled, {@code SamplingMechanism.APPSEC} is exempt from the sampling priority lock, + * so an extra direct call would both consume a slot of the one-per-minute trickle and overwrite + * the priority the publish path went on to assign. + */ + private static List samplingPrioritiesOfTraces(Sampler sampler, int traceCount) { + ListWriter writer = new ListWriter(); + CoreTracer tracer = CoreTracer.builder().writer(writer).sampler(sampler).build(); + try { + for (int i = 0; i < traceCount; i++) { + tracer.buildSpan("datadog", "test").start().finish(); + } + writer.waitForTraces(traceCount); + + List priorities = new ArrayList<>(); + for (List trace : writer) { + priorities.add(trace.get(0).getSamplingPriority()); + } + return priorities; + } catch (InterruptedException | TimeoutException e) { + throw new AssertionError("the traces were never written", e); + } finally { + tracer.close(); + } + } + + /** + * Runs a single trace through a tracer using {@code sampler}, applying {@code tagger} to its root + * span before it is finished, and returns the priority it was written with. + */ + private static int samplingPriorityOfTrace(Sampler sampler, Consumer tagger) { + ListWriter writer = new ListWriter(); + CoreTracer tracer = CoreTracer.builder().writer(writer).sampler(sampler).build(); + try { + DDSpan span = (DDSpan) tracer.buildSpan("datadog", "test").start(); + tagger.accept(span); + span.finish(); + writer.waitForTraces(1); + + return writer.firstTrace().get(0).getSamplingPriority(); + } catch (InterruptedException | TimeoutException e) { + throw new AssertionError("the trace was never written", e); + } finally { + tracer.close(); + } + } + + /** + * Asserts the trace is marked dropped but still written. Products that ride the tracer and ship + * their spans elsewhere — LLM Observability sends them to its own intake — depend on the spans + * still reaching the writers, so dropping APM traces must mean a drop priority, not a discarded + * span. + */ + private static void assertApmTracesDropped() { + Sampler sampler = Sampler.Builder.forConfig(Config.get(), null); + + assertInstanceOf(ApmTracingDisabledSampler.class, sampler); + + ListWriter writer = new ListWriter(); + CoreTracer tracer = CoreTracer.builder().writer(writer).sampler(sampler).build(); + try { + DDSpan span = (DDSpan) tracer.buildSpan("datadog", "test").start(); + ((PrioritySampler) sampler).setSamplingPriority(span); + + assertEquals(SAMPLER_DROP, (int) span.getSamplingPriority()); + assertTrue(sampler.sample(span)); + + span.finish(); + writer.waitForTraces(1); + assertEquals(singletonList(span), writer.firstTrace()); + } catch (InterruptedException | TimeoutException e) { + throw new AssertionError("the dropped trace was never written", e); + } finally { + tracer.close(); + } + } } diff --git a/dd-trace-core/src/test/java/datadog/trace/common/writer/DDIntakeWriterCombinedTest.java b/dd-trace-core/src/test/java/datadog/trace/common/writer/DDIntakeWriterCombinedTest.java index c53286cb106..bd63f037cec 100644 --- a/dd-trace-core/src/test/java/datadog/trace/common/writer/DDIntakeWriterCombinedTest.java +++ b/dd-trace-core/src/test/java/datadog/trace/common/writer/DDIntakeWriterCombinedTest.java @@ -1,6 +1,7 @@ package datadog.trace.common.writer; import static datadog.trace.common.writer.ddagent.Prioritization.ENSURE_TRACE; +import static datadog.trace.common.writer.ddagent.Prioritization.FAST_LANE; import static org.junit.jupiter.api.Assertions.assertEquals; import static org.junit.jupiter.api.Assertions.assertTrue; import static org.mockito.ArgumentMatchers.any; @@ -16,6 +17,7 @@ import static org.mockito.Mockito.verify; import static org.mockito.Mockito.verifyNoMoreInteractions; +import com.fasterxml.jackson.databind.ObjectMapper; import datadog.communication.http.OkHttpUtils; import datadog.communication.serialization.FlushingBuffer; import datadog.communication.serialization.msgpack.MsgPackWriter; @@ -24,6 +26,10 @@ import datadog.trace.agent.test.server.http.JavaTestHttpServer; import datadog.trace.api.civisibility.CiVisibilityWellKnownTags; import datadog.trace.api.intake.TrackType; +import datadog.trace.api.sampling.PrioritySampling; +import datadog.trace.api.sampling.SamplingMechanism; +import datadog.trace.bootstrap.instrumentation.api.InternalSpanTypes; +import datadog.trace.bootstrap.instrumentation.api.Tags; import datadog.trace.common.writer.ddintake.DDIntakeApi; import datadog.trace.common.writer.ddintake.DDIntakeMapperDiscovery; import datadog.trace.core.DDCoreJavaSpecification; @@ -31,20 +37,26 @@ import datadog.trace.core.monitor.HealthMetrics; import datadog.trace.core.monitor.TracerHealthMetrics; import datadog.trace.test.util.Flaky; +import java.io.ByteArrayOutputStream; import java.io.IOException; +import java.nio.channels.Channels; +import java.nio.channels.WritableByteChannel; import java.util.Collections; import java.util.List; +import java.util.Map; import java.util.concurrent.CountDownLatch; import java.util.concurrent.Phaser; import java.util.concurrent.Semaphore; import java.util.concurrent.TimeUnit; import java.util.concurrent.atomic.AtomicInteger; +import java.util.concurrent.atomic.AtomicReference; import okhttp3.HttpUrl; import okhttp3.OkHttpClient; import org.junit.jupiter.api.AfterEach; import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.api.Test; import org.junit.jupiter.api.Timeout; +import org.msgpack.jackson.dataformat.MessagePackFactory; @Timeout(value = 10, unit = TimeUnit.SECONDS) class DDIntakeWriterCombinedTest extends DDCoreJavaSpecification { @@ -808,6 +820,64 @@ void statsdCommFailure() throws Exception { healthMetrics.close(); } + /** + * With APM tracing disabled every APM trace is sampled out, so the LLM Observability spans only + * reach their intake if the writer still serializes p0 traces. Built like the LLMObs writer in + * {@link WriterFactory}: no dropping policy and FAST_LANE prioritization. + */ + @SuppressWarnings("unchecked") + @Test + void sampledOutTraceStillDeliversLlmObsSpans() throws Exception { + DDIntakeApi api = mock(DDIntakeApi.class); + DDIntakeWriter writer = + DDIntakeWriter.builder() + .addTrack(TrackType.LLMOBS, api) + .prioritization(FAST_LANE) + .monitoring(monitoring) + .flushIntervalMilliseconds(-1) + .alwaysFlush(false) + .build(); + writer.start(); + clearInvocations(api); + + AtomicReference sent = new AtomicReference<>(); + doAnswer( + invocation -> { + sent.set(writeTo(invocation.getArgument(0))); + return RemoteApi.Response.success(200); + }) + .when(api) + .sendSerializedTraces(any()); + + DDSpan span = + (DDSpan) + dummyTracer + .buildSpan("datadog", "openai.request") + .withTag("_ml_obs_tag.span.kind", Tags.LLMOBS_LLM_SPAN_KIND) + .start(); + span.setSpanType(InternalSpanTypes.LLMOBS); + span.setSamplingPriority(PrioritySampling.SAMPLER_DROP, SamplingMechanism.DEFAULT); + + writer.write(Collections.singletonList(span)); + writer.flush(); + + verify(api, times(1)).sendSerializedTraces(argThat(payload -> payload.traceCount() == 1)); + Map result = + new ObjectMapper(new MessagePackFactory()).readValue(sent.get(), Map.class); + List> spans = (List>) result.get("spans"); + assertEquals(1, spans.size()); + assertEquals(String.valueOf(span.getSpanId()), spans.get(0).get("span_id")); + + writer.close(); + } + + private static byte[] writeTo(Payload payload) throws IOException { + ByteArrayOutputStream out = new ByteArrayOutputStream(); + WritableByteChannel channel = Channels.newChannel(out); + payload.writeTo(channel); + return out.toByteArray(); + } + static String buildIntakePath(TrackType trackType, String apiVersion) { return String.format("/api/%s/%s", apiVersion, trackType.name().toLowerCase()); } diff --git a/dd-trace-core/src/test/java/datadog/trace/common/writer/TraceProcessingWorkerTest.java b/dd-trace-core/src/test/java/datadog/trace/common/writer/TraceProcessingWorkerTest.java index 86dbfa61b22..98878865b95 100644 --- a/dd-trace-core/src/test/java/datadog/trace/common/writer/TraceProcessingWorkerTest.java +++ b/dd-trace-core/src/test/java/datadog/trace/common/writer/TraceProcessingWorkerTest.java @@ -23,6 +23,7 @@ import datadog.trace.test.util.DDJavaSpecification; import java.lang.reflect.Method; import java.util.ArrayList; +import java.util.Arrays; import java.util.Collections; import java.util.List; import java.util.concurrent.TimeUnit; @@ -148,6 +149,57 @@ void testAFlushShouldClearThePrimaryQueue() { } } + // ------------------------------------------------------------------------- + // Test 3b: a flush should dispatch the secondary queue before completing + // ------------------------------------------------------------------------- + + /** + * The flush marker is only offered to the primary queue, so a trace on the secondary queue has to + * be swept into the dispatcher before the flush completes — otherwise {@code close()} interrupts + * the serializer thread and the trace is never dispatched. Sampled-out APM traces can afford + * that, but with APM tracing disabled the dropped trace also carries the LLM Observability spans. + */ + @Test + void testAFlushShouldDispatchTheSecondaryQueueBeforeCompleting() { + List calls = Collections.synchronizedList(new ArrayList<>()); + PayloadDispatcherImpl dispatcher = mock(PayloadDispatcherImpl.class); + doAnswer( + inv -> { + calls.add("addTrace"); + return null; + }) + .when(dispatcher) + .addTrace(any()); + doAnswer( + inv -> { + calls.add("flush"); + return null; + }) + .when(dispatcher) + .flush(); + + try (TraceProcessingWorker worker = + new TraceProcessingWorker( + 10, + mock(HealthMetrics.class), + dispatcher, + () -> false, + FAST_LANE, + 100, + TimeUnit.SECONDS, // prevent heartbeats from helping the flush happen + null)) { + List trace = Collections.singletonList(mock(DDSpan.class)); + worker.getSecondaryQueue().offer(trace); + worker.start(); + boolean flushed = worker.flush(10, TimeUnit.SECONDS); + + assertTrue(flushed); + assertTrue(worker.getSecondaryQueue().isEmpty()); + // the trace reached the dispatcher, and it did so before the flush was reported complete + assertEquals(Arrays.asList("addTrace", "flush"), calls); + } + } + // ------------------------------------------------------------------------- // Test 4: should report failure if serialization fails // ------------------------------------------------------------------------- diff --git a/dd-trace-core/src/test/java/datadog/trace/core/propagation/DatadogHttpExtractorTest.java b/dd-trace-core/src/test/java/datadog/trace/core/propagation/DatadogHttpExtractorTest.java index 462848c57a0..59372fa0459 100644 --- a/dd-trace-core/src/test/java/datadog/trace/core/propagation/DatadogHttpExtractorTest.java +++ b/dd-trace-core/src/test/java/datadog/trace/core/propagation/DatadogHttpExtractorTest.java @@ -1,7 +1,9 @@ package datadog.trace.core.propagation; +import static datadog.trace.api.config.GeneralConfig.APM_TRACING_ENABLED; import static datadog.trace.api.config.TracerConfig.REQUEST_HEADER_TAGS_COMMA_ALLOWED; import static datadog.trace.api.config.TracerConfig.TRACE_BAGGAGE_MAX_ITEMS; +import static datadog.trace.api.sampling.PrioritySampling.SAMPLER_KEEP; import static datadog.trace.api.sampling.PrioritySampling.UNSET; import static datadog.trace.bootstrap.instrumentation.api.ContextVisitors.stringValuesMap; import static datadog.trace.core.propagation.DatadogHttpCodec.DATADOG_TAGS_KEY; @@ -381,6 +383,52 @@ void extractMappedBaggageIsSubjectToTheItemLimit() { assertEquals(singletonMap(SOME_BAGGAGE, "mappedBaggageValue"), context.getBaggage()); } + /** + * With APM tracing disabled, an upstream sampling decision is only kept when {@code _dd.p.ts} + * marks the trace for ASM or AI Guard. Otherwise the local sampler sets the priority. + */ + @Nested + @WithConfig(key = APM_TRACING_ENABLED, value = "false") + class ApmTracingDisabled { + @Test + void upstreamPriorityDroppedWhenNoProductMarkedTheTrace() { + Map headers = new HashMap<>(); + headers.put(TRACE_ID_KEY, "1"); + headers.put(SPAN_ID_KEY, "2"); + headers.put(SAMPLING_PRIORITY_KEY, String.valueOf(SAMPLER_KEEP)); + + ExtractedContext context = (ExtractedContext) extractor.extract(headers, stringValuesMap()); + + assertEquals(UNSET, context.getSamplingPriority()); + } + + @Test + void upstreamPriorityKeptWhenProductMarkedTheTrace() { + Map headers = new HashMap<>(); + headers.put(TRACE_ID_KEY, "1"); + headers.put(SPAN_ID_KEY, "2"); + headers.put(SAMPLING_PRIORITY_KEY, String.valueOf(SAMPLER_KEEP)); + headers.put(DATADOG_TAGS_KEY, "_dd.p.ts=02"); + + ExtractedContext context = (ExtractedContext) extractor.extract(headers, stringValuesMap()); + + assertEquals(SAMPLER_KEEP, context.getSamplingPriority()); + } + + @Test + void upstreamPriorityDroppedWhenOnlyDsmMarkedTheTrace() { + Map headers = new HashMap<>(); + headers.put(TRACE_ID_KEY, "1"); + headers.put(SPAN_ID_KEY, "2"); + headers.put(SAMPLING_PRIORITY_KEY, String.valueOf(SAMPLER_KEEP)); + headers.put(DATADOG_TAGS_KEY, "_dd.p.ts=04"); + + ExtractedContext context = (ExtractedContext) extractor.extract(headers, stringValuesMap()); + + assertEquals(UNSET, context.getSamplingPriority()); + } + } + @Nested class BaggageLimits extends AbstractOTBaggageTest { @Override diff --git a/dd-trace-core/src/test/java/datadog/trace/core/propagation/W3COtelTraceStateContinuationTest.java b/dd-trace-core/src/test/java/datadog/trace/core/propagation/W3COtelTraceStateContinuationTest.java index 3d28facc067..336a121015b 100644 --- a/dd-trace-core/src/test/java/datadog/trace/core/propagation/W3COtelTraceStateContinuationTest.java +++ b/dd-trace-core/src/test/java/datadog/trace/core/propagation/W3COtelTraceStateContinuationTest.java @@ -91,6 +91,7 @@ void compoundExtractionKeepsFirstPriorityAndReconcilesOtelState( when(config.getTracePropagationStylesToExtract()) .thenReturn(new LinkedHashSet<>(asList(DATADOG, TRACECONTEXT))); when(config.getxDatadogTagsMaxLength()).thenReturn(DEFAULT_TRACE_X_DATADOG_TAGS_MAX_LENGTH); + when(config.isApmTracingEnabled()).thenReturn(true); CoreTracer tracer = tracerBuilder().build(); try { @@ -134,6 +135,7 @@ void compoundExtractionWithoutDatadogPriorityReconcilesFirstLocalDecision() { when(config.getTracePropagationStylesToExtract()) .thenReturn(new LinkedHashSet<>(asList(DATADOG, TRACECONTEXT))); when(config.getxDatadogTagsMaxLength()).thenReturn(DEFAULT_TRACE_X_DATADOG_TAGS_MAX_LENGTH); + when(config.isApmTracingEnabled()).thenReturn(true); CoreTracer tracer = tracerBuilder().build(); try { diff --git a/internal-api/src/main/java/datadog/trace/api/sampling/SamplingMechanism.java b/internal-api/src/main/java/datadog/trace/api/sampling/SamplingMechanism.java index 6b0cd397626..813125feccd 100644 --- a/internal-api/src/main/java/datadog/trace/api/sampling/SamplingMechanism.java +++ b/internal-api/src/main/java/datadog/trace/api/sampling/SamplingMechanism.java @@ -77,12 +77,21 @@ public static boolean validateWithSamplingPriority(int mechanism, int priority) /** * Returns true if sampling priority lock can be avoided for the given mechanism and priority * + *

With APM tracing disabled, the {@link #DEFAULT} drop overrides an already-set priority (e.g. + * {@code manual.keep}), so the trace isn't kept as APM. Traces marked for ASM or AI Guard never + * reach the sampler. The drop is not exempt when Data Jobs is enabled, so it can't overwrite a + * {@link #DATA_JOBS} keep. + * * @param priority the sampling priority * @param mechanism the sampling mechanism * @return {@code true} if the sampling priority lock can be avoided, {@code false} otherwise */ public static boolean canAvoidSamplingPriorityLock(int priority, int mechanism) { - return (!Config.get().isApmTracingEnabled() && mechanism == SamplingMechanism.APPSEC) + return (!Config.get().isApmTracingEnabled() + && (mechanism == SamplingMechanism.APPSEC + || (mechanism == DEFAULT + && priority == SAMPLER_DROP + && !Config.get().isDataJobsEnabled()))) || (Config.get().isDataJobsEnabled() && mechanism == DATA_JOBS); } diff --git a/internal-api/src/test/groovy/datadog/trace/api/sampling/SamplingMechanismTest.groovy b/internal-api/src/test/groovy/datadog/trace/api/sampling/SamplingMechanismTest.groovy index 9ad58eb1e34..96bee6ec4e4 100644 --- a/internal-api/src/test/groovy/datadog/trace/api/sampling/SamplingMechanismTest.groovy +++ b/internal-api/src/test/groovy/datadog/trace/api/sampling/SamplingMechanismTest.groovy @@ -108,6 +108,7 @@ class SamplingMechanismTest extends DDSpecification { APPSEC | UNSET | true APPSEC | SAMPLER_KEEP | true UNKNOWN | SAMPLER_KEEP | false + DEFAULT | SAMPLER_DROP | true DEFAULT | SAMPLER_KEEP | false AGENT_RATE | SAMPLER_KEEP | false REMOTE_AUTO_RATE | SAMPLER_KEEP | false @@ -117,4 +118,29 @@ class SamplingMechanismTest extends DDSpecification { DATA_JOBS | SAMPLER_KEEP | false EXTERNAL_OVERRIDE | SAMPLER_KEEP | false } + + void 'Test canAvoidSamplingPriorityLock with Data Jobs enabled'(){ + setup: + injectSysConfig("dd.apm.tracing.enabled", "false") + injectSysConfig("dd.data.jobs.enabled", "true") + + expect: + canAvoidSamplingPriorityLock(priority, mechanism) == valid + + where: + mechanism | priority | valid + DEFAULT | SAMPLER_DROP | false + APPSEC | SAMPLER_KEEP | true + DATA_JOBS | USER_KEEP | true + } + + void 'Test canAvoidSamplingPriorityLock with APM tracing enabled'(){ + expect: + canAvoidSamplingPriorityLock(priority, mechanism) == valid + + where: + mechanism | priority | valid + DEFAULT | SAMPLER_DROP | false + APPSEC | SAMPLER_KEEP | false + } }