Repository: getsentry/sentry-python

Issue:
asyncio.gather tracing produces nested spans
I've got the following code:

```python
import asyncio
import sentry_sdk

async def foo():
    with sentry_sdk.start_span(op="foo"):
        await asyncio.sleep(1)

async def bar():
    with sentry_sdk.start_span(op="bar"):
        await asyncio.sleep(1)

with sentry_sdk.start_span(op="root"):
    await asyncio.gather(foo(), bar(), return_exceptions=True)
```

I see that `bar` is nested under `foo`, and `foo` is nested under `root` in the Performance UI.

```
root
    |- foo
         |- bar
```

Expected behavior: both `foo` and `bar` are nested under `root` and are on the same level.

```
root
    |- foo
    |- bar
```


Discussion:
yes, our asyncio support is somewhat limited at the moment. As a workaround you can do:

```python
import asyncio
import sentry_sdk

async def foo():
    with sentry_sdk.Hub(sentry_sdk.Hub.current):
        with sentry_sdk.start_span(op="foo"):
            await asyncio.sleep(1)

async def bar():
    with sentry_sdk.Hub(sentry_sdk.Hub.current):
        with sentry_sdk.start_span(op="bar"):
            await asyncio.sleep(1)

with sentry_sdk.start_span(op="root"):
    await asyncio.gather(asyncio.create_task(foo()), asyncio.create_task(bar()), return_exceptions=True)
```

I am not sure if `create_task` is strictly necessary or if `gather` does that for you. If `gather` already creates independent tasks then we can potentially handle this case transparently.
Yes, `asyncio.gather` launches new `Task`-s, via https://docs.python.org/3/library/asyncio-task.html#asyncio.gather

> If any awaitable in aws is a coroutine, it is automatically scheduled as a Task.

All right, thanks for the workaround!
@untitaker do you have any plans to add this as first-class support? we are heavy users of this asyncio and this creates unreadable traces
@yoav-orca not at the moment. It would also require heavy monkeypatching of the asyncio runtime afaik
We're currently looking into this and as far as I can tell we can tell asyncio to use a custom task factory which will wrap all tasks with hub cloning and that should fix this issue.

https://docs.python.org/3/library/asyncio-eventloop.html#asyncio.loop.set_task_factory

We just have to be careful to take into account if there's another factory override and be good citizens on our end in both cases (default/existing custom factory).


Resolving patch:
diff --git a/sentry_sdk/consts.py b/sentry_sdk/consts.py
--- a/sentry_sdk/consts.py
+++ b/sentry_sdk/consts.py
@@ -111,6 +111,7 @@ class OP:
     DB = "db"
     DB_REDIS = "db.redis"
     EVENT_DJANGO = "event.django"
+    FUNCTION = "function"
     FUNCTION_AWS = "function.aws"
     FUNCTION_GCP = "function.gcp"
     HTTP_CLIENT = "http.client"
diff --git a/sentry_sdk/integrations/asyncio.py b/sentry_sdk/integrations/asyncio.py
new file mode 100644
--- /dev/null
+++ b/sentry_sdk/integrations/asyncio.py
@@ -0,0 +1,64 @@
+from __future__ import absolute_import
+
+from sentry_sdk.consts import OP
+from sentry_sdk.hub import Hub
+from sentry_sdk.integrations import Integration, DidNotEnable
+from sentry_sdk._types import MYPY
+
+try:
+    import asyncio
+    from asyncio.tasks import Task
+except ImportError:
+    raise DidNotEnable("asyncio not available")
+
+
+if MYPY:
+    from typing import Any
+
+
+def _sentry_task_factory(loop, coro):
+    # type: (Any, Any) -> Task[None]
+
+    async def _coro_creating_hub_and_span():
+        # type: () -> None
+        hub = Hub(Hub.current)
+        with hub:
+            with hub.start_span(op=OP.FUNCTION, description=coro.__qualname__):
+                await coro
+
+    # Trying to use user set task factory (if there is one)
+    orig_factory = loop.get_task_factory()
+    if orig_factory:
+        return orig_factory(loop, _coro_creating_hub_and_span)
+
+    # The default task factory in `asyncio` does not have its own function
+    # but is just a couple of lines in `asyncio.base_events.create_task()`
+    # Those lines are copied here.
+
+    # WARNING:
+    # If the default behavior of the task creation in asyncio changes,
+    # this will break!
+    task = Task(_coro_creating_hub_and_span, loop=loop)  # type: ignore
+    if task._source_traceback:  # type: ignore
+        del task._source_traceback[-1]  # type: ignore
+
+    return task
+
+
+def patch_asyncio():
+    # type: () -> None
+    try:
+        loop = asyncio.get_running_loop()
+        loop.set_task_factory(_sentry_task_factory)
+    except RuntimeError:
+        # When there is no running loop, we have nothing to patch.
+        pass
+
+
+class AsyncioIntegration(Integration):
+    identifier = "asyncio"
+
+    @staticmethod
+    def setup_once():
+        # type: () -> None
+        patch_asyncio()


Test patch:
diff --git a/tests/integrations/asyncio/__init__.py b/tests/integrations/asyncio/__init__.py
new file mode 100644
diff --git a/tests/integrations/asyncio/test_asyncio.py b/tests/integrations/asyncio/test_asyncio.py
new file mode 100644
--- /dev/null
+++ b/tests/integrations/asyncio/test_asyncio.py
@@ -0,0 +1,118 @@
+import asyncio
+import sys
+
+import pytest
+import pytest_asyncio
+
+import sentry_sdk
+from sentry_sdk.consts import OP
+from sentry_sdk.integrations.asyncio import AsyncioIntegration
+
+
+minimum_python_36 = pytest.mark.skipif(
+    sys.version_info < (3, 6), reason="ASGI is only supported in Python >= 3.6"
+)
+
+
+async def foo():
+    await asyncio.sleep(0.01)
+
+
+async def bar():
+    await asyncio.sleep(0.01)
+
+
+@pytest_asyncio.fixture(scope="session")
+def event_loop(request):
+    """Create an instance of the default event loop for each test case."""
+    loop = asyncio.get_event_loop_policy().new_event_loop()
+    yield loop
+    loop.close()
+
+
+@minimum_python_36
+@pytest.mark.asyncio
+async def test_create_task(
+    sentry_init,
+    capture_events,
+    event_loop,
+):
+    sentry_init(
+        traces_sample_rate=1.0,
+        send_default_pii=True,
+        debug=True,
+        integrations=[
+            AsyncioIntegration(),
+        ],
+    )
+
+    events = capture_events()
+
+    with sentry_sdk.start_transaction(name="test_transaction_for_create_task"):
+        with sentry_sdk.start_span(op="root", description="not so important"):
+            tasks = [event_loop.create_task(foo()), event_loop.create_task(bar())]
+            await asyncio.wait(tasks, return_when=asyncio.FIRST_EXCEPTION)
+
+            sentry_sdk.flush()
+
+    (transaction_event,) = events
+
+    assert transaction_event["spans"][0]["op"] == "root"
+    assert transaction_event["spans"][0]["description"] == "not so important"
+
+    assert transaction_event["spans"][1]["op"] == OP.FUNCTION
+    assert transaction_event["spans"][1]["description"] == "foo"
+    assert (
+        transaction_event["spans"][1]["parent_span_id"]
+        == transaction_event["spans"][0]["span_id"]
+    )
+
+    assert transaction_event["spans"][2]["op"] == OP.FUNCTION
+    assert transaction_event["spans"][2]["description"] == "bar"
+    assert (
+        transaction_event["spans"][2]["parent_span_id"]
+        == transaction_event["spans"][0]["span_id"]
+    )
+
+
+@minimum_python_36
+@pytest.mark.asyncio
+async def test_gather(
+    sentry_init,
+    capture_events,
+):
+    sentry_init(
+        traces_sample_rate=1.0,
+        send_default_pii=True,
+        debug=True,
+        integrations=[
+            AsyncioIntegration(),
+        ],
+    )
+
+    events = capture_events()
+
+    with sentry_sdk.start_transaction(name="test_transaction_for_gather"):
+        with sentry_sdk.start_span(op="root", description="not so important"):
+            await asyncio.gather(foo(), bar(), return_exceptions=True)
+
+        sentry_sdk.flush()
+
+    (transaction_event,) = events
+
+    assert transaction_event["spans"][0]["op"] == "root"
+    assert transaction_event["spans"][0]["description"] == "not so important"
+
+    assert transaction_event["spans"][1]["op"] == OP.FUNCTION
+    assert transaction_event["spans"][1]["description"] == "foo"
+    assert (
+        transaction_event["spans"][1]["parent_span_id"]
+        == transaction_event["spans"][0]["span_id"]
+    )
+
+    assert transaction_event["spans"][2]["op"] == OP.FUNCTION
+    assert transaction_event["spans"][2]["description"] == "bar"
+    assert (
+        transaction_event["spans"][2]["parent_span_id"]
+        == transaction_event["spans"][0]["span_id"]
+    )