diff --git a/ddtrace/profiling/collector/memalloc.py b/ddtrace/profiling/collector/memalloc.py index 8254cc13f3f..8a52a749a4d 100644 --- a/ddtrace/profiling/collector/memalloc.py +++ b/ddtrace/profiling/collector/memalloc.py @@ -9,6 +9,8 @@ from typing import Type from typing import cast +from typing_extensions import Self + try: from ddtrace.profiling.collector import _memalloc @@ -55,8 +57,9 @@ def start(self) -> None: _memalloc.stop() _memalloc.start(self.max_nframe, self.heap_sample_size) - def __enter__(self) -> None: + def __enter__(self) -> Self: self.start() + return self def __exit__( self, @@ -66,7 +69,7 @@ def __exit__( ) -> None: self.stop() - def join(self) -> None: + def join(self, timeout: Optional[float] = None) -> None: pass def stop(self) -> None: @@ -76,7 +79,8 @@ def stop(self) -> None: except RuntimeError: LOG.debug("Failed to stop memalloc profiling on shutdown", exc_info=True) - def _get_thread_id_ignore_set(self) -> Set[int]: + @staticmethod + def _get_thread_id_ignore_set() -> Set[int]: # This method is not perfect and prone to race condition in theory, but very little in practice. # Anyhow it's not a big deal — it's a best effort feature. return { @@ -86,7 +90,8 @@ def _get_thread_id_ignore_set(self) -> Set[int]: } def snapshot(self) -> None: - thread_id_ignore_set = self._get_thread_id_ignore_set() + """Take a snapshot of collected data, to be exported.""" + thread_id_ignore_set = MemoryCollector._get_thread_id_ignore_set() try: if _memalloc is None: diff --git a/ddtrace/profiling/profiler.py b/ddtrace/profiling/profiler.py index 7d5163421d7..900dae72834 100644 --- a/ddtrace/profiling/profiler.py +++ b/ddtrace/profiling/profiler.py @@ -142,7 +142,9 @@ def __init__( self.endpoint_collection_enabled: bool = endpoint_collection_enabled # Non-user-supplied values - self._collectors: List[collector.Collector] = [] + # Note: memalloc.MemoryCollector is not a subclass of collector.Collector, so we need to use a union type. + # This is because its snapshot method cannot be static. + self._collectors: List[collector.Collector | memalloc.MemoryCollector] = [] self._collectors_on_import: Optional[List[tuple[str, Callable[[Any], None]]]] = None self._scheduler: Optional[Union[scheduler.Scheduler, scheduler.ServerlessScheduler]] = None self._lambda_function_name: Optional[str] = os.environ.get("AWS_LAMBDA_FUNCTION_NAME") @@ -251,7 +253,7 @@ def start_collector(collector_class: Type[collector.Collector]) -> None: ModuleWatchdog.register_module_hook(module, hook) if self._memory_collector_enabled: - self._collectors.append(memalloc.MemoryCollector()) # type: ignore[arg-type] + self._collectors.append(memalloc.MemoryCollector()) self._build_default_exporters()