Skip to content

Commit 37416c1

Browse files
P403n1x87taegyunkim
P403n1x87
authored and
taegyunkim
committed
fix(pytest): disable module cloning (#14996)
We prevent module cloning from triggering in pytest sessions as it is likely to happen at the wrong time, potentially causing module import issues.
1 parent 65db18b commit 37416c1

File tree

4 files changed

+132
-110
lines changed

4 files changed

+132
-110
lines changed

‎ddtrace/bootstrap/cloning.py‎

Lines changed: 120 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,120 @@
1+
import logging
2+
import os
3+
import sys
4+
import warnings
5+
6+
from ddtrace.internal.module import ModuleWatchdog
7+
from ddtrace.internal.module import is_module_installed
8+
from ddtrace.internal.utils.formats import asbool # noqa:F401
9+
10+
11+
MODULES_REQUIRING_CLEANUP = ("gevent",)
12+
13+
14+
enabled = (
15+
any(is_module_installed(m) for m in MODULES_REQUIRING_CLEANUP)
16+
if (_unload_modules := os.getenv("DD_UNLOAD_MODULES_FROM_SITECUSTOMIZE", default="auto").lower()) == "auto"
17+
else asbool(_unload_modules)
18+
)
19+
20+
21+
if "gevent" in sys.modules or "gevent.monkey" in sys.modules:
22+
import gevent.monkey # noqa:F401
23+
24+
if gevent.monkey.is_module_patched("threading"):
25+
warnings.warn( # noqa: B028
26+
"Loading ddtrace after gevent.monkey.patch_all() is not supported and is "
27+
"likely to break the application. Use ddtrace-run to fix this, or "
28+
"import ddtrace.auto before calling gevent.monkey.patch_all().",
29+
RuntimeWarning,
30+
)
31+
32+
33+
def cleanup_loaded_modules() -> None:
34+
global enabled
35+
36+
if not enabled:
37+
return
38+
39+
def drop(module_name: str) -> None:
40+
module = sys.modules.get(module_name)
41+
# Don't delete modules that are currently being imported (they might be None or incomplete)
42+
# or that don't exist. This can happen when pytest's assertion rewriter is importing modules
43+
# that themselves import ddtrace.auto, which triggers this cleanup during the import process.
44+
if module is None:
45+
return
46+
# Skip modules that don't have a __spec__ attribute yet (still being imported)
47+
if not hasattr(module, "__spec__"):
48+
return
49+
# Check if the module is currently being initialized
50+
# During import, __spec__._initializing is True
51+
spec = getattr(module, "__spec__", None)
52+
if spec is not None and getattr(spec, "_initializing", False):
53+
return
54+
del sys.modules[module_name]
55+
56+
# We need to import these modules to make sure they grab references to the
57+
# right modules before we start unloading stuff.
58+
import ddtrace.internal.http # noqa
59+
import ddtrace.internal.uds # noqa
60+
61+
# Unload all the modules that we have imported, except for the ddtrace one.
62+
# NB: this means that every `import threading` anywhere in `ddtrace/` code
63+
# uses a copy of that module that is distinct from the copy that user code
64+
# gets when it does `import threading`. The same applies to every module
65+
# not in `KEEP_MODULES`.
66+
KEEP_MODULES = frozenset(
67+
[
68+
"atexit",
69+
"copyreg", # pickling issues for tracebacks with gevent
70+
"ddtrace",
71+
"concurrent",
72+
"typing",
73+
"_operator", # pickling issues with typing module
74+
"re", # referenced by the typing module
75+
"sre_constants", # imported by re at runtime
76+
"logging",
77+
"attr",
78+
"google",
79+
"google.protobuf", # the upb backend in >= 4.21 does not like being unloaded
80+
"wrapt",
81+
"bytecode", # needed by before-fork hooks
82+
]
83+
)
84+
for m in list(_ for _ in sys.modules if _ not in ddtrace.LOADED_MODULES):
85+
if any(m == _ or m.startswith(_ + ".") for _ in KEEP_MODULES):
86+
continue
87+
88+
drop(m)
89+
90+
# TODO: The better strategy is to identify the core modules in LOADED_MODULES
91+
# that should not be unloaded, and then unload as much as possible.
92+
UNLOAD_MODULES = frozenset(
93+
[
94+
# imported in Python >= 3.10 and patched by gevent
95+
"time",
96+
# we cannot unload the whole concurrent hierarchy, but this
97+
# submodule makes use of threading so it is critical to unload when
98+
# gevent is used.
99+
"concurrent.futures",
100+
# We unload the threading module in case it was imported by
101+
# CPython on boot.
102+
"threading",
103+
"_thread",
104+
]
105+
)
106+
for u in UNLOAD_MODULES:
107+
for m in list(sys.modules):
108+
if m == u or m.startswith(u + "."):
109+
drop(m)
110+
111+
# Because we are not unloading it, the logging module requires a reference
112+
# to the newly imported threading module to allow it to retrieve the correct
113+
# thread object information, like the thread name. We register a post-import
114+
# hook on the threading module to perform this update.
115+
@ModuleWatchdog.after_module_imported("threading")
116+
def _(threading):
117+
logging.threading = threading
118+
119+
# Do module cloning only once
120+
enabled = False

‎ddtrace/bootstrap/sitecustomize.py‎

Lines changed: 2 additions & 110 deletions
Original file line numberDiff line numberDiff line change
@@ -15,129 +15,21 @@
1515
# perform the correct initialisation for the library. All the actual
1616
# initialisation logic should be placed in preload.py.
1717
import ddtrace # isort:skip
18-
import logging # noqa:I001
1918
import os # noqa:F401
2019
import sys
21-
import warnings # noqa:F401
2220

23-
from ddtrace import config # noqa:F401
21+
import ddtrace.bootstrap.cloning as cloning
2422
from ddtrace.internal.logger import get_logger # noqa:F401
25-
from ddtrace.internal.module import ModuleWatchdog # noqa:F401
26-
from ddtrace.internal.module import is_module_installed
2723
from ddtrace.internal.telemetry import telemetry_writer
28-
from ddtrace.internal.utils.formats import asbool # noqa:F401
2924

3025

3126
log = get_logger(__name__)
3227

3328

34-
if "gevent" in sys.modules or "gevent.monkey" in sys.modules:
35-
import gevent.monkey # noqa:F401
36-
37-
if gevent.monkey.is_module_patched("threading"):
38-
warnings.warn( # noqa: B028
39-
"Loading ddtrace after gevent.monkey.patch_all() is not supported and is "
40-
"likely to break the application. Use ddtrace-run to fix this, or "
41-
"import ddtrace.auto before calling gevent.monkey.patch_all().",
42-
RuntimeWarning,
43-
)
44-
45-
46-
def cleanup_loaded_modules():
47-
def drop(module_name):
48-
# type: (str) -> None
49-
module = sys.modules.get(module_name)
50-
# Don't delete modules that are currently being imported (they might be None or incomplete)
51-
# or that don't exist. This can happen when pytest's assertion rewriter is importing modules
52-
# that themselves import ddtrace.auto, which triggers this cleanup during the import process.
53-
if module is None:
54-
return
55-
# Skip modules that don't have a __spec__ attribute yet (still being imported)
56-
if not hasattr(module, "__spec__"):
57-
return
58-
# Check if the module is currently being initialized
59-
# During import, __spec__._initializing is True
60-
spec = getattr(module, "__spec__", None)
61-
if spec is not None and getattr(spec, "_initializing", False):
62-
return
63-
del sys.modules[module_name]
64-
65-
MODULES_REQUIRING_CLEANUP = ("gevent",)
66-
do_cleanup = os.getenv("DD_UNLOAD_MODULES_FROM_SITECUSTOMIZE", default="auto").lower()
67-
if do_cleanup == "auto":
68-
do_cleanup = any(is_module_installed(m) for m in MODULES_REQUIRING_CLEANUP)
69-
70-
if not asbool(do_cleanup):
71-
return
72-
73-
# We need to import these modules to make sure they grab references to the
74-
# right modules before we start unloading stuff.
75-
import ddtrace.internal.http # noqa
76-
import ddtrace.internal.uds # noqa
77-
78-
# Unload all the modules that we have imported, except for the ddtrace one.
79-
# NB: this means that every `import threading` anywhere in `ddtrace/` code
80-
# uses a copy of that module that is distinct from the copy that user code
81-
# gets when it does `import threading`. The same applies to every module
82-
# not in `KEEP_MODULES`.
83-
KEEP_MODULES = frozenset(
84-
[
85-
"atexit",
86-
"copyreg", # pickling issues for tracebacks with gevent
87-
"ddtrace",
88-
"concurrent",
89-
"typing",
90-
"_operator", # pickling issues with typing module
91-
"re", # referenced by the typing module
92-
"sre_constants", # imported by re at runtime
93-
"logging",
94-
"attr",
95-
"google",
96-
"google.protobuf", # the upb backend in >= 4.21 does not like being unloaded
97-
"wrapt",
98-
"bytecode", # needed by before-fork hooks
99-
]
100-
)
101-
for m in list(_ for _ in sys.modules if _ not in ddtrace.LOADED_MODULES):
102-
if any(m == _ or m.startswith(_ + ".") for _ in KEEP_MODULES):
103-
continue
104-
105-
drop(m)
106-
107-
# TODO: The better strategy is to identify the core modules in LOADED_MODULES
108-
# that should not be unloaded, and then unload as much as possible.
109-
UNLOAD_MODULES = frozenset(
110-
[
111-
# imported in Python >= 3.10 and patched by gevent
112-
"time",
113-
# we cannot unload the whole concurrent hierarchy, but this
114-
# submodule makes use of threading so it is critical to unload when
115-
# gevent is used.
116-
"concurrent.futures",
117-
# We unload the threading module in case it was imported by
118-
# CPython on boot.
119-
"threading",
120-
"_thread",
121-
]
122-
)
123-
for u in UNLOAD_MODULES:
124-
for m in list(sys.modules):
125-
if m == u or m.startswith(u + "."):
126-
drop(m)
127-
128-
# Because we are not unloading it, the logging module requires a reference
129-
# to the newly imported threading module to allow it to retrieve the correct
130-
# thread object information, like the thread name. We register a post-import
131-
# hook on the threading module to perform this update.
132-
@ModuleWatchdog.after_module_imported("threading")
133-
def _(threading):
134-
logging.threading = threading
135-
136-
13729
try:
13830
import ddtrace.bootstrap.preload as preload # Perform the actual initialisation
13931

140-
cleanup_loaded_modules()
32+
cloning.cleanup_loaded_modules()
14133

14234
# Check for and import any sitecustomize that would have normally been used
14335
# had ddtrace-run not been used.

‎ddtrace/contrib/internal/pytest/_plugin_v2.py‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -294,6 +294,11 @@ def _pytest_load_initial_conftests_pre_yield(early_config, parser, args):
294294
ModuleCodeCollector has a tangible impact on the time it takes to load modules, so it should only be installed if
295295
coverage collection is requested by the backend.
296296
"""
297+
# Disable module cloning in test runs as early as possible
298+
import ddtrace.bootstrap.cloning as cloning
299+
300+
cloning.enabled = False
301+
297302
take_over_logger_stream_handler()
298303

299304
# Log early initialization details

‎releasenotes/notes/fix-pytest-plugin-disable-module-cloning-b86dce8ee693e02d.yaml‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,5 @@
1+
---
2+
fixes:
3+
- |
4+
pytest plugin: fix for potential ``KeyError`` exceptions in test runs when
5+
gevent is detected within the environment.

0 commit comments

Comments
 (0)