feat: task's idle execution

Author: Anton

taskiq/brokers/inmemory_queue_broker.py

@@ -0,0 +1,64 @@
+import asyncio
+from typing import AsyncGenerator
+
+from taskiq import BrokerMessage
+from taskiq.brokers.inmemory_broker import InMemoryBroker
+
+
+class InMemoryQueueBroker(InMemoryBroker):
+    """In memory Broker based on asyncio.Queue."""
+
+    def __init__(
+        self,
+        sync_tasks_pool_size: int = 4,
+        max_stored_results: int = 100,
+        cast_types: bool = True,
+        max_async_tasks: int = 30,
+    ) -> None:
+        super().__init__(
+            sync_tasks_pool_size,
+            max_stored_results,
+            cast_types,
+            max_async_tasks,
+        )
+        self.queue: asyncio.Queue[BrokerMessage] = asyncio.Queue()
+
+    async def kick(self, message: BrokerMessage) -> None:
+        """
+        Kicking task.
+
+        :param message: incoming message
+        """
+        await self.queue.put(message)
+
+    async def listen(self) -> AsyncGenerator[bytes, None]:
+        """
+        Listening for messages.
+
+        :yields: message's raw data
+        """
+        running = asyncio.create_task(self.running.wait())
+
+        while not self.running.is_set():
+            message = asyncio.create_task(self.queue.get())
+            await asyncio.wait(
+                [running, message],
+                return_when=asyncio.FIRST_COMPLETED,
+            )
+
+            if message.done():
+                yield (await message).message
+                continue
+
+            message.cancel()
+            await running
+
+    async def startup(self) -> None:
+        """Runs startup events for client and worker side."""
+        await super().startup()
+        self.running = asyncio.Event()
+
+    async def shutdown(self) -> None:
+        """Runs shutdown events for client and worker side."""
+        await super().shutdown()
+        self.running.set()

taskiq/cli/worker/args.py

@@ -41,6 +41,7 @@ class WorkerArgs:
     receiver: str = "taskiq.receiver:Receiver"
     receiver_arg: List[Tuple[str, str]] = field(default_factory=list)
     max_prefetch: int = 0
+    max_idle_tasks: Optional[int] = None
 
     @classmethod
     def from_cli(  # noqa: WPS213
@@ -176,6 +177,13 @@ def from_cli(  # noqa: WPS213
             default=0,
             help="Maximum prefetched tasks per worker process. ",
         )
+        parser.add_argument(
+            "--max-idle-tasks",
+            type=int,
+            dest="max_idle_tasks",
+            default=None,
+            help="Maximum idle tasks per worker process. ",
+        )
 
         namespace = parser.parse_args(args)
         return WorkerArgs(**namespace.__dict__)

taskiq/cli/worker/run.py

@@ -133,6 +133,7 @@ def interrupt_handler(signum: int, _frame: Any) -> None:
                 validate_params=not args.no_parse,
                 max_async_tasks=args.max_async_tasks,
                 max_prefetch=args.max_prefetch,
+                max_idle_tasks=args.max_idle_tasks,
                 **receiver_args,
             )
             loop.run_until_complete(receiver.listen())

taskiq/context.py

@@ -1,4 +1,4 @@
-from typing import TYPE_CHECKING
+from typing import TYPE_CHECKING, Awaitable, Callable
 
 from taskiq.abc.broker import AsyncBroker
 from taskiq.message import TaskiqMessage
@@ -10,8 +10,14 @@
 class Context:
     """Context class."""
 
-    def __init__(self, message: TaskiqMessage, broker: AsyncBroker) -> None:
+    def __init__(
+        self,
+        message: TaskiqMessage,
+        broker: AsyncBroker,
+        task_idler: Callable[[float], Awaitable[None]],
+    ) -> None:
         self.message = message
         self.broker = broker
         self.state: "TaskiqState" = None  # type: ignore
         self.state = broker.state
+        self.task_idler = task_idler

taskiq/receiver/receiver.py

@@ -15,10 +15,11 @@
 from taskiq.receiver.params_parser import parse_params
 from taskiq.result import TaskiqResult
 from taskiq.state import TaskiqState
-from taskiq.utils import maybe_awaitable
+from taskiq.utils import DequeQueue, DequeSemaphore, maybe_awaitable
 
 logger = getLogger(__name__)
 QUEUE_DONE = b"-1"
+QUEUE_SKIP = b"-2"
 
 
 def _run_sync(target: Callable[..., Any], message: TaskiqMessage) -> Any:
@@ -45,6 +46,7 @@ def __init__(  # noqa: WPS211
         validate_params: bool = True,
         max_async_tasks: "Optional[int]" = None,
         max_prefetch: int = 0,
+        max_idle_tasks: Optional[int] = None,
     ) -> None:
         self.broker = broker
         self.executor = executor
@@ -56,15 +58,20 @@ def __init__(  # noqa: WPS211
             self.task_signatures[task.task_name] = inspect.signature(task.original_func)
             self.task_hints[task.task_name] = get_type_hints(task.original_func)
             self.dependency_graphs[task.task_name] = DependencyGraph(task.original_func)
-        self.sem: "Optional[asyncio.Semaphore]" = None
+        self.sem: "Optional[DequeSemaphore]" = None
         if max_async_tasks is not None and max_async_tasks > 0:
-            self.sem = asyncio.Semaphore(max_async_tasks)
+            self.sem = DequeSemaphore(max_async_tasks)
         else:
             logger.warning(
                 "Setting unlimited number of async tasks "
                 + "can result in undefined behavior",
             )
-        self.sem_prefetch = asyncio.Semaphore(max_prefetch)
+        self.sem_prefetch = DequeSemaphore(max_prefetch)
+        self.queue: DequeQueue[bytes] = DequeQueue()
+
+        self.sem_idle: Optional[asyncio.Semaphore] = None
+        if max_idle_tasks and max_idle_tasks > 0:
+            self.sem_idle = asyncio.Semaphore(max_idle_tasks)
 
     async def callback(  # noqa: C901, WPS213
         self,
@@ -176,7 +183,7 @@ async def run_task(  # noqa: C901, WPS210
             broker_ctx = self.broker.custom_dependency_context
             broker_ctx.update(
                 {
-                    Context: Context(message, self.broker),
+                    Context: Context(message, self.broker, self.task_idler),
                     TaskiqState: self.broker.state,
                 },
             )
@@ -243,11 +250,10 @@ async def listen(self) -> None:  # pragma: no cover
         """
         await self.broker.startup()
         logger.info("Listening started.")
-        queue: asyncio.Queue[bytes] = asyncio.Queue()
 
         async with anyio.create_task_group() as gr:
-            gr.start_soon(self.prefetcher, queue)
-            gr.start_soon(self.runner, queue)
+            gr.start_soon(self.prefetcher, self.queue)
+            gr.start_soon(self.runner, self.queue)
 
     async def prefetcher(self, queue: "asyncio.Queue[Any]") -> None:
         """
@@ -268,7 +274,7 @@ async def prefetcher(self, queue: "asyncio.Queue[Any]") -> None:
 
         await queue.put(QUEUE_DONE)
 
-    async def runner(self, queue: "asyncio.Queue[bytes]") -> None:
+    async def runner(self, queue: "asyncio.Queue[bytes]") -> None:  # noqa: C901, WPS213
         """
         Run tasks.
 
@@ -299,6 +305,15 @@ def task_cb(task: "asyncio.Task[Any]") -> None:
             message = await queue.get()
             if message is QUEUE_DONE:
                 break
+            if message is QUEUE_SKIP:
+                # Decrease max_prefetch
+                prefetch_dec = asyncio.create_task(self.sem_prefetch.acquire_first())
+                prefetch_dec.add_done_callback(tasks.discard)
+                tasks.add(prefetch_dec)
+
+                if self.sem is not None:
+                    self.sem.release()
+                continue
 
             task = asyncio.create_task(
                 self.callback(message=message, raise_err=False),
@@ -311,3 +326,32 @@ def task_cb(task: "asyncio.Task[Any]") -> None:
             # and it considered to be Hisenbug.
             # https://textual.textualize.io/blog/2023/02/11/the-heisenbug-lurking-in-your-async-code/
             task.add_done_callback(task_cb)
+
+    async def task_idler(self, wait: float) -> None:
+        """
+        Temporary increasing `max_async_tasks` for at least `wait` amount of time.
+
+        :param wait: time
+        """
+        if not self.sem:
+            await asyncio.sleep(wait)
+            return
+
+        if not self.sem_idle:
+            logger.warning("`max_idle_tasks` is undefined. Idle is unavailable.")
+            await asyncio.sleep(wait)
+            return
+
+        start_time = time()
+        async with self.sem_idle:
+            # Increase max_tasks
+            # Increase max_prefetch in runner
+            self.sem.release()
+
+            # Wait
+            await asyncio.sleep(wait - (time() - start_time))
+
+            # Decrease max_prefetch in runner
+            await self.queue.put_first(QUEUE_SKIP)
+            # Decrease max_tasks
+            await self.sem.acquire_first()

taskiq/utils.py

@@ -1,5 +1,8 @@
+import asyncio
 import inspect
-from typing import Any, Coroutine, TypeVar, Union
+from typing import TYPE_CHECKING, Any, Coroutine, Deque, Generic, TypeVar, Union
+
+from typing_extensions import Literal
 
 _T = TypeVar("_T")  # noqa: WPS111
 
@@ -35,3 +38,103 @@ def remove_suffix(text: str, suffix: str) -> str:
     if text.endswith(suffix):
         return text[: -len(suffix)]
     return text
+
+
+class DequeSemaphore(asyncio.Semaphore):
+    """Deque based Semaphore."""
+
+    if TYPE_CHECKING:  # noqa: WPS604
+        _loop: asyncio.BaseEventLoop
+
+    async def acquire_first(self) -> Literal[True]:
+        """
+        Acquire as soon as possible. LIFO style.
+
+        :raises BaseException: exception
+        :return: true
+        """
+        self._value -= 1
+
+        while self._value < 0:
+            fut: asyncio.Future[Any] = self._loop.create_future()
+            self._waiters.appendleft(fut)
+            try:
+                await fut
+            except BaseException:  # noqa: WPS424
+                self._value += 1
+
+                fut.cancel()
+                if not self.locked() and not fut.cancelled():
+                    self._wake_up_next()
+                raise
+
+        return True
+
+
+class DequeQueue(
+    asyncio.Queue,  # type: ignore
+    Generic[_T],
+):
+    """Deque based Queue."""
+
+    if TYPE_CHECKING:  # noqa: WPS604
+        _loop: asyncio.BaseEventLoop
+        _queue: Deque[_T]
+        _putters: Deque[Any]
+        _getters: Deque[Any]
+        _unfinished_tasks: int
+        _finished: asyncio.Event
+        _wakeup_next: Any
+
+    async def put_first(self, item: _T) -> None:
+        """
+        Wait till queue is not full. Put item in Queue as soon as possible. LIFO style.
+
+        :param item: value to prepend
+        :raises BaseException: something goes wrong
+        :returns: nothing
+        """
+        while self.full():
+            putter = self._loop.create_future()
+            self._putters.appendleft(putter)
+            try:
+                await putter
+            except BaseException:  # noqa: WPS424
+                putter.cancel()  # Just in case putter is not done yet.
+                try:  # noqa: WPS505
+                    # Clean self._putters from canceled putters.
+                    self._putters.remove(putter)
+                except ValueError:
+                    # The putter could be removed from self._putters by a
+                    # previous get_nowait call.
+                    pass  # noqa: WPS420
+                if not self.full() and not putter.cancelled():
+                    # We were woken up by get_nowait(), but can't take
+                    # the call.  Wake up the next in line.
+                    self._wakeup_next(self._putters)
+                raise
+
+        return self.put_first_nowait(item)
+
+    def put_first_nowait(self, item: _T) -> None:
+        """
+        Put item in Queue as soon as possible. LIFO style.
+
+        :param item: value to prepend
+        :raises QueueFull: queue is full
+        """
+        if self.full():
+            raise asyncio.QueueFull()
+
+        self._put_first(item)
+        self._unfinished_tasks += 1
+        self._finished.clear()
+        self._wakeup_next(self._getters)
+
+    def _put_first(self, item: _T) -> None:
+        """
+        Prepend item.
+
+        :param item: value to prepend
+        """
+        self._queue.appendleft(item)