Skip to content

Commit 5b6d387

Browse files
committed
Added __aexit__ and __aenter__ support. (#614)
Fixed process-pool.
1 parent 3d83293 commit 5b6d387

4 files changed

Lines changed: 258 additions & 8 deletions

File tree

‎taskiq/abc/broker.py‎

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -533,3 +533,11 @@ def _register_task(
533533
if task.broker != self:
534534
raise TaskBrokerMismatchError(broker=task.broker)
535535
self.local_task_registry[task_name] = task
536+
537+
async def __aenter__(self) -> None:
538+
"""Starts the broker as ctx manager."""
539+
await self.startup()
540+
541+
async def __aexit__(self, *args: object, **kwargs: Any) -> None:
542+
"""Shuts down the broker as ctx manager."""
543+
await self.shutdown()

‎taskiq/receiver/receiver.py‎

Lines changed: 50 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -4,7 +4,7 @@
44
import inspect
55
import sys
66
from collections.abc import Callable
7-
from concurrent.futures import Executor
7+
from concurrent.futures import Executor, ProcessPoolExecutor
88
from logging import getLogger
99
from time import time
1010
from typing import Any, get_type_hints
@@ -28,6 +28,33 @@
2828
QUEUE_DONE = b"-1"
2929

3030

31+
def _execute_sync_task_in_executor(
32+
target: Callable[..., Any],
33+
args: tuple[Any, ...],
34+
kwargs: dict[str, Any],
35+
) -> Any:
36+
"""Execute a sync task.
37+
38+
This is a wrapper to ensure we pass the target function directly
39+
to the executor, avoiding issues with pickling bound methods like ctx.run.
40+
41+
:param target: function to execute
42+
:param args: positional arguments
43+
:param kwargs: keyword arguments
44+
:return: result of the function call
45+
"""
46+
return target(*args, **kwargs)
47+
48+
49+
def _is_process_pool_executor(executor: Executor) -> bool:
50+
"""Check if the executor is a ProcessPoolExecutor.
51+
52+
:param executor: executor instance
53+
:return: True if it's a ProcessPoolExecutor
54+
"""
55+
return isinstance(executor, ProcessPoolExecutor)
56+
57+
3158
class Receiver:
3259
"""Class that uses as a callback handler."""
3360

@@ -247,13 +274,28 @@ async def run_task( # noqa: C901, PLR0912, PLR0915
247274
is_coroutine = False
248275
# If this is a synchronous function, we
249276
# run it in executor and preserve the context.
250-
ctx = contextvars.copy_context()
251-
func = functools.partial(target, *message.args, **kwargs)
252-
target_future = loop.run_in_executor(
253-
self.executor,
254-
ctx.run,
255-
func,
256-
)
277+
if _is_process_pool_executor(self.executor):
278+
# For ProcessPoolExecutor, we can't use ctx.run because it contains
279+
# a reference to contextvars.Context which cannot be pickled.
280+
# Instead, we call the target function directly in the executor.
281+
# Each worker process starts with its own context, so we don't need
282+
# to preserve the parent context.
283+
target_future = loop.run_in_executor(
284+
self.executor,
285+
_execute_sync_task_in_executor,
286+
target,
287+
message.args,
288+
kwargs,
289+
)
290+
else:
291+
# For ThreadPoolExecutor, we can use ctx.run with functools.partial
292+
ctx = contextvars.copy_context()
293+
func = functools.partial(target, *message.args, **kwargs)
294+
target_future = loop.run_in_executor(
295+
self.executor,
296+
ctx.run,
297+
func,
298+
)
257299
timeout = message.labels.get("timeout")
258300
if timeout is not None:
259301
if not is_coroutine:

‎tests/abc/test_broker.py‎

Lines changed: 107 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,9 +1,13 @@
11
from collections.abc import AsyncGenerator
22
from copy import copy
33

4+
import pytest
5+
46
from taskiq.abc.broker import AsyncBroker
57
from taskiq.decor import AsyncTaskiqDecoratedTask
8+
from taskiq.events import TaskiqEvents
69
from taskiq.message import BrokerMessage
10+
from taskiq.state import TaskiqState
711

812

913
class _TestBroker(AsyncBroker):
@@ -76,3 +80,106 @@ async def test_task() -> None: ...
7680
assert "another_label" in test_kicker.labels
7781

7882
assert test_task.labels == old_labels
83+
84+
85+
@pytest.mark.anyio
86+
async def test_async_context_manager_enter() -> None:
87+
"""Test that __aenter__ calls startup."""
88+
broker = _TestBroker()
89+
startup_called = False
90+
91+
@broker.on_event(TaskiqEvents.CLIENT_STARTUP)
92+
async def track_startup(state: TaskiqState) -> None:
93+
nonlocal startup_called
94+
startup_called = True
95+
96+
async with broker:
97+
assert startup_called is True
98+
99+
100+
@pytest.mark.anyio
101+
async def test_async_context_manager_exit() -> None:
102+
"""Test that __aexit__ calls shutdown."""
103+
broker = _TestBroker()
104+
shutdown_called = False
105+
106+
@broker.on_event(TaskiqEvents.CLIENT_SHUTDOWN)
107+
async def track_shutdown(state: TaskiqState) -> None:
108+
nonlocal shutdown_called
109+
shutdown_called = True
110+
111+
async with broker:
112+
pass
113+
114+
assert shutdown_called is True
115+
116+
117+
@pytest.mark.anyio
118+
async def test_async_context_manager_enter_worker() -> None:
119+
"""Test that __aenter__ calls worker startup when is_worker_process is True."""
120+
broker = _TestBroker()
121+
broker.is_worker_process = True
122+
startup_called = False
123+
124+
@broker.on_event(TaskiqEvents.WORKER_STARTUP)
125+
async def track_startup(state: TaskiqState) -> None:
126+
nonlocal startup_called
127+
startup_called = True
128+
129+
async with broker:
130+
assert startup_called is True
131+
132+
133+
@pytest.mark.anyio
134+
async def test_async_context_manager_exit_worker() -> None:
135+
"""Test that __aexit__ calls worker shutdown when is_worker_process is True."""
136+
broker = _TestBroker()
137+
broker.is_worker_process = True
138+
shutdown_called = False
139+
140+
@broker.on_event(TaskiqEvents.WORKER_SHUTDOWN)
141+
async def track_shutdown(state: TaskiqState) -> None:
142+
nonlocal shutdown_called
143+
shutdown_called = True
144+
145+
async with broker:
146+
pass
147+
148+
assert shutdown_called is True
149+
150+
151+
@pytest.mark.anyio
152+
async def test_async_context_manager_exit_on_exception() -> None:
153+
"""Test that __aexit__ calls shutdown even if exception is raised."""
154+
broker = _TestBroker()
155+
shutdown_called = False
156+
157+
@broker.on_event(TaskiqEvents.CLIENT_SHUTDOWN)
158+
async def track_shutdown(state: TaskiqState) -> None:
159+
nonlocal shutdown_called
160+
shutdown_called = True
161+
162+
with pytest.raises(ValueError, match="Test exception"):
163+
async with broker:
164+
raise ValueError("Test exception")
165+
166+
assert shutdown_called is True
167+
168+
169+
@pytest.mark.anyio
170+
async def test_async_context_manager_exit_worker_on_exception() -> None:
171+
"""Test that __aexit__ calls worker shutdown even if exception is raised."""
172+
broker = _TestBroker()
173+
broker.is_worker_process = True
174+
shutdown_called = False
175+
176+
@broker.on_event(TaskiqEvents.WORKER_SHUTDOWN)
177+
async def track_shutdown(state: TaskiqState) -> None:
178+
nonlocal shutdown_called
179+
shutdown_called = True
180+
181+
with pytest.raises(ValueError, match="Test exception"):
182+
async with broker:
183+
raise ValueError("Test exception")
184+
185+
assert shutdown_called is True
Lines changed: 93 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,93 @@
1+
"""Test for ProcessPoolExecutor sync task execution."""
2+
import asyncio
3+
from concurrent.futures import ProcessPoolExecutor
4+
5+
import pytest
6+
from taskiq.message import TaskiqMessage
7+
from taskiq.receiver import Receiver
8+
from tests.utils import AsyncQueueBroker
9+
10+
11+
def get_process_pool_receiver(broker) -> Receiver:
12+
"""Get a receiver with ProcessPoolExecutor."""
13+
return Receiver(
14+
broker,
15+
executor=ProcessPoolExecutor(max_workers=2),
16+
validate_params=True,
17+
)
18+
19+
20+
# Module-level functions that can be pickled
21+
def my_test_func(param: int) -> int:
22+
"""A simple test function for ProcessPoolExecutor."""
23+
return param
24+
25+
26+
def my_test_func_with_error() -> None:
27+
"""A test function that raises an error."""
28+
raise ValueError("test error")
29+
30+
31+
def my_test_func_with_args_and_kwargs(a: int, b: str, c: float = 3.14) -> str:
32+
"""A test function with args and kwargs."""
33+
return f"a={a}, b={b}, c={c}"
34+
35+
36+
async def test_run_task_successful_sync_with_process_pool():
37+
"""Tests that run_task can run sync tasks with ProcessPoolExecutor."""
38+
39+
receiver = get_process_pool_receiver(AsyncQueueBroker())
40+
41+
result = await receiver.run_task(
42+
my_test_func,
43+
TaskiqMessage(
44+
task_id="",
45+
task_name="",
46+
labels={},
47+
args=[42],
48+
kwargs={},
49+
),
50+
)
51+
assert result.return_value == 42
52+
53+
54+
async def test_run_task_exception_with_process_pool():
55+
"""Tests that run_task can handle exceptions with ProcessPoolExecutor."""
56+
57+
receiver = get_process_pool_receiver(AsyncQueueBroker())
58+
59+
result = await receiver.run_task(
60+
my_test_func_with_error,
61+
TaskiqMessage(
62+
task_id="",
63+
task_name="",
64+
labels={},
65+
args=[],
66+
kwargs={},
67+
),
68+
)
69+
assert result.return_value is None
70+
assert result.is_err
71+
assert isinstance(result.error, ValueError)
72+
73+
74+
async def test_sync_task_with_args_and_kwargs_process_pool():
75+
"""Tests sync task with both args and kwargs using ProcessPoolExecutor."""
76+
77+
receiver = get_process_pool_receiver(AsyncQueueBroker())
78+
79+
result = await receiver.run_task(
80+
my_test_func_with_args_and_kwargs,
81+
TaskiqMessage(
82+
task_id="",
83+
task_name="",
84+
labels={},
85+
args=[10, "hello"],
86+
kwargs={"c": 2.71},
87+
),
88+
)
89+
assert result.return_value == "a=10, b=hello, c=2.71"
90+
91+
92+
if __name__ == "__main__":
93+
pytest.main([__file__, "-v"])

0 commit comments

Comments
 (0)