[WIP] Use weak reference in langchain instrumentation span map

instrumentation-genai/opentelemetry-instrumentation-langchain/pyproject.toml

@@ -27,7 +27,8 @@ classifiers = [
 dependencies = [
   "opentelemetry-api >= 1.31.0",
   "opentelemetry-instrumentation ~= 0.57b0",
-  "opentelemetry-semantic-conventions ~= 0.57b0"
+  "opentelemetry-semantic-conventions ~= 0.57b0",
+  "cachetools ~= 5.0"
 ]
 
 [project.optional-dependencies]

instrumentation-genai/opentelemetry-instrumentation-langchain/src/opentelemetry/instrumentation/langchain/span_manager.py

@@ -16,6 +16,7 @@
 from typing import Dict, List, Optional
 from uuid import UUID
 
+from cachetools import TTLCache
 from opentelemetry.semconv._incubating.attributes import (
     gen_ai_attributes as GenAI,
 )
@@ -42,8 +43,8 @@ def __init__(
         self._tracer = tracer
 
         # Map from run_id -> _SpanState, to keep track of spans and parent/child relationships
-        # TODO: Use weak references or a TTL cache to avoid memory leaks in long-running processes. See #3735
-        self.spans: Dict[UUID, _SpanState] = {}
+        # Using a TTL cache to avoid memory leaks in long-running processes where end_span might not be called.
+        self.spans: TTLCache[UUID, _SpanState] = TTLCache(maxsize=1024, ttl=3600)
 
     def _create_span(
         self,
@@ -92,12 +93,16 @@ def create_chat_span(
         return span
 
     def end_span(self, run_id: UUID) -> None:
-        state = self.spans[run_id]
+        state = self.spans.get(run_id)
+        if not state:
+            return
+
         for child_id in state.children:
             child_state = self.spans.get(child_id)
             if child_state:
                 child_state.span.end()
                 del self.spans[child_id]
+
         state.span.end()
         del self.spans[run_id]
 

instrumentation-genai/opentelemetry-instrumentation-langchain/tests/test_span_manager.py

@@ -1,5 +1,6 @@
 import unittest.mock
 import uuid
+import time
 
 import pytest
 
@@ -98,3 +99,28 @@ def test_end_span(self, handler):
         child_mock_span.end.assert_called_once()
         assert run_id not in handler.spans
         assert child_run_id not in handler.spans
+
+    @unittest.mock.patch(
+            "opentelemetry.instrumentation.langchain.span_manager.TTLCache",
+        )
+    def test_span_ttl_expiration(self, mock_ttlc_class, tracer):
+        # Arrange
+        from cachetools import TTLCache
+
+        mock_ttlc_class.side_effect = lambda maxsize, ttl: TTLCache(
+            maxsize=1024, ttl=0.01
+        )
+        handler = _SpanManager(tracer=tracer)
+        run_id = uuid.uuid4()
+
+        # Act
+        handler._create_span(run_id, None, "test_span")
+
+        # Assert: Span is present immediately after creation
+        assert run_id in handler.spans
+
+        # Wait
+        time.sleep(0.02)
+
+        # Assert: Span is gone after TTL expiration
+        assert run_id not in handler.spans