fix(receiver): harden prefetch and listener lifecycle (#645)

* fix(receiver): harden prefetch and listener lifecycle

Track delivery and prefetch capacity ownership across shutdown paths.
Close pending listener reads, preserve late deliveries, and restore
fail-fast transport error propagation.

Keep middleware accounting balanced, report discarded deliveries, and
retain deferred lifecycle diagnostics when task group failures overlap.

* fix(receiver): tighten prefetch ownership checks

Reject invalid prefetch values before receiver retries and retain delivery accounting until execution capacity is acquired.

Add deterministic coverage for CLI defaults, cancellation races, and listener lifecycle cleanup.

Refs #528

* fix(receiver): simplify prefetch listener lifecycle

- remove undocumented prefetch middleware hooks and metrics
- preserve listener errors during task group failures
- restore capacity when callback handoff fails
- strengthen shutdown and cleanup coverage

Refs: #528

* refactor(receiver): simplify prefetch queue cleanup

Author: GefMar

docs/guide/cli.md

@@ -162,7 +162,7 @@ The number of signals before a hard kill can be configured with the `--hardkill-
 * `--log-format` is used to set a log format (default `%(asctime)s][%(name)s][%(levelname)-7s][%(processName)s] %(message)s`).
 * `--max-async-tasks` - maximum number of simultaneously running async tasks.
 * `--max-async-tasks-jitter` – Randomly varies the max async task limit between --max-async-tasks and a jittered value, helping prevent simultaneous worker restarts.
-* `--max-prefetch` - number of tasks to be prefetched before execution. (Useful for systems with high message rates, but brokers should support acknowledgements).
+* `--max-prefetch` - maximum number of deliveries allowed to wait for execution beyond the available async execution capacity. With a finite `--max-async-tasks` limit, a worker admits at most its effective async execution limit plus `max_prefetch` running or waiting deliveries. The default `0` disables additional buffering, and the value must be non-negative. This option is useful for systems with high message rates, but brokers should support acknowledgements.
 * `--max-threadpool-threads` - number of threads for sync function execution.
 * `--no-propagate-errors` - if this parameter is enabled, exceptions won't be thrown in generator dependencies.
 * `--receiver` - python path to custom receiver class.

taskiq/api/receiver.py

@@ -45,13 +45,18 @@ async def run_receiver_task(
     :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 max_prefetch: maximum number of deliveries allowed to wait beyond
+        the available async execution capacity. Must be non-negative.
     :param propagate_exceptions: whether to propagate exceptions in generators or not.
     :param run_startup: whether to run startup function or not.
     :param ack_time: acknowledge type to use.
     :param use_process_pool: whether to use process pool or threadpool.
     :raises asyncio.CancelledError: if the task was cancelled.
+    :raises ValueError: if max_prefetch is negative.
     """
+    if max_prefetch < 0:
+        raise ValueError("max_prefetch cannot be negative.")
+
     finish_event = asyncio.Event()
 
     def on_exit(_: Receiver) -> None:

taskiq/cli/worker/args.py

@@ -224,7 +224,10 @@ def from_cli(
             type=int,
             dest="max_prefetch",
             default=0,
-            help="Maximum prefetched tasks per worker process. ",
+            help=(
+                "Maximum deliveries waiting beyond async execution capacity. "
+                "Must be non-negative. "
+            ),
         )
         parser.add_argument(
             "--no-configure-logging",
@@ -283,6 +286,8 @@ def from_cli(
             args,
             namespace=None if defaults is None else Namespace(**defaults),
         )
+        if namespace.max_prefetch < 0:
+            parser.error("argument --max-prefetch: max_prefetch cannot be negative.")
         # If there are any patterns specified, remove default.
         # This is an argparse limitation.
         if len(namespace.tasks_pattern) > 1:

taskiq/middlewares/opentelemetry_middleware.py

@@ -230,12 +230,6 @@ def __init__(
             unit="1",
             description="Number of tasks currently executing in the worker.",
         )
-        # 9- Number of tasks executing
-        self.number_of_broker_prefetched_tasks = self._meter.create_up_down_counter(
-            "worker_prefetched_tasks",
-            unit="1",
-            description="Number of tasks currently prefetched in the worker.",
-        )
 
     def _observe_memory(self, options: Any) -> Generator[Observation, None, None]:
         if self.broker and self.broker.is_worker_process:
@@ -447,11 +441,3 @@ def post_execute(
             -1,
             attributes={"task_name": message.task_name},
         )
-
-    def on_prefetch_queue_add(self) -> None:
-        """This hook is called after task is added to the worker prefetch queue."""
-        self.number_of_broker_prefetched_tasks.add(1)
-
-    def on_prefetch_queue_remove(self) -> None:
-        """This hook is called after task is removed from the worker prefetch queue."""
-        self.number_of_broker_prefetched_tasks.add(-1)