Use WeakKeyDictionary to safely track instrumented objects

Author: allen-k1m

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

@@ -5,13 +5,12 @@
 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.
+Some Python objects (like coroutines or functions) may not support setting
+custom attributes or weak references. To avoid memory leaks and runtime
+errors, this module uses a WeakKeyDictionary to safely track instrumented
+objects.
 
-If an object cannot be weak-referenced, we skip tracking it to avoid
-runtime errors.
+If an object cannot be weak-referenced, it is silently skipped.
 
 Usage:
     if not _is_instrumented(obj):
@@ -21,54 +20,41 @@
 
 import weakref
 
-# A global dictionary to track whether an object has been instrumented.
-# Keys are weak references to avoid preventing garbage collection.
-_instrumented_tasks = {}
+# A global WeakKeyDictionary to track instrumented objects.
+# Entries are automatically removed when the objects are garbage collected.
+_instrumented_tasks = weakref.WeakKeyDictionary()
 
 
-def _get_weak_key(obj):
+def _is_instrumented(obj) -> bool:
     """
-    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.
+    Check whether the given object has already been marked as instrumented.
 
     Args:
-        obj: The object to generate a weak reference for.
+        obj: A coroutine, function, or Future.
 
     Returns:
-        A weakref.ref to the object if supported, otherwise None.
+        True if the object has already been marked as instrumented.
+        False if it has not, or if it does not support weak references.
     """
     try:
-        return weakref.ref(obj)
+        return _instrumented_tasks.get(obj, False)
     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
+        # The object does not support weak references.
+        return False
 
 
 def _mark_instrumented(obj):
     """
-    Mark the object as instrumented to avoid double-instrumentation.
+    Mark the given object as instrumented.
 
-    Only objects that support weak references are tracked. Unsupported
-    objects are silently skipped.
+    This function only tracks objects that support weak references.
+    Unsupported objects are ignored without raising errors.
 
     Args:
-        obj: The coroutine, function, or Future to mark.
+        obj: A coroutine, function, or Future.
     """
-    key = _get_weak_key(obj)
-    if key:
-        _instrumented_tasks[key] = True
+    try:
+        _instrumented_tasks[obj] = True
+    except TypeError:
+        # The object does not support weak references.
+        pass