Skip to content

Commit 5aa2436

Browse files
authored
Fix asyncio task factory
* Make sure the correct co-routine object is used. * Make sure that if a users task factory is set, it is used.
1 parent c471331 commit 5aa2436

File tree

1 file changed

+27
-26
lines changed

1 file changed

+27
-26
lines changed

sentry_sdk/integrations/asyncio.py

Lines changed: 27 additions & 26 deletions
Original file line numberDiff line numberDiff line change
@@ -16,39 +16,40 @@
1616
from typing import Any
1717

1818

19-
def _sentry_task_factory(loop, coro):
20-
# type: (Any, Any) -> Task[None]
19+
def patch_asyncio():
20+
# type: () -> None
21+
orig_task_factory = None
22+
try:
23+
loop = asyncio.get_running_loop()
24+
orig_task_factory = loop.get_task_factory()
2125

22-
async def _coro_creating_hub_and_span():
23-
# type: () -> None
24-
hub = Hub(Hub.current)
25-
with hub:
26-
with hub.start_span(op=OP.FUNCTION, description=coro.__qualname__):
27-
await coro
26+
def _sentry_task_factory(loop, coro):
27+
# type: (Any, Any) -> Any
2828

29-
# Trying to use user set task factory (if there is one)
30-
orig_factory = loop.get_task_factory()
31-
if orig_factory:
32-
return orig_factory(loop, _coro_creating_hub_and_span)
29+
async def _coro_creating_hub_and_span():
30+
# type: () -> None
31+
hub = Hub(Hub.current)
32+
with hub:
33+
with hub.start_span(op=OP.FUNCTION, description=coro.__qualname__):
34+
await coro
3335

34-
# The default task factory in `asyncio` does not have its own function
35-
# but is just a couple of lines in `asyncio.base_events.create_task()`
36-
# Those lines are copied here.
36+
# Trying to use user set task factory (if there is one)
37+
if orig_task_factory:
38+
return orig_task_factory(loop, _coro_creating_hub_and_span()) # type: ignore
3739

38-
# WARNING:
39-
# If the default behavior of the task creation in asyncio changes,
40-
# this will break!
41-
task = Task(_coro_creating_hub_and_span, loop=loop) # type: ignore
42-
if task._source_traceback: # type: ignore
43-
del task._source_traceback[-1] # type: ignore
40+
# The default task factory in `asyncio` does not have its own function
41+
# but is just a couple of lines in `asyncio.base_events.create_task()`
42+
# Those lines are copied here.
4443

45-
return task
44+
# WARNING:
45+
# If the default behavior of the task creation in asyncio changes,
46+
# this will break!
47+
task = Task(_coro_creating_hub_and_span(), loop=loop)
48+
if task._source_traceback: # type: ignore
49+
del task._source_traceback[-1] # type: ignore
4650

51+
return task
4752

48-
def patch_asyncio():
49-
# type: () -> None
50-
try:
51-
loop = asyncio.get_running_loop()
5253
loop.set_task_factory(_sentry_task_factory)
5354
except RuntimeError:
5455
# When there is no running loop, we have nothing to patch.

0 commit comments

Comments
 (0)