fix(internal): do not replace sys.modules [backport 2.0] (#7068)

Backport 2bd2c2a from #6925 to 2.0.

We refactor the ModuleWatchdog class to avoid replacing the sys.modules
dictionary with a proxy. Instead, we use a WeakValueDictionary to
prevent adding strong references to modules that might be unloaded, and
also preventing the origin mapping from growing unchecked over time.

Fixes #6923.

## Testing Strategy

The existing test suites should be good enough to catch any regressions
from this substantial change. As for the mentioned issue, the change was
tested manually on the reproducer provided in the description provided
by the reporter.

## Risk

This is a substantial change in behaviour for `ModuleWatchdog`. However,
we expect things to run smoother now that `sys.modules` is not _hacked_,
at the cost of a potentially slightly delayed memory clean-up.

## Checklist

- [x] Change(s) are motivated and described in the PR description.
- [x] Testing strategy is described if automated tests are not included
in the PR.
- [x] Risk is outlined (performance impact, potential for breakage,
maintainability, etc).
- [x] Change is maintainable (easy to change, telemetry, documentation).
- [x] [Library release note
guidelines](https://ddtrace.readthedocs.io/en/stable/releasenotes.html)
are followed. If no release note is required, add label
`changelog/no-changelog`.
- [x] Documentation is included (in-code, generated user docs, [public
corp docs](https://github.com/DataDog/documentation/)).
- [x] Backport labels are set (if
[applicable](https://ddtrace.readthedocs.io/en/latest/contributing.html#backporting))

## Reviewer Checklist

- [x] Title is accurate.
- [x] No unnecessary changes are introduced.
- [x] Description motivates each change.
- [x] Avoids breaking
[API](https://ddtrace.readthedocs.io/en/stable/versioning.html#interfaces)
changes unless absolutely necessary.
- [x] Testing strategy adequately addresses listed risk(s).
- [x] Change is maintainable (easy to change, telemetry, documentation).
- [x] Release note makes sense to a user of the library.
- [x] Reviewer has explicitly acknowledged and discussed the performance
implications of this PR as reported in the benchmarks PR comment.
- [x] Backport labels are set in a manner that is consistent with the
[release branch maintenance
policy](https://ddtrace.readthedocs.io/en/latest/contributing.html#backporting)
- [x] If this PR touches code that signs or publishes builds or
packages, or handles credentials of any kind, I've requested a review
from `@DataDog/security-design-and-guidance`.
- [x] This PR doesn't touch any of that.

Co-authored-by: Gabriele N. Tornetta <[email protected]>

Author: github-actions[bot]

ddtrace/internal/module.py

@@ -7,6 +7,7 @@
 import sys
 import typing
 from typing import cast
+from weakref import WeakValueDictionary as wvdict
 
 
 if typing.TYPE_CHECKING:
@@ -15,7 +16,6 @@
     from typing import Callable
     from typing import DefaultDict
     from typing import Dict
-    from typing import Iterator
     from typing import List
     from typing import Optional
     from typing import Set
@@ -223,12 +223,11 @@ def _exec_module(self, module):
             callback(module)
 
 
-class ModuleWatchdog(dict):
+class ModuleWatchdog(object):
     """Module watchdog.
 
-    Replace the standard ``sys.modules`` dictionary to detect when modules are
-    loaded/unloaded. This is also responsible for triggering any registered
-    import hooks.
+    Hooks into the import machinery to detect when modules are loaded/unloaded.
+    This is also responsible for triggering any registered import hooks.
 
     Subclasses might customize the default behavior by overriding the
     ``after_import`` method, which is triggered on every module import, once
@@ -240,31 +239,22 @@ class ModuleWatchdog(dict):
     def __init__(self):
         # type: () -> None
         self._hook_map = defaultdict(list)  # type: DefaultDict[str, List[ModuleHookType]]
-        self._om = None  # type: Optional[Dict[str, ModuleType]]
-        self._modules = sys.modules  # type: Union[dict, ModuleWatchdog]
+        self._om = None  # type: Optional[wvdict[str, ModuleType]]
         self._finding = set()  # type: Set[str]
         self._pre_exec_module_hooks = []  # type: List[Tuple[PreExecHookCond, PreExecHookType]]
 
-    def __getitem__(self, item):
-        # type: (str) -> ModuleType
-        return self._modules.__getitem__(item)
-
-    def __setitem__(self, name, module):
-        # type: (str, ModuleType) -> None
-        self._modules.__setitem__(name, module)
-
     @property
     def _origin_map(self):
-        # type: () -> Dict[str, ModuleType]
+        # type: () -> wvdict[str, ModuleType]
         if self._om is None:
             try:
-                self._om = {origin(module): module for module in sys.modules.values()}
+                self._om = wvdict({origin(module): module for module in sys.modules.values()})
             except RuntimeError:
                 # The state of sys.modules might have been mutated by another
                 # thread. We try to build the full mapping at the next occasion.
                 # For now we take the more expensive route of building a list of
                 # the current values, which might be incomplete.
-                return {origin(module): module for module in list(sys.modules.values())}
+                return wvdict({origin(module): module for module in list(sys.modules.values())})
 
         return self._om
 
@@ -284,8 +274,11 @@ def _find_in_meta_path(cls):
     def _remove_from_meta_path(cls):
         # type: () -> None
         i = cls._find_in_meta_path()
-        if i is not None:
-            sys.meta_path.pop(i)
+
+        if i is None:
+            raise RuntimeError("%s is not installed" % cls.__name__)
+
+        sys.meta_path.pop(i)
 
     def after_import(self, module):
         # type: (ModuleType) -> None
@@ -328,48 +321,6 @@ def get_by_origin(cls, _origin):
 
         return None
 
-    def __delitem__(self, name):
-        # type: (str) -> None
-        try:
-            path = origin(sys.modules[name])
-            # Drop the module reference to reclaim memory
-            del self._origin_map[path]
-        except KeyError:
-            pass
-
-        self._modules.__delitem__(name)
-
-    def __getattribute__(self, name):
-        # type: (str) -> Any
-        if LEGACY_DICT_COPY and name == "keys":
-            # This is a potential attempt to make a copy of sys.modules using
-            # dict(sys.modules) on a Python version that uses the C API to
-            # perform the operation. Since we are an instance of a dict, this
-            # means that we will end up looking like the empty dict, so we take
-            # this chance to actually look like sys.modules.
-            # NOTE: This is a potential source of memory leaks. However, we
-            # expect this to occur only on defunct Python versions, and only
-            # during special code executions, like test runs.
-            super(ModuleWatchdog, self).clear()
-            super(ModuleWatchdog, self).update(self._modules)
-
-        try:
-            return super(ModuleWatchdog, self).__getattribute__("_modules").__getattribute__(name)
-        except AttributeError:
-            return super(ModuleWatchdog, self).__getattribute__(name)
-
-    def __contains__(self, name):
-        # type: (object) -> bool
-        return self._modules.__contains__(name)
-
-    def __len__(self):
-        # type: () -> int
-        return self._modules.__len__()
-
-    def __iter__(self):
-        # type: () -> Iterator
-        return self._modules.__iter__()
-
     def find_module(self, fullname, path=None):
         # type: (str, Optional[str]) -> Union[Loader, None]
         if fullname in self._finding:
@@ -452,6 +403,13 @@ def register_origin_hook(cls, origin, hook):
         instance._hook_map[path].append(hook)
         try:
             module = instance._origin_map[path]
+            # Sanity check: the module might have been removed from sys.modules
+            # but not yet garbage collected.
+            try:
+                sys.modules[module.__name__]
+            except KeyError:
+                del instance._origin_map[path]
+                raise
         except KeyError:
             # The module is not loaded yet. Nothing more we can do.
             return
@@ -499,7 +457,7 @@ def register_module_hook(cls, module, hook):
         instance = cast(ModuleWatchdog, cls._instance)
         instance._hook_map[module].append(hook)
         try:
-            module_object = instance[module]
+            module_object = sys.modules[module]
         except KeyError:
             # The module is not loaded yet. Nothing more we can do.
             return
@@ -567,8 +525,8 @@ def install(cls):
         if cls.is_installed():
             raise RuntimeError("%s is already installed" % cls.__name__)
 
-        cls._instance = sys.modules = cls()
-        sys.modules._add_to_meta_path()
+        cls._instance = cls()
+        cls._instance._add_to_meta_path()
         log.debug("%s installed", cls)
 
     @classmethod
@@ -585,17 +543,8 @@ def uninstall(cls):
         class.
         """
         cls._check_installed()
+        cls._remove_from_meta_path()
 
-        parent, current = None, sys.modules
-        while isinstance(current, ModuleWatchdog):
-            if type(current) is cls:
-                cls._remove_from_meta_path()
-                if parent is not None:
-                    parent._modules = current._modules
-                else:
-                    sys.modules = current._modules
-                cls._instance = None
-                log.debug("%s uninstalled", cls)
-                return
-            parent = current
-            current = current._modules
+        cls._instance = None
+
+        log.debug("%s uninstalled", cls)

releasenotes/notes/fix-module-watchdog-no-sys-module-replacement-0e019c59ad745244.yaml

@@ -0,0 +1,5 @@
+---
+fixes:
+  - |
+    Resolves issues with the import machinery that might have caused the
+    garbage collector to clean up weak references in an unexpected order.

tests/contrib/patch.py

@@ -695,6 +695,8 @@ def test_ddtrace_run_patch_on_import(self):
                         """
                         import sys
 
+                        from ddtrace.internal.module import ModuleWatchdog
+
                         from ddtrace.vendor.wrapt import wrap_function_wrapper as wrap
 
                         patched = False
@@ -710,7 +712,7 @@ def patch_wrapper(wrapped, _, args, kwrags):
 
                             wrap(module.__name__, module.patch.__name__, patch_wrapper)
 
-                        sys.modules.register_module_hook("ddtrace.contrib.%s.patch", patch_hook)
+                        ModuleWatchdog.register_module_hook("ddtrace.contrib.%s.patch", patch_hook)
 
                         sys.stdout.write("O")
 

tests/debugging/test_debugger.py

@@ -128,7 +128,7 @@ def test_debugger_probe_new_delete(probe, trigger):
         d.add_probes(probe)
 
         assert probe in d._probe_registry
-        assert _get_probe_location(probe) in sys.modules._locations
+        assert _get_probe_location(probe) in d.__watchdog__._instance._locations
 
         trigger()
 
@@ -137,7 +137,7 @@ def test_debugger_probe_new_delete(probe, trigger):
         # Test that the probe was ejected
         assert probe not in d._probe_registry
 
-        assert _get_probe_location(probe) not in sys.modules._locations
+        assert _get_probe_location(probe) not in d.__watchdog__._instance._locations
 
         trigger()
 

tests/internal/test_module.py

@@ -1,5 +1,7 @@
+import json
 import os
 import sys
+from warnings import warn
 
 import mock
 import pytest
@@ -35,7 +37,13 @@ def ensure_no_module_watchdog():
         yield
     finally:
         if was_installed:
-            ModuleWatchdog.install()
+            if ModuleWatchdog.is_installed():
+                warn(
+                    "ModuleWatchdog still installed after test run. This might also be caused by a test that failed "
+                    "while a ModuleWatchdog was installed."
+                )
+            else:
+                ModuleWatchdog.install()
 
 
 @pytest.fixture
@@ -50,11 +58,18 @@ def module_watchdog():
 
 
 def test_watchdog_install_uninstall():
-    assert not isinstance(sys.modules, ModuleWatchdog)
+    assert not ModuleWatchdog.is_installed()
+    assert not any(isinstance(m, ModuleWatchdog) for m in sys.meta_path)
+
     ModuleWatchdog.install()
-    assert isinstance(sys.modules, ModuleWatchdog)
+
+    assert ModuleWatchdog.is_installed()
+    assert isinstance(sys.meta_path[0], ModuleWatchdog)
+
     ModuleWatchdog.uninstall()
-    assert not isinstance(sys.modules, ModuleWatchdog)
+
+    assert not ModuleWatchdog.is_installed()
+    assert not any(isinstance(m, ModuleWatchdog) for m in sys.meta_path)
 
 
 def test_import_origin_hook_for_imported_module(module_watchdog):
@@ -156,39 +171,48 @@ def test_import_module_hook_for_module_not_yet_imported():
     ModuleWatchdog.uninstall()
 
 
-@pytest.mark.subprocess(env=dict(MODULE_ORIGIN=origin(tests.test_module)))
+@pytest.mark.subprocess(env=dict(MODULE_ORIGIN=origin(json)))
 def test_module_deleted():
+    import gc
     import os
     import sys
 
-    from mock import mock
-
     from ddtrace.internal.module import ModuleWatchdog
 
-    name = "tests.test_module"
+    if "json" in sys.modules:
+        del sys.modules["json"]
+        gc.collect()
+
+    name = "json"
     path = os.getenv("MODULE_ORIGIN")
-    hook = mock.Mock()
+
+    class Counter(object):
+        count = 0
+
+        def __call__(self, _):
+            self.count += 1
+
+    hook = Counter()
 
     ModuleWatchdog.register_origin_hook(path, hook)
     ModuleWatchdog.register_module_hook(name, hook)
 
     __import__(name)
 
-    calls = [mock.call(sys.modules[name])] * 2
-    hook.assert_has_calls(calls)
+    assert hook.count == 2, hook.count
 
     assert path in ModuleWatchdog._instance._origin_map
 
     del sys.modules[name]
+    gc.collect()
 
     assert path not in ModuleWatchdog._instance._origin_map
 
     # We are not deleting the registered hooks, so if we re-import the module
     # new hook calls are triggered
     __import__(name)
 
-    calls.extend([mock.call(sys.modules[name])] * 2)
-    hook.assert_has_calls(calls)
+    assert hook.count == 4
 
     ModuleWatchdog.uninstall()
 
@@ -343,25 +367,6 @@ class Bob(BaseCollector):
     Alice.uninstall()
 
 
-def test_module_watchdog_dict_shallow_copy():
-    # Save original reference to sys.modules
-    original_sys_modules = sys.modules
-
-    ModuleWatchdog.install()
-
-    # Ensure that we have replaced sys.modules
-    assert original_sys_modules is not sys.modules
-
-    # Make a shallow copy of both using the dict constructor
-    original_modules = set(dict(original_sys_modules).keys())
-    new_modules = set(dict(sys.modules).keys())
-
-    # Ensure that they match
-    assert original_modules == new_modules
-
-    ModuleWatchdog.uninstall()
-
-
 @pytest.mark.skipif(sys.version_info < (3, 5), reason="LazyLoader was introduced in Python 3.5")
 @pytest.mark.subprocess(out="ddtrace imported\naccessing lazy module\nlazy loaded\n")
 def test_module_watchdog_no_lazy_force_load():