Add support for jitter so that we can ensure evenly distributed tasks don't cause all workers to restart at same time

Author: adiberk

taskiq/api/receiver.py

@@ -15,6 +15,7 @@ async def run_receiver_task(
     sync_workers: int | None = None,
     validate_params: bool = True,
     max_async_tasks: int = 100,
+    max_async_tasks_jitter: int = 0,
     max_prefetch: int = 0,
     propagate_exceptions: bool = True,
     run_startup: bool = False,
@@ -43,6 +44,7 @@ async def run_receiver_task(
         or processes in processpool that runs sync tasks.
     :param validate_params: whether to validate params or not.
     :param max_async_tasks: maximum number of simultaneous async tasks.
+    :param max_async_tasks_jitter: random jitter to add to max_async_tasks.
     :param max_prefetch: maximum number of tasks to prefetch.
     :param propagate_exceptions: whether to propagate exceptions in generators or not.
     :param run_startup: whether to run startup function or not.
@@ -79,6 +81,7 @@ def on_exit(_: Receiver) -> None:
                     run_startup=run_startup,
                     validate_params=validate_params,
                     max_async_tasks=max_async_tasks,
+                    max_async_tasks_jitter=max_async_tasks_jitter,
                     max_prefetch=max_prefetch,
                     propagate_exceptions=propagate_exceptions,
                     on_exit=on_exit,

taskiq/brokers/inmemory_broker.py

@@ -127,6 +127,7 @@ def __init__(
         max_stored_results: int = 100,
         cast_types: bool = True,
         max_async_tasks: int = 30,
+        max_async_tasks_jitter: int = 0,
         propagate_exceptions: bool = True,
         await_inplace: bool = False,
     ) -> None:
@@ -140,6 +141,7 @@ def __init__(
             executor=self.executor,
             validate_params=cast_types,
             max_async_tasks=max_async_tasks,
+            max_async_tasks_jitter=max_async_tasks_jitter,
             propagate_exceptions=propagate_exceptions,
         )
         self.await_inplace = await_inplace

taskiq/cli/worker/args.py

@@ -44,6 +44,7 @@ class WorkerArgs:
     reload_dirs: list[str] = field(default_factory=list)
     no_gitignore: bool = False
     max_async_tasks: int = 100
+    max_async_tasks_jitter: int = 0
     receiver: str = "taskiq.receiver:Receiver"
     receiver_arg: list[tuple[str, str]] = field(default_factory=list)
     max_prefetch: int = 0
@@ -210,6 +211,14 @@ def from_cli(
             default=100,
             help="Maximum simultaneous async tasks per worker process. ",
         )
+        parser.add_argument(
+            "--max-async-tasks-jitter",
+            type=int,
+            dest="max_async_tasks_jitter",
+            default=0,
+            help="Add random jitter (0 to this value) to max-async-tasks to prevent "
+            "all workers from closing at the same time. ",
+        )
         parser.add_argument(
             "--max-prefetch",
             type=int,

taskiq/cli/worker/run.py

@@ -165,6 +165,7 @@ def interrupt_handler(signum: int, _frame: Any) -> None:
                 executor=pool,
                 validate_params=not args.no_parse,
                 max_async_tasks=args.max_async_tasks,
+                max_async_tasks_jitter=args.max_async_tasks_jitter,
                 max_prefetch=args.max_prefetch,
                 propagate_exceptions=not args.no_propagate_errors,
                 ack_type=args.ack_type,

taskiq/receiver/receiver.py

@@ -2,6 +2,7 @@
 import contextvars
 import functools
 import inspect
+import random
 import sys
 from collections.abc import Callable
 from concurrent.futures import Executor
@@ -37,6 +38,7 @@ def __init__(
         executor: Executor | None = None,
         validate_params: bool = True,
         max_async_tasks: "int | None" = None,
+        max_async_tasks_jitter: int = 0,
         max_prefetch: int = 0,
         propagate_exceptions: bool = True,
         run_startup: bool = True,
@@ -62,7 +64,11 @@ def __init__(
             self._prepare_task(task.task_name, task.original_func)
         self.sem: asyncio.Semaphore | None = None
         if max_async_tasks is not None and max_async_tasks > 0:
-            self.sem = asyncio.Semaphore(max_async_tasks)
+            # Apply jitter to prevent all workers from hitting the limit simultaneously
+            actual_limit = max_async_tasks
+            if max_async_tasks_jitter > 0:
+                actual_limit = max_async_tasks + random.randint(0, max_async_tasks_jitter)
+            self.sem = asyncio.Semaphore(actual_limit)
         else:
             logger.warning(
                 "Setting unlimited number of async tasks "