From 6c1e2cc953fca98dd822e8b0d75f09032e968adb Mon Sep 17 00:00:00 2001 From: Yaroslav <22854425+fil1n@users.noreply.github.com> Date: Fri, 1 Aug 2025 15:10:04 +0500 Subject: [PATCH 01/14] feat: asyncio task supervisor --- faststream/_internal/_compat.py | 2 +- .../_internal/endpoint/subscriber/mixins.py | 20 ++++++-- .../endpoint/subscriber/supervisor.py | 50 +++++++++++++++++++ faststream/confluent/broker/broker.py | 2 +- faststream/confluent/broker/router.py | 2 +- faststream/confluent/helpers/client.py | 6 +-- faststream/confluent/message.py | 6 +-- faststream/confluent/subscriber/usecase.py | 2 +- faststream/kafka/subscriber/usecase.py | 4 +- .../usecases/key_value_subscriber.py | 2 +- .../usecases/object_storage_subscriber.py | 2 +- .../usecases/stream_pull_subscriber.py | 6 +-- faststream/redis/subscriber/usecases/basic.py | 2 +- 13 files changed, 84 insertions(+), 22 deletions(-) create mode 100644 faststream/_internal/endpoint/subscriber/supervisor.py diff --git a/faststream/_internal/_compat.py b/faststream/_internal/_compat.py index fe1d224ff64..a0de06faa48 100644 --- a/faststream/_internal/_compat.py +++ b/faststream/_internal/_compat.py @@ -41,7 +41,7 @@ ujson: Any try: - import orjson # type: ignore[no-redef] + import orjson except ImportError: orjson = None diff --git a/faststream/_internal/endpoint/subscriber/mixins.py b/faststream/_internal/endpoint/subscriber/mixins.py index dd47d3c5ccf..0ecea942b3e 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: tuple[Any] | tuple[()] = func_args or () + kwargs: dict[str, Any] = 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..86d892a26e8 --- /dev/null +++ b/faststream/_internal/endpoint/subscriber/supervisor.py @@ -0,0 +1,50 @@ +from __future__ import annotations + +from asyncio import CancelledError +from logging import getLogger +from typing import TYPE_CHECKING, Any, ClassVar + +if TYPE_CHECKING: + from asyncio import Task + from collections.abc import Callable, Coroutine + + from faststream._internal.endpoint.subscriber.mixins import TasksMixin + +# stores how many times each coroutine has been retried +_attempts_counter: dict[Callable[..., Coroutine[Any, Any, Any]], int] = {} + + +class TaskCallbackSupervisor: + __ignored_exceptions: ClassVar = [CancelledError] + + def __init__( + self, + func: Callable[..., Coroutine[Any, Any, Any]], + func_args: tuple[Any] | None, + func_kwargs: dict[str, Any] | None, + subscriber: TasksMixin, + *, + max_attempts: int = 3, + ) -> None: + self.subscriber = subscriber + self.func = func + self.args: tuple[Any] | tuple[()] = func_args or () + self.kwargs: dict[str, Any] = func_kwargs or {} + self.max_attempts = max_attempts + + def _register_task(self) -> None: + attempts: int = _attempts_counter.get(self.func, 1) + if attempts < self.max_attempts: + self.subscriber.add_task(self.func, self.args, self.kwargs) + _attempts_counter[self.func] = attempts + 1 + + def __call__(self, task: Task[Any]) -> None: + if task.cancelled(): + return + + if (exc := task.exception()) and exc not in self.__ignored_exceptions: + logger = getattr(self.subscriber, "logger", getLogger(__name__)) + logger.error( + f"{task.get_name()} raised an exception, retrying...", exc_info=exc + ) + self._register_task() diff --git a/faststream/confluent/broker/broker.py b/faststream/confluent/broker/broker.py index 54ffade3d09..ff36ebd87be 100644 --- a/faststream/confluent/broker/broker.py +++ b/faststream/confluent/broker/broker.py @@ -260,7 +260,7 @@ def __init__( ) super().__init__( - routers=routers, + routers=routers, # type: ignore[arg-type] config=KafkaBrokerConfig( connection_config=connection_config, client_id=client_id, diff --git a/faststream/confluent/broker/router.py b/faststream/confluent/broker/router.py index 7083c8fc500..87d54c4571e 100644 --- a/faststream/confluent/broker/router.py +++ b/faststream/confluent/broker/router.py @@ -372,5 +372,5 @@ def __init__( include_in_schema=include_in_schema, prefix=prefix, ), - routers=routers, + routers=routers, # type: ignore[arg-type] ) diff --git a/faststream/confluent/helpers/client.py b/faststream/confluent/helpers/client.py index b15ea4fd1c6..c9d87580580 100644 --- a/faststream/confluent/helpers/client.py +++ b/faststream/confluent/helpers/client.py @@ -45,7 +45,7 @@ def __init__( self.config = config.producer_config self.producer = Producer( self.config, - logger=self.logger_state.logger.logger, + logger=self.logger_state.logger.logger, # type: ignore[call-arg] ) self.__running = True @@ -233,7 +233,7 @@ def __init__( } | config.consumer_config self.config = config_from_params - self.consumer = Consumer(self.config, logger=self.logger_state.logger.logger) + self.consumer = Consumer(self.config, logger=self.logger_state.logger.logger) # type: ignore[call-arg] # A pool with single thread is used in order to execute the commands of the consumer sequentially: # https://github.com/ag2ai/faststream/issues/1904#issuecomment-2506990895 @@ -334,7 +334,7 @@ async def getmany( """Consumes a batch of messages from Kafka and groups them by topic and partition.""" raw_messages: list[Message | None] = await run_in_executor( self._thread_pool, - self.consumer.consume, + self.consumer.consume, # type: ignore[arg-type] num_messages=max_records or 10, timeout=timeout, ) diff --git a/faststream/confluent/message.py b/faststream/confluent/message.py index 52aa2485c79..55dbb943325 100644 --- a/faststream/confluent/message.py +++ b/faststream/confluent/message.py @@ -74,8 +74,8 @@ async def nack(self) -> None: else self.raw_message ) await self.consumer.seek( - topic=raw_message.topic(), - partition=raw_message.partition(), - offset=raw_message.offset(), + topic=raw_message.topic(), # type: ignore[arg-type] + partition=raw_message.partition(), # type: ignore[arg-type] + offset=raw_message.offset(), # type: ignore[arg-type] ) await super().nack() 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() From a65485acad81133580c853e2e0189f7dbdb143c7 Mon Sep 17 00:00:00 2001 From: Yaroslav <22854425+fil1n@users.noreply.github.com> Date: Fri, 1 Aug 2025 15:10:37 +0500 Subject: [PATCH 02/14] test: supervisor --- tests/supervisor/__init__.py | 0 tests/supervisor/conftest.py | 14 +++++++++ tests/supervisor/test_supervisor.py | 47 +++++++++++++++++++++++++++++ 3 files changed, 61 insertions(+) create mode 100644 tests/supervisor/__init__.py create mode 100644 tests/supervisor/conftest.py create mode 100644 tests/supervisor/test_supervisor.py diff --git a/tests/supervisor/__init__.py b/tests/supervisor/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/supervisor/conftest.py b/tests/supervisor/conftest.py new file mode 100644 index 00000000000..8c166057cab --- /dev/null +++ b/tests/supervisor/conftest.py @@ -0,0 +1,14 @@ +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 diff --git a/tests/supervisor/test_supervisor.py b/tests/supervisor/test_supervisor.py new file mode 100644 index 00000000000..9637f479145 --- /dev/null +++ b/tests/supervisor/test_supervisor.py @@ -0,0 +1,47 @@ +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 + + # wait until tasks are done + await asyncio.sleep(5) + assert len(subscriber_with_task_mixin.tasks) == 3 + + +@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 From fc4db8ece2acd5ed1906d4f1fa2a4bd6373fa8fa Mon Sep 17 00:00:00 2001 From: Yaroslav <22854425+fil1n@users.noreply.github.com> Date: Fri, 1 Aug 2025 17:48:45 +0500 Subject: [PATCH 03/14] refactor: moved tests to broker dir --- tests/{ => brokers}/supervisor/__init__.py | 0 tests/{ => brokers}/supervisor/conftest.py | 0 tests/{ => brokers}/supervisor/test_supervisor.py | 0 3 files changed, 0 insertions(+), 0 deletions(-) rename tests/{ => brokers}/supervisor/__init__.py (100%) rename tests/{ => brokers}/supervisor/conftest.py (100%) rename tests/{ => brokers}/supervisor/test_supervisor.py (100%) diff --git a/tests/supervisor/__init__.py b/tests/brokers/supervisor/__init__.py similarity index 100% rename from tests/supervisor/__init__.py rename to tests/brokers/supervisor/__init__.py diff --git a/tests/supervisor/conftest.py b/tests/brokers/supervisor/conftest.py similarity index 100% rename from tests/supervisor/conftest.py rename to tests/brokers/supervisor/conftest.py diff --git a/tests/supervisor/test_supervisor.py b/tests/brokers/supervisor/test_supervisor.py similarity index 100% rename from tests/supervisor/test_supervisor.py rename to tests/brokers/supervisor/test_supervisor.py From 2b9fa03a822e58023d86d4b3fb44237a3af1900f Mon Sep 17 00:00:00 2001 From: Yaroslav <22854425+fil1n@users.noreply.github.com> Date: Fri, 1 Aug 2025 18:57:35 +0500 Subject: [PATCH 04/14] refactor: slots and ignored exceptions check --- .../_internal/endpoint/subscriber/supervisor.py | 16 +++++++++++++--- 1 file changed, 13 insertions(+), 3 deletions(-) diff --git a/faststream/_internal/endpoint/subscriber/supervisor.py b/faststream/_internal/endpoint/subscriber/supervisor.py index 86d892a26e8..4fcd7426491 100644 --- a/faststream/_internal/endpoint/subscriber/supervisor.py +++ b/faststream/_internal/endpoint/subscriber/supervisor.py @@ -2,7 +2,7 @@ from asyncio import CancelledError from logging import getLogger -from typing import TYPE_CHECKING, Any, ClassVar +from typing import TYPE_CHECKING, Any if TYPE_CHECKING: from asyncio import Task @@ -15,7 +15,17 @@ class TaskCallbackSupervisor: - __ignored_exceptions: ClassVar = [CancelledError] + """Supervisor for asyncio.Task spawned in TaskMixin implemented via task callback.""" + + __ignored_exceptions: tuple[type[BaseException], ...] = (CancelledError,) + + __slots__ = ( + "args", + "func", + "kwargs", + "max_attempts", + "subscriber", + ) def __init__( self, @@ -42,7 +52,7 @@ def __call__(self, task: Task[Any]) -> None: if task.cancelled(): return - if (exc := task.exception()) and exc not in self.__ignored_exceptions: + if (exc := task.exception()) and not isinstance(exc, self.__ignored_exceptions): logger = getattr(self.subscriber, "logger", getLogger(__name__)) logger.error( f"{task.get_name()} raised an exception, retrying...", exc_info=exc From 2e4f3fb79ce035e75c709c083eb8c5e07956cb0e Mon Sep 17 00:00:00 2001 From: fil1n <22854425+fil1n@users.noreply.github.com> Date: Sun, 3 Aug 2025 20:51:18 +0500 Subject: [PATCH 05/14] chore: ignored mypy error --- faststream/_internal/_compat.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/faststream/_internal/_compat.py b/faststream/_internal/_compat.py index a0de06faa48..fe1d224ff64 100644 --- a/faststream/_internal/_compat.py +++ b/faststream/_internal/_compat.py @@ -41,7 +41,7 @@ ujson: Any try: - import orjson + import orjson # type: ignore[no-redef] except ImportError: orjson = None From 8a862439d663b1c8b2a7fa83f5c84a30de7f813b Mon Sep 17 00:00:00 2001 From: fil1n <22854425+fil1n@users.noreply.github.com> Date: Sun, 3 Aug 2025 21:46:56 +0500 Subject: [PATCH 06/14] feat: supervisor killswitch --- .../endpoint/subscriber/supervisor.py | 18 ++++++++++++++++-- 1 file changed, 16 insertions(+), 2 deletions(-) diff --git a/faststream/_internal/endpoint/subscriber/supervisor.py b/faststream/_internal/endpoint/subscriber/supervisor.py index 4fcd7426491..85be205ccb5 100644 --- a/faststream/_internal/endpoint/subscriber/supervisor.py +++ b/faststream/_internal/endpoint/subscriber/supervisor.py @@ -1,8 +1,9 @@ from __future__ import annotations +import os from asyncio import CancelledError from logging import getLogger -from typing import TYPE_CHECKING, Any +from typing import TYPE_CHECKING, Any, Final if TYPE_CHECKING: from asyncio import Task @@ -13,6 +14,9 @@ # stores how many times each coroutine has been retried _attempts_counter: dict[Callable[..., Coroutine[Any, Any, Any]], int] = {} +# supervisor can affect some test cases, so it might be useful to have global killswitch +SUPERVISOR_DISABLING_ENV_NAME: Final[str] = "FASTSTREAM_SUPERVISOR_DISABLED" + class TaskCallbackSupervisor: """Supervisor for asyncio.Task spawned in TaskMixin implemented via task callback.""" @@ -48,8 +52,18 @@ def _register_task(self) -> None: self.subscriber.add_task(self.func, self.args, self.kwargs) _attempts_counter[self.func] = attempts + 1 + @property + def _is_disabled(self) -> bool: + """Checks if supervisor is disabled globally.""" + try: + integer: int = int(os.getenv(SUPERVISOR_DISABLING_ENV_NAME, default="0")) + except (ValueError, TypeError): + return False + + return bool(integer) + def __call__(self, task: Task[Any]) -> None: - if task.cancelled(): + if task.cancelled() or self._is_disabled: return if (exc := task.exception()) and not isinstance(exc, self.__ignored_exceptions): From 42442261684d8a960a4fe77b753d594e9f13501d Mon Sep 17 00:00:00 2001 From: fil1n <22854425+fil1n@users.noreply.github.com> Date: Sun, 3 Aug 2025 21:48:47 +0500 Subject: [PATCH 07/14] test: disable supervisor --- tests/brokers/supervisor/conftest.py | 8 ++++++++ tests/conftest.py | 8 ++++++++ 2 files changed, 16 insertions(+) diff --git a/tests/brokers/supervisor/conftest.py b/tests/brokers/supervisor/conftest.py index 8c166057cab..e2f6007b7ef 100644 --- a/tests/brokers/supervisor/conftest.py +++ b/tests/brokers/supervisor/conftest.py @@ -4,6 +4,9 @@ import pytest from faststream._internal.endpoint.subscriber.mixins import TasksMixin +from faststream._internal.endpoint.subscriber.supervisor import ( + SUPERVISOR_DISABLING_ENV_NAME, +) @pytest.fixture() @@ -12,3 +15,8 @@ def subscriber_with_task_mixin(): mock.tasks = [] mock.add_task = MethodType(TasksMixin.add_task, mock) return mock + + +@pytest.fixture(autouse=True) +def disable_supervisor(monkeypatch): + monkeypatch.setenv(SUPERVISOR_DISABLING_ENV_NAME, "0") diff --git a/tests/conftest.py b/tests/conftest.py index e7c3079b36d..b42999afabe 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -9,6 +9,9 @@ from faststream.__about__ import __version__ from faststream._internal.context import ContextRepo +from faststream._internal.endpoint.subscriber.supervisor import ( + SUPERVISOR_DISABLING_ENV_NAME, +) @pytest.hookimpl(tryfirst=True) @@ -72,3 +75,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(SUPERVISOR_DISABLING_ENV_NAME, "1") From 19af06f1c6d230009eedcaff3005d8491b51901b Mon Sep 17 00:00:00 2001 From: Yaroslav <22854425+fil1n@users.noreply.github.com> Date: Tue, 12 Aug 2025 18:34:43 +0500 Subject: [PATCH 08/14] refactor: changes requested by the review --- .../_internal/endpoint/subscriber/mixins.py | 4 +- .../endpoint/subscriber/supervisor.py | 42 ++++++++----------- 2 files changed, 19 insertions(+), 27 deletions(-) diff --git a/faststream/_internal/endpoint/subscriber/mixins.py b/faststream/_internal/endpoint/subscriber/mixins.py index 0ecea942b3e..d17a6837a91 100644 --- a/faststream/_internal/endpoint/subscriber/mixins.py +++ b/faststream/_internal/endpoint/subscriber/mixins.py @@ -24,8 +24,8 @@ def add_task( func_args: tuple[Any, ...] | None = None, func_kwargs: dict[str, Any] | None = None, ) -> asyncio.Task[Any]: - args: tuple[Any] | tuple[()] = func_args or () - kwargs: dict[str, Any] = func_kwargs or {} + 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) diff --git a/faststream/_internal/endpoint/subscriber/supervisor.py b/faststream/_internal/endpoint/subscriber/supervisor.py index 85be205ccb5..525b970b087 100644 --- a/faststream/_internal/endpoint/subscriber/supervisor.py +++ b/faststream/_internal/endpoint/subscriber/supervisor.py @@ -1,12 +1,12 @@ from __future__ import annotations import os -from asyncio import CancelledError -from logging import getLogger -from typing import TYPE_CHECKING, Any, Final +from asyncio import CancelledError, Task +from typing import TYPE_CHECKING, Any + +from faststream._internal.logger import logger if TYPE_CHECKING: - from asyncio import Task from collections.abc import Callable, Coroutine from faststream._internal.endpoint.subscriber.mixins import TasksMixin @@ -14,18 +14,14 @@ # stores how many times each coroutine has been retried _attempts_counter: dict[Callable[..., Coroutine[Any, Any, Any]], int] = {} -# supervisor can affect some test cases, so it might be useful to have global killswitch -SUPERVISOR_DISABLING_ENV_NAME: Final[str] = "FASTSTREAM_SUPERVISOR_DISABLED" - class TaskCallbackSupervisor: """Supervisor for asyncio.Task spawned in TaskMixin implemented via task callback.""" - __ignored_exceptions: tuple[type[BaseException], ...] = (CancelledError,) - __slots__ = ( "args", "func", + "ignored_exceptions", "kwargs", "max_attempts", "subscriber", @@ -39,35 +35,31 @@ def __init__( subscriber: TasksMixin, *, max_attempts: int = 3, + ignored_exceptions: tuple[type[BaseException], ...] = (CancelledError,), ) -> None: self.subscriber = subscriber self.func = func - self.args: tuple[Any] | tuple[()] = func_args or () - self.kwargs: dict[str, Any] = func_kwargs or {} + self.args = func_args or () + self.kwargs = func_kwargs or {} self.max_attempts = max_attempts + 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 _register_task(self) -> None: - attempts: int = _attempts_counter.get(self.func, 1) + attempts = _attempts_counter.get(self.func, 1) if attempts < self.max_attempts: self.subscriber.add_task(self.func, self.args, self.kwargs) _attempts_counter[self.func] = attempts + 1 - @property - def _is_disabled(self) -> bool: - """Checks if supervisor is disabled globally.""" - try: - integer: int = int(os.getenv(SUPERVISOR_DISABLING_ENV_NAME, default="0")) - except (ValueError, TypeError): - return False - - return bool(integer) - def __call__(self, task: Task[Any]) -> None: - if task.cancelled() or self._is_disabled: + if task.cancelled() or self.is_disabled: return - if (exc := task.exception()) and not isinstance(exc, self.__ignored_exceptions): - logger = getattr(self.subscriber, "logger", getLogger(__name__)) + if (exc := task.exception()) and not isinstance(exc, self.ignored_exceptions): logger.error( f"{task.get_name()} raised an exception, retrying...", exc_info=exc ) From 9a35cf29fbcc709e6f2629f57d636075f8a6e599 Mon Sep 17 00:00:00 2001 From: Yaroslav <22854425+fil1n@users.noreply.github.com> Date: Tue, 12 Aug 2025 18:34:56 +0500 Subject: [PATCH 09/14] refactor: changed var name --- tests/brokers/supervisor/conftest.py | 5 +---- tests/conftest.py | 5 +---- 2 files changed, 2 insertions(+), 8 deletions(-) diff --git a/tests/brokers/supervisor/conftest.py b/tests/brokers/supervisor/conftest.py index e2f6007b7ef..9f4aae1bca4 100644 --- a/tests/brokers/supervisor/conftest.py +++ b/tests/brokers/supervisor/conftest.py @@ -4,9 +4,6 @@ import pytest from faststream._internal.endpoint.subscriber.mixins import TasksMixin -from faststream._internal.endpoint.subscriber.supervisor import ( - SUPERVISOR_DISABLING_ENV_NAME, -) @pytest.fixture() @@ -19,4 +16,4 @@ def subscriber_with_task_mixin(): @pytest.fixture(autouse=True) def disable_supervisor(monkeypatch): - monkeypatch.setenv(SUPERVISOR_DISABLING_ENV_NAME, "0") + monkeypatch.setenv("FASTSTREAM_SUPERVISOR_DISABLED", "0") diff --git a/tests/conftest.py b/tests/conftest.py index b42999afabe..0549cd2be51 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -9,9 +9,6 @@ from faststream.__about__ import __version__ from faststream._internal.context import ContextRepo -from faststream._internal.endpoint.subscriber.supervisor import ( - SUPERVISOR_DISABLING_ENV_NAME, -) @pytest.hookimpl(tryfirst=True) @@ -79,4 +76,4 @@ def kafka_ascynapi_project() -> str: @pytest.fixture(autouse=True) def disable_supervisor(monkeypatch): - monkeypatch.setenv(SUPERVISOR_DISABLING_ENV_NAME, "1") + monkeypatch.setenv("FASTSTREAM_SUPERVISOR_DISABLED", "1") From 992d238f2e55905ba13087c44d2d83cb817ce508 Mon Sep 17 00:00:00 2001 From: fil1n <22854425+fil1n@users.noreply.github.com> Date: Mon, 1 Sep 2025 01:53:14 +0500 Subject: [PATCH 10/14] refactor: fixed mypy errors --- faststream/confluent/broker/broker.py | 2 +- faststream/confluent/broker/router.py | 2 +- faststream/confluent/helpers/client.py | 6 +++--- faststream/confluent/message.py | 6 +++--- 4 files changed, 8 insertions(+), 8 deletions(-) diff --git a/faststream/confluent/broker/broker.py b/faststream/confluent/broker/broker.py index ff36ebd87be..54ffade3d09 100644 --- a/faststream/confluent/broker/broker.py +++ b/faststream/confluent/broker/broker.py @@ -260,7 +260,7 @@ def __init__( ) super().__init__( - routers=routers, # type: ignore[arg-type] + routers=routers, config=KafkaBrokerConfig( connection_config=connection_config, client_id=client_id, diff --git a/faststream/confluent/broker/router.py b/faststream/confluent/broker/router.py index 87d54c4571e..7083c8fc500 100644 --- a/faststream/confluent/broker/router.py +++ b/faststream/confluent/broker/router.py @@ -372,5 +372,5 @@ def __init__( include_in_schema=include_in_schema, prefix=prefix, ), - routers=routers, # type: ignore[arg-type] + routers=routers, ) diff --git a/faststream/confluent/helpers/client.py b/faststream/confluent/helpers/client.py index c9d87580580..b15ea4fd1c6 100644 --- a/faststream/confluent/helpers/client.py +++ b/faststream/confluent/helpers/client.py @@ -45,7 +45,7 @@ def __init__( self.config = config.producer_config self.producer = Producer( self.config, - logger=self.logger_state.logger.logger, # type: ignore[call-arg] + logger=self.logger_state.logger.logger, ) self.__running = True @@ -233,7 +233,7 @@ def __init__( } | config.consumer_config self.config = config_from_params - self.consumer = Consumer(self.config, logger=self.logger_state.logger.logger) # type: ignore[call-arg] + self.consumer = Consumer(self.config, logger=self.logger_state.logger.logger) # A pool with single thread is used in order to execute the commands of the consumer sequentially: # https://github.com/ag2ai/faststream/issues/1904#issuecomment-2506990895 @@ -334,7 +334,7 @@ async def getmany( """Consumes a batch of messages from Kafka and groups them by topic and partition.""" raw_messages: list[Message | None] = await run_in_executor( self._thread_pool, - self.consumer.consume, # type: ignore[arg-type] + self.consumer.consume, num_messages=max_records or 10, timeout=timeout, ) diff --git a/faststream/confluent/message.py b/faststream/confluent/message.py index 55dbb943325..52aa2485c79 100644 --- a/faststream/confluent/message.py +++ b/faststream/confluent/message.py @@ -74,8 +74,8 @@ async def nack(self) -> None: else self.raw_message ) await self.consumer.seek( - topic=raw_message.topic(), # type: ignore[arg-type] - partition=raw_message.partition(), # type: ignore[arg-type] - offset=raw_message.offset(), # type: ignore[arg-type] + topic=raw_message.topic(), + partition=raw_message.partition(), + offset=raw_message.offset(), ) await super().nack() From a927af04ea44e6e2b8eeaee2a0b070ae0df78cb6 Mon Sep 17 00:00:00 2001 From: fil1n <22854425+fil1n@users.noreply.github.com> Date: Sun, 7 Sep 2025 23:04:18 +0500 Subject: [PATCH 11/14] refactor: removed retries --- .../_internal/endpoint/subscriber/supervisor.py | 13 +------------ 1 file changed, 1 insertion(+), 12 deletions(-) diff --git a/faststream/_internal/endpoint/subscriber/supervisor.py b/faststream/_internal/endpoint/subscriber/supervisor.py index 525b970b087..9ca7a4bcdf9 100644 --- a/faststream/_internal/endpoint/subscriber/supervisor.py +++ b/faststream/_internal/endpoint/subscriber/supervisor.py @@ -11,9 +11,6 @@ from faststream._internal.endpoint.subscriber.mixins import TasksMixin -# stores how many times each coroutine has been retried -_attempts_counter: dict[Callable[..., Coroutine[Any, Any, Any]], int] = {} - class TaskCallbackSupervisor: """Supervisor for asyncio.Task spawned in TaskMixin implemented via task callback.""" @@ -34,14 +31,12 @@ def __init__( func_kwargs: dict[str, Any] | None, subscriber: TasksMixin, *, - max_attempts: int = 3, ignored_exceptions: tuple[type[BaseException], ...] = (CancelledError,), ) -> None: self.subscriber = subscriber self.func = func self.args = func_args or () self.kwargs = func_kwargs or {} - self.max_attempts = max_attempts self.ignored_exceptions = ignored_exceptions @property @@ -49,12 +44,6 @@ 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 _register_task(self) -> None: - attempts = _attempts_counter.get(self.func, 1) - if attempts < self.max_attempts: - self.subscriber.add_task(self.func, self.args, self.kwargs) - _attempts_counter[self.func] = attempts + 1 - def __call__(self, task: Task[Any]) -> None: if task.cancelled() or self.is_disabled: return @@ -63,4 +52,4 @@ def __call__(self, task: Task[Any]) -> None: logger.error( f"{task.get_name()} raised an exception, retrying...", exc_info=exc ) - self._register_task() + self.subscriber.add_task(self.func, self.args, self.kwargs) From f5ad63c86934fa197d5e0fc4de0f18abede349e0 Mon Sep 17 00:00:00 2001 From: fil1n <22854425+fil1n@users.noreply.github.com> Date: Mon, 8 Sep 2025 01:50:54 +0500 Subject: [PATCH 12/14] test: removed retry check --- tests/brokers/supervisor/test_supervisor.py | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/tests/brokers/supervisor/test_supervisor.py b/tests/brokers/supervisor/test_supervisor.py index 9637f479145..d1aad89da9c 100644 --- a/tests/brokers/supervisor/test_supervisor.py +++ b/tests/brokers/supervisor/test_supervisor.py @@ -16,9 +16,7 @@ async def failing_task(): with suppress(ValueError): await task - # wait until tasks are done - await asyncio.sleep(5) - assert len(subscriber_with_task_mixin.tasks) == 3 + assert len(subscriber_with_task_mixin.tasks) > 1 @pytest.mark.asyncio() From 888bb58dd45e488ba20839f90bd4cdcde9219f98 Mon Sep 17 00:00:00 2001 From: fil1n <22854425+fil1n@users.noreply.github.com> Date: Mon, 8 Sep 2025 02:40:43 +0500 Subject: [PATCH 13/14] feat: more logging --- faststream/_internal/endpoint/subscriber/supervisor.py | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/faststream/_internal/endpoint/subscriber/supervisor.py b/faststream/_internal/endpoint/subscriber/supervisor.py index 9ca7a4bcdf9..2c4b242826b 100644 --- a/faststream/_internal/endpoint/subscriber/supervisor.py +++ b/faststream/_internal/endpoint/subscriber/supervisor.py @@ -45,11 +45,15 @@ def is_disabled(self) -> bool: 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...", exc_info=exc + 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) From f178a024f27371f7e386812c19c73aeb17d43016 Mon Sep 17 00:00:00 2001 From: fil1n <22854425+fil1n@users.noreply.github.com> Date: Mon, 8 Sep 2025 03:26:21 +0500 Subject: [PATCH 14/14] test: fix logging test --- tests/log/test_logging.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) 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"