refactor(profiling): build a generic _ThreadLink class (#3060)

This class implements the basics needed for the _ThreadSpanLinks class but
allow it to be re-usable for other purpose.

Co-authored-by: Brett Langdon <[email protected]>
Co-authored-by: Munir Abdinur <[email protected]>

Author: 3 people

ddtrace/profiling/collector/stack.pyx

@@ -403,29 +403,41 @@ cdef stack_collect(ignore_profiler, thread_time, max_nframes, interval, wall_tim
     return stack_events, exc_events
 
 
+# cython does not play well with mypy
+if typing.TYPE_CHECKING:
+    _T = typing.TypeVar("_T")
+    _thread_link_base = typing.Generic[_T]
+    _weakref_type = weakref.ReferenceType[_T]
+else:
+    _thread_link_base = object
+    _weakref_type = typing.Any
+
+
 @attr.s(slots=True, eq=False)
-class _ThreadSpanLinks(object):
+class _ThreadLink(_thread_link_base):
+    """Link a thread with an object.
+
+    Object is removed when the thread disappears.
+    """
 
     # Key is a thread_id
-    # Value is a weakref to latest active span
-    _thread_id_to_spans = attr.ib(factory=dict, repr=False, init=False, type=typing.Dict[int, ddspan.Span])
+    # Value is a weakref to an object
+    _thread_id_to_object = attr.ib(factory=dict, repr=False, init=False, type=typing.Dict[int, _weakref_type])
     _lock = attr.ib(factory=nogevent.Lock, repr=False, init=False, type=nogevent.Lock)
 
-    def link_span(
+    def link_object(
             self,
-            span # type: typing.Optional[typing.Union[context.Context, ddspan.Span]]
+            obj  # type: _T
     ):
         # type: (...) -> None
-        """Link a span to its running environment.
-
-        Track threads, tasks, etc.
-        """
+        """Link an object to the current running thread."""
         # Since we're going to iterate over the set, make sure it's locked
-        if isinstance(span, ddspan.Span):
-            with self._lock:
-                self._thread_id_to_spans[nogevent.thread_get_ident()] = weakref.ref(span)
+        with self._lock:
+            self._thread_id_to_object[nogevent.thread_get_ident()] = weakref.ref(obj)
 
-    def clear_threads(self, existing_thread_ids):
+    def clear_threads(self,
+                      existing_thread_ids,  # type: typing.Set[int]
+                      ):
         """Clear the stored list of threads based on the list of existing thread ids.
 
         If any thread that is part of this list was stored, its data will be deleted.
@@ -434,9 +446,48 @@ class _ThreadSpanLinks(object):
         """
         with self._lock:
             # Iterate over a copy of the list of keys since it's mutated during our iteration.
-            for thread_id in list(self._thread_id_to_spans.keys()):
+            for thread_id in list(self._thread_id_to_object.keys()):
                 if thread_id not in existing_thread_ids:
-                    del self._thread_id_to_spans[thread_id]
+                    del self._thread_id_to_object[thread_id]
+
+    def get_object(
+            self,
+            thread_id  # type: int
+    ):
+        # type: (...) -> _T
+        """Return the object attached to thread.
+
+        :param thread_id: The thread id.
+        :return: The attached object.
+        """
+
+        with self._lock:
+            obj_ref = self._thread_id_to_object.get(thread_id)
+            if obj_ref is not None:
+                return obj_ref()
+
+
+if typing.TYPE_CHECKING:
+    _thread_span_links_base = _ThreadLink[ddspan.Span]
+else:
+    _thread_span_links_base = _ThreadLink
+
+
+@attr.s(slots=True, eq=False)
+class _ThreadSpanLinks(_thread_span_links_base):
+
+    def link_span(
+            self,
+            span # type: typing.Optional[typing.Union[context.Context, ddspan.Span]]
+    ):
+        # type: (...) -> None
+        """Link a span to its running environment.
+
+        Track threads, tasks, etc.
+        """
+        # Since we're going to iterate over the set, make sure it's locked
+        if isinstance(span, ddspan.Span):
+            self.link_object(span)
 
     def get_active_span_from_thread_id(
             self,
@@ -448,14 +499,10 @@ class _ThreadSpanLinks(object):
         :param thread_id: The thread id.
         :return: A set with the active spans.
         """
-
-        with self._lock:
-            active_span_ref = self._thread_id_to_spans.get(thread_id)
-            if active_span_ref is not None:
-                active_span = active_span_ref()
-                if active_span is not None and not active_span.finished:
-                    return active_span
-                return None
+        active_span = self.get_object(thread_id)
+        if active_span is not None and not active_span.finished:
+            return active_span
+        return None
 
 
 def _default_min_interval_time():