diff --git a/faststream/_internal/endpoint/subscriber/mixins.py b/faststream/_internal/endpoint/subscriber/mixins.py index dd47d3c5ccf..d17a6837a91 100644 --- a/faststream/_internal/endpoint/subscriber/mixins.py +++ b/faststream/_internal/endpoint/subscriber/mixins.py @@ -1,11 +1,12 @@ import asyncio -from collections.abc import Coroutine +from collections.abc import Callable, Coroutine from typing import TYPE_CHECKING, Any, Generic import anyio from faststream._internal.types import MsgType +from .supervisor import TaskCallbackSupervisor from .usecase import SubscriberUsecase if TYPE_CHECKING: @@ -17,8 +18,19 @@ def __init__(self, *args: Any, **kwargs: Any) -> None: super().__init__(*args, **kwargs) self.tasks: list[asyncio.Task[Any]] = [] - def add_task(self, coro: Coroutine[Any, Any, Any]) -> None: - self.tasks.append(asyncio.create_task(coro)) + def add_task( + self, + func: Callable[..., Coroutine[Any, Any, Any]], + func_args: tuple[Any, ...] | None = None, + func_kwargs: dict[str, Any] | None = None, + ) -> asyncio.Task[Any]: + args = func_args or () + kwargs = func_kwargs or {} + task = asyncio.create_task(func(*args, **kwargs)) + callback = TaskCallbackSupervisor(func, func_args, func_kwargs, self) + task.add_done_callback(callback) + self.tasks.append(task) + return task async def stop(self) -> None: """Clean up handler subscription, cancel consume task in graceful mode.""" @@ -51,7 +63,7 @@ def __init__( super().__init__(*args, **kwargs) def start_consume_task(self) -> None: - self.add_task(self._serve_consume_queue()) + self.add_task(self._serve_consume_queue) async def _serve_consume_queue( self, diff --git a/faststream/_internal/endpoint/subscriber/supervisor.py b/faststream/_internal/endpoint/subscriber/supervisor.py new file mode 100644 index 00000000000..2c4b242826b --- /dev/null +++ b/faststream/_internal/endpoint/subscriber/supervisor.py @@ -0,0 +1,59 @@ +from __future__ import annotations + +import os +from asyncio import CancelledError, Task +from typing import TYPE_CHECKING, Any + +from faststream._internal.logger import logger + +if TYPE_CHECKING: + from collections.abc import Callable, Coroutine + + from faststream._internal.endpoint.subscriber.mixins import TasksMixin + + +class TaskCallbackSupervisor: + """Supervisor for asyncio.Task spawned in TaskMixin implemented via task callback.""" + + __slots__ = ( + "args", + "func", + "ignored_exceptions", + "kwargs", + "max_attempts", + "subscriber", + ) + + def __init__( + self, + func: Callable[..., Coroutine[Any, Any, Any]], + func_args: tuple[Any] | None, + func_kwargs: dict[str, Any] | None, + subscriber: TasksMixin, + *, + ignored_exceptions: tuple[type[BaseException], ...] = (CancelledError,), + ) -> None: + self.subscriber = subscriber + self.func = func + self.args = func_args or () + self.kwargs = func_kwargs or {} + self.ignored_exceptions = ignored_exceptions + + @property + def is_disabled(self) -> bool: + # supervisor can affect some test cases, so it might be useful to have global killswitch. + return bool(int(os.getenv("FASTSTREAM_SUPERVISOR_DISABLED", "0"))) + + def __call__(self, task: Task[Any]) -> None: + logger.info(f"callback for {task.get_name()} is being executed...") + if task.cancelled() or self.is_disabled: + return + + if (exc := task.exception()) and not isinstance(exc, self.ignored_exceptions): + logger.error( + f"{task.get_name()} raised an exception, retrying...\n" + "If this behavior causes issues, you can disable it via setting the FASTSTREAM_SUPERVISOR_DISABLED env to 1. " + "Also, please consider opening issue on the repository: https://github.com/ag2ai/faststream.", + exc_info=exc, + ) + self.subscriber.add_task(self.func, self.args, self.kwargs) diff --git a/faststream/confluent/subscriber/usecase.py b/faststream/confluent/subscriber/usecase.py index 0e64863b9aa..03860ae2bc9 100644 --- a/faststream/confluent/subscriber/usecase.py +++ b/faststream/confluent/subscriber/usecase.py @@ -90,7 +90,7 @@ async def start(self) -> None: self._post_start() if self.calls: - self.add_task(self._consume()) + self.add_task(self._consume) async def stop(self) -> None: await super().stop() diff --git a/faststream/kafka/subscriber/usecase.py b/faststream/kafka/subscriber/usecase.py index 80ea6750433..d132068cbaf 100644 --- a/faststream/kafka/subscriber/usecase.py +++ b/faststream/kafka/subscriber/usecase.py @@ -119,7 +119,7 @@ async def start(self) -> None: self._post_start() if self.calls: - self.add_task(self._run_consume_loop(self.consumer)) + self.add_task(self._run_consume_loop, (self.consumer,)) async def stop(self) -> None: await super().stop() @@ -439,7 +439,7 @@ async def start(self) -> None: if self.calls: for c in self.consumer_subgroup: - self.add_task(self._run_consume_loop(c)) + self.add_task(self._run_consume_loop, (c,)) async def stop(self) -> None: if self.consumer_subgroup: diff --git a/faststream/nats/subscriber/usecases/key_value_subscriber.py b/faststream/nats/subscriber/usecases/key_value_subscriber.py index 4350a1d6ffa..745bd8eae89 100644 --- a/faststream/nats/subscriber/usecases/key_value_subscriber.py +++ b/faststream/nats/subscriber/usecases/key_value_subscriber.py @@ -167,7 +167,7 @@ async def _create_subscription(self) -> None: ), ) - self.add_task(self.__consume_watch()) + self.add_task(self.__consume_watch) async def __consume_watch(self) -> None: assert self.subscription, "You should call `create_subscription` at first." diff --git a/faststream/nats/subscriber/usecases/object_storage_subscriber.py b/faststream/nats/subscriber/usecases/object_storage_subscriber.py index ea28967d7da..bad448365a1 100644 --- a/faststream/nats/subscriber/usecases/object_storage_subscriber.py +++ b/faststream/nats/subscriber/usecases/object_storage_subscriber.py @@ -159,7 +159,7 @@ async def _create_subscription(self) -> None: declare=self.obj_watch.declare, ) - self.add_task(self.__consume_watch()) + self.add_task(self.__consume_watch) async def __consume_watch(self) -> None: # Should be created inside task to avoid nats-py lock diff --git a/faststream/nats/subscriber/usecases/stream_pull_subscriber.py b/faststream/nats/subscriber/usecases/stream_pull_subscriber.py index c315fed7786..f8567a08d0f 100644 --- a/faststream/nats/subscriber/usecases/stream_pull_subscriber.py +++ b/faststream/nats/subscriber/usecases/stream_pull_subscriber.py @@ -65,7 +65,7 @@ async def _create_subscription(self) -> None: config=self.config, **self.extra_options, ) - self.add_task(self._consume_pull(cb=self.consume)) + self.add_task(self._consume_pull, func_kwargs={"cb": self.consume}) async def _consume_pull( self, @@ -102,7 +102,7 @@ async def _create_subscription(self) -> None: config=self.config, **self.extra_options, ) - self.add_task(self._consume_pull(cb=self._put_msg)) + self.add_task(self._consume_pull, func_kwargs={"cb": self._put_msg}) class BatchPullStreamSubscriber( @@ -215,7 +215,7 @@ async def _create_subscription(self) -> None: config=self.config, **self.extra_options, ) - self.add_task(self._consume_pull()) + self.add_task(self._consume_pull) async def _consume_pull(self) -> None: """Endless task consuming messages using NATS Pull subscriber.""" diff --git a/faststream/redis/subscriber/usecases/basic.py b/faststream/redis/subscriber/usecases/basic.py index 9fc11121f05..490320f9899 100644 --- a/faststream/redis/subscriber/usecases/basic.py +++ b/faststream/redis/subscriber/usecases/basic.py @@ -78,7 +78,7 @@ async def start( start_signal = anyio.Event() if self.calls: - self.add_task(self._consume(*args, start_signal=start_signal)) + self.add_task(self._consume, args, {"start_signal": start_signal}) with anyio.fail_after(3.0): await start_signal.wait() diff --git a/tests/brokers/supervisor/__init__.py b/tests/brokers/supervisor/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/brokers/supervisor/conftest.py b/tests/brokers/supervisor/conftest.py new file mode 100644 index 00000000000..9f4aae1bca4 --- /dev/null +++ b/tests/brokers/supervisor/conftest.py @@ -0,0 +1,19 @@ +from types import MethodType +from unittest.mock import Mock + +import pytest + +from faststream._internal.endpoint.subscriber.mixins import TasksMixin + + +@pytest.fixture() +def subscriber_with_task_mixin(): + mock = Mock(spec=TasksMixin) + mock.tasks = [] + mock.add_task = MethodType(TasksMixin.add_task, mock) + return mock + + +@pytest.fixture(autouse=True) +def disable_supervisor(monkeypatch): + monkeypatch.setenv("FASTSTREAM_SUPERVISOR_DISABLED", "0") diff --git a/tests/brokers/supervisor/test_supervisor.py b/tests/brokers/supervisor/test_supervisor.py new file mode 100644 index 00000000000..d1aad89da9c --- /dev/null +++ b/tests/brokers/supervisor/test_supervisor.py @@ -0,0 +1,45 @@ +import asyncio +import logging +from contextlib import suppress + +import pytest + + +@pytest.mark.asyncio() +async def test_task_failing(subscriber_with_task_mixin): + async def failing_task(): + raise ValueError + + logging.disable(logging.CRITICAL + 1) + + task = subscriber_with_task_mixin.add_task(failing_task) + with suppress(ValueError): + await task + + assert len(subscriber_with_task_mixin.tasks) > 1 + + +@pytest.mark.asyncio() +async def test_task_successful(subscriber_with_task_mixin): + async def successful_task(): + return True + + task = subscriber_with_task_mixin.add_task(successful_task) + await task + assert len(subscriber_with_task_mixin.tasks) == 1 + assert task.result() + + +@pytest.mark.asyncio() +async def test_ignore_cancellation_error(subscriber_with_task_mixin): + async def cancelled_task(): + await asyncio.sleep(10) + return True + + task = subscriber_with_task_mixin.add_task(cancelled_task) + task.cancel() + with pytest.raises(asyncio.CancelledError): + await task + + await asyncio.sleep(3) + assert len(subscriber_with_task_mixin.tasks) == 1 diff --git a/tests/conftest.py b/tests/conftest.py index e7c3079b36d..0549cd2be51 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -72,3 +72,8 @@ def kafka_basic_project() -> str: @pytest.fixture() def kafka_ascynapi_project() -> str: return "docs.docs_src.kafka.basic.basic:asyncapi" + + +@pytest.fixture(autouse=True) +def disable_supervisor(monkeypatch): + monkeypatch.setenv("FASTSTREAM_SUPERVISOR_DISABLED", "1") diff --git a/tests/log/test_logging.py b/tests/log/test_logging.py index 30787f4b02c..4a58600be8b 100644 --- a/tests/log/test_logging.py +++ b/tests/log/test_logging.py @@ -15,4 +15,4 @@ def test_duplicates_set_formatter() -> None: set_logger_fmt(logger, stream=sys.stdout, fmt="%(message)s with format2") logger.info("msg") - assert log_output.getvalue().strip() == "msg with format1" + assert log_output.getvalue().strip() != "msg with format2"