feat: task supervisor

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,

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)

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()

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:

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."

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

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."""

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()

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")

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

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")