Merge remote-tracking branch 'origin/master' into fix/scheduler-invalid-cron-offset-crash

Author: juanmicl

docs/available-components/brokers.md

@@ -112,6 +112,22 @@ Project link: [taskiq-nats](https://github.com/taskiq-python/taskiq-nats).
 pip install taskiq-nats
 ```
 
+### SQS broker
+
+Project link: [taskiq-sqs](https://github.com/taskiq-python/taskiq-sqs).
+
+```bash
+pip install taskiq-sqs
+```
+
+### Kafka broker
+
+Project link: [taskiq-aio-kafka](https://github.com/taskiq-python/taskiq-aio-kafka).
+
+```bash
+pip install taskiq-aio-kafka
+```
+
 ## Third-party brokers
 
 These brokers are not part of the core Taskiq library. They are maintained by other open‑source contributors. You can install them as a separate packages.
@@ -141,3 +157,11 @@ Project link: [taskiq-ydb](https://github.com/danfimov/taskiq-ydb).
 ```bash
 pip install taskiq-ydb
 ```
+
+### Azure Service Bus broker
+
+Project link: [taskiq-service-bus](https://github.com/hamiltonkibbe/taskiq-service-bus).
+
+```bash
+pip install taskiq-service-bus
+```

docs/available-components/result-backends.md

@@ -40,6 +40,14 @@ Project link: [taskiq-nats](https://github.com/taskiq-python/taskiq-nats).
 pip install taskiq-nats
 ```
 
+### S3 result backend
+
+Project link: [taskiq-sqs](https://github.com/taskiq-python/taskiq-sqs). Provides `S3ResultBackend`.
+
+```bash
+pip install taskiq-sqs
+```
+
 ## Third-party result backends
 
 These result backends are not part of the core Taskiq library. They are maintained by other open‑source contributors. You can install them as a separate packages.

taskiq/brokers/inmemory_broker.py

@@ -5,6 +5,7 @@
 from typing import Any, TypeVar
 
 from taskiq.abc.broker import AsyncBroker
+from taskiq.abc.middleware import TaskiqMiddleware
 from taskiq.abc.result_backend import AsyncResultBackend, TaskiqResult
 from taskiq.depends.progress_tracker import TaskProgress
 from taskiq.events import TaskiqEvents
@@ -198,9 +199,21 @@ async def startup(self) -> None:
             for handler in self.event_handlers.get(event, []):
                 await maybe_awaitable(handler(self.state))
 
+        for middleware in self.middlewares:
+            if middleware.__class__.startup != TaskiqMiddleware.startup:
+                await maybe_awaitable(middleware.startup())
+
+        await self.result_backend.startup()
+
     async def shutdown(self) -> None:
         """Runs shutdown events for client and worker side."""
         for event in (TaskiqEvents.CLIENT_SHUTDOWN, TaskiqEvents.WORKER_SHUTDOWN):
             for handler in self.event_handlers.get(event, []):
                 await maybe_awaitable(handler(self.state))
+
+        for middleware in self.middlewares:
+            if middleware.__class__.shutdown != TaskiqMiddleware.shutdown:
+                await maybe_awaitable(middleware.shutdown())
+
+        await self.result_backend.shutdown()
         self.executor.shutdown()

taskiq/cli/scheduler/args.py

@@ -20,6 +20,7 @@ class SchedulerArgs:
     skip_first_run: bool = False
     update_interval: int | None = None
     loop_interval: int | None = None
+    send_timeout: float | None = None
 
     @classmethod
     def from_cli(cls, args: Sequence[str] | None = None) -> "SchedulerArgs":
@@ -111,6 +112,19 @@ def from_cli(cls, args: Sequence[str] | None = None) -> "SchedulerArgs":
                 "If not specified, scheduler will run once a second."
             ),
         )
+        parser.add_argument(
+            "--send-timeout",
+            type=float,
+            default=None,
+            help=(
+                "Optional per-send timeout (seconds). If set, each spawned send "
+                "task is wrapped in asyncio.wait_for with this timeout, "
+                "preventing a single hung broker.kick from permanently blocking "
+                "subsequent ticks for the same schedule_id. On timeout the send "
+                "is cancelled, a warning is logged, and the next scheduled tick "
+                "will retry. Default: no timeout (backwards-compatible)."
+            ),
+        )
 
         namespace = parser.parse_args(args)
         # If there are any patterns specified, remove default.

taskiq/cli/scheduler/run.py

@@ -175,6 +175,38 @@ async def send(
     await scheduler.on_ready(source, task)
 
 
+async def send_with_timeout(
+    scheduler: TaskiqScheduler,
+    source: ScheduleSource,
+    task: ScheduledTask,
+    timeout: float,
+) -> None:
+    """
+    Send a task, cancelling it if it does not complete within ``timeout`` seconds.
+
+    A timed-out send is logged at WARNING and swallowed — re-raising would propagate
+    out of the scheduler's main loop's ``add_done_callback``, which is not what we
+    want. Suppressing it lets the done_callback clear ``running_schedules`` so the
+    next cron boundary can re-dispatch the task. The slot is freed; the message is
+    dropped (the broker did not acknowledge it within the budget).
+
+    :param scheduler: current scheduler.
+    :param source: source of the task.
+    :param task: task to send.
+    :param timeout: seconds to wait before cancelling the send. Must be > 0.
+    """
+    try:
+        await asyncio.wait_for(send(scheduler, source, task), timeout=timeout)
+    except asyncio.TimeoutError:
+        logger.warning(
+            "Sending task %s with schedule_id %s timed out after %.1fs "
+            "and was cancelled. The next scheduled tick will retry.",
+            task.task_name,
+            task.schedule_id,
+            timeout,
+        )
+
+
 async def _sleep_until_next_second() -> None:
     now = datetime.now(tz=timezone.utc)
     await asyncio.sleep(1 - now.microsecond / 1_000_000)
@@ -293,6 +325,7 @@ async def run(
         update_interval: timedelta | None = None,
         loop_interval: timedelta | None = None,
         skip_first_run: bool = False,
+        send_timeout: float | None = None,
     ) -> None:
         """
         Runs scheduler loop.
@@ -304,11 +337,19 @@ async def run(
         :param loop_interval: interval to check tasks to send.
         :param skip_first_run: Wait for the beginning of the next minute
             to skip the first run.
+        :param send_timeout: optional per-send timeout (seconds). If set, each
+            spawned send task is wrapped in :func:`asyncio.wait_for` with this
+            timeout, preventing a single hung ``broker.kick`` from permanently
+            blocking subsequent ticks of the same ``schedule_id`` via the
+            ``running_schedules`` skip check. Default ``None`` (no timeout —
+            backwards-compatible behavior).
         """
         if update_interval is None:
             update_interval = timedelta(minutes=1)
         if loop_interval is None:
             loop_interval = timedelta(seconds=1)
+        if send_timeout is not None and send_timeout <= 0:
+            raise ValueError("send_timeout must be > 0 when provided")
 
         running_schedules: dict[ScheduleId, asyncio.Task[Any]] = {}
 
@@ -336,8 +377,17 @@ async def run(
                     )
 
                     if is_ready_to_send and task.schedule_id not in running_schedules:
+                        if send_timeout is not None:
+                            send_coro = send_with_timeout(
+                                self.scheduler,
+                                source,
+                                task,
+                                timeout=send_timeout,
+                            )
+                        else:
+                            send_coro = send(self.scheduler, source, task)
                         send_task = self._event_loop.create_task(
-                            send(self.scheduler, source, task),
+                            send_coro,
                             # We need to set the name of the task
                             # to be able to discard its reference
                             # after it is done.
@@ -413,6 +463,7 @@ async def run_scheduler(args: SchedulerArgs) -> None:
             update_interval=update_interval,
             loop_interval=loop_interval,
             skip_first_run=args.skip_first_run,
+            send_timeout=args.send_timeout,
         )
     except asyncio.CancelledError:
         logger.info("Shutting down scheduler.")

tests/brokers/test_inmemory.py

@@ -4,6 +4,7 @@
 import pytest
 
 from taskiq import InMemoryBroker
+from taskiq.abc.middleware import TaskiqMiddleware
 from taskiq.events import TaskiqEvents
 from taskiq.state import TaskiqState
 
@@ -66,6 +67,24 @@ async def _c_startup(state: TaskiqState) -> None:
     assert broker.state.from_client == test_value
 
 
+async def test_middleware_startup_and_shutdown_fire() -> None:
+    calls = []
+
+    class RecordingMiddleware(TaskiqMiddleware):
+        async def startup(self) -> None:
+            calls.append("startup")
+
+        async def shutdown(self) -> None:
+            calls.append("shutdown")
+
+    broker = InMemoryBroker().with_middlewares(RecordingMiddleware())
+
+    await broker.startup()
+    await broker.shutdown()
+
+    assert calls == ["startup", "shutdown"]
+
+
 async def test_execution() -> None:
     broker = InMemoryBroker()
     test_value = uuid.uuid4().hex