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 a57dd858b45..b38dffeb0f8 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 @@ -30,6 +30,11 @@ public static void start(Instrumentation inst, SharedCommunicationObjects sco) { return; } + if (!config.isTraceEnabled()) { + LOGGER.debug("LLM Observability is disabled: tracing is disabled"); + return; + } + sco.createRemaining(config); String mlApp = config.getLlmObsMlApp(); diff --git a/dd-java-agent/agent-llmobs/src/main/java/datadog/trace/llmobs/domain/DDLLMObsSpan.java b/dd-java-agent/agent-llmobs/src/main/java/datadog/trace/llmobs/domain/DDLLMObsSpan.java index 6532829cfa6..996590514b3 100644 --- a/dd-java-agent/agent-llmobs/src/main/java/datadog/trace/llmobs/domain/DDLLMObsSpan.java +++ b/dd-java-agent/agent-llmobs/src/main/java/datadog/trace/llmobs/domain/DDLLMObsSpan.java @@ -5,7 +5,9 @@ import datadog.trace.api.DDSpanTypes; import datadog.trace.api.DDTraceApiInfo; import datadog.trace.api.DDTraceId; +import datadog.trace.api.ProductTraceSource; import datadog.trace.api.WellKnownTags; +import datadog.trace.api.internal.TraceSegment; import datadog.trace.api.llmobs.LLMObs; import datadog.trace.api.llmobs.LLMObsContext; import datadog.trace.api.llmobs.LLMObsSpan; @@ -110,6 +112,12 @@ public DDLLMObsSpan( } span.setTag(LLMOBS_TAG_PREFIX + PARENT_ID_TAG_INTERNAL, parentSpanID); scope = LLMObsContext.attach(span.context()); + + // Mark this span as originating from LLM Observability product + TraceSegment segment = AgentTracer.get().getTraceSegment(); + if (segment != null) { + segment.setTagTop(Tags.PROPAGATED_TRACE_SOURCE, ProductTraceSource.LLMOBS); + } } @Override diff --git a/dd-java-agent/agent-llmobs/src/test/groovy/datadog/trace/llmobs/LLMObsSystemTest.groovy b/dd-java-agent/agent-llmobs/src/test/groovy/datadog/trace/llmobs/LLMObsSystemTest.groovy new file mode 100644 index 00000000000..ef2f6c82dd4 --- /dev/null +++ b/dd-java-agent/agent-llmobs/src/test/groovy/datadog/trace/llmobs/LLMObsSystemTest.groovy @@ -0,0 +1,53 @@ +package datadog.trace.llmobs + +import datadog.communication.ddagent.SharedCommunicationObjects +import datadog.trace.test.util.DDSpecification +import okhttp3.HttpUrl + +class LLMObsSystemTest extends DDSpecification { + + void 'start disabled when llmobs is disabled'() { + setup: + injectSysConfig('llmobs.enabled', 'false') + rebuildConfig() + final inst = Mock(java.lang.instrument.Instrumentation) + final sco = Mock(SharedCommunicationObjects) + + when: + LLMObsSystem.start(inst, sco) + + then: + 0 * sco._ + } + + void 'start disabled when trace is disabled'() { + setup: + injectSysConfig('llmobs.enabled', 'true') + injectSysConfig('trace.enabled', 'false') + rebuildConfig() + final inst = Mock(java.lang.instrument.Instrumentation) + final sco = Mock(SharedCommunicationObjects) + + when: + LLMObsSystem.start(inst, sco) + + then: + 0 * sco._ + } + + void 'start enabled when apm tracing disabled but llmobs enabled'() { + setup: + injectSysConfig('llmobs.enabled', 'true') + injectSysConfig('apm.tracing.enabled', 'false') + rebuildConfig() + final inst = Mock(java.lang.instrument.Instrumentation) + final sco = Mock(SharedCommunicationObjects) + sco.agentUrl = HttpUrl.parse('http://localhost:8126') + + when: + LLMObsSystem.start(inst, sco) + + then: + 1 * sco.createRemaining(_) + } +} diff --git a/dd-smoke-tests/apm-tracing-disabled/src/main/java/datadog/smoketest/apmtracingdisabled/Controller.java b/dd-smoke-tests/apm-tracing-disabled/src/main/java/datadog/smoketest/apmtracingdisabled/Controller.java index 3bb55197614..61f1e4ffef0 100644 --- a/dd-smoke-tests/apm-tracing-disabled/src/main/java/datadog/smoketest/apmtracingdisabled/Controller.java +++ b/dd-smoke-tests/apm-tracing-disabled/src/main/java/datadog/smoketest/apmtracingdisabled/Controller.java @@ -1,5 +1,7 @@ package datadog.smoketest.apmtracingdisabled; +import datadog.trace.api.llmobs.LLMObs; +import datadog.trace.api.llmobs.LLMObsSpan; import edu.umd.cs.findbugs.annotations.SuppressFBWarnings; import io.opentracing.Span; import io.opentracing.util.GlobalTracer; @@ -73,6 +75,17 @@ public void write( } } + @GetMapping("/llmobs/test") + public String llmobsTest() { + // Create LLMObs span using public API + LLMObsSpan llmSpan = + LLMObs.startLLMSpan("llmobs-test-operation", "gpt-4", "openai", null, null); + llmSpan.annotateIO("test input", "test output"); + llmSpan.finish(); + + return "LLMObs test completed"; + } + private String forceKeepSpan() { final Span span = GlobalTracer.get().activeSpan(); if (span != null) { diff --git a/dd-smoke-tests/apm-tracing-disabled/src/test/groovy/datadog/smoketest/apmtracingdisabled/LlmObsApmDisabledSmokeTest.groovy b/dd-smoke-tests/apm-tracing-disabled/src/test/groovy/datadog/smoketest/apmtracingdisabled/LlmObsApmDisabledSmokeTest.groovy new file mode 100644 index 00000000000..ef1c6b8f7aa --- /dev/null +++ b/dd-smoke-tests/apm-tracing-disabled/src/test/groovy/datadog/smoketest/apmtracingdisabled/LlmObsApmDisabledSmokeTest.groovy @@ -0,0 +1,92 @@ +package datadog.smoketest.apmtracingdisabled + +import datadog.trace.api.sampling.PrioritySampling +import okhttp3.Request + +class LlmObsApmDisabledSmokeTest extends AbstractApmTracingDisabledSmokeTest { + + static final String LLMOBS_SERVICE_NAME = "llmobs-apm-disabled-test" + + static final String[] LLMOBS_APM_DISABLED_PROPERTIES = [ + "-Ddd.apm.tracing.enabled=false", + "-Ddd.llmobs.enabled=true", + "-Ddd.llmobs.ml-app=test-app", + "-Ddd.service.name=${LLMOBS_SERVICE_NAME}", + ] + + @Override + ProcessBuilder createProcessBuilder() { + return createProcess(LLMOBS_APM_DISABLED_PROPERTIES) + } + + void 'When APM disabled and LLMObs enabled, LLMObs spans should be kept and APM spans should be dropped'() { + setup: + final llmobsUrl = "http://localhost:${httpPort}/rest-api/llmobs/test" + final llmobsRequest = new Request.Builder().url(llmobsUrl).get().build() + + final apmUrl = "http://localhost:${httpPort}/rest-api/greetings" + final apmRequest = new Request.Builder().url(apmUrl).get().build() + + when: "Create LLMObs span" + final llmobsResponse = client.newCall(llmobsRequest).execute() + + then: "LLMObs request should succeed" + llmobsResponse.successful + + when: "Create regular APM span" + final apmResponse = client.newCall(apmRequest).execute() + + then: "APM request should succeed" + apmResponse.successful + + and: "Wait for traces" + waitForTraceCount(2) + + and: "LLMObs trace should be kept (SAMPLER_KEEP)" + def llmobsTrace = traces.find { trace -> + trace.spans.find { span -> + span.meta["http.url"] == llmobsUrl + } + } + assert llmobsTrace != null + // The LLMObs child span should have LLMObs tags + def llmobsChildSpan = llmobsTrace.spans.find { span -> + span.meta["_ml_obs_tag.model_name"] == "gpt-4" + } + assert llmobsChildSpan != null : "LLMObs child span with model_name=gpt-4 should exist" + + and: "Regular APM trace should be dropped (SAMPLER_DROP)" + def apmTrace = traces.find { trace -> + trace.spans.find { span -> + span.meta["http.url"] == apmUrl + } + } + assert apmTrace != null + checkRootSpanPrioritySampling(apmTrace, PrioritySampling.SAMPLER_DROP) + + and: "No NPE or errors in logs" + !isLogPresent { it.contains("NullPointerException") } + !isLogPresent { it.contains("ERROR") } + } + + void 'LLMObs spans should have PROPAGATED_TRACE_SOURCE tag set'() { + setup: + final llmobsUrl = "http://localhost:${httpPort}/rest-api/llmobs/test" + final llmobsRequest = new Request.Builder().url(llmobsUrl).get().build() + + when: + final response = client.newCall(llmobsRequest).execute() + + then: + response.successful + waitForTraceCount(1) + + and: "LLMObs span should be created successfully" + def trace = traces[0] + assert trace != null + def llmobsSpan = trace.spans.find { span -> + span.meta["_ml_obs_tag.model_name"] == "gpt-4" + } + assert llmobsSpan != null : "LLMObs span with model_name should exist" + } +} diff --git a/dd-smoke-tests/apm-tracing-disabled/src/test/groovy/datadog/smoketest/apmtracingdisabled/LlmObsTraceDisabledSmokeTest.groovy b/dd-smoke-tests/apm-tracing-disabled/src/test/groovy/datadog/smoketest/apmtracingdisabled/LlmObsTraceDisabledSmokeTest.groovy new file mode 100644 index 00000000000..5501a9cc845 --- /dev/null +++ b/dd-smoke-tests/apm-tracing-disabled/src/test/groovy/datadog/smoketest/apmtracingdisabled/LlmObsTraceDisabledSmokeTest.groovy @@ -0,0 +1,34 @@ +package datadog.smoketest.apmtracingdisabled + +import okhttp3.Request + +class LlmObsTraceDisabledSmokeTest extends AbstractApmTracingDisabledSmokeTest { + + static final String[] LLMOBS_TRACE_DISABLED_PROPERTIES = [ + "-Ddd.trace.enabled=false", + "-Ddd.llmobs.enabled=true", + "-Ddd.llmobs.ml-app=test-app", + "-Ddd.service.name=llmobs-trace-disabled-test", + ] + + @Override + ProcessBuilder createProcessBuilder() { + return createProcess(LLMOBS_TRACE_DISABLED_PROPERTIES) + } + + void 'DD_TRACE_ENABLED=false with DD_LLMOBS_ENABLED=true should disable LLMObs gracefully'() { + setup: + final llmobsUrl = "http://localhost:${httpPort}/rest-api/llmobs/test" + final llmobsRequest = new Request.Builder().url(llmobsUrl).get().build() + + when: "Call LLMObs endpoint" + final response = client.newCall(llmobsRequest).execute() + + then: "Request should succeed" + response.successful + response.code() == 200 + + and: "LLMObs disabled message in logs" + isLogPresent { it.contains("LLM Observability is disabled: tracing is disabled") } + } +} diff --git a/dd-trace-core/src/main/java/datadog/trace/common/sampling/LlmObsAndAsmStandaloneSampler.java b/dd-trace-core/src/main/java/datadog/trace/common/sampling/LlmObsAndAsmStandaloneSampler.java new file mode 100644 index 00000000000..23185c8dc00 --- /dev/null +++ b/dd-trace-core/src/main/java/datadog/trace/common/sampling/LlmObsAndAsmStandaloneSampler.java @@ -0,0 +1,73 @@ +package datadog.trace.common.sampling; + +import static datadog.trace.api.sampling.PrioritySampling.SAMPLER_DROP; +import static datadog.trace.api.sampling.PrioritySampling.SAMPLER_KEEP; + +import datadog.trace.api.ProductTraceSource; +import datadog.trace.api.sampling.SamplingMechanism; +import datadog.trace.core.CoreSpan; +import datadog.trace.core.DDSpan; +import java.time.Clock; +import java.util.concurrent.atomic.AtomicLong; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +/** + * This sampler is used when APM tracing is disabled but both LLM Observability and ASM are enabled. + * It keeps all LLMObs and ASM traces, and allows 1 APM trace per minute for billing/service catalog + * purposes. + */ +public class LlmObsAndAsmStandaloneSampler implements Sampler, PrioritySampler { + + private static final Logger log = LoggerFactory.getLogger(LlmObsAndAsmStandaloneSampler.class); + private static final int RATE_IN_MILLISECONDS = 60000; // 1 minute + + private final AtomicLong lastSampleTime; + private final Clock clock; + + public LlmObsAndAsmStandaloneSampler(final Clock clock) { + this.clock = clock; + this.lastSampleTime = new AtomicLong(clock.millis() - RATE_IN_MILLISECONDS); + } + + @Override + public > boolean sample(final T span) { + // Priority sampling sends all traces to the core agent, including traces marked dropped. + // This allows the core agent to collect stats on all traces. + return true; + } + + @Override + public > void setSamplingPriority(final T span) { + T rootSpan = span.getLocalRootSpan(); + if (rootSpan instanceof DDSpan) { + DDSpan ddRootSpan = (DDSpan) rootSpan; + int traceSource = ddRootSpan.context().getPropagationTags().getTraceSource(); + if (ProductTraceSource.isProductMarked(traceSource, ProductTraceSource.LLMOBS)) { + log.debug("Set SAMPLER_KEEP for LLMObs span {}", span.getSpanId()); + span.setSamplingPriority(SAMPLER_KEEP, SamplingMechanism.DEFAULT); + return; + } + if (ProductTraceSource.isProductMarked(traceSource, ProductTraceSource.ASM)) { + log.debug("Set SAMPLER_KEEP for ASM span {}", span.getSpanId()); + span.setSamplingPriority(SAMPLER_KEEP, SamplingMechanism.APPSEC); + return; + } + } + // For APM-only traces, allow 1 per minute for billing/catalog purposes + if (shouldSample()) { + log.debug("Set SAMPLER_KEEP for APM span {}", span.getSpanId()); + span.setSamplingPriority(SAMPLER_KEEP, SamplingMechanism.APPSEC); + } else { + log.debug("Set SAMPLER_DROP for APM span {}", span.getSpanId()); + span.setSamplingPriority(SAMPLER_DROP, SamplingMechanism.APPSEC); + } + } + + private boolean shouldSample() { + long now = clock.millis(); + return lastSampleTime.updateAndGet( + lastTime -> now - lastTime >= RATE_IN_MILLISECONDS ? now : lastTime) + == now; + } +} diff --git a/dd-trace-core/src/main/java/datadog/trace/common/sampling/LlmObsStandaloneSampler.java b/dd-trace-core/src/main/java/datadog/trace/common/sampling/LlmObsStandaloneSampler.java new file mode 100644 index 00000000000..8f1b30a72df --- /dev/null +++ b/dd-trace-core/src/main/java/datadog/trace/common/sampling/LlmObsStandaloneSampler.java @@ -0,0 +1,47 @@ +package datadog.trace.common.sampling; + +import static datadog.trace.api.sampling.PrioritySampling.SAMPLER_DROP; +import static datadog.trace.api.sampling.PrioritySampling.SAMPLER_KEEP; + +import datadog.trace.api.ProductTraceSource; +import datadog.trace.api.sampling.SamplingMechanism; +import datadog.trace.core.CoreSpan; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +/** + * This sampler is used when APM tracing is disabled but LLM Observability is enabled. Unlike ASM + * standalone mode which only needs 1 trace per minute for billing/catalog purposes, LLM + * Observability needs to capture all LLM interactions to track costs, latency, and quality metrics. + * Therefore, this sampler keeps all LLMOBS traces and drops all APM-only traces. + */ +public class LlmObsStandaloneSampler implements Sampler, PrioritySampler { + + private static final Logger log = LoggerFactory.getLogger(LlmObsStandaloneSampler.class); + + @Override + public > boolean sample(final T span) { + // Priority sampling sends all traces to the core agent, including traces marked dropped. + // This allows the core agent to collect stats on all traces. + return true; + } + + @Override + public > void setSamplingPriority(final T span) { + // Only keep traces that have the LLMOBS product flag + // Drop regular APM traces when APM tracing is disabled + T rootSpan = span.getLocalRootSpan(); + if (rootSpan instanceof datadog.trace.core.DDSpan) { + datadog.trace.core.DDSpan ddRootSpan = (datadog.trace.core.DDSpan) rootSpan; + int traceSource = ddRootSpan.context().getPropagationTags().getTraceSource(); + if (ProductTraceSource.isProductMarked(traceSource, ProductTraceSource.LLMOBS)) { + log.debug("Set SAMPLER_KEEP for LLMObs span {}", span.getSpanId()); + span.setSamplingPriority(SAMPLER_KEEP, SamplingMechanism.DEFAULT); + return; + } + } + // Drop APM-only traces when APM tracing is disabled + log.debug("Set SAMPLER_DROP for APM-only span {}", span.getSpanId()); + span.setSamplingPriority(SAMPLER_DROP, SamplingMechanism.DEFAULT); + } +} 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 af1045e39df..ff1f61f9bca 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 @@ -36,9 +36,22 @@ final class Builder { public static Sampler forConfig(final Config config, final TraceConfig traceConfig) { 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 (config.isLlmObsEnabled() && isAsmEnabled(config)) { + log.debug( + "APM is disabled, but both LLMObs and ASM are enabled. All LLMObs and ASM traces will be kept, only 1 APM trace per minute will be sent."); + return new LlmObsAndAsmStandaloneSampler(Clock.systemUTC()); + } else if (config.isLlmObsEnabled()) { + log.debug("APM is disabled, but LLMObs is enabled. All LLMObs traces will be kept."); + return new LlmObsStandaloneSampler(); + } else if (isAsmEnabled(config)) { + log.debug( + "APM is disabled, but ASM is enabled. Only 1 APM trace per minute will be sent, all ASM traces will be kept."); + return new AsmStandaloneSampler(Clock.systemUTC()); + } + // APM disabled and no other products enabled - drop all APM traces + log.debug("APM is disabled. All APM traces will be dropped."); + return new ForcePrioritySampler(PrioritySampling.SAMPLER_DROP, SamplingMechanism.DEFAULT); } final Map serviceRules = config.getTraceSamplingServiceRules(); final Map operationRules = config.getTraceSamplingOperationRules(); 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 777fc3889bf..0da18680295 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 @@ -65,10 +65,13 @@ public void setSamplingPriorityIfNecessary() { DDSpan rootSpan = getRootSpan(); if (traceConfig.sampler instanceof PrioritySampler && rootSpan != null) { // Ignore the force-keep priority in the absence of propagated _dd.p.ts span tag marked for - // ASM. + // ASM or LLMOBS. if ((!Config.get().isApmTracingEnabled() && !ProductTraceSource.isProductMarked( - rootSpan.context().getPropagationTags().getTraceSource(), ProductTraceSource.ASM)) + rootSpan.context().getPropagationTags().getTraceSource(), ProductTraceSource.ASM) + && !ProductTraceSource.isProductMarked( + rootSpan.context().getPropagationTags().getTraceSource(), + ProductTraceSource.LLMOBS)) || rootSpan.context().getSamplingPriority() == PrioritySampling.UNSET) { ((PrioritySampler) traceConfig.sampler).setSamplingPriority(rootSpan); } diff --git a/dd-trace-core/src/test/groovy/datadog/trace/common/sampling/LlmObsAndAsmStandaloneSamplerTest.groovy b/dd-trace-core/src/test/groovy/datadog/trace/common/sampling/LlmObsAndAsmStandaloneSamplerTest.groovy new file mode 100644 index 00000000000..489b44450c6 --- /dev/null +++ b/dd-trace-core/src/test/groovy/datadog/trace/common/sampling/LlmObsAndAsmStandaloneSamplerTest.groovy @@ -0,0 +1,98 @@ +package datadog.trace.common.sampling + +import datadog.trace.api.ProductTraceSource +import datadog.trace.api.sampling.PrioritySampling +import datadog.trace.bootstrap.instrumentation.api.Tags +import datadog.trace.common.writer.ListWriter +import datadog.trace.core.test.DDCoreSpecification + +import java.time.Clock +import java.util.concurrent.atomic.AtomicLong + +class LlmObsAndAsmStandaloneSamplerTest extends DDCoreSpecification { + + def writer = new ListWriter() + + void "test LLMObs spans are kept"() { + setup: + def sampler = new LlmObsAndAsmStandaloneSampler(Clock.systemUTC()) + def tracer = tracerBuilder().writer(writer).sampler(sampler).build() + + when: + def span = tracer.buildSpan("testInstrumentation", "llm-call").start() + def scope = tracer.activateSpan(span) + tracer.getTraceSegment().setTagTop(Tags.PROPAGATED_TRACE_SOURCE, ProductTraceSource.LLMOBS) + sampler.setSamplingPriority(span) + scope.close() + + then: + span.getSamplingPriority() == PrioritySampling.SAMPLER_KEEP + + cleanup: + tracer.close() + } + + void "test ASM spans are kept"() { + setup: + def sampler = new LlmObsAndAsmStandaloneSampler(Clock.systemUTC()) + def tracer = tracerBuilder().writer(writer).sampler(sampler).build() + + when: + def span = tracer.buildSpan("testInstrumentation", "http-request").start() + def scope = tracer.activateSpan(span) + tracer.getTraceSegment().setTagTop(Tags.PROPAGATED_TRACE_SOURCE, ProductTraceSource.ASM) + sampler.setSamplingPriority(span) + scope.close() + + then: + span.getSamplingPriority() == PrioritySampling.SAMPLER_KEEP + + cleanup: + tracer.close() + } + + void "test APM-only spans are rate-limited to 1 per minute"() { + setup: + def current = new AtomicLong(System.currentTimeMillis()) + final Clock clock = Mock(Clock) { + millis() >> { + current.get() + } + } + def sampler = new LlmObsAndAsmStandaloneSampler(clock) + def tracer = tracerBuilder().writer(writer).sampler(sampler).build() + + when: "first APM span" + def span1 = tracer.buildSpan("testInstrumentation", "apm-request").start() + sampler.setSamplingPriority(span1) + + then: + 1 * clock.millis() >> { + current.updateAndGet(v -> v + 1000) + } + span1.getSamplingPriority() == PrioritySampling.SAMPLER_KEEP + + when: "second APM span within the same minute" + def span2 = tracer.buildSpan("testInstrumentation", "apm-request2").start() + sampler.setSamplingPriority(span2) + + then: + 1 * clock.millis() >> { + current.updateAndGet(v -> v + 1000) + } + span2.getSamplingPriority() == PrioritySampling.SAMPLER_DROP + + when: "third APM span after 1 minute" + def span3 = tracer.buildSpan("testInstrumentation", "apm-request3").start() + sampler.setSamplingPriority(span3) + + then: + clock.millis() >> { + current.updateAndGet(v -> v + 60000) + } + span3.getSamplingPriority() == PrioritySampling.SAMPLER_KEEP + + cleanup: + tracer.close() + } +} diff --git a/dd-trace-core/src/test/groovy/datadog/trace/common/sampling/LlmObsStandaloneSamplerTest.groovy b/dd-trace-core/src/test/groovy/datadog/trace/common/sampling/LlmObsStandaloneSamplerTest.groovy new file mode 100644 index 00000000000..86be1207cc7 --- /dev/null +++ b/dd-trace-core/src/test/groovy/datadog/trace/common/sampling/LlmObsStandaloneSamplerTest.groovy @@ -0,0 +1,62 @@ +package datadog.trace.common.sampling + +import datadog.trace.api.ProductTraceSource +import datadog.trace.bootstrap.instrumentation.api.Tags +import datadog.trace.common.writer.ListWriter +import datadog.trace.core.test.DDCoreSpecification +import datadog.trace.api.sampling.PrioritySampling + +class LlmObsStandaloneSamplerTest extends DDCoreSpecification { + + def writer = new ListWriter() + + void "test LLMOBS spans are kept"() { + setup: + def sampler = new LlmObsStandaloneSampler() + def tracer = tracerBuilder().writer(writer).sampler(sampler).build() + + when: "LLMOBS span" + def span = tracer.buildSpan("llm-call").start() + def scope = tracer.activateSpan(span) + tracer.getTraceSegment().setTagTop(Tags.PROPAGATED_TRACE_SOURCE, ProductTraceSource.LLMOBS) + sampler.setSamplingPriority(span) + scope.close() + + then: + span.getSamplingPriority() == PrioritySampling.SAMPLER_KEEP + + cleanup: + tracer.close() + } + + void "test APM-only spans are dropped"() { + setup: + def sampler = new LlmObsStandaloneSampler() + def tracer = tracerBuilder().writer(writer).sampler(sampler).build() + + when: "APM-only span (no LLMOBS flag)" + def span = tracer.buildSpan("http-request").start() + sampler.setSamplingPriority(span) + + then: + span.getSamplingPriority() == PrioritySampling.SAMPLER_DROP + + cleanup: + tracer.close() + } + + void "test sample method always returns true"() { + setup: + def sampler = new LlmObsStandaloneSampler() + def tracer = tracerBuilder().writer(writer).sampler(sampler).build() + + when: + def span = tracer.buildSpan("test").start() + + then: + sampler.sample(span) == true + + cleanup: + tracer.close() + } +} \ No newline at end of file diff --git a/dd-trace-core/src/test/groovy/datadog/trace/common/sampling/SamplerTest.groovy b/dd-trace-core/src/test/groovy/datadog/trace/common/sampling/SamplerTest.groovy index 706f45ea43f..0bcaeb87b48 100644 --- a/dd-trace-core/src/test/groovy/datadog/trace/common/sampling/SamplerTest.groovy +++ b/dd-trace-core/src/test/groovy/datadog/trace/common/sampling/SamplerTest.groovy @@ -44,16 +44,29 @@ class SamplerTest extends DDSpecification{ sampler instanceof AsmStandaloneSampler } - void "test that AsmStandaloneSampler is not selected when apm tracing and asm not enabled"() { + void "test that LlmObsStandaloneSampler is selected when apm tracing disabled and llmobs enabled"() { setup: System.setProperty("dd.apm.tracing.enabled", "false") + System.setProperty("dd.llmobs.enabled", "true") Config config = new Config() when: Sampler sampler = Sampler.Builder.forConfig(config, null) then: - !(sampler instanceof AsmStandaloneSampler) + sampler instanceof LlmObsStandaloneSampler + } + + void "test that ForcePrioritySampler with SAMPLER_DROP is selected when apm tracing disabled and no other products enabled"() { + setup: + System.setProperty("dd.apm.tracing.enabled", "false") + Config config = new Config() + + when: + Sampler sampler = Sampler.Builder.forConfig(config, null) + + then: + sampler instanceof ForcePrioritySampler } void "test that AsmStandaloneSampler is not selected when apm tracing enabled and asm not enabled"() { diff --git a/internal-api/src/main/java/datadog/trace/api/ProductTraceSource.java b/internal-api/src/main/java/datadog/trace/api/ProductTraceSource.java index 792829dde3a..df95be22538 100644 --- a/internal-api/src/main/java/datadog/trace/api/ProductTraceSource.java +++ b/internal-api/src/main/java/datadog/trace/api/ProductTraceSource.java @@ -22,6 +22,7 @@ public class ProductTraceSource { public static final int DSM = 0x04; public static final int DJM = 0x08; public static final int DBM = 0x10; + public static final int LLMOBS = 0x20; /** Updates the bitfield by setting the bit corresponding to a specific product. */ public static int updateProduct(int bitfield, int product) {