sys.monitoring gradual deinstrumentation

ddtrace/internal/coverage/code.py

@@ -3,6 +3,7 @@
 from copy import deepcopy
 from inspect import getmodule
 import os
+import sys
 from types import CodeType
 from types import ModuleType
 import typing as t
@@ -22,6 +23,9 @@
 from ddtrace.internal.utils.inspection import resolved_code_origin
 
 
+if sys.version_info >= (3, 12):
+    from ddtrace.internal.coverage.instrumentation_py3_12 import reset_monitoring_for_new_context
+
 log = get_logger(__name__)
 
 _original_exec = exec
@@ -231,6 +235,11 @@ def __enter__(self):
             if self.is_import_coverage:
                 ctx_is_import_coverage.set(self.is_import_coverage)
 
+            # For Python 3.12+, re-enable monitoring that was disabled by previous contexts
+            # This ensures each test/suite gets accurate coverage data
+            if sys.version_info >= (3, 12):
+                reset_monitoring_for_new_context()
+
             return self
 
         def __exit__(self, *args, **kwargs):

ddtrace/internal/coverage/instrumentation_py3_12.py

@@ -21,10 +21,30 @@
 RETURN_CONST = dis.opmap["RETURN_CONST"]
 EMPTY_MODULE_BYTES = bytes([RESUME, 0, RETURN_CONST, 0])
 
+# Store: (hook, path, import_names_by_line)
 _CODE_HOOKS: t.Dict[CodeType, t.Tuple[HookType, str, t.Dict[int, t.Tuple[str, t.Optional[t.Tuple[str]]]]]] = {}
 
+# Track all instrumented code objects so we can re-enable monitoring between tests/suites
+_DEINSTRUMENTED_CODE_OBJECTS: t.Set[CodeType] = set()
+
 
 def instrument_all_lines(code: CodeType, hook: HookType, path: str, package: str) -> t.Tuple[CodeType, CoverageLines]:
+    """
+    Instrument code for coverage tracking using Python 3.12's monitoring API.
+
+    Args:
+        code: The code object to instrument
+        hook: The hook function to call
+        path: The file path
+        package: The package name
+
+    Note: Python 3.12+ uses an optimized approach where each line callback returns DISABLE
+    after recording. This means:
+    - Each line is only reported once per coverage context (test/suite)
+    - No overhead for repeated line executions (e.g., in loops)
+    - Full line-by-line coverage data is captured
+    - reset_monitoring_for_new_context() re-enables monitoring between contexts
+    """
     coverage_tool = sys.monitoring.get_tool(sys.monitoring.COVERAGE_ID)
     if coverage_tool is not None and coverage_tool != "datadog":
         log.debug("Coverage tool '%s' already registered, not gathering coverage", coverage_tool)
@@ -38,9 +58,24 @@ def instrument_all_lines(code: CodeType, hook: HookType, path: str, package: str
 
 
 def _line_event_handler(code: CodeType, line: int) -> t.Any:
-    hook, path, import_names = _CODE_HOOKS[code]
+    hook_data = _CODE_HOOKS.get(code)
+    if hook_data is None:
+        # Track this code object so we can re-enable monitoring for it later
+        _DEINSTRUMENTED_CODE_OBJECTS.add(code)
+        return sys.monitoring.DISABLE
+
+    hook, path, import_names = hook_data
+
+    # Report the line and then disable monitoring for this specific line
+    # This ensures each line is only reported once per context, even if executed multiple times (e.g., in loops)
     import_name = import_names.get(line, None)
-    return hook((line, path, import_name))
+    hook((line, path, import_name))
+
+    # Track this code object so we can re-enable monitoring for it later
+    _DEINSTRUMENTED_CODE_OBJECTS.add(code)
+    # Return DISABLE to prevent future callbacks for this specific line
+    # This provides full line coverage with minimal overhead
+    return sys.monitoring.DISABLE
 
 
 def _register_monitoring():
@@ -55,6 +90,23 @@ def _register_monitoring():
     )  # noqa
 
 
+def reset_monitoring_for_new_context():
+    """
+    Re-enable monitoring for all instrumented code objects.
+
+    This should be called when starting a new coverage context (e.g., per-test or per-suite).
+    It re-enables monitoring that was disabled by previous DISABLE returns.
+    """
+    # restart_events() re-enables all events that were disabled by returning DISABLE
+    # This resets the per-line disable state across all code objects
+    sys.monitoring.restart_events()
+
+    # Then re-enable local events for all instrumented code objects
+    # This ensures monitoring is active for the new context
+    for code in _DEINSTRUMENTED_CODE_OBJECTS:
+        sys.monitoring.set_local_events(sys.monitoring.COVERAGE_ID, code, sys.monitoring.events.LINE)  # noqa
+
+
 def _instrument_all_lines_with_monitoring(
     code: CodeType, hook: HookType, path: str, package: str
 ) -> t.Tuple[CodeType, CoverageLines]:

tests/coverage/included_path/reinstrumentation_test_module.py

@@ -0,0 +1,40 @@
+"""
+Simple test module for testing coverage re-instrumentation across contexts.
+
+This module provides simple, predictable functions with known line numbers
+to help test that coverage collection works correctly across multiple contexts.
+"""
+
+
+def simple_function(x, y):
+    """A simple function with a few lines."""
+    result = x + y
+    return result
+
+
+def function_with_loop(n):
+    """A function with a loop to test repeated line execution."""
+    total = 0
+    for i in range(n):
+        total += i
+    return total
+
+
+def function_with_branches(condition):
+    """A function with branches to test different code paths."""
+    if condition:
+        result = "true_branch"
+    else:
+        result = "false_branch"
+    return result
+
+
+def multi_line_function(a, b, c):
+    """A function with multiple lines to test comprehensive coverage."""
+    step1 = a + b
+    step2 = step1 * c
+    step3 = step2 - a
+    step4 = step3 / (b if b != 0 else 1)
+    result = step4 ** 2
+    return result
+