ref(asyncio integration): Refactor asyncio context filtering

GH-4699

Author: srothh

sentry_sdk/integrations/asyncio.py

@@ -8,7 +8,7 @@
     event_from_exception,
     logger,
     reraise,
-    is_sentry_internal_task,
+    is_internal_task,
 )
 from sentry_sdk.transport import AsyncHttpTransport
 
@@ -89,16 +89,18 @@ def _sentry_task_factory(
         ) -> asyncio.Future[Any]:
 
             # Check if this is an internal Sentry task
-            is_internal = is_sentry_internal_task.get()
+            is_internal = is_internal_task()
 
             if is_internal:
+                task = None
                 if orig_task_factory:
-                    return orig_task_factory(loop, coro, **kwargs)
-                else:
+                    task = orig_task_factory(loop, coro, **kwargs)
+                if task is None:
                     task = Task(coro, loop=loop, **kwargs)
                     if task._source_traceback:  # type: ignore
                         del task._source_traceback[-1]  # type: ignore
-                    return task
+
+                return task
 
             async def _task_with_sentry_span_creation() -> Any:
                 result = None

sentry_sdk/utils.py

@@ -74,19 +74,23 @@
 
 _installed_modules = None
 
-is_sentry_internal_task = contextvars.ContextVar(
+_is_sentry_internal_task = contextvars.ContextVar(
     "is_sentry_internal_task", default=False
 )
 
 
+def is_internal_task():
+    return _is_sentry_internal_task.get()
+
+
 @contextmanager
 def mark_sentry_task_internal():
     """Context manager to mark a task as Sentry internal."""
-    token = is_sentry_internal_task.set(True)
+    token = _is_sentry_internal_task.set(True)
     try:
         yield
     finally:
-        is_sentry_internal_task.reset(token)
+        _is_sentry_internal_task.reset(token)
 
 
 BASE64_ALPHABET = re.compile(r"^[a-zA-Z0-9/+=]*$")

tests/integrations/asyncio/test_asyncio.py

@@ -433,57 +433,47 @@ def test_loop_close_flushes_async_transport(sentry_init):
 @minimum_python_38
 @pytest.mark.asyncio
 async def test_internal_tasks_not_wrapped(sentry_init, capture_events):
-    """Test that internal Sentry tasks marked with context manager are not wrapped."""
 
-    # Get the event loop and save original task factory
-    loop = asyncio.get_running_loop()
-    original_task_factory = loop.get_task_factory()
-
-    try:
-        sentry_init(integrations=[AsyncioIntegration()], traces_sample_rate=1.0)
-        events = capture_events()
-
-        # Create a user task that should be wrapped
-        async def user_task():
-            await asyncio.sleep(0.01)
-            return "user_result"
+    sentry_init(integrations=[AsyncioIntegration()], traces_sample_rate=1.0)
+    events = capture_events()
 
-        # Create an internal task that should NOT be wrapped
-        async def internal_task():
-            await asyncio.sleep(0.01)
-            return "internal_result"
+    # Create a user task that should be wrapped
+    async def user_task():
+        await asyncio.sleep(0.01)
+        return "user_result"
 
-        with sentry_sdk.start_transaction(name="test_transaction"):
-            user_task_obj = asyncio.create_task(user_task())
+    # Create an internal task that should NOT be wrapped
+    async def internal_task():
+        await asyncio.sleep(0.01)
+        return "internal_result"
 
-            with mark_sentry_task_internal():
-                internal_task_obj = asyncio.create_task(internal_task())
+    with sentry_sdk.start_transaction(name="test_transaction"):
+        user_task_obj = asyncio.create_task(user_task())
 
-            user_result = await user_task_obj
-            internal_result = await internal_task_obj
+        with mark_sentry_task_internal():
+            internal_task_obj = asyncio.create_task(internal_task())
 
-        assert user_result == "user_result"
-        assert internal_result == "internal_result"
+        user_result = await user_task_obj
+        internal_result = await internal_task_obj
 
-        assert len(events) == 1
-        transaction = events[0]
+    assert user_result == "user_result"
+    assert internal_result == "internal_result"
 
-        user_spans = []
-        internal_spans = []
+    assert len(events) == 1
+    transaction = events[0]
 
-        for span in transaction.get("spans", []):
-            if "user_task" in span.get("description", ""):
-                user_spans.append(span)
-            elif "internal_task" in span.get("description", ""):
-                internal_spans.append(span)
+    user_spans = []
+    internal_spans = []
 
-        assert (
-            len(user_spans) > 0
-        ), f"User task should have been traced. All spans: {[s.get('description') for s in transaction.get('spans', [])]}"
-        assert (
-            len(internal_spans) == 0
-        ), f"Internal task should NOT have been traced. All spans: {[s.get('description') for s in transaction.get('spans', [])]}"
+    for span in transaction.get("spans", []):
+        if "user_task" in span.get("description", ""):
+            user_spans.append(span)
+        elif "internal_task" in span.get("description", ""):
+            internal_spans.append(span)
 
-    finally:
-        # Restore original task factory to avoid interfering with other tests
-        loop.set_task_factory(original_task_factory)
+    assert (
+        len(user_spans) > 0
+    ), f"User task should have been traced. All spans: {[s.get('description') for s in transaction.get('spans', [])]}"
+    assert (
+        len(internal_spans) == 0
+    ), f"Internal task should NOT have been traced. All spans: {[s.get('description') for s in transaction.get('spans', [])]}"