Merge pull request #2 from d3nbr0/v1.1

aiocarrot task scheduler

Author: d3nbr0

aiocarrot/__meta__.py

@@ -1 +1 @@
-__version__ = '1.0.7'
+__version__ = '1.1.0'

aiocarrot/carrot.py

@@ -2,9 +2,11 @@
 
 from loguru import logger
 
+from .scheduler import Scheduler
+
 from typing import Optional, TYPE_CHECKING
 
-import asyncio, aio_pika, ujson, uuid, copy, signal
+import asyncio, aio_pika, ujson, uuid, copy, signal, aiormq
 
 if TYPE_CHECKING:
     from aiormq.abc import ConfirmationFrameType
@@ -23,6 +25,7 @@ class Carrot:
     _connection: Optional['aio_pika.abc.AbstractConnection'] = None
     _channel: Optional['aio_pika.abc.AbstractChannel'] = None
     _queue: Optional['aio_pika.abc.AbstractQueue'] = None
+    _scheduler: Optional['Scheduler'] = None
 
     def __init__(self, url: str, queue_name: str) -> None:
         """
@@ -35,6 +38,7 @@ def __init__(self, url: str, queue_name: str) -> None:
         self._url = url
         self._tasks = []
         self._queue_name = queue_name
+        self._scheduler = Scheduler(carrot=self)
 
     async def send(self, _cnm: str, **kwargs) -> 'ConfirmationFrameType':
         """
@@ -76,6 +80,19 @@ def setup_consumer(self, consumer: 'Consumer') -> None:
 
         self._consumer = consumer
 
+        self._scheduler.clear()
+        self._scheduler.stop()
+
+        for _, message in self._consumer._messages.items():
+            if not message.schedule:
+                continue
+
+            self._scheduler.add_task(message)
+
+        if self._is_consumer_alive and self._scheduler.has_tasks:
+            scheduler_task = asyncio.create_task(self._scheduler.reload())
+            self._tasks.append(scheduler_task)
+
     async def run(self) -> None:
         """
         Starts the main loop of the Carrot new message listener
@@ -99,6 +116,10 @@ async def run(self) -> None:
             logger.info(f'  * {message_name}')
 
         logger.info('')
+
+        if self._scheduler.has_tasks:
+            asyncio.create_task(self._scheduler.start())
+
         logger.info('Starting listener loop...')
 
         signal.signal(signal.SIGINT, self._exit_signal_handler)
@@ -131,6 +152,7 @@ async def shutdown(self, silent: bool = False) -> None:
         :return:
         """
 
+        self._scheduler.stop()
         pending_tasks = [x for x in self._tasks if not x.done()]
 
         if len(pending_tasks) > 0:
@@ -163,51 +185,62 @@ async def _consumer_loop(self) -> None:
         logger.info('Consumer is successfully connected to queue')
 
         async with queue.iterator() as queue_iterator:
-            async for message in queue_iterator:
-                for task in copy.copy(self._tasks):
-                    if task.done():
-                        self._tasks.remove(task)
+            if not self._is_consumer_alive:
+                return
+
+            try:
+                await self._iterate_queue(queue_iterator)
+            except aiormq.ChannelClosed:
+                return
+
+    async def _iterate_queue(self, queue_iterator: 'aio_pika.abc.AbstractQueueIterator') -> None:
+        """ Iterates over the queue iterator and passes the message on to the handler """
 
-                async with message.process():
-                    decoded_message: str = message.body.decode()
+        async for message in queue_iterator:
+            for task in copy.copy(self._tasks):
+                if task.done():
+                    self._tasks.remove(task)
 
-                    try:
-                        message_payload = ujson.loads(decoded_message)
+            async with message.process():
+                decoded_message: str = message.body.decode()
 
-                        assert isinstance(message_payload, dict)
-                    except ujson.JSONDecodeError:
-                        logger.error(f'Error receiving the message (failed to receive JSON): {decoded_message}')
-                        continue
+                try:
+                    message_payload = ujson.loads(decoded_message)
 
-                    message_id = message_payload.get('_cid')
-                    message_name = message_payload.get('_cnm')
+                    assert isinstance(message_payload, dict)
+                except ujson.JSONDecodeError:
+                    logger.error(f'Error receiving the message (failed to receive JSON): {decoded_message}')
+                    continue
 
-                    if not message_id:
-                        logger.error(
-                            'The message format could not be determined (identifier is missing): '
-                            f'{message_payload}'
-                        )
+                message_id = message_payload.get('_cid')
+                message_name = message_payload.get('_cnm')
 
-                        continue
+                if not message_id:
+                    logger.error(
+                        'The message format could not be determined (identifier is missing): '
+                        f'{message_payload}'
+                    )
+
+                    continue
 
-                    if not message_name:
-                        logger.error(
-                            'The message format could not be determined (message name is missing): '
-                            f'{message_payload}'
-                        )
+                if not message_name:
+                    logger.error(
+                        'The message format could not be determined (message name is missing): '
+                        f'{message_payload}'
+                    )
 
-                        continue
+                    continue
 
-                    del message_payload['_cid']
-                    del message_payload['_cnm']
+                del message_payload['_cid']
+                del message_payload['_cnm']
 
-                    task = asyncio.create_task(self._consumer.on_message(
-                        message_id,
-                        message_name,
-                        **message_payload,
-                    ))
+                task = asyncio.create_task(self._consumer.on_message(
+                    message_id,
+                    message_name,
+                    **message_payload,
+                ))
 
-                    self._tasks.append(task)
+                self._tasks.append(task)
 
     async def _get_queue(self) -> 'aio_pika.abc.AbstractQueue':
         """

aiocarrot/consumer/consumer.py

@@ -6,7 +6,7 @@
 from loguru import logger
 
 from abc import ABC
-from typing import TYPE_CHECKING
+from typing import Optional, TYPE_CHECKING
 from copy import deepcopy
 
 if TYPE_CHECKING:
@@ -52,7 +52,7 @@ class Consumer(AbstractConsumer):
     def __init__(self) -> None:
         self._messages = {}
 
-    def create_message(self, name: str, handler: 'Callable') -> None:
+    def create_message(self, name: str, handler: 'Callable', schedule: Optional[str] = None) -> None:
         split_name = name.split()
 
         if len(split_name) != 1:
@@ -62,19 +62,20 @@ def create_message(self, name: str, handler: 'Callable') -> None:
             name=name,
             handler=handler,
             dependant=get_dependant(handler=handler),
+            schedule=schedule,
         )
 
         self._messages[name] = message
 
-    def message(self, name: str | list[str]) -> 'Callable':
+    def message(self, name: str | list[str], schedule: Optional[str] = None) -> 'Callable':
         def decorator(func: 'Callable') -> 'Callable':
             if not isinstance(name, list):
                 message_names = [name]
             else:
                 message_names = name
 
             for message_name in message_names:
-                self.create_message(name=message_name, handler=func)
+                self.create_message(name=message_name, handler=func, schedule=schedule)
 
             return func
 

aiocarrot/consumer/types.py

@@ -1,10 +1,10 @@
-from pydantic import TypeAdapter, ValidationError, ConfigDict
+from typing import Optional, Callable, NamedTuple, TypeVar, Annotated
+from dataclasses import dataclass
+
+from pydantic import TypeAdapter, ValidationError
 from pydantic.fields import FieldInfo
 from pydantic_core import PydanticUndefined as Undefined
 
-from typing import Callable, NamedTuple, TypeVar, Annotated
-from dataclasses import dataclass
-
 
 T = TypeVar('T')
 
@@ -33,7 +33,7 @@ def __post_init__(self) -> None:
             Annotated[self.field_info.annotation, self.field_info],
         )
 
-    def validate(self, value: any) -> tuple[T | None, str | None]:
+    def validate(self, value: any) -> tuple[Optional[T], Optional[str]]:
         try:
             return self._type_adapter.validate_python(value, from_attributes=True), None
         except ValidationError as exc:
@@ -48,9 +48,10 @@ def __init__(self, *, params: list[Field] = None):
 
 
 class Message(NamedTuple):
-    name: str | None
+    name: Optional[str]
     handler: Callable
     dependant: Dependant
+    schedule: Optional[str] = None
 
 
 __all__ = (

aiocarrot/scheduler/__init__.py

@@ -0,0 +1 @@
+from .scheduler import Scheduler

aiocarrot/scheduler/scheduler.py

@@ -0,0 +1,93 @@
+import asyncio
+
+from typing import TYPE_CHECKING
+from datetime import datetime, timezone
+
+from loguru import logger
+
+from .types import Task
+
+if TYPE_CHECKING:
+    from ..carrot import Carrot
+    from ..consumer.types import Message
+
+
+class Scheduler:
+    """ Common Carrot task scheduler """
+
+    _carrot: 'Carrot'
+    _tasks: list[Task] = None
+    _is_alive: bool = False
+    _max_interval: int = 10
+    _adjust: float = 0.001
+
+    def __init__(self, carrot: 'Carrot') -> None:
+        """ The basic and basic task scheduler for the aiocarrot framework """
+
+        self._carrot = carrot
+        self._tasks = []
+
+    def add_task(self, message: 'Message') -> None:
+        if not message.schedule:
+            raise AttributeError(f'Message <{message.name}> has no schedule attribute')
+
+        try:
+            task = Task(message=message, carrot=self._carrot)
+        except (ValueError, TypeError):
+            return logger.error(f'Issue cron format error for message <{message.name}>')
+
+        self._tasks.append(task)
+
+    async def start(self) -> None:
+        if self._is_alive:
+            return logger.error('Scheduler already running')
+
+        if len(self._tasks) == 0:
+            return logger.warning('Scheduler cannot be started because there are no scheduled tasks')
+
+        logger.info(f'Scheduler registered {len(self._tasks)} messages')
+        self._is_alive = True
+
+        while self._is_alive:
+            next_task_at = self._get_next_time()
+            now = datetime.now(timezone.utc)
+
+            delta = (next_task_at - now).total_seconds()
+            if delta < 0:
+                delta = self._max_interval
+
+            interval = min(delta, self._max_interval) + self._adjust
+
+            await asyncio.sleep(interval)
+
+            now = datetime.now(timezone.utc)
+
+            for task in self._tasks:
+                if not task.is_due(now):
+                    continue
+
+                await task.run()
+
+    async def reload(self) -> None:
+        logger.info(f'Reloading scheduler after changing consumer...')
+
+        await asyncio.sleep(5)
+        await self.start()
+
+    def stop(self) -> None:
+        self._is_alive = False
+
+    def clear(self) -> None:
+        self._tasks = []
+
+    @property
+    def has_tasks(self) -> bool:
+        return self._tasks and len(self._tasks) > 0
+
+    def _get_next_time(self) -> datetime:
+        return min(self._tasks, key=lambda x: x.next_run).next_run
+
+
+__all__ = (
+    'Scheduler',
+)

aiocarrot/scheduler/types.py

@@ -0,0 +1,53 @@
+from typing import TYPE_CHECKING
+from datetime import datetime, timezone
+
+from loguru import logger
+
+from croniter import croniter
+
+if TYPE_CHECKING:
+    from ..carrot import Carrot
+    from ..consumer.types import Message
+
+
+class Task:
+    """ The basic class of the task for the scheduler """
+
+    _carrot: 'Carrot'
+    _message: 'Message'
+    _cron: croniter
+    _next_run: datetime
+    _resync_at_next_tick: bool = False
+
+    def __init__(self, carrot: 'Carrot', message: 'Message') -> None:
+        self._carrot = carrot
+        self._message = message
+
+        now = datetime.now(timezone.utc)
+
+        self._cron = croniter(message.schedule, now)
+        self._sync(now)
+
+    async def run(self):
+        logger.info(f'Scheduled task <{self._message.name}> has been queued')
+        await self._carrot.send(self._message.name)
+        self._resync_at_next_tick = True
+
+    def is_due(self, now: datetime) -> bool:
+        if self._resync_at_next_tick:
+            self._sync(now)
+
+        return now >= self._next_run
+
+    @property
+    def next_run(self) -> datetime:
+        return self._next_run
+
+    def _sync(self, start_datetime: datetime) -> None:
+        self._next_run = self._cron.get_next(datetime, start_time=start_datetime)
+        self._resync_at_next_tick = False
+
+
+__all__ = (
+    'Task',
+)