feat: Introduce max span limit to ApiServerSpanExporter

It defaults to previous behavior but config allows to set hard limit on total number of spans in memory

PiperOrigin-RevId: 921329049

Author: damianmomotgoogle

dev/src/main/java/com/google/adk/web/config/OpenTelemetryConfig.java

@@ -17,15 +17,18 @@
 package com.google.adk.web.config;
 
 import com.google.adk.web.service.ApiServerSpanExporter;
+import com.google.adk.web.service.ApiServerSpanExporterConfig;
 import io.opentelemetry.api.OpenTelemetry;
 import io.opentelemetry.api.common.AttributeKey;
 import io.opentelemetry.api.common.Attributes;
 import io.opentelemetry.sdk.OpenTelemetrySdk;
 import io.opentelemetry.sdk.resources.Resource;
 import io.opentelemetry.sdk.trace.SdkTracerProvider;
 import io.opentelemetry.sdk.trace.export.SimpleSpanProcessor;
+import java.util.Optional;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
+import org.springframework.beans.factory.annotation.Value;
 import org.springframework.context.annotation.Bean;
 import org.springframework.context.annotation.Configuration;
 
@@ -35,8 +38,14 @@ public class OpenTelemetryConfig {
   private static final Logger otelLog = LoggerFactory.getLogger(OpenTelemetryConfig.class);
 
   @Bean
-  public ApiServerSpanExporter apiServerSpanExporter() {
-    return new ApiServerSpanExporter();
+  public ApiServerSpanExporterConfig apiServerSpanExporterConfig(
+      @Value("${adk.debug.trace.max-spans:#{null}}") Optional<Integer> maxSpansToKeep) {
+    return ApiServerSpanExporterConfig.builder().maxSpansToKeep(maxSpansToKeep).build();
+  }
+
+  @Bean
+  public ApiServerSpanExporter apiServerSpanExporter(ApiServerSpanExporterConfig config) {
+    return new ApiServerSpanExporter(config);
   }
 
   @Bean(destroyMethod = "shutdown")

dev/src/main/java/com/google/adk/web/service/ApiServerSpanExporter.java

@@ -20,13 +20,13 @@
 import io.opentelemetry.sdk.common.CompletableResultCode;
 import io.opentelemetry.sdk.trace.data.SpanData;
 import io.opentelemetry.sdk.trace.export.SpanExporter;
+import java.util.ArrayDeque;
 import java.util.ArrayList;
 import java.util.Collection;
-import java.util.Collections;
+import java.util.Deque;
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
-import java.util.concurrent.ConcurrentHashMap;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
@@ -41,74 +41,131 @@
 public class ApiServerSpanExporter implements SpanExporter {
   private static final Logger exporterLog = LoggerFactory.getLogger(ApiServerSpanExporter.class);
 
-  private final Map<String, Map<String, Object>> eventIdTraceStorage = new ConcurrentHashMap<>();
+  private final ApiServerSpanExporterConfig config;
+
+  private final Map<String, Integer> eventIdRefCount = new HashMap<>();
+  private final Map<String, Map<String, Object>> eventIdTraceStorage = new HashMap<>();
 
   // Session ID -> Trace IDs -> Trace Object
-  private final Map<String, List<String>> sessionToTraceIdsMap = new ConcurrentHashMap<>();
+  private final Map<String, List<String>> sessionToTraceIdsMap = new HashMap<>();
+
+  private final Deque<SpanData> allExportedSpans = new ArrayDeque<>();
 
-  private final List<SpanData> allExportedSpans = Collections.synchronizedList(new ArrayList<>());
+  public ApiServerSpanExporter() {
+    this(ApiServerSpanExporterConfig.builder().build());
+  }
 
-  public ApiServerSpanExporter() {}
+  public ApiServerSpanExporter(ApiServerSpanExporterConfig config) {
+    this.config = config;
+  }
 
   public Map<String, Object> getEventTraceAttributes(String eventId) {
-    return this.eventIdTraceStorage.get(eventId);
+    synchronized (allExportedSpans) {
+      return this.eventIdTraceStorage.get(eventId);
+    }
   }
 
   public Map<String, List<String>> getSessionToTraceIdsMap() {
-    return this.sessionToTraceIdsMap;
+    synchronized (allExportedSpans) {
+      return new HashMap<>(this.sessionToTraceIdsMap);
+    }
   }
 
   public List<SpanData> getAllExportedSpans() {
-    return this.allExportedSpans;
+    synchronized (allExportedSpans) {
+      return new ArrayList<>(this.allExportedSpans);
+    }
   }
 
   @Override
   public CompletableResultCode export(Collection<SpanData> spans) {
     exporterLog.debug("ApiServerSpanExporter received {} spans to export.", spans.size());
-    List<SpanData> currentBatch = new ArrayList<>(spans);
-    allExportedSpans.addAll(currentBatch);
-
-    for (SpanData span : currentBatch) {
-      String spanName = span.getName();
-      if ("call_llm".equals(spanName)
-          || "send_data".equals(spanName)
-          || (spanName != null && spanName.startsWith("tool_response"))) {
-        String eventId =
-            span.getAttributes().get(AttributeKey.stringKey("gcp.vertex.agent.event_id"));
-        if (eventId != null && !eventId.isEmpty()) {
-          Map<String, Object> attributesMap = new HashMap<>();
-          span.getAttributes().forEach((key, value) -> attributesMap.put(key.getKey(), value));
-          attributesMap.put("trace_id", span.getSpanContext().getTraceId());
-          attributesMap.put("span_id", span.getSpanContext().getSpanId());
-          attributesMap.putIfAbsent("gcp.vertex.agent.event_id", eventId);
-          exporterLog.debug("Storing event-based trace attributes for event_id: {}", eventId);
-          this.eventIdTraceStorage.put(eventId, attributesMap); // Use internal storage
-        } else {
-          exporterLog.trace(
-              "Span {} for event-based trace did not have 'gcp.vertex.agent.event_id'"
-                  + " attribute or it was empty.",
-              spanName);
+
+    synchronized (allExportedSpans) {
+      for (SpanData span : spans) {
+        if (config.maxSpansToKeep().isPresent()
+            && allExportedSpans.size() >= config.maxSpansToKeep().get()) {
+          SpanData evicted = allExportedSpans.pollFirst();
+          if (evicted != null) {
+            handleEviction(evicted);
+          }
         }
+        allExportedSpans.addLast(span);
+        handleAddition(span);
       }
+    }
+    return CompletableResultCode.ofSuccess();
+  }
 
-      if ("call_llm".equals(spanName)) {
-        String sessionId =
-            span.getAttributes().get(AttributeKey.stringKey("gcp.vertex.agent.session_id"));
-        if (sessionId != null && !sessionId.isEmpty()) {
-          String traceId = span.getSpanContext().getTraceId();
-          sessionToTraceIdsMap
-              .computeIfAbsent(sessionId, k -> Collections.synchronizedList(new ArrayList<>()))
-              .add(traceId);
-          exporterLog.trace(
-              "Associated trace_id {} with session_id {} for session tracing", traceId, sessionId);
+  private void handleAddition(SpanData span) {
+    String spanName = span.getName();
+    String eventId = span.getAttributes().get(AttributeKey.stringKey("gcp.vertex.agent.event_id"));
+    boolean isEventTraceSpan =
+        "call_llm".equals(spanName)
+            || "send_data".equals(spanName)
+            || (spanName != null && spanName.startsWith("tool_response"));
+    if (eventId != null && !eventId.isEmpty()) {
+      eventIdRefCount.merge(eventId, 1, Integer::sum);
+      if (isEventTraceSpan) {
+        Map<String, Object> attributesMap = new HashMap<>();
+        span.getAttributes().forEach((key, value) -> attributesMap.put(key.getKey(), value));
+        attributesMap.put("trace_id", span.getSpanContext().getTraceId());
+        attributesMap.put("span_id", span.getSpanContext().getSpanId());
+        attributesMap.putIfAbsent("gcp.vertex.agent.event_id", eventId);
+        exporterLog.debug("Storing event-based trace attributes for event_id: {}", eventId);
+        eventIdTraceStorage.put(eventId, attributesMap);
+      }
+    } else if (isEventTraceSpan) {
+      exporterLog.trace(
+          "Span {} for event-based trace did not have 'gcp.vertex.agent.event_id'"
+              + " attribute or it was empty.",
+          spanName);
+    }
+
+    if ("call_llm".equals(spanName)) {
+      String sessionId =
+          span.getAttributes().get(AttributeKey.stringKey("gcp.vertex.agent.session_id"));
+      if (sessionId != null && !sessionId.isEmpty()) {
+        String traceId = span.getSpanContext().getTraceId();
+        sessionToTraceIdsMap.computeIfAbsent(sessionId, k -> new ArrayList<>()).add(traceId);
+        exporterLog.trace(
+            "Associated trace_id {} with session_id {} for session tracing", traceId, sessionId);
+      } else {
+        exporterLog.trace(
+            "Span {} for session trace did not have 'gcp.vertex.agent.session_id' attribute.",
+            spanName);
+      }
+    }
+  }
+
+  private void handleEviction(SpanData span) {
+    String spanName = span.getName();
+    String eventId = span.getAttributes().get(AttributeKey.stringKey("gcp.vertex.agent.event_id"));
+    if (eventId != null && !eventId.isEmpty()) {
+      Integer count = eventIdRefCount.get(eventId);
+      if (count != null) {
+        if (count <= 1) {
+          eventIdRefCount.remove(eventId);
+          eventIdTraceStorage.remove(eventId);
         } else {
-          exporterLog.trace(
-              "Span {} for session trace did not have 'gcp.vertex.agent.session_id' attribute.",
-              spanName);
+          eventIdRefCount.put(eventId, count - 1);
+        }
+      }
+    }
+
+    if ("call_llm".equals(spanName)) {
+      String sessionId =
+          span.getAttributes().get(AttributeKey.stringKey("gcp.vertex.agent.session_id"));
+      if (sessionId != null && !sessionId.isEmpty()) {
+        List<String> traceIds = sessionToTraceIdsMap.get(sessionId);
+        if (traceIds != null) {
+          traceIds.remove(span.getSpanContext().getTraceId());
+          if (traceIds.isEmpty()) {
+            sessionToTraceIdsMap.remove(sessionId);
+          }
         }
       }
     }
-    return CompletableResultCode.ofSuccess();
   }
 
   @Override
@@ -119,7 +176,12 @@ public CompletableResultCode flush() {
   @Override
   public CompletableResultCode shutdown() {
     exporterLog.debug("Shutting down ApiServerSpanExporter.");
-    // no need to clear storage on shutdown, as everything is currently stored in memory.
+    synchronized (allExportedSpans) {
+      allExportedSpans.clear();
+      eventIdRefCount.clear();
+      eventIdTraceStorage.clear();
+      sessionToTraceIdsMap.clear();
+    }
     return CompletableResultCode.ofSuccess();
   }
 }

dev/src/main/java/com/google/adk/web/service/ApiServerSpanExporterConfig.java

@@ -0,0 +1,59 @@
+/*
+ * Copyright 2025 Google LLC
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.google.adk.web.service;
+
+import com.google.auto.value.AutoValue;
+import java.util.Optional;
+
+/** Configuration for {@link ApiServerSpanExporter}. */
+@AutoValue
+public abstract class ApiServerSpanExporterConfig {
+
+  /**
+   * The maximum number of spans to keep in memory. When the limit is reached, the oldest spans are
+   * evicted (FIFO). If empty, no limit is enforced and spans accumulate without bound.
+   *
+   * <p>When set, the value must be a positive integer ({@code >= 1}).
+   */
+  public abstract Optional<Integer> maxSpansToKeep();
+
+  public static Builder builder() {
+    return new AutoValue_ApiServerSpanExporterConfig.Builder();
+  }
+
+  /** Builder for {@link ApiServerSpanExporterConfig}. */
+  @AutoValue.Builder
+  public abstract static class Builder {
+    public abstract Builder maxSpansToKeep(Optional<Integer> maxSpansToKeep);
+
+    abstract ApiServerSpanExporterConfig autoBuild();
+
+    public final ApiServerSpanExporterConfig build() {
+      ApiServerSpanExporterConfig config = autoBuild();
+      config
+          .maxSpansToKeep()
+          .ifPresent(
+              max -> {
+                if (max < 1) {
+                  throw new IllegalArgumentException(
+                      "maxSpansToKeep must be >= 1 when set, got: " + max);
+                }
+              });
+      return config;
+    }
+  }
+}