fix(ci_visibility): fix coverage crash and data loss with copied/restored execution contexts

ddtrace/internal/coverage/code.py

@@ -48,10 +48,29 @@ def _is_site_packages_path(path: Path) -> bool:
     return not _SITE_PACKAGES_DIRNAMES.isdisjoint(path.parts)
 
 
+class _ContextLinesEntry(defaultdict[str, CoverageLines]):
+    """Coverage lines container for one collector, tracking whether it completed.
+
+    Execution contexts copied while a collector is active keep a reference to a
+    stack that still contains its entry (eg. a task scheduled by a module that is
+    imported inside a test). Marking the entry closed when the collector finishes
+    lets the context resolvers skip completed collectors and attribute new
+    coverage to the nearest collector that is still active.
+    """
+
+    closed: bool = False
+
+
+class _ContextFilesEntry(set[str]):
+    """File-level counterpart of _ContextLinesEntry."""
+
+    closed: bool = False
+
+
 # NOTE: A mutable ContextVar default would be shared across threads until set() is called.
-# Keep None so CollectInContext initializes a separate coverage stack in each context.
-ctx_covered: ContextVar[t.Optional[list[defaultdict[str, CoverageLines]]]] = ContextVar("ctx_covered", default=None)
-ctx_covered_files: ContextVar[t.Optional[list[set[str]]]] = ContextVar("ctx_covered_files", default=None)
+# Keep None so CollectInContext starts a separate coverage stack in each context.
+ctx_covered: ContextVar[t.Optional[list[_ContextLinesEntry]]] = ContextVar("ctx_covered", default=None)
+ctx_covered_files: ContextVar[t.Optional[list[_ContextFilesEntry]]] = ContextVar("ctx_covered_files", default=None)
 ctx_is_import_coverage = ContextVar("ctx_is_import_coverage", default=False)
 ctx_coverage_enabled = ContextVar("ctx_coverage_enabled", default=False)
 
@@ -63,13 +82,20 @@ def _is_site_packages_path(path: Path) -> bool:
 def _get_ctx_covered_lines() -> defaultdict[str, CoverageLines]:
     if ctx_coverage_enabled.get():
         if context_stack := ctx_covered.get():
-            return context_stack[-1]
-        log.debug("_get_ctx_covered_lines() called but ctx_covered stack is empty")
+            for entry in reversed(context_stack):
+                if not entry.closed:
+                    return entry
+            # Every entry on this stack belongs to a collector that has completed:
+            # this context inherited the stack before they exited. Fall through to
+            # the TLS fallback / an empty container instead of attributing new
+            # coverage to a completed collector.
+        else:
+            log.debug("_get_ctx_covered_lines() called but ctx_covered stack is empty")
 
     # Fallback for Python 3.14+ where sys.monitoring callbacks can't see ContextVars
     if _PY_GE_314:
         tls_covered = getattr(_tls_coverage, "covered", None)
-        if tls_covered is not None:
+        if tls_covered is not None and not tls_covered.closed:
             return tls_covered
 
     return defaultdict(CoverageLines)
@@ -78,12 +104,15 @@ def _get_ctx_covered_lines() -> defaultdict[str, CoverageLines]:
 def _get_ctx_covered_files() -> set[str]:
     if ctx_coverage_enabled.get():
         if context_stack := ctx_covered_files.get():
-            return context_stack[-1]
-        log.debug("_get_ctx_covered_files() called but ctx_covered_files stack is empty")
+            for entry in reversed(context_stack):
+                if not entry.closed:
+                    return entry
+        else:
+            log.debug("_get_ctx_covered_files() called but ctx_covered_files stack is empty")
 
     if _PY_GE_314:
         tls_covered_files = getattr(_tls_coverage, "covered_files", None)
-        if tls_covered_files is not None:
+        if tls_covered_files is not None and not tls_covered_files.closed:
             return tls_covered_files
 
     return set()
@@ -370,17 +399,41 @@ def _get_covered_file_paths_with_imports(self, covered_file_paths: set[str]) ->
             self._file_level_covered_paths_cache.popitem(last=False)
         return paths
 
+    class _ContextStack(list[t.Any]):
+        """Per-context stack of coverage data that compares by identity, not value.
+
+        Context-propagation helpers (e.g. asgiref's _restore_context, used by Django's
+        async test support via async_to_sync/sync_to_async) restore context variables
+        by comparing the current value with the incoming one using !=. A plain list
+        compares by value, which both silently masks legitimate stack swaps (when
+        two distinct stacks happen to contain equal entries) and allows one context's
+        stack to be replaced by another context's stack object. Comparing stacks by
+        identity makes such propagation respect stack ownership: restores only
+        propagate a stack reference into a context that does not already hold that
+        exact stack object, keeping the coverage data attributed to the right context.
+        """
+
+        __slots__ = ()
+
+        def __eq__(self, other: object) -> bool:  # noqa: D105
+            return self is other
+
+        def __ne__(self, other: object) -> bool:  # noqa: D105
+            return self is not other
+
     class CollectInContext:
         def __init__(self, is_import_coverage: bool = False):
             self.is_import_coverage = is_import_coverage
-            if ctx_covered.get() is None:
-                ctx_covered.set([])
-            if ctx_covered_files.get() is None:
-                ctx_covered_files.set([])
 
         def __enter__(self):
-            ctx_covered.get().append(defaultdict(CoverageLines))
-            ctx_covered_files.get().append(set())
+            # ContextVar values are copied by reference into new execution contexts.
+            # Replace the stacks so a nested collector cannot mutate its parent's stack.
+            self._covered_lines = _ContextLinesEntry(CoverageLines)
+            self._covered_files = _ContextFilesEntry()
+            ctx_covered.set(ModuleCodeCollector._ContextStack((ctx_covered.get() or []) + [self._covered_lines]))
+            ctx_covered_files.set(
+                ModuleCodeCollector._ContextStack((ctx_covered_files.get() or []) + [self._covered_files])
+            )
             ctx_coverage_enabled.set(True)
 
             if self.is_import_coverage:
@@ -412,13 +465,41 @@ def __enter__(self):
             return self
 
         def __exit__(self, *args, **kwargs):
-            covered_lines_stack = ctx_covered.get()
-            covered_files_stack = ctx_covered_files.get()
-            covered_lines_stack.pop()
-            covered_files_stack.pop()
+            # The collector is completing. Stacks inherited by other contexts (copied
+            # while this collector was active) still reference its entry: mark it
+            # closed so the resolvers skip it and attribute new coverage to the
+            # nearest active collector instead of this completed one.
+            self._covered_lines.closed = True
+            self._covered_files.closed = True
+
+            covered_lines_stack = ctx_covered.get() or []
+            covered_files_stack = ctx_covered_files.get() or []
+            if (
+                covered_lines_stack
+                and covered_files_stack
+                and covered_lines_stack[-1] is self._covered_lines
+                and covered_files_stack[-1] is self._covered_files
+            ):
+                covered_lines_stack = ModuleCodeCollector._ContextStack(covered_lines_stack[:-1])
+                covered_files_stack = ModuleCodeCollector._ContextStack(covered_files_stack[:-1])
+                ctx_covered.set(covered_lines_stack)
+                ctx_covered_files.set(covered_files_stack)
+            else:
+                # A copied context may finish a collector that was entered elsewhere.
+                # Leave this context's collector intact instead of popping the wrong one.
+                if _PY_GE_314:
+                    # The exited collector may still be this thread's TLS fallback: the
+                    # sys.monitoring callbacks fall back to this thread-local state when
+                    # their snapshot context cannot observe ContextVar changes. Re-sync
+                    # the fallback to the collector that is actually active in this
+                    # context (or clear it) so coverage keeps being attributed to the
+                    # right entry instead of the completed collector.
+                    _tls_coverage.covered = covered_lines_stack[-1] if covered_lines_stack else None
+                    _tls_coverage.covered_files = covered_files_stack[-1] if covered_files_stack else None
+                return
 
             # Stop coverage if we're exiting the last context
-            if len(covered_lines_stack) == 0:
+            if not covered_lines_stack:
                 ctx_coverage_enabled.set(False)
                 if _PY_GE_314:
                     _tls_coverage.covered = None
@@ -428,7 +509,7 @@ def __exit__(self, *args, **kwargs):
                 _tls_coverage.covered_files = covered_files_stack[-1]
 
         def get_covered_lines(self) -> dict[str, CoverageLines]:
-            covered_lines = _get_ctx_covered_lines()
+            covered_lines = self._covered_lines
             if global_instance := ModuleCodeCollector._instance:
                 global_instance._add_import_time_lines(covered_lines)
             return covered_lines
@@ -437,8 +518,8 @@ def get_covered_file_paths(self) -> t.AbstractSet[str]:
             # Python < 3.12 and injected child-process coverage may only update the line-oriented
             # context data. Merge those keys into the file set so file-level uploads still include
             # every file that would have been emitted by get_covered_lines().
-            covered_file_paths = set(_get_ctx_covered_files())
-            covered_file_paths.update(_get_ctx_covered_lines())
+            covered_file_paths = set(self._covered_files)
+            covered_file_paths.update(self._covered_lines)
             if global_instance := ModuleCodeCollector._instance:
                 return global_instance._get_covered_file_paths_with_imports(covered_file_paths)
             return covered_file_paths

releasenotes/notes/fix-ci-visibility-copied-context-coverage-f8c173bd.yaml

@@ -0,0 +1,7 @@
+---
+fixes:
+  - |
+    CI Visibility: Fixes an issue where pytest runs using copied or restored execution
+    contexts (e.g. Django async tests via asgiref's ``async_to_sync``) could end with an
+    internal coverage error, or silently lose the coverage data of the tests that use
+    them.

tests/coverage/test_coverage.py

@@ -12,6 +12,137 @@
 import pytest
 
 
+def test_coverage_stacks_are_isolated_across_copied_contexts():
+    from contextvars import copy_context
+
+    from ddtrace.internal.coverage.code import ModuleCodeCollector
+    from ddtrace.internal.coverage.code import ctx_covered
+    from ddtrace.internal.coverage.code import ctx_covered_files
+
+    with ModuleCodeCollector.CollectInContext():
+        parent_lines_stack = ctx_covered.get()
+        parent_files_stack = ctx_covered_files.get()
+        parent_depth = len(parent_lines_stack)
+        child_context = copy_context()
+
+        def collect_in_child_context():
+            with ModuleCodeCollector.CollectInContext():
+                assert len(ctx_covered.get()) == len(ctx_covered_files.get()) == parent_depth + 1
+                assert ctx_covered.get()[-1] is not parent_lines_stack[-1]
+                assert ctx_covered_files.get()[-1] is not parent_files_stack[-1]
+
+        child_context.run(collect_in_child_context)
+        assert ctx_covered.get() is parent_lines_stack
+        assert ctx_covered_files.get() is parent_files_stack
+        assert len(parent_lines_stack) == len(parent_files_stack) == parent_depth
+
+
+def test_exiting_collector_in_another_context_preserves_active_coverage():
+    from contextvars import copy_context
+
+    from ddtrace.internal.coverage.code import ModuleCodeCollector
+    from ddtrace.internal.coverage.code import ctx_covered
+    from ddtrace.internal.coverage.code import ctx_covered_files
+
+    with ModuleCodeCollector.CollectInContext():
+        parent_depth = len(ctx_covered.get())
+        parent_lines = ctx_covered.get()[-1]
+        parent_files = ctx_covered_files.get()[-1]
+        child_context = copy_context()
+        child = ModuleCodeCollector.CollectInContext()
+        child_context.run(child.__enter__)
+
+        child.__exit__()
+        assert ctx_covered.get()[-1] is parent_lines
+        assert ctx_covered_files.get()[-1] is parent_files
+
+        child_context.run(child.__exit__)
+        assert len(ctx_covered.get()) == len(ctx_covered_files.get()) == parent_depth
+
+
+def test_completed_collector_entries_do_not_capture_inherited_context_coverage():
+    """Contexts that inherit a stack still holding a completed entry must not write to it.
+
+    A module imported inside a test may create an asyncio task before finishing its
+    import collector. The task inherits the test's context, whose stack still
+    references the (now completed) import entry. New coverage in the task must be
+    attributed to the enclosing live collector instead of the orphaned entry.
+    """
+    from contextvars import copy_context
+
+    import ddtrace.internal.coverage.code as coverage_code
+    from ddtrace.internal.coverage.code import ModuleCodeCollector
+
+    with ModuleCodeCollector.CollectInContext() as test_collector:
+        with ModuleCodeCollector.CollectInContext() as import_collector:
+            task_context = copy_context()
+
+        assert import_collector._covered_lines.closed
+
+        # The task context still sees the completed import entry atop its stack.
+        task_stack = task_context.run(coverage_code.ctx_covered.get)
+        assert task_stack[-1] is import_collector._covered_lines
+
+        # Resolution inside the task context must skip the completed entry and
+        # attribute coverage to the still-active test collector.
+        assert task_context.run(coverage_code._get_ctx_covered_lines) is test_collector._covered_lines
+        assert task_context.run(coverage_code._get_ctx_covered_files) is test_collector._covered_files
+
+        # A live collector entered in the task context takes precedence even though
+        # the completed import entry remains buried beneath it on the stack.
+        nested = ModuleCodeCollector.CollectInContext()
+        task_context.run(nested.__enter__)
+        assert task_context.run(coverage_code._get_ctx_covered_lines) is nested._covered_lines
+        task_context.run(nested.__exit__)
+        assert task_context.run(coverage_code._get_ctx_covered_lines) is test_collector._covered_lines
+
+    # Once every collector the task inherited has completed, new coverage lands in
+    # a fresh container rather than in any of the completed entries.
+    stale = task_context.run(coverage_code._get_ctx_covered_lines)
+    assert stale is not import_collector._covered_lines
+    assert stale is not test_collector._covered_lines
+
+
+def test_mismatched_exit_resyncs_tls_fallback():
+    """A mismatched exit must not leave the TLS fallback on the completed collector.
+
+    On Python 3.14+, sys.monitoring callbacks run in a snapshot context and fall back
+    to the thread-local coverage state when they cannot observe ContextVar changes.
+    When a collector entered in a copied context is exited from a different context
+    (a mismatched exit), that thread-local fallback must be re-synced to the active
+    collector instead of pointing at the collector that just completed.
+    """
+    from contextvars import copy_context
+
+    import ddtrace.internal.coverage.code as coverage_code
+    from ddtrace.internal.coverage.code import ModuleCodeCollector
+    from ddtrace.internal.coverage.code import ctx_covered
+
+    original_flag = coverage_code._PY_GE_314
+    coverage_code._PY_GE_314 = True
+    try:
+        with ModuleCodeCollector.CollectInContext():
+            parent_lines = ctx_covered.get()[-1]
+            parent_files = coverage_code.ctx_covered_files.get()[-1]
+            child_context = copy_context()
+            child = ModuleCodeCollector.CollectInContext()
+            child_context.run(child.__enter__)
+
+            # The thread-local fallback tracks the most recent collector entered in
+            # this thread, which is the child's.
+            assert coverage_code._tls_coverage.covered is child._covered_lines
+
+            # Exiting the child from the parent context is a mismatched exit.
+            child.__exit__()
+
+            # The fallback must be re-synced to the parent's active entries rather
+            # than left pointing at the completed child collector.
+            assert coverage_code._tls_coverage.covered is parent_lines
+            assert coverage_code._tls_coverage.covered_files is parent_files
+    finally:
+        coverage_code._PY_GE_314 = original_flag
+
+
 @pytest.mark.skipif(sys.version_info < (3, 12), reason="Test specific to Python 3.12+ monitoring API")
 @pytest.mark.subprocess()
 def test_coverage_defaults_to_file_level_when_env_unset():