Handle both LLMObs and ASM enabled in standalone sampler

When APM tracing is disabled but both LLMObs and ASM are enabled,
the previous code returned LlmObsStandaloneSampler which never sent
the 1 APM trace/minute required by ASM for billing and service catalog.

Introduce LlmObsAndAsmStandaloneSampler that keeps all LLMObs and
ASM traces while rate-limiting plain APM traces to 1 per minute.

Also clarify the log message for the ASM-only standalone case.

Signed-off-by: matsumo-and <yh134.toisanda@gmail.com>

Author: matsumo-and

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 <T extends CoreSpan<T>> 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 <T extends CoreSpan<T>> 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;
+  }
+}

dd-trace-core/src/main/java/datadog/trace/common/sampling/Sampler.java

@@ -37,11 +37,16 @@ public static Sampler forConfig(final Config config, final TraceConfig traceConf
       Sampler sampler;
       if (config != null) {
         if (!config.isApmTracingEnabled()) {
-          if (config.isLlmObsEnabled()) {
+          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 trace per minute will be sent.");
+            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

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()
+  }
+}