diff --git a/ddtrace/internal/module.py b/ddtrace/internal/module.py index 576987aa975..8861104ce61 100644 --- a/ddtrace/internal/module.py +++ b/ddtrace/internal/module.py @@ -230,8 +230,11 @@ def call_back(self, module: ModuleType) -> None: # loader type module.register_loader_type(_ImportHookChainedLoader, module.DefaultProvider) - for callback in self.callbacks.values(): - callback(module) + for key, callback in self.callbacks.items(): + try: + callback(module) + except Exception: + log.exception("Exception ignored in after_import hook %r for module %s", key, module.__name__) def load_module(self, fullname: str) -> t.Optional[ModuleType]: if self.loader is None: @@ -260,20 +263,10 @@ def _create_module(self, spec): def _find_first_hook( self, module: ModuleType, hooks_attr: str ) -> t.Optional[t.Callable[[t.Any, ModuleType], None]]: - for _ in sys.meta_path: - if isinstance(_, ModuleWatchdog): - try: - for ( - cond, - hook, - ) in getattr(_, hooks_attr, []): - if (isinstance(cond, str) and cond == module.__name__) or ( - callable(cond) and cond(module.__name__) - ): - return hook - except Exception: - log.debug("Exception happened while processing %s", hooks_attr, exc_info=True) - return None + universal = _UniversalModuleWatchdog._instance + if universal is None: + return None + return universal._find_first_hook(module, hooks_attr) def _find_first_exception_hook(self, module: ModuleType) -> t.Optional[t.Callable[[t.Any, ModuleType], None]]: return self._find_first_hook(module, "_import_exception_hooks") @@ -330,15 +323,21 @@ def get_code(_loader, fullname): log.exception("Failed to call back on module %s", module) -class BaseModuleWatchdog(abc.ABC): - """Base module watchdog. +class _UniversalModuleWatchdog: + """The single finder ever inserted into ``sys.meta_path`` for module watchdog purposes. - Invokes ``after_import`` every time a new module is imported. + Concrete ``BaseModuleWatchdog`` subclasses register their instances here as + virtual participants, instead of each inserting itself into ``sys.meta_path``. + This collapses what would otherwise be O(N) redundant, mutually-recursive + ``find_spec`` scans (one per installed watchdog, each re-scanning the rest of + ``sys.meta_path`` to find the real underlying finder) into a single O(1) lookup + per import, regardless of how many watchdogs are installed. """ - _instance: t.Optional["BaseModuleWatchdog"] = None + _instance: t.Optional["_UniversalModuleWatchdog"] = None def __init__(self) -> None: + self._watchdogs: list["BaseModuleWatchdog"] = [] self._finding: set[str] = set() # DEV: pkg_resources support to prevent errors such as @@ -363,31 +362,14 @@ def __init__(self) -> None: else: log.warning("Cannot ensure correct support with pkg_resources") - def _add_to_meta_path(self) -> None: - sys.meta_path.insert(0, self) # type: ignore[arg-type] - - @classmethod - def _find_in_meta_path(cls) -> t.Optional[int]: - for i, meta_path in enumerate(sys.meta_path): - if type(meta_path) is cls: - return i - return None - - @classmethod - def _remove_from_meta_path(cls) -> None: - i = cls._find_in_meta_path() - - if i is None: - log.warning("%s is not installed", cls.__name__) - return - - sys.meta_path.pop(i) - - def after_import(self, module: ModuleType) -> None: - raise NotImplementedError() + def register(self, participant: "BaseModuleWatchdog") -> None: + self._watchdogs.append(participant) - def transform(self, code: CodeType, _module: ModuleType) -> CodeType: - return code + def unregister(self, participant: "BaseModuleWatchdog") -> None: + try: + self._watchdogs.remove(participant) + except ValueError: + pass def find_module(self, fullname: str, path: t.Optional[str] = None) -> t.Optional["Loader"]: if fullname in self._finding: @@ -404,8 +386,9 @@ def find_module(self, fullname: str, path: t.Optional[str] = None) -> t.Optional else original_loader ) - loader.add_callback(type(self), self.after_import) - loader.add_transformer(type(self), self.transform) + for watchdog in list(self._watchdogs): + loader.add_callback(type(watchdog), watchdog.after_import) + loader.add_transformer(type(watchdog), watchdog.transform) return t.cast("Loader", loader) @@ -452,21 +435,73 @@ def find_spec( if not isinstance(loader, _ImportHookChainedLoader): spec.loader = t.cast("Loader", _ImportHookChainedLoader(loader, spec)) - t.cast(_ImportHookChainedLoader, spec.loader).add_callback(type(self), self.after_import) - t.cast(_ImportHookChainedLoader, spec.loader).add_transformer(type(self), self.transform) + for watchdog in list(self._watchdogs): + t.cast(_ImportHookChainedLoader, spec.loader).add_callback(type(watchdog), watchdog.after_import) + t.cast(_ImportHookChainedLoader, spec.loader).add_transformer(type(watchdog), watchdog.transform) return spec finally: self._finding.remove(fullname) + def _find_first_hook( + self, module: ModuleType, hooks_attr: str + ) -> t.Optional[t.Callable[[t.Any, ModuleType], None]]: + for watchdog in list(self._watchdogs): + try: + for ( + cond, + hook, + ) in getattr(watchdog, hooks_attr, []): + if (isinstance(cond, str) and cond == module.__name__) or ( + callable(cond) and cond(module.__name__) + ): + return hook + except Exception: + log.debug("Exception happened while processing %s", hooks_attr, exc_info=True) + return None + + @classmethod + def _register(cls, participant: "BaseModuleWatchdog") -> None: + if cls._instance is None: + cls._instance = cls() + sys.meta_path.insert(0, cls._instance) # type: ignore[arg-type] + cls._instance.register(participant) + + @classmethod + def _unregister(cls, participant: "BaseModuleWatchdog") -> None: + if cls._instance is None: + return + cls._instance.unregister(participant) + if not cls._instance._watchdogs: + try: + sys.meta_path.remove(cls._instance) # type: ignore[arg-type] + except ValueError: + log.warning("%s is not installed", cls.__name__) + cls._instance = None + + +class BaseModuleWatchdog(abc.ABC): + """Base module watchdog. + + Invokes ``after_import`` every time a new module is imported. + """ + + _instance: t.Optional["BaseModuleWatchdog"] = None + + def after_import(self, module: ModuleType) -> None: + raise NotImplementedError() + + def transform(self, code: CodeType, _module: ModuleType) -> CodeType: + return code + @classmethod def install(cls) -> None: """Install the module watchdog.""" if cls.is_installed(): return cls._instance = cls() - cls._instance._add_to_meta_path() + _UniversalModuleWatchdog._register(cls._instance) log.debug("%s installed", cls) @classmethod @@ -484,7 +519,7 @@ def uninstall(cls) -> None: if not cls.is_installed(): return - cls._remove_from_meta_path() + _UniversalModuleWatchdog._unregister(t.cast("BaseModuleWatchdog", cls._instance)) cls._instance = None diff --git a/tests/internal/test_module.py b/tests/internal/test_module.py index d13fb95f8a1..ac5941a1d9c 100644 --- a/tests/internal/test_module.py +++ b/tests/internal/test_module.py @@ -56,22 +56,24 @@ def test_watchdog_install_uninstall(): import sys from ddtrace.internal.module import ModuleWatchdog + from ddtrace.internal.module import _UniversalModuleWatchdog if ModuleWatchdog.is_installed(): ModuleWatchdog.uninstall() assert not ModuleWatchdog.is_installed() - assert not any(isinstance(m, ModuleWatchdog) for m in sys.meta_path) + assert not any(isinstance(m, _UniversalModuleWatchdog) for m in sys.meta_path) ModuleWatchdog.install() assert ModuleWatchdog.is_installed() - assert isinstance(sys.meta_path[0], ModuleWatchdog) + assert isinstance(sys.meta_path[0], _UniversalModuleWatchdog) + assert ModuleWatchdog._instance in sys.meta_path[0]._watchdogs ModuleWatchdog.uninstall() assert not ModuleWatchdog.is_installed() - assert not any(isinstance(m, ModuleWatchdog) for m in sys.meta_path) + assert not any(isinstance(m, _UniversalModuleWatchdog) for m in sys.meta_path) def test_import_origin_hook_for_imported_module(module_watchdog): @@ -616,6 +618,7 @@ def test_module_watchdog_find_spec_no_cross_thread_deadlock(): import threading from ddtrace.internal.module import ModuleWatchdog + from ddtrace.internal.module import _UniversalModuleWatchdog bootstrap = sys.modules.get("importlib._bootstrap") if bootstrap is None or not hasattr(bootstrap, "_get_module_lock"): @@ -634,7 +637,7 @@ def test_module_watchdog_find_spec_no_cross_thread_deadlock(): def background(): for finder in sys.meta_path: - if isinstance(finder, ModuleWatchdog): + if isinstance(finder, _UniversalModuleWatchdog): finder.find_spec("tests._ddtrace_regression_nonexistent", None, None) break background_done.set() @@ -649,3 +652,96 @@ def background(): ) finally: lock.release() + + +def test_universal_module_watchdog_single_finder_invariant(): + from ddtrace.internal.module import _UniversalModuleWatchdog + + classes = [type(f"Watchdog{i}", (ModuleWatchdog,), {}) for i in range(5)] + for cls in classes: + cls.install() + try: + universal_finders = [m for m in sys.meta_path if isinstance(m, _UniversalModuleWatchdog)] + assert len(universal_finders) == 1 + assert not any(cls._instance in sys.meta_path for cls in classes) + finally: + for cls in classes: + cls.uninstall() + + +def test_universal_module_watchdog_constant_find_spec_calls(): + from ddtrace.internal.module import _UniversalModuleWatchdog + + # Ensure the parent package is already imported so that importing the + # child module below triggers exactly one find_spec call, rather than one + # per not-yet-imported ancestor package. + import tests.submod.stuff # noqa:F401 + + for num_watchdogs in (1, 3, 6): + classes = [type(f"Watchdog{i}", (ModuleWatchdog,), {}) for i in range(num_watchdogs)] + for cls in classes: + cls.install() + try: + universal = _UniversalModuleWatchdog._instance + call_count = 0 + original_find_spec = universal.find_spec + + def counting_find_spec(*args, **kwargs): + nonlocal call_count + call_count += 1 + return original_find_spec(*args, **kwargs) + + universal.find_spec = counting_find_spec + try: + sys.modules.pop("tests.submod.stuff", None) + import tests.submod.stuff # noqa:F401,F811 + finally: + del universal.find_spec + + assert call_count == 1, f"expected 1 find_spec call for {num_watchdogs} watchdogs, got {call_count}" + finally: + for cls in classes: + cls.uninstall() + sys.modules.pop("tests.submod.stuff", None) + + +def test_universal_module_watchdog_first_registered_wins(): + calls = [] + + class First(ModuleWatchdog): + pass + + class Second(ModuleWatchdog): + pass + + First.install() + Second.install() + try: + First.register_pre_exec_module_hook( + lambda name: name == "tests.submod.stuff", lambda loader, module: calls.append("first") + ) + Second.register_pre_exec_module_hook( + lambda name: name == "tests.submod.stuff", lambda loader, module: calls.append("second") + ) + + import tests.submod.stuff # noqa:F401 + + assert calls == ["first"] + finally: + sys.modules.pop("tests.submod.stuff", None) + Second.uninstall() + First.uninstall() + + +def test_universal_module_watchdog_teardown(): + from ddtrace.internal.module import _UniversalModuleWatchdog + + classes = [type(f"Watchdog{i}", (ModuleWatchdog,), {}) for i in range(3)] + for cls in classes: + cls.install() + + for cls in classes: + cls.uninstall() + + assert _UniversalModuleWatchdog._instance is None + assert not any(isinstance(m, _UniversalModuleWatchdog) for m in sys.meta_path)