refactor(internal): use a single universal finder for module watchdogs

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
 

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)