Merge branch 'master' into on-send-error-middleware-hook

pyproject.toml

@@ -37,7 +37,7 @@ packaging = ">=19"
 # For prometheus metrics
 prometheus_client = { version = "^0", optional = true }
 # For ZMQBroker
-pyzmq = { version = "^23.2.0", optional = true, markers = "python_version < '3.12'"}
+pyzmq = { version = "^26", optional = true }
 # For speed
 uvloop = { version = ">=0.16.0,<1", optional = true, markers = "sys_platform != 'win32'" }
 # For hot-reload.

taskiq/abc/broker.py

@@ -215,7 +215,7 @@ async def shutdown(self) -> None:
 
         for middleware in self.middlewares:
             if middleware.__class__.shutdown != TaskiqMiddleware.shutdown:
-                await maybe_awaitable(middleware.shutdown)
+                await maybe_awaitable(middleware.shutdown())
 
         await self.result_backend.shutdown()
 

taskiq/cli/worker/process_manager.py

@@ -1,4 +1,5 @@
 import logging
+import os
 import signal
 import sys
 from contextlib import suppress
@@ -265,7 +266,10 @@ def start(self) -> Optional[int]:  # noqa: C901
                     action.handle(self.workers, self.args, self.worker_function)
                     reloaded_workers.add(action.worker_num)
                 elif isinstance(action, ShutdownAction):
-                    logger.debug("Process manager closed.")
+                    logger.debug("Process manager closed, killing workers.")
+                    for worker in self.workers:
+                        if worker.pid:
+                            os.kill(worker.pid, signal.SIGINT)
                     return None
 
             for worker_num, worker in enumerate(self.workers):

taskiq/kicker.py

@@ -50,6 +50,7 @@ def __init__(
         self.broker = broker
         self.labels = labels
         self.custom_task_id: Optional[str] = None
+        self.custom_schedule_id: Optional[str] = None
 
     def with_labels(
         self,
@@ -77,6 +78,19 @@ def with_task_id(self, task_id: str) -> "AsyncKicker[_FuncParams, _ReturnType]":
         self.custom_task_id = task_id
         return self
 
+    def with_schedule_id(
+        self,
+        schedule_id: str,
+    ) -> "AsyncKicker[_FuncParams, _ReturnType]":
+        """
+        Set schedule_id for current execution.
+
+        :param schedule_id: custom schedule id.
+        :return: kicker with custom schedule id.
+        """
+        self.custom_schedule_id = schedule_id
+        return self
+
     def with_broker(
         self,
         broker: "AsyncBroker",
@@ -184,7 +198,9 @@ async def schedule_by_cron(
 
         :return: schedule id.
         """
-        schedule_id = self.broker.id_generator()
+        schedule_id = self.custom_schedule_id
+        if schedule_id is None:
+            schedule_id = self.broker.id_generator()
         message = self._prepare_message(*args, **kwargs)
         cron_offset = None
         if isinstance(cron, CronSpec):
@@ -219,7 +235,9 @@ async def schedule_by_time(
         :param args: function's args.
         :param kwargs: function's kwargs.
         """
-        schedule_id = self.broker.id_generator()
+        schedule_id = self.custom_schedule_id
+        if schedule_id is None:
+            schedule_id = self.broker.id_generator()
         message = self._prepare_message(*args, **kwargs)
         scheduled = ScheduledTask(
             schedule_id=schedule_id,

taskiq/receiver/receiver.py

@@ -1,5 +1,6 @@
 import asyncio
 import inspect
+import signal
 from concurrent.futures import Executor
 from logging import getLogger
 from time import time
@@ -334,6 +335,12 @@ async def listen(self) -> None:  # pragma: no cover
             gr.start_soon(self.prefetcher, queue)
             gr.start_soon(self.runner, queue)
 
+            # Propagate cancellation to the prefetcher & runner
+            def _cancel(*_: Any) -> None:
+                gr.cancel_scope.cancel()
+
+            signal.signal(signal.SIGINT, _cancel)
+
         if self.on_exit is not None:
             self.on_exit(self)
 
@@ -361,9 +368,7 @@ async def prefetcher(
                 message = await iterator.__anext__()
                 fetched_tasks += 1
                 await queue.put(message)
-            except asyncio.CancelledError:
-                break
-            except StopAsyncIteration:
+            except (asyncio.CancelledError, StopAsyncIteration):
                 break
 
         await queue.put(QUEUE_DONE)
@@ -394,31 +399,35 @@ def task_cb(task: "asyncio.Task[Any]") -> None:
                 self.sem.release()
 
         while True:
-            # Waits for semaphore to be released.
-            if self.sem is not None:
-                await self.sem.acquire()
-
-            self.sem_prefetch.release()
-            message = await queue.get()
-            if message is QUEUE_DONE:
-                # asyncio.wait will throw an error if there is nothing to wait for
-                if tasks:
-                    logger.info("Waiting for running tasks to complete.")
-                    await asyncio.wait(tasks, timeout=self.wait_tasks_timeout)
-                break
+            try:
+                # Waits for semaphore to be released.
+                if self.sem is not None:
+                    await self.sem.acquire()
+
+                self.sem_prefetch.release()
+                message = await queue.get()
+                if message is QUEUE_DONE:
+                    # asyncio.wait will throw an error if there is nothing to wait for
+                    if tasks:
+                        logger.info("Waiting for running tasks to complete.")
+                        await asyncio.wait(tasks, timeout=self.wait_tasks_timeout)
+                    break
 
-            task = asyncio.create_task(
-                self.callback(message=message, raise_err=False),
-            )
-            tasks.add(task)
-
-            # We want the task to remove itself from the set when it's done.
-            #
-            # Because if we won't save it anywhere,
-            # python's GC can silently cancel task
-            # and this behaviour considered to be a Hisenbug.
-            # https://textual.textualize.io/blog/2023/02/11/the-heisenbug-lurking-in-your-async-code/
-            task.add_done_callback(task_cb)
+                task = asyncio.create_task(
+                    self.callback(message=message, raise_err=False),
+                )
+                tasks.add(task)
+
+                # We want the task to remove itself from the set when it's done.
+                #
+                # Because if we won't save it anywhere,
+                # python's GC can silently cancel task
+                # and this behaviour considered to be a Hisenbug.
+                # https://textual.textualize.io/blog/2023/02/11/the-heisenbug-lurking-in-your-async-code/
+                task.add_done_callback(task_cb)
+
+            except asyncio.CancelledError:
+                break
 
     def _prepare_task(self, name: str, handler: Callable[..., Any]) -> None:
         """