feat(asyncio): add weakref-based tracking for instrumented objects

Author: allen-k1m

instrumentation/opentelemetry-instrumentation-asyncio/src/opentelemetry/instrumentation/asyncio/__init__.py

@@ -93,6 +93,10 @@ def func():
 
 from wrapt import wrap_function_wrapper as _wrap
 
+from opentelemetry.instrumentation.asyncio.instrumentation_state import (
+    _is_instrumented,
+    _mark_instrumented,
+)
 from opentelemetry.instrumentation.asyncio.package import _instruments
 from opentelemetry.instrumentation.asyncio.utils import (
     get_coros_to_trace,
@@ -107,7 +111,7 @@ def func():
 from opentelemetry.trace.status import Status, StatusCode
 
 ASYNCIO_PREFIX = "asyncio"
-_IS_ASYNCIO_INSTRUMENTED_ATTRIBUTE = "_otel_asyncio_instrumented"
+
 
 class AsyncioInstrumentor(BaseInstrumentor):
     """
@@ -240,10 +244,10 @@ def trace_to_thread(self, func: callable):
         """
         Trace a function, but if already instrumented, skip double-wrapping.
         """
-        if getattr(func, _IS_ASYNCIO_INSTRUMENTED_ATTRIBUTE, False):
+        if _is_instrumented(func):
             return func
 
-        setattr(func, _IS_ASYNCIO_INSTRUMENTED_ATTRIBUTE, True)
+        _mark_instrumented(func)
 
         start = default_timer()
         func_name = getattr(func, "__name__", None)
@@ -281,10 +285,10 @@ async def trace_coroutine(self, coro):
         Wrap a coroutine so that we measure its duration, metrics, etc.
         If already instrumented, simply 'await coro' to preserve call behavior.
         """
-        if getattr(coro, _IS_ASYNCIO_INSTRUMENTED_ATTRIBUTE, False):
+        if _is_instrumented(coro):
             return await coro
 
-        setattr(coro, _IS_ASYNCIO_INSTRUMENTED_ATTRIBUTE, True)
+        _mark_instrumented(coro)
 
         if not hasattr(coro, "__name__"):
             return await coro
@@ -322,10 +326,10 @@ def trace_future(self, future):
         """
         Wrap a Future's done callback. If already instrumented, skip re-wrapping.
         """
-        if getattr(future, _IS_ASYNCIO_INSTRUMENTED_ATTRIBUTE, False):
+        if _is_instrumented(future):
             return future
 
-        setattr(future, _IS_ASYNCIO_INSTRUMENTED_ATTRIBUTE, True)
+        _mark_instrumented(future)
 
         start = default_timer()
         span = (

instrumentation/opentelemetry-instrumentation-asyncio/src/opentelemetry/instrumentation/asyncio/instrumentation_state.py

@@ -0,0 +1,74 @@
+"""
+Instrumentation State Tracker
+
+This module provides helper functions to safely track whether a coroutine,
+Future, or function has already been instrumented by the OpenTelemetry
+asyncio instrumentation layer.
+
+Because some Python objects (like coroutines or functions) may not support
+adding custom attributes or may not be weak-referenceable, we use a
+weak-reference-based dictionary to track instrumented objects safely and
+efficiently without causing memory leaks.
+
+If an object cannot be weak-referenced, we skip tracking it to avoid
+runtime errors.
+
+Usage:
+    if not _is_instrumented(obj):
+        _mark_instrumented(obj)
+        # instrument the object...
+"""
+
+import weakref
+
+# A global dictionary to track whether an object has been instrumented.
+# Keys are weak references to avoid preventing garbage collection.
+_instrumented_tasks = {}
+
+
+def _get_weak_key(obj):
+    """
+    Attempt to create a weak reference key for the given object.
+
+    Some object types (e.g., built-in functions or async_generator_asend)
+    do not support weak references. In those cases, return None.
+
+    Args:
+        obj: The object to generate a weak reference for.
+
+    Returns:
+        A weakref.ref to the object if supported, otherwise None.
+    """
+    try:
+        return weakref.ref(obj)
+    except TypeError:
+        return None
+
+
+def _is_instrumented(obj) -> bool:
+    """
+    Check if the object has already been instrumented.
+
+    Args:
+        obj: The coroutine, function, or Future to check.
+
+    Returns:
+        True if the object is already marked as instrumented, False otherwise.
+    """
+    key = _get_weak_key(obj)
+    return key in _instrumented_tasks if key else False
+
+
+def _mark_instrumented(obj):
+    """
+    Mark the object as instrumented to avoid double-instrumentation.
+
+    Only objects that support weak references are tracked. Unsupported
+    objects are silently skipped.
+
+    Args:
+        obj: The coroutine, function, or Future to mark.
+    """
+    key = _get_weak_key(obj)
+    if key:
+        _instrumented_tasks[key] = True

instrumentation/opentelemetry-instrumentation-asyncio/tests/test_asyncio_duplicate_instrument.py

@@ -21,16 +21,16 @@
 
 import asyncio
 
-from opentelemetry.test.test_base import TestBase
-
 from opentelemetry.instrumentation.asyncio import AsyncioInstrumentor
+from opentelemetry.test.test_base import TestBase
 
 
 class MockSubscription:
     """
     Example class holding an unsubscribe_future, similar to something like
     aiokafka's subscription.
     """
+
     def __init__(self):
         self.unsubscribe_future = asyncio.Future()
 
@@ -39,6 +39,7 @@ class MockGroupCoordinator:
     """
     Example class modeling repeated instrumentation of the same Future objects.
     """
+
     def __init__(self):
         self._closing = asyncio.Future()
         self.subscription = MockSubscription()
@@ -62,11 +63,16 @@ class TestAsyncioDuplicateInstrument(TestBase):
 
     def setUp(self):
         super().setUp()
+        self.loop = asyncio.new_event_loop()
+        asyncio.set_event_loop(self.loop)
+
         self.instrumentor = AsyncioInstrumentor()
         self.instrumentor.instrument()
 
     def tearDown(self):
         self.instrumentor.uninstrument()
+        self.loop.close()
+        asyncio.set_event_loop(None)
         super().tearDown()
 
     def test_duplicate_instrumentation_of_futures(self):
@@ -79,11 +85,15 @@ def test_duplicate_instrumentation_of_futures(self):
         # Simulate calling the routine multiple times
         num_iterations = 10
         for _ in range(num_iterations):
-            asyncio.run(coordinator.run_routine(self.instrumentor))
+            self.loop.run_until_complete(
+                coordinator.run_routine(self.instrumentor)
+            )
 
         # Check for callback accumulation
         closing_cb_count = len(coordinator._closing._callbacks)
-        unsub_cb_count = len(coordinator.subscription.unsubscribe_future._callbacks)
+        unsub_cb_count = len(
+            coordinator.subscription.unsubscribe_future._callbacks
+        )
         rejoin_cb_count = len(coordinator._rejoin_needed_fut._callbacks)
 
         # If instrumentation is properly deduplicated, each Future might have ~1-2 callbacks.