Skip to content

Commit 1e2bd37

Browse files
damianmomotgooglecopybara-github
damianmomotgoogle
authored and
copybara-github
committed
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: 920891402
1 parent cc3b9ce commit 1e2bd37

4 files changed

Lines changed: 331 additions & 45 deletions

File tree

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

Lines changed: 11 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -17,15 +17,18 @@
1717
package com.google.adk.web.config;
1818

1919
import com.google.adk.web.service.ApiServerSpanExporter;
20+
import com.google.adk.web.service.ApiServerSpanExporterConfig;
2021
import io.opentelemetry.api.OpenTelemetry;
2122
import io.opentelemetry.api.common.AttributeKey;
2223
import io.opentelemetry.api.common.Attributes;
2324
import io.opentelemetry.sdk.OpenTelemetrySdk;
2425
import io.opentelemetry.sdk.resources.Resource;
2526
import io.opentelemetry.sdk.trace.SdkTracerProvider;
2627
import io.opentelemetry.sdk.trace.export.SimpleSpanProcessor;
28+
import java.util.Optional;
2729
import org.slf4j.Logger;
2830
import org.slf4j.LoggerFactory;
31+
import org.springframework.beans.factory.annotation.Value;
2932
import org.springframework.context.annotation.Bean;
3033
import org.springframework.context.annotation.Configuration;
3134

@@ -35,8 +38,14 @@ public class OpenTelemetryConfig {
3538
private static final Logger otelLog = LoggerFactory.getLogger(OpenTelemetryConfig.class);
3639

3740
@Bean
38-
public ApiServerSpanExporter apiServerSpanExporter() {
39-
return new ApiServerSpanExporter();
41+
public ApiServerSpanExporterConfig apiServerSpanExporterConfig(
42+
@Value("${adk.debug.trace.max-spans:#{null}}") Optional<Integer> maxSpansToKeep) {
43+
return ApiServerSpanExporterConfig.builder().maxSpansToKeep(maxSpansToKeep).build();
44+
}
45+
46+
@Bean
47+
public ApiServerSpanExporter apiServerSpanExporter(ApiServerSpanExporterConfig config) {
48+
return new ApiServerSpanExporter(config);
4049
}
4150

4251
@Bean(destroyMethod = "shutdown")

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

Lines changed: 95 additions & 43 deletions
Original file line numberDiff line numberDiff line change
@@ -16,17 +16,18 @@
1616

1717
package com.google.adk.web.service;
1818

19+
import com.google.errorprone.annotations.CanIgnoreReturnValue;
1920
import io.opentelemetry.api.common.AttributeKey;
2021
import io.opentelemetry.sdk.common.CompletableResultCode;
2122
import io.opentelemetry.sdk.trace.data.SpanData;
2223
import io.opentelemetry.sdk.trace.export.SpanExporter;
24+
import java.util.ArrayDeque;
2325
import java.util.ArrayList;
2426
import java.util.Collection;
25-
import java.util.Collections;
27+
import java.util.Deque;
2628
import java.util.HashMap;
2729
import java.util.List;
2830
import java.util.Map;
29-
import java.util.concurrent.ConcurrentHashMap;
3031
import org.slf4j.Logger;
3132
import org.slf4j.LoggerFactory;
3233

@@ -41,74 +42,119 @@
4142
public class ApiServerSpanExporter implements SpanExporter {
4243
private static final Logger exporterLog = LoggerFactory.getLogger(ApiServerSpanExporter.class);
4344

44-
private final Map<String, Map<String, Object>> eventIdTraceStorage = new ConcurrentHashMap<>();
45+
private final ApiServerSpanExporterConfig config;
46+
47+
private final Map<String, Integer> eventIdRefCount = new HashMap<>();
48+
private final Map<String, Map<String, Object>> eventIdTraceStorage = new HashMap<>();
4549

4650
// Session ID -> Trace IDs -> Trace Object
47-
private final Map<String, List<String>> sessionToTraceIdsMap = new ConcurrentHashMap<>();
51+
private final Map<String, List<String>> sessionToTraceIdsMap = new HashMap<>();
52+
53+
private final Deque<SpanData> allExportedSpans = new ArrayDeque<>();
4854

49-
private final List<SpanData> allExportedSpans = Collections.synchronizedList(new ArrayList<>());
55+
public ApiServerSpanExporter() {
56+
this(ApiServerSpanExporterConfig.builder().build());
57+
}
5058

51-
public ApiServerSpanExporter() {}
59+
public ApiServerSpanExporter(ApiServerSpanExporterConfig config) {
60+
this.config = config;
61+
}
5262

5363
public Map<String, Object> getEventTraceAttributes(String eventId) {
54-
return this.eventIdTraceStorage.get(eventId);
64+
synchronized (allExportedSpans) {
65+
return this.eventIdTraceStorage.get(eventId);
66+
}
5567
}
5668

5769
public Map<String, List<String>> getSessionToTraceIdsMap() {
58-
return this.sessionToTraceIdsMap;
70+
synchronized (allExportedSpans) {
71+
return new HashMap<>(this.sessionToTraceIdsMap);
72+
}
5973
}
6074

6175
public List<SpanData> getAllExportedSpans() {
62-
return this.allExportedSpans;
76+
synchronized (allExportedSpans) {
77+
return new ArrayList<>(this.allExportedSpans);
78+
}
6379
}
6480

6581
@Override
82+
@CanIgnoreReturnValue
6683
public CompletableResultCode export(Collection<SpanData> spans) {
6784
exporterLog.debug("ApiServerSpanExporter received {} spans to export.", spans.size());
68-
List<SpanData> currentBatch = new ArrayList<>(spans);
69-
allExportedSpans.addAll(currentBatch);
7085

71-
for (SpanData span : currentBatch) {
72-
String spanName = span.getName();
86+
synchronized (allExportedSpans) {
87+
for (SpanData span : spans) {
88+
if (config.maxSpansToKeep().isPresent()
89+
&& allExportedSpans.size() >= config.maxSpansToKeep().get()) {
90+
SpanData evicted = allExportedSpans.pollFirst();
91+
if (evicted != null) {
92+
handleEviction(evicted);
93+
}
94+
}
95+
allExportedSpans.addLast(span);
96+
handleAddition(span);
97+
}
98+
}
99+
return CompletableResultCode.ofSuccess();
100+
}
101+
102+
private void handleAddition(SpanData span) {
103+
String spanName = span.getName();
104+
String eventId = span.getAttributes().get(AttributeKey.stringKey("gcp.vertex.agent.event_id"));
105+
if (eventId != null && !eventId.isEmpty()) {
106+
eventIdRefCount.merge(eventId, 1, Integer::sum);
73107
if ("call_llm".equals(spanName)
74108
|| "send_data".equals(spanName)
75109
|| (spanName != null && spanName.startsWith("tool_response"))) {
76-
String eventId =
77-
span.getAttributes().get(AttributeKey.stringKey("gcp.vertex.agent.event_id"));
78-
if (eventId != null && !eventId.isEmpty()) {
79-
Map<String, Object> attributesMap = new HashMap<>();
80-
span.getAttributes().forEach((key, value) -> attributesMap.put(key.getKey(), value));
81-
attributesMap.put("trace_id", span.getSpanContext().getTraceId());
82-
attributesMap.put("span_id", span.getSpanContext().getSpanId());
83-
attributesMap.putIfAbsent("gcp.vertex.agent.event_id", eventId);
84-
exporterLog.debug("Storing event-based trace attributes for event_id: {}", eventId);
85-
this.eventIdTraceStorage.put(eventId, attributesMap); // Use internal storage
110+
Map<String, Object> attributesMap = new HashMap<>();
111+
span.getAttributes().forEach((key, value) -> attributesMap.put(key.getKey(), value));
112+
attributesMap.put("trace_id", span.getSpanContext().getTraceId());
113+
attributesMap.put("span_id", span.getSpanContext().getSpanId());
114+
attributesMap.putIfAbsent("gcp.vertex.agent.event_id", eventId);
115+
eventIdTraceStorage.put(eventId, attributesMap);
116+
}
117+
}
118+
119+
if ("call_llm".equals(spanName)) {
120+
String sessionId =
121+
span.getAttributes().get(AttributeKey.stringKey("gcp.vertex.agent.session_id"));
122+
if (sessionId != null && !sessionId.isEmpty()) {
123+
sessionToTraceIdsMap
124+
.computeIfAbsent(sessionId, k -> new ArrayList<>())
125+
.add(span.getSpanContext().getTraceId());
126+
}
127+
}
128+
}
129+
130+
private void handleEviction(SpanData span) {
131+
String spanName = span.getName();
132+
String eventId = span.getAttributes().get(AttributeKey.stringKey("gcp.vertex.agent.event_id"));
133+
if (eventId != null && !eventId.isEmpty()) {
134+
Integer count = eventIdRefCount.get(eventId);
135+
if (count != null) {
136+
if (count <= 1) {
137+
eventIdRefCount.remove(eventId);
138+
eventIdTraceStorage.remove(eventId);
86139
} else {
87-
exporterLog.trace(
88-
"Span {} for event-based trace did not have 'gcp.vertex.agent.event_id'"
89-
+ " attribute or it was empty.",
90-
spanName);
140+
eventIdRefCount.put(eventId, count - 1);
91141
}
92142
}
143+
}
93144

94-
if ("call_llm".equals(spanName)) {
95-
String sessionId =
96-
span.getAttributes().get(AttributeKey.stringKey("gcp.vertex.agent.session_id"));
97-
if (sessionId != null && !sessionId.isEmpty()) {
98-
String traceId = span.getSpanContext().getTraceId();
99-
sessionToTraceIdsMap
100-
.computeIfAbsent(sessionId, k -> Collections.synchronizedList(new ArrayList<>()))
101-
.add(traceId);
102-
exporterLog.trace(
103-
"Associated trace_id {} with session_id {} for session tracing", traceId, sessionId);
104-
} else {
105-
exporterLog.trace(
106-
"Span {} for session trace did not have 'gcp.vertex.agent.session_id' attribute.",
107-
spanName);
145+
if ("call_llm".equals(spanName)) {
146+
String sessionId =
147+
span.getAttributes().get(AttributeKey.stringKey("gcp.vertex.agent.session_id"));
148+
if (sessionId != null && !sessionId.isEmpty()) {
149+
List<String> traceIds = sessionToTraceIdsMap.get(sessionId);
150+
if (traceIds != null) {
151+
traceIds.remove(span.getSpanContext().getTraceId());
152+
if (traceIds.isEmpty()) {
153+
sessionToTraceIdsMap.remove(sessionId);
154+
}
108155
}
109156
}
110157
}
111-
return CompletableResultCode.ofSuccess();
112158
}
113159

114160
@Override
@@ -117,9 +163,15 @@ public CompletableResultCode flush() {
117163
}
118164

119165
@Override
166+
@CanIgnoreReturnValue
120167
public CompletableResultCode shutdown() {
121168
exporterLog.debug("Shutting down ApiServerSpanExporter.");
122-
// no need to clear storage on shutdown, as everything is currently stored in memory.
169+
synchronized (allExportedSpans) {
170+
allExportedSpans.clear();
171+
eventIdRefCount.clear();
172+
eventIdTraceStorage.clear();
173+
sessionToTraceIdsMap.clear();
174+
}
123175
return CompletableResultCode.ofSuccess();
124176
}
125177
}

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

Lines changed: 39 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,39 @@
1+
/*
2+
* Copyright 2025 Google LLC
3+
*
4+
* Licensed under the Apache License, Version 2.0 (the "License");
5+
* you may not use this file except in compliance with the License.
6+
* You may obtain a copy of the License at
7+
*
8+
* http://www.apache.org/licenses/LICENSE-2.0
9+
*
10+
* Unless required by applicable law or agreed to in writing, software
11+
* distributed under the License is distributed on an "AS IS" BASIS,
12+
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13+
* See the License for the specific language governing permissions and
14+
* limitations under the License.
15+
*/
16+
17+
package com.google.adk.web.service;
18+
19+
import com.google.auto.value.AutoValue;
20+
import java.util.Optional;
21+
22+
/** Configuration for {@link ApiServerSpanExporter}. */
23+
@AutoValue
24+
public abstract class ApiServerSpanExporterConfig {
25+
26+
public abstract Optional<Integer> maxSpansToKeep();
27+
28+
public static Builder builder() {
29+
return new AutoValue_ApiServerSpanExporterConfig.Builder();
30+
}
31+
32+
/** Builder for {@link ApiServerSpanExporterConfig}. */
33+
@AutoValue.Builder
34+
public abstract static class Builder {
35+
public abstract Builder maxSpansToKeep(Optional<Integer> maxSpansToKeep);
36+
37+
public abstract ApiServerSpanExporterConfig build();
38+
}
39+
}

0 commit comments

Comments
 (0)