chore(ci_visibility): introduce experimental internal coverage collector

ddtrace/contrib/pytest/plugin.py

@@ -13,6 +13,7 @@
 """
 from doctest import DocTest
 import json
+import os
 import re
 from typing import Dict  # noqa:F401
 
@@ -49,6 +50,7 @@
 from ddtrace.internal.ci_visibility.constants import SUITE_ID as _SUITE_ID
 from ddtrace.internal.ci_visibility.constants import SUITE_TYPE as _SUITE_TYPE
 from ddtrace.internal.ci_visibility.constants import TEST
+from ddtrace.internal.ci_visibility.coverage import USE_DD_COVERAGE
 from ddtrace.internal.ci_visibility.coverage import _module_has_dd_coverage_enabled
 from ddtrace.internal.ci_visibility.coverage import _report_coverage_to_span
 from ddtrace.internal.ci_visibility.coverage import _start_coverage
@@ -61,7 +63,9 @@
 from ddtrace.internal.ci_visibility.utils import get_relative_or_absolute_path_for_path
 from ddtrace.internal.ci_visibility.utils import take_over_logger_stream_handler
 from ddtrace.internal.constants import COMPONENT
+from ddtrace.internal.coverage.code import ModuleCodeCollector
 from ddtrace.internal.logger import get_logger
+from ddtrace.internal.utils.formats import asbool
 
 
 PATCH_ALL_HELP_MSG = "Call ddtrace.patch_all before running tests."
@@ -70,6 +74,10 @@
 
 _global_skipped_elements = 0
 
+# COVER_SESSION is an experimental feature flag that provides full coverage (similar to coverage run), and is an
+# experimental feature. It currently significantly increases test import time and should not be used.
+COVER_SESSION = asbool(os.environ.get("_DD_COVER_SESSION", "false"))
+
 
 def _is_pytest_8_or_later():
     if hasattr(pytest, "version_tuple"):
@@ -821,6 +829,20 @@ def pytest_runtest_protocol(item, nextitem):
             _stop_coverage(pytest)
 
 
+def pytest_load_initial_conftests(early_config, parser, args):
+    # Enables experimental use of ModuleCodeCollector for coverage collection.
+    if USE_DD_COVERAGE:
+        if not ModuleCodeCollector.is_installed():
+            ModuleCodeCollector.install()
+        if COVER_SESSION:
+            ModuleCodeCollector.start_coverage()
+    else:
+        if COVER_SESSION:
+            log.warning(
+                "_DD_COVER_SESSION must be used with _DD_USE_INTERNAL_COVERAGE but not DD_CIVISIBILITY_ITR_ENABLED"
+            )
+
+
 @pytest.hookimpl(hookwrapper=True)
 def pytest_runtest_makereport(item, call):
     """Store outcome for tracing."""
@@ -936,3 +958,10 @@ def pytest_ddtrace_get_item_test_name(item):
         if item.config.getoption("ddtrace-include-class-name") or item.config.getini("ddtrace-include-class-name"):
             return "%s.%s" % (item.cls.__name__, item.name)
     return item.name
+
+
+@pytest.hookimpl(trylast=True)
+def pytest_terminal_summary(terminalreporter, exitstatus, config):
+    # Reports coverage if experimental session-level coverage is enabled.
+    if USE_DD_COVERAGE and COVER_SESSION:
+        ModuleCodeCollector.report()

ddtrace/internal/ci_visibility/coverage.py

@@ -1,20 +1,28 @@
 from itertools import groupby
 import json
+import os
 from typing import Dict  # noqa:F401
 from typing import Iterable  # noqa:F401
 from typing import List  # noqa:F401
 from typing import Optional  # noqa:F401
 from typing import Tuple  # noqa:F401
+from typing import Union  # noqa:F401
 
 import ddtrace
 from ddtrace.internal.ci_visibility.constants import COVERAGE_TAG_NAME
 from ddtrace.internal.ci_visibility.utils import get_relative_or_absolute_path_for_path
+from ddtrace.internal.coverage.code import ModuleCodeCollector
 from ddtrace.internal.logger import get_logger
+from ddtrace.internal.utils.formats import asbool
 
 
 log = get_logger(__name__)
 _global_relative_file_paths_for_cov: Dict[str, Dict[str, str]] = {}
 
+# This feature-flags experimental collection of code coverage via our internal ModuleCodeCollector.
+# It is disabled by default because it is not production-ready.
+USE_DD_COVERAGE = asbool(os.environ.get("_DD_USE_INTERNAL_COVERAGE", "false"))
+
 try:
     from coverage import Coverage
     from coverage import version_info as coverage_version
@@ -45,19 +53,30 @@ def _initialize_coverage(root_dir):
 
 
 def _start_coverage(root_dir: str):
+    # Experimental feature to use internal coverage collection
+    if USE_DD_COVERAGE:
+        ctx = ModuleCodeCollector.CollectInContext()
+        return ctx
     coverage = _initialize_coverage(root_dir)
     coverage.start()
     return coverage
 
 
 def _stop_coverage(module):
+    # Experimental feature to use internal coverage collection
+    if USE_DD_COVERAGE:
+        module._dd_coverage.__exit__()
+        return
     if _module_has_dd_coverage_enabled(module):
         module._dd_coverage.stop()
         module._dd_coverage.erase()
         del module._dd_coverage
 
 
 def _module_has_dd_coverage_enabled(module, silent_mode: bool = False) -> bool:
+    # Experimental feature to use internal coverage collection
+    if USE_DD_COVERAGE:
+        return hasattr(module, "_dd_coverage")
     if not hasattr(module, "_dd_coverage"):
         if not silent_mode:
             log.warning("Datadog Coverage has not been initiated")
@@ -73,7 +92,15 @@ def _coverage_has_valid_data(coverage_data: Coverage, silent_mode: bool = False)
     return True
 
 
-def _switch_coverage_context(coverage_data: Coverage, unique_test_name: str):
+def _switch_coverage_context(
+    coverage_data: Union[Coverage, ModuleCodeCollector.CollectInContext], unique_test_name: str
+):
+    # Experimental feature to use internal coverage collection
+    if isinstance(coverage_data, ModuleCodeCollector.CollectInContext):
+        if USE_DD_COVERAGE:
+            # In this case, coverage_data is the context manager supplied by ModuleCodeCollector.CollectInContext
+            coverage_data.__enter__()
+        return
     if not _coverage_has_valid_data(coverage_data, silent_mode=True):
         return
     coverage_data._collector.data.clear()  # type: ignore[union-attr]
@@ -83,7 +110,24 @@ def _switch_coverage_context(coverage_data: Coverage, unique_test_name: str):
         log.warning(err)
 
 
-def _report_coverage_to_span(coverage_data: Coverage, span: ddtrace.Span, root_dir: str):
+def _report_coverage_to_span(
+    coverage_data: Union[Coverage, ModuleCodeCollector.CollectInContext], span: ddtrace.Span, root_dir: str
+):
+    # Experimental feature to use internal coverage collection
+    if isinstance(coverage_data, ModuleCodeCollector.CollectInContext):
+        if USE_DD_COVERAGE:
+            # In this case, coverage_data is the context manager supplied by ModuleCodeCollector.CollectInContext
+            files = ModuleCodeCollector.report_seen_lines()
+            if not files:
+                return
+            span.set_tag_str(
+                COVERAGE_TAG_NAME,
+                json.dumps({"files": files}),
+            )
+            coverage_data.__exit__(None, None, None)
+
+        return
+
     span_id = str(span.trace_id)
     if not _coverage_has_valid_data(coverage_data):
         return

ddtrace/internal/coverage/_native.c

@@ -0,0 +1,61 @@
+#define PY_SSIZE_T_CLEAN
+#include <Python.h>
+
+#if PY_VERSION_HEX < 0x030c0000
+#if defined __GNUC__ && defined HAVE_STD_ATOMIC
+#undef HAVE_STD_ATOMIC
+#endif
+#endif
+
+// ----------------------------------------------------------------------------
+static PyObject*
+replace_in_tuple(PyObject* m, PyObject* args)
+{
+    PyObject* tuple = NULL;
+    PyObject* item = NULL;
+    PyObject* replacement = NULL;
+
+    if (!PyArg_ParseTuple(args, "O!OO", &PyTuple_Type, &tuple, &item, &replacement))
+        return NULL;
+
+    for (Py_ssize_t i = 0; i < PyTuple_Size(tuple); i++) {
+        PyObject* current = PyTuple_GetItem(tuple, i);
+        if (current == item) {
+            Py_DECREF(current);
+            // !!! DANGER !!!
+            PyTuple_SET_ITEM(tuple, i, replacement);
+            Py_INCREF(replacement);
+        }
+    }
+
+    Py_RETURN_NONE;
+}
+
+// ----------------------------------------------------------------------------
+static PyMethodDef native_methods[] = {
+    { "replace_in_tuple", replace_in_tuple, METH_VARARGS, "Replace an item in a tuple." },
+    { NULL, NULL, 0, NULL } /* Sentinel */
+};
+
+// ----------------------------------------------------------------------------
+static struct PyModuleDef nativemodule = {
+    PyModuleDef_HEAD_INIT,
+    "_native", /* name of module */
+    NULL,      /* module documentation, may be NULL */
+    -1,        /* size of per-interpreter state of the module,
+                  or -1 if the module keeps state in global variables. */
+    native_methods,
+};
+
+// ----------------------------------------------------------------------------
+PyMODINIT_FUNC
+PyInit__native(void)
+{
+    PyObject* m;
+
+    m = PyModule_Create(&nativemodule);
+    if (m == NULL)
+        return NULL;
+
+    return m;
+}

ddtrace/internal/coverage/_native.pyi

@@ -0,0 +1,3 @@
+import typing as t
+
+def replace_in_tuple(tup: tuple, item: t.Any, replacement: t.Any) -> None: ...