refactor the logic touching the trace ids for event spans

Author: pawel-maciejczek

server/adkrest/controllers/debug_test.go

@@ -92,7 +92,7 @@ func TestSessionSpansHandler(t *testing.T) {
 		t.Run(tt.name, func(t *testing.T) {
 			eventID := "test-event"
 			opName := semconv.GenAIOperationNameExecuteTool.Value.AsString()
-			testTelemetry := setupTestTelemetry()
+			testTelemetry := setupTestTelemetry(t)
 
 			apiController := controllers.NewDebugAPIController(nil, nil, testTelemetry.dt)
 			req, err := http.NewRequest(http.MethodGet, "/debug/sessions/"+tt.reqSessionID+"/spans", nil)
@@ -204,7 +204,7 @@ func TestEventSpanHandler(t *testing.T) {
 	for _, tt := range tc {
 		t.Run(tt.name, func(t *testing.T) {
 			sessionID := "test-session"
-			testTelemetry := setupTestTelemetry()
+			testTelemetry := setupTestTelemetry(t)
 
 			apiController := controllers.NewDebugAPIController(nil, nil, testTelemetry.dt)
 			req, err := http.NewRequest(http.MethodGet, "/debug/events/"+tt.reqEventID+"/span", nil)
@@ -258,8 +258,11 @@ type testTelemetry struct {
 	lp     *sdklog.LoggerProvider
 }
 
-func setupTestTelemetry() *testTelemetry {
-	dt := services.NewDebugTelemetry(nil)
+func setupTestTelemetry(t *testing.T) *testTelemetry {
+	dt, err := services.NewDebugTelemetryWithConfig(nil)
+	if err != nil {
+		t.Fatalf("failed to create debug telemetry: %v", err)
+	}
 
 	tp := sdktrace.NewTracerProvider(sdktrace.WithSpanProcessor(dt.SpanProcessor()))
 	lp := sdklog.NewLoggerProvider(sdklog.WithProcessor(dt.LogProcessor()))

server/adkrest/internal/services/debugtelemetry.go

@@ -156,15 +156,22 @@ func (s *spanStore) getSpansByEventID(id string) []DebugSpan {
 	defer s.mu.RUnlock()
 	// Create a copy of the slice to avoid race conditions.
 	records := slices.Clone(s.recordsByEventID[id])
-	seenTraces := make(map[string]bool)
+	s.touchTraces(records)
+	return convertRecords(records)
+}
+
+// touchTraces marks traces as recently used. Required because fetching by event ID bypasses the trace LRU cache.
+func (s *spanStore) touchTraces(records []*spanRecord) {
+	// touchedTraces is used to avoid touching the same trace multiple times.
+	touchedTraces := make(map[string]bool)
 	for _, r := range records {
 		traceIDStr := r.Context.TraceID().String()
-		if traceIDStr != "" && !seenTraces[traceIDStr] {
+		if traceIDStr != "" && !touchedTraces[traceIDStr] {
+			touchedTraces[traceIDStr] = true
+			// Get the trace to update its access time in the LRU cache, ignore the result.
 			s.recordsByTraceID.Get(traceIDStr)
-			seenTraces[traceIDStr] = true
 		}
 	}
-	return convertRecords(records)
 }
 
 func (s *spanStore) getSpansBySessionID(sessionID string) []DebugSpan {
@@ -304,7 +311,7 @@ func (s *spanStore) evict(traceID string, spans []*spanRecord) {
 func (s *spanStore) evictRecordsByEventID(eventID string, span *spanRecord) {
 	records := s.recordsByEventID[eventID]
 	records = slices.DeleteFunc(records, func(r *spanRecord) bool {
-		return r == span
+		return r.Context.SpanID() == span.Context.SpanID()
 	})
 	if len(records) == 0 {
 		delete(s.recordsByEventID, eventID)