from types import FrameType import typing as t import gevent from gevent import thread import gevent.greenlet from gevent.greenlet import Greenlet as _Greenlet import gevent.hub from greenlet import greenlet from greenlet import settrace from ddtrace.internal.datadog.profiling import stack # Original objects _gevent_hub_spawn_raw: t.Callable[..., _Greenlet] = gevent.hub.spawn_raw _gevent_joinall: t.Callable[..., t.Sequence[_Greenlet]] = gevent.joinall _gevent_wait: t.Callable[..., t.Any] = gevent.wait _gevent_iwait: t.Callable[..., t.Any] = gevent.iwait # Global package state _tracked_greenlets: set[int] = set() _original_greenlet_tracer: t.Optional[t.Callable[[str, t.Any], None]] = None _greenlet_parent_map: dict[int, int] = {} _parent_greenlet_count: dict[int, int] = {} FRAME_NOT_SET: bool = False # Sentinel for when the frame is not set class GreenletTrackingError(Exception): """Exception raised when a greenlet cannot be tracked.""" pass def track_gevent_greenlet(greenlet: _Greenlet) -> _Greenlet: greenlet_id: int = thread.get_ident(greenlet) frame: t.Union[FrameType, bool, None] = FRAME_NOT_SET try: stack.track_greenlet(greenlet_id, greenlet.name or type(greenlet).__qualname__, frame) except AttributeError as e: raise GreenletTrackingError("Cannot track greenlet with no name attribute") from e except Exception as e: raise GreenletTrackingError("Cannot track greenlet") from e # Untrack on completion try: greenlet.rawlink(untrack_greenlet) except AttributeError: # This greenlet cannot be linked (e.g. the Hub) pass except Exception as e: raise GreenletTrackingError("Cannot link greenlet for untracking") from e _tracked_greenlets.add(greenlet_id) return greenlet def update_greenlet_frame(greenlet_id: int, frame: t.Union[FrameType, bool, None]) -> None: _tracked_greenlets.add(greenlet_id) stack.update_greenlet_frame(greenlet_id, frame) def greenlet_tracer(event: str, args: t.Any) -> None: # Greenlets that already exist when profiling is enabled are discovered lazily. # We only start tracking them once a post-patch "switch"/"throw" event is observed. # A greenlet that exits before switching again may not be tracked. if event in {"switch", "throw"}: # This tracer function runs in the context of the target origin, target = t.cast(tuple[_Greenlet, _Greenlet], args) if (origin_id := thread.get_ident(origin)) not in _tracked_greenlets: try: track_gevent_greenlet(origin) except GreenletTrackingError: # Not something that we can track pass if (target_id := thread.get_ident(target)) not in _tracked_greenlets: # This is likely the hub. We take this chance to track it. try: track_gevent_greenlet(target) except GreenletTrackingError: # Not something that we can track pass try: # If this is being set to None, it means the greenlet is likely # finished. We use the sentinel again to signal this. update_greenlet_frame( origin_id, t.cast(t.Optional[FrameType], origin.gr_frame) or FRAME_NOT_SET, ) if target_id not in _parent_greenlet_count: # We don't want to wipe the frame of a parent greenlet because # we need to unwind it. We definitely know it is still running # so if we allow the tracer to set its tracked frame to None, # we won't be able to unwind the full stack. update_greenlet_frame(target_id, target.gr_frame) # this *is* None except KeyError: # TODO: Log missing greenlet pass if _original_greenlet_tracer is not None: _original_greenlet_tracer(event, args) def untrack_greenlet(greenlet: _Greenlet) -> None: greenlet_id: int = thread.get_ident(greenlet) stack.untrack_greenlet(greenlet_id) _tracked_greenlets.discard(greenlet_id) _parent_greenlet_count.pop(greenlet_id, None) if (parent_id := _greenlet_parent_map.pop(greenlet_id, None)) is not None: _parent_greenlet_count[parent_id] -= 1 if _parent_greenlet_count[parent_id] <= 0: del _parent_greenlet_count[parent_id] def link_greenlets(greenlet_id: int, parent_id: int) -> None: stack.link_greenlets(greenlet_id, parent_id) _parent_greenlet_count[parent_id] = _parent_greenlet_count.get(parent_id, 0) + 1 _greenlet_parent_map[greenlet_id] = parent_id class Greenlet(_Greenlet): @classmethod def spawn(cls, *args: t.Any, **kwargs: t.Any) -> _Greenlet: greenlet = super().spawn(*args, **kwargs) try: return track_gevent_greenlet(greenlet) except GreenletTrackingError: # If we cannot track the greenlet, we just return it as is. return greenlet @classmethod def spawn_later(cls, *args: t.Any, **kwargs: t.Any) -> _Greenlet: greenlet = super().spawn_later(*args, **kwargs) try: return track_gevent_greenlet(greenlet) except GreenletTrackingError: return greenlet def join(self, *args: t.Any, **kwargs: t.Any) -> None: target_id: int = thread.get_ident(self) origin_id: int = thread.get_ident(gevent.getcurrent()) link_greenlets(target_id, origin_id) super().join(*args, **kwargs) def wrap_spawn(original: t.Callable[..., _Greenlet]) -> t.Callable[..., _Greenlet]: def _(*args: t.Any, **kwargs: t.Any) -> _Greenlet: greenlet = original(*args, **kwargs) try: return track_gevent_greenlet(greenlet) except GreenletTrackingError: return greenlet return _ def joinall(greenlets: t.Sequence[_Greenlet], *args: t.Any, **kwargs: t.Any) -> t.Sequence[_Greenlet]: # This is a wrapper around gevent.joinall to track the greenlets # that are being joined. current_greenlet = gevent.getcurrent() if isinstance(current_greenlet, greenlet): current_greenlet = gevent.hub.get_hub() current_greenlet_id: int = thread.get_ident(current_greenlet) for g in greenlets: link_greenlets(thread.get_ident(g), current_greenlet_id) return _gevent_joinall(greenlets, *args, **kwargs) def wait_wrapper(original: t.Callable[..., t.Any]) -> t.Callable[..., t.Any]: def _(*args: t.Any, **kwargs: t.Any) -> t.Any: try: objects = args[0] except IndexError: objects = kwargs.get("args", []) if greenlets := [_ for _ in objects if isinstance(_, (greenlet, gevent.Greenlet))]: current_greenlet = gevent.getcurrent() if isinstance(current_greenlet, greenlet): current_greenlet = gevent.hub.get_hub() current_greenlet_id: int = thread.get_ident(current_greenlet) for g in greenlets: link_greenlets(thread.get_ident(g), current_greenlet_id) return original(*args, **kwargs) return _ def get_current_greenlet_task() -> tuple[t.Optional[int], t.Optional[str], t.Optional[FrameType]]: current_greenlet = gevent.getcurrent() task_id = thread.get_ident(current_greenlet) # Import locally to avoid eager import order interactions. from ddtrace.profiling import _threading task_name = _threading.get_thread_name(task_id) frame = t.cast(t.Optional[FrameType], current_greenlet.gr_frame) return task_id, task_name, frame def patch() -> None: global _original_greenlet_tracer # Patch the spawn method to track greenlets. gevent.Greenlet = gevent.greenlet.Greenlet = Greenlet gevent.spawn = Greenlet.spawn gevent.spawn_later = Greenlet.spawn_later gevent.joinall = joinall gevent.wait = wait_wrapper(_gevent_wait) gevent.iwait = wait_wrapper(_gevent_iwait) gevent.hub.spawn_raw = wrap_spawn(_gevent_hub_spawn_raw) _original_greenlet_tracer = t.cast(t.Callable[[str, t.Any], None], settrace(greenlet_tracer)) def unpatch() -> None: # Unpatch the spawn method to stop tracking greenlets. gevent.Greenlet = gevent.greenlet.Greenlet = _Greenlet gevent.spawn = _Greenlet.spawn gevent.spawn_later = _Greenlet.spawn_later gevent.joinall = _gevent_joinall gevent.wait = _gevent_wait gevent.iwait = _gevent_iwait gevent.hub.spawn_raw = _gevent_hub_spawn_raw settrace(_original_greenlet_tracer)