Aggregate LLMObs span.finished telemetry counts per interval

internal-api/src/main/java/datadog/trace/api/telemetry/LLMObsMetricCollector.java

@@ -7,11 +7,24 @@
 import java.util.Collection;
 import java.util.Collections;
 import java.util.List;
+import java.util.Map;
 import java.util.concurrent.ArrayBlockingQueue;
 import java.util.concurrent.BlockingQueue;
+import java.util.concurrent.ConcurrentHashMap;
+import java.util.concurrent.atomic.LongAdder;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
+/**
+ * Collects telemetry metrics for LLM Observability spans.
+ *
+ * <p>Counts are aggregated per tag combination in-process and emitted as one point per metrics
+ * interval. Emitting one point of value 1 per span instead would under-report badly: points are
+ * timestamped at second granularity, so every point a series produces within the same second
+ * collapses to a single value at the metrics intake, capping the reported rate at roughly 1/s per
+ * series regardless of the real span rate. This matches how {@link CoreMetricCollector} and the
+ * other tracers (dd-trace-py, dd-trace-js) report counts.
+ */
 public final class LLMObsMetricCollector
     implements MetricCollector<LLMObsMetricCollector.LLMObsMetric> {
   private static final String METRIC_NAMESPACE = "mlobs";
@@ -35,19 +48,38 @@ public static LLMObsMetricCollector get() {
   private static final String HAS_SESSION_ID_TRUE = "has_session_id:1";
   private static final String HAS_SESSION_ID_FALSE = "has_session_id:0";
 
+  /**
+   * Upper bound on the number of distinct tag combinations tracked concurrently. Tag values are
+   * drawn from bounded sets (a handful of integrations and span kinds, plus four booleans), so this
+   * is only a guard against an unexpected high-cardinality source. It is also kept low enough that
+   * several {@link #prepareMetrics()} intervals can be staged without overflowing {@link
+   * MetricCollector#RAW_QUEUE_SIZE} before the next {@link #drain()}.
+   */
+  static final int MAX_TAG_COMBINATIONS = 128;
+
   private final BlockingQueue<LLMObsMetric> metricsQueue;
   private final DDCache<String, String> integrationTagCache;
   private final DDCache<String, String> spanKindTagCache;
+  /**
+   * Counter per tag combination, aggregated in-process and flushed once per metrics interval by
+   * {@link #prepareMetrics()}. Counting here rather than enqueuing one entry per span is what keeps
+   * the reported count accurate at high span rates.
+   */
+  private final ConcurrentHashMap<List<String>, LongAdder> spanFinishedCounters;
 
   private LLMObsMetricCollector() {
     this.metricsQueue = new ArrayBlockingQueue<>(RAW_QUEUE_SIZE);
     this.integrationTagCache = DDCaches.newFixedSizeCache(8);
     this.spanKindTagCache = DDCaches.newFixedSizeCache(8);
+    this.spanFinishedCounters = new ConcurrentHashMap<>();
   }
 
   /**
    * Record a span finished metric for LLMObs telemetry.
    *
+   * <p>This only increments an in-process counter. The counter is converted into a single telemetry
+   * metric per tag combination by {@link #prepareMetrics()}, once per metrics interval.
+   *
    * @param integration the integration name (e.g., "openai")
    * @param spanKind the span kind (e.g., "llm", "embedding")
    * @param isRootSpan whether this is a root span
@@ -74,16 +106,49 @@ public void recordSpanFinished(
             isAutoInstrumented ? AUTOINSTRUMENTED_TRUE : AUTOINSTRUMENTED_FALSE,
             hasError ? ERROR_TRUE : ERROR_FALSE,
             hasSessionId ? HAS_SESSION_ID_TRUE : HAS_SESSION_ID_FALSE);
-    LLMObsMetric metric =
-        new LLMObsMetric(METRIC_NAMESPACE, true, SPAN_FINISHED_METRIC, COUNT_METRIC_TYPE, 1L, tags);
-    if (!metricsQueue.offer(metric)) {
-      log.debug("Unable to add telemetry metric {} for {}", SPAN_FINISHED_METRIC, integration);
+
+    LongAdder counter = spanFinishedCounters.get(tags);
+    if (counter == null) {
+      // Soft bound: concurrent recorders may overshoot slightly, which is fine for a guard.
+      if (spanFinishedCounters.size() >= MAX_TAG_COMBINATIONS) {
+        log.debug(
+            "Dropping telemetry metric {} for {}: tag combination limit ({}) reached",
+            SPAN_FINISHED_METRIC,
+            integration,
+            MAX_TAG_COMBINATIONS);
+        return;
+      }
+      counter = spanFinishedCounters.computeIfAbsent(tags, key -> new LongAdder());
     }
+    counter.increment();
   }
 
   @Override
   public void prepareMetrics() {
-    // metrics are added directly via recordSpanFinished; no additional preparation needed
+    // Entries are never removed: a recorder thread may already hold a reference to a LongAdder, so
+    // removing it here would silently drop a concurrent increment. Tag values come from bounded
+    // sets, so retaining idle combinations costs at most MAX_TAG_COMBINATIONS entries.
+    for (Map.Entry<List<String>, LongAdder> entry : spanFinishedCounters.entrySet()) {
+      long value = entry.getValue().sumThenReset();
+      if (value == 0) {
+        continue;
+      }
+      LLMObsMetric metric =
+          new LLMObsMetric(
+              METRIC_NAMESPACE,
+              true,
+              SPAN_FINISHED_METRIC,
+              COUNT_METRIC_TYPE,
+              value,
+              entry.getKey());
+      if (!metricsQueue.offer(metric)) {
+        // Queue is full; give the count back to the counter so it is reported in a later interval
+        // instead of being lost, and stop staging for now.
+        entry.getValue().add(value);
+        log.debug("Unable to add telemetry metric {}: queue is full", SPAN_FINISHED_METRIC);
+        break;
+      }
+    }
   }
 
   @Override
@@ -96,6 +161,12 @@ public Collection<LLMObsMetric> drain() {
     return drained;
   }
 
+  /** Clears all staged counters and metrics. Visible for testing only. */
+  public void resetForTesting() {
+    spanFinishedCounters.clear();
+    metricsQueue.clear();
+  }
+
   public static class LLMObsMetric extends MetricCollector.Metric {
     public LLMObsMetric(
         String namespace,

internal-api/src/test/java/datadog/trace/api/telemetry/LLMObsMetricCollectorTest.java

@@ -0,0 +1,167 @@
+package datadog.trace.api.telemetry;
+
+import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertTrue;
+
+import java.util.ArrayList;
+import java.util.Arrays;
+import java.util.Collection;
+import java.util.HashSet;
+import java.util.List;
+import org.junit.jupiter.api.AfterEach;
+import org.junit.jupiter.api.BeforeEach;
+import org.junit.jupiter.api.Test;
+
+class LLMObsMetricCollectorTest {
+  private final LLMObsMetricCollector collector = LLMObsMetricCollector.get();
+
+  @BeforeEach
+  void clearStaleMetrics() {
+    collector.resetForTesting();
+  }
+
+  @AfterEach
+  void clearMetrics() {
+    collector.resetForTesting();
+  }
+
+  @Test
+  void noMetricsDrainsEmptyList() {
+    collector.prepareMetrics();
+
+    assertTrue(collector.drain().isEmpty());
+  }
+
+  @Test
+  void recordsOneMetricPerDistinctTagCombination() {
+    collector.recordSpanFinished("openai", "llm", true, true, false, false);
+    collector.recordSpanFinished("openai", "llm", false, true, false, true);
+    collector.recordSpanFinished("anthropic", "embedding", true, false, true, false);
+    collector.prepareMetrics();
+
+    List<LLMObsMetricCollector.LLMObsMetric> metrics = sorted(collector.drain());
+
+    assertEquals(3, metrics.size());
+    for (LLMObsMetricCollector.LLMObsMetric metric : metrics) {
+      assertEquals("mlobs", metric.namespace);
+      assertEquals("span.finished", metric.metricName);
+      assertEquals("count", metric.type);
+      assertEquals(1L, metric.value);
+    }
+    assertEquals(
+        new HashSet<>(
+            Arrays.asList(
+                new HashSet<>(
+                    Arrays.asList(
+                        "integration:openai",
+                        "span_kind:llm",
+                        "is_root_span:1",
+                        "autoinstrumented:1",
+                        "error:0",
+                        "has_session_id:0")),
+                new HashSet<>(
+                    Arrays.asList(
+                        "integration:openai",
+                        "span_kind:llm",
+                        "is_root_span:0",
+                        "autoinstrumented:1",
+                        "error:0",
+                        "has_session_id:1")),
+                new HashSet<>(
+                    Arrays.asList(
+                        "integration:anthropic",
+                        "span_kind:embedding",
+                        "is_root_span:1",
+                        "autoinstrumented:0",
+                        "error:1",
+                        "has_session_id:0")))),
+        tagSets(metrics));
+  }
+
+  @Test
+  void aggregatesIdenticalTagCombinationsIntoASingleCount() {
+    collector.recordSpanFinished("openai", "llm", true, true, false, false);
+    collector.recordSpanFinished("openai", "llm", true, true, false, false);
+    collector.recordSpanFinished("openai", "llm", true, true, false, false);
+    collector.prepareMetrics();
+
+    Collection<LLMObsMetricCollector.LLMObsMetric> metrics = collector.drain();
+
+    assertEquals(1, metrics.size());
+    LLMObsMetricCollector.LLMObsMetric metric = metrics.iterator().next();
+    assertEquals("count", metric.type);
+    assertEquals(3L, metric.value);
+  }
+
+  /**
+   * Regression test for the reported ~100x under-count: emitting one raw metric per span both
+   * overflowed the bounded raw queue and collapsed to ~1/s per series at the metrics intake, since
+   * points are timestamped at second granularity. The count must survive well past {@link
+   * MetricCollector#RAW_QUEUE_SIZE} spans in a single interval.
+   */
+  @Test
+  void reportsExactCountWellBeyondRawQueueSize() {
+    int spans = MetricCollector.RAW_QUEUE_SIZE * 10;
+    for (int i = 0; i < spans; i++) {
+      collector.recordSpanFinished("openai", "llm", false, true, false, false);
+    }
+    collector.prepareMetrics();
+
+    Collection<LLMObsMetricCollector.LLMObsMetric> metrics = collector.drain();
+
+    assertEquals(1, metrics.size());
+    assertEquals((long) spans, metrics.iterator().next().value);
+  }
+
+  @Test
+  void countersResetBetweenIntervals() {
+    collector.recordSpanFinished("openai", "llm", true, true, false, false);
+    collector.prepareMetrics();
+    assertEquals(1, collector.drain().size());
+
+    collector.prepareMetrics();
+
+    assertTrue(collector.drain().isEmpty(), "an idle interval must not re-report a stale count");
+  }
+
+  @Test
+  void reportsCountsAccumulatedAcrossIntervalsWithoutADrain() {
+    collector.recordSpanFinished("openai", "llm", true, true, false, false);
+    collector.prepareMetrics();
+    collector.recordSpanFinished("openai", "llm", true, true, false, false);
+    collector.recordSpanFinished("openai", "llm", true, true, false, false);
+    collector.prepareMetrics();
+
+    List<LLMObsMetricCollector.LLMObsMetric> metrics = sorted(collector.drain());
+
+    assertEquals(2, metrics.size());
+    assertEquals(1L, metrics.get(0).value);
+    assertEquals(2L, metrics.get(1).value);
+  }
+
+  @Test
+  void boundsTheNumberOfTrackedTagCombinations() {
+    for (int i = 0; i < LLMObsMetricCollector.MAX_TAG_COMBINATIONS * 2; i++) {
+      collector.recordSpanFinished("integration-" + i, "llm", true, true, false, false);
+    }
+    collector.prepareMetrics();
+
+    assertEquals(LLMObsMetricCollector.MAX_TAG_COMBINATIONS, collector.drain().size());
+  }
+
+  private static List<LLMObsMetricCollector.LLMObsMetric> sorted(
+      Collection<LLMObsMetricCollector.LLMObsMetric> metrics) {
+    List<LLMObsMetricCollector.LLMObsMetric> sorted = new ArrayList<>(metrics);
+    sorted.sort((a, b) -> Long.compare(a.value.longValue(), b.value.longValue()));
+    return sorted;
+  }
+
+  private static HashSet<HashSet<String>> tagSets(
+      Collection<LLMObsMetricCollector.LLMObsMetric> metrics) {
+    HashSet<HashSet<String>> tagSets = new HashSet<>();
+    for (LLMObsMetricCollector.LLMObsMetric metric : metrics) {
+      tagSets.add(new HashSet<>(metric.tags));
+    }
+    return tagSets;
+  }
+}