jararaca 0.3.10__py3-none-any.whl → 0.3.11__py3-none-any.whl

jararaca/messagebus/worker.py

@@ -1,16 +1,33 @@
 import asyncio
 import inspect
 import logging
+import random
 import signal
+import time
+import uuid
+from abc import ABC
 from contextlib import asynccontextmanager, suppress
-from dataclasses import dataclass
-from typing import Any, AsyncContextManager, AsyncGenerator, Type, get_origin
+from dataclasses import dataclass, field
+from datetime import UTC, datetime
+from typing import (
+    Any,
+    AsyncContextManager,
+    AsyncGenerator,
+    Awaitable,
+    Optional,
+    Type,
+    get_origin,
+)
+from urllib.parse import parse_qs, urlparse
 
 import aio_pika
 import aio_pika.abc
 import uvloop
+from aio_pika.exceptions import AMQPError, ChannelClosed, ChannelNotFoundEntity
 from pydantic import BaseModel
 
+from jararaca.broker_backend import MessageBrokerBackend
+from jararaca.broker_backend.mapper import get_message_broker_backend_from_url
 from jararaca.core.uow import UnitOfWorkContextProvider
 from jararaca.di import Container
 from jararaca.lifecycle import AppLifecycle
@@ -20,13 +37,22 @@ from jararaca.messagebus.bus_message_controller import (
 )
 from jararaca.messagebus.decorators import (
     MESSAGE_HANDLER_DATA_SET,
+    SCHEDULED_ACTION_DATA_SET,
     MessageBusController,
     MessageHandler,
     MessageHandlerData,
+    ScheduleDispatchData,
 )
 from jararaca.messagebus.message import Message, MessageOf
-from jararaca.microservice import MessageBusAppContext, Microservice
+from jararaca.microservice import (
+    AppTransactionContext,
+    MessageBusTransactionData,
+    Microservice,
+    SchedulerTransactionData,
+)
+from jararaca.scheduler.decorators import ScheduledActionData
 from jararaca.utils.rabbitmq_utils import RabbitmqUtils
+from jararaca.utils.retry import RetryConfig, retry_with_backoff
 
 logger = logging.getLogger(__name__)
 
@@ -36,6 +62,22 @@ class AioPikaWorkerConfig:
     url: str
     exchange: str
     prefetch_count: int
+    connection_retry_config: RetryConfig = field(
+        default_factory=lambda: RetryConfig(
+            max_retries=15,
+            initial_delay=1.0,
+            max_delay=60.0,
+            backoff_factor=2.0,
+        )
+    )
+    consumer_retry_config: RetryConfig = field(
+        default_factory=lambda: RetryConfig(
+            max_retries=15,
+            initial_delay=0.5,
+            max_delay=40.0,
+            backoff_factor=2.0,
+        )
+    )
 
 
 class AioPikaMessage(MessageOf[Message]):
@@ -76,78 +118,732 @@ class MessageProcessingLocker:
         await asyncio.gather(*self.current_processing_messages_set)
 
 
-class AioPikaMicroserviceConsumer:
+class MessageBusConsumer(ABC):
+
+    async def consume(self) -> None:
+        raise NotImplementedError("consume method not implemented")
+
+    def shutdown(self) -> None: ...
+
+    async def close(self) -> None:
+        """Close all resources related to the consumer"""
+
+
+class AioPikaMicroserviceConsumer(MessageBusConsumer):
     def __init__(
         self,
+        broker_backend: MessageBrokerBackend,
         config: AioPikaWorkerConfig,
         message_handler_set: MESSAGE_HANDLER_DATA_SET,
+        scheduled_actions: SCHEDULED_ACTION_DATA_SET,
         uow_context_provider: UnitOfWorkContextProvider,
     ):
+
+        self.broker_backend = broker_backend
         self.config = config
         self.message_handler_set = message_handler_set
+        self.scheduled_actions = scheduled_actions
         self.incoming_map: dict[str, MessageHandlerData] = {}
         self.uow_context_provider = uow_context_provider
         self.shutdown_event = asyncio.Event()
         self.lock = asyncio.Lock()
         self.tasks: set[asyncio.Task[Any]] = set()
+        self.connection: aio_pika.abc.AbstractConnection | None = None
+        self.channels: dict[str, aio_pika.abc.AbstractChannel] = {}
+
+    async def _verify_infrastructure(self) -> bool:
+        """
+        Verify that the required RabbitMQ infrastructure (exchanges, queues) exists.
+        Returns True if all required infrastructure is in place.
+        """
+        try:
+            async with self.connect() as connection:
+                # Create a main channel just for checking infrastructure
+                async with connection.channel() as main_channel:
+                    # Get existing exchange and queues to verify infrastructure is in place
+                    await RabbitmqUtils.get_main_exchange(
+                        channel=main_channel,
+                        exchange_name=self.config.exchange,
+                    )
+                    await RabbitmqUtils.get_dl_exchange(channel=main_channel)
+                    await RabbitmqUtils.get_dl_queue(channel=main_channel)
+                    return True
+        except (ChannelNotFoundEntity, ChannelClosed, AMQPError) as e:
+            logger.critical(
+                f"Required exchange or queue infrastructure not found. "
+                f"Please use the declare command first to create the required infrastructure. Error: {e}"
+            )
+            return False
+
+    async def _setup_message_handler_consumer(
+        self, handler: MessageHandlerData
+    ) -> bool:
+        """
+        Set up a consumer for a message handler with retry mechanism.
+        Returns True if successful, False otherwise.
+        """
+        queue_name = f"{handler.message_type.MESSAGE_TOPIC}.{handler.instance_callable.__module__}.{handler.instance_callable.__qualname__}"
+        routing_key = f"{handler.message_type.MESSAGE_TOPIC}.#"
+
+        async def setup_consumer() -> None:
+            # Create a channel using the context manager
+            async with self.create_channel(queue_name) as channel:
+                queue = await RabbitmqUtils.get_queue(
+                    channel=channel, queue_name=queue_name
+                )
+
+                # Configure consumer right away while in the context
+                await queue.consume(
+                    callback=MessageHandlerCallback(
+                        consumer=self,
+                        queue_name=queue_name,
+                        routing_key=routing_key,
+                        message_handler=handler,
+                    ),
+                    no_ack=handler.spec.auto_ack,
+                )
+
+                logger.info(
+                    f"Consuming message handler {queue_name} on dedicated channel"
+                )
+
+        try:
+            # Setup with retry
+            await retry_with_backoff(
+                setup_consumer,
+                retry_config=self.config.consumer_retry_config,
+                retry_exceptions=(ChannelNotFoundEntity, ChannelClosed, AMQPError),
+            )
+            return True
+        except Exception as e:
+            logger.error(
+                f"Failed to setup consumer for queue '{queue_name}' after retries: {e}"
+            )
+            return False
+
+    async def _setup_scheduled_action_consumer(
+        self, scheduled_action: ScheduledActionData
+    ) -> bool:
+        """
+        Set up a consumer for a scheduled action with retry mechanism.
+        Returns True if successful, False otherwise.
+        """
+        queue_name = f"{scheduled_action.callable.__module__}.{scheduled_action.callable.__qualname__}"
+        routing_key = queue_name
+
+        async def setup_consumer() -> None:
+            # Create a channel using the context manager
+            async with self.create_channel(queue_name) as channel:
+                queue = await RabbitmqUtils.get_queue(
+                    channel=channel, queue_name=queue_name
+                )
+
+                # Configure consumer right away while in the context
+                await queue.consume(
+                    callback=ScheduledMessageHandlerCallback(
+                        consumer=self,
+                        queue_name=queue_name,
+                        routing_key=routing_key,
+                        scheduled_action=scheduled_action,
+                    ),
+                    no_ack=True,
+                )
+
+                logger.info(f"Consuming scheduler {queue_name} on dedicated channel")
+
+        try:
+            # Setup with retry
+            await retry_with_backoff(
+                setup_consumer,
+                retry_config=self.config.consumer_retry_config,
+                retry_exceptions=(ChannelNotFoundEntity, ChannelClosed, AMQPError),
+            )
+            return True
+        except Exception as e:
+            logger.error(
+                f"Failed to setup consumer for scheduler queue '{queue_name}' after retries: {e}"
+            )
+            return False
+
+    async def consume(self) -> None:
+        """
+        Main consume method that sets up all message handlers and scheduled actions with retry mechanisms.
+        """
+        # Verify infrastructure with retry
+        infra_check_success = await retry_with_backoff(
+            self._verify_infrastructure,
+            retry_config=self.config.connection_retry_config,
+            retry_exceptions=(Exception,),
+        )
 
-    async def consume(self, passive_declare: bool) -> None:
+        if not infra_check_success:
+            logger.critical("Failed to verify RabbitMQ infrastructure. Shutting down.")
+            self.shutdown_event.set()
+            return
 
-        connection = await aio_pika.connect(self.config.url)
+        async def wait_for(
+            type: str, name: str, coroutine: Awaitable[bool]
+        ) -> tuple[str, str, bool]:
+            return type, name, await coroutine
 
-        channel = await connection.channel()
+        tasks: set[asyncio.Task[tuple[str, str, bool]]] = set()
+
+        # Setup message handlers
+        for handler in self.message_handler_set:
+            queue_name = f"{handler.message_type.MESSAGE_TOPIC}.{handler.instance_callable.__module__}.{handler.instance_callable.__qualname__}"
+            self.incoming_map[queue_name] = handler
 
+            tasks.add(
+                task := asyncio.create_task(
+                    wait_for(
+                        "message_handler",
+                        queue_name,
+                        self._setup_message_handler_consumer(handler),
+                    )
+                )
+            )
+            # task.add_done_callback(tasks.discard)
+            # success = await self._setup_message_handler_consumer(handler)
+            # if not success:
+            #     logger.warning(
+            #         f"Failed to set up consumer for {queue_name}, will not process messages from this queue"
+            #     )
+
+        # Setup scheduled actions
+        for scheduled_action in self.scheduled_actions:
+
+            queue_name = f"{scheduled_action.callable.__module__}.{scheduled_action.callable.__qualname__}"
+            tasks.add(
+                task := asyncio.create_task(
+                    wait_for(
+                        "scheduled_action",
+                        queue_name,
+                        self._setup_scheduled_action_consumer(scheduled_action),
+                    )
+                )
+            )
+            # task.add_done_callback(tasks.discard)
+
+            # success = await self._setup_scheduled_action_consumer(scheduled_action)
+            # if not success:
+            #     queue_name = f"{scheduled_action.callable.__module__}.{scheduled_action.callable.__qualname__}"
+            #     logger.warning(
+            #         f"Failed to set up consumer for scheduled action {queue_name}, will not process scheduled tasks from this queue"
+            #     )
+
+        async def handle_task_results() -> None:
+            for task in asyncio.as_completed(tasks):
+                type, name, success = await task
+                if success:
+                    logger.info(f"Successfully set up {type} consumer for {name}")
+                else:
+                    logger.warning(
+                        f"Failed to set up {type} consumer for {name}, will not process messages from this queue"
+                    )
+
+        handle_task_results_task = asyncio.create_task(handle_task_results())
+
+        # Wait for shutdown signal
+        await self.shutdown_event.wait()
+        logger.info("Shutdown event received, stopping consumers")
+        handle_task_results_task.cancel()
+        with suppress(asyncio.CancelledError):
+            await handle_task_results_task
+        for task in tasks:
+            if not task.done():
+                task.cancel()
+                with suppress(asyncio.CancelledError):
+                    await task
+        logger.info("Worker shutting down")
+
+        # Wait for all tasks to complete
+        await self.wait_all_tasks_done()
+
+        # Close all channels and the connection
+        await self.close_channels_and_connection()
+
+    async def wait_all_tasks_done(self) -> None:
+        if not self.tasks:
+            return
+
+        logger.info(f"Waiting for {len(self.tasks)} in-flight tasks to complete")
+        async with self.lock:
+            # Use gather with return_exceptions=True to ensure all tasks are awaited
+            # even if some raise exceptions
+            results = await asyncio.gather(*self.tasks, return_exceptions=True)
+
+            # Log any exceptions that occurred
+            for result in results:
+                if isinstance(result, Exception):
+                    logger.error(f"Task raised an exception during shutdown: {result}")
+
+    async def close_channels_and_connection(self) -> None:
+        """Close all channels and then the connection"""
+        # Close all channels
+        channel_close_tasks = []
+        for queue_name, channel in self.channels.items():
+            try:
+                if not channel.is_closed:
+                    logger.info(f"Closing channel for queue {queue_name}")
+                    channel_close_tasks.append(channel.close())
+                else:
+                    logger.info(f"Channel for queue {queue_name} already closed")
+            except Exception as e:
+                logger.error(
+                    f"Error preparing to close channel for queue {queue_name}: {e}"
+                )
+
+        # Wait for all channels to close (if any)
+        if channel_close_tasks:
+            try:
+                await asyncio.gather(*channel_close_tasks, return_exceptions=True)
+            except Exception as e:
+                logger.error(f"Error during channel closures: {e}")
+
+        # Clear channels dictionary
+        self.channels.clear()
+
+        # Close the connection
+        if self.connection:
+            try:
+                if not self.connection.is_closed:
+                    logger.info("Closing RabbitMQ connection")
+                    await self.connection.close()
+                else:
+                    logger.info("RabbitMQ connection already closed")
+            except Exception as e:
+                logger.error(f"Error closing RabbitMQ connection: {e}")
+            self.connection = None
+
+    def shutdown(self) -> None:
+        """Signal for shutdown"""
+        logger.info("Initiating graceful shutdown")
+        self.shutdown_event.set()
+
+    async def close(self) -> None:
+        """Implement MessageBusConsumer.close for cleanup"""
+        self.shutdown()
+        await self.wait_all_tasks_done()
+        await self.close_channels_and_connection()
+
+    async def get_channel(self, queue_name: str) -> aio_pika.abc.AbstractChannel | None:
+        """
+        Get the channel for a specific queue, or None if not found.
+        This helps with error handling when a channel might have been closed.
+        """
+        if queue_name not in self.channels:
+            logger.warning(f"No channel found for queue {queue_name}")
+            return None
+
+        try:
+            channel = self.channels[queue_name]
+            if channel.is_closed:
+                logger.warning(f"Channel for queue {queue_name} is closed")
+                # Attempt to recreate the channel if needed
+                if self.connection and not self.connection.is_closed:
+                    logger.info(f"Creating new channel for {queue_name}")
+                    self.channels[queue_name] = await self.connection.channel()
+                    await self.channels[queue_name].set_qos(
+                        prefetch_count=self.config.prefetch_count
+                    )
+                    return self.channels[queue_name]
+                return None
+            return channel
+        except Exception as e:
+            logger.error(f"Error accessing channel for queue {queue_name}: {e}")
+            return None
+
+    async def _establish_channel(self, queue_name: str) -> aio_pika.abc.AbstractChannel:
+        """
+        Creates a new channel for the specified queue with proper QoS settings.
+        """
+        if self.connection is None or self.connection.is_closed:
+            logger.warning(
+                f"Cannot create channel for {queue_name}: connection is not available"
+            )
+            raise RuntimeError("Connection is not available")
+
+        logger.debug(f"Creating channel for queue {queue_name}")
+        channel = await self.connection.channel()
         await channel.set_qos(prefetch_count=self.config.prefetch_count)
+        logger.debug(f"Created channel for queue {queue_name}")
+        return channel
+
+    @asynccontextmanager
+    async def create_channel(
+        self, queue_name: str
+    ) -> AsyncGenerator[aio_pika.abc.AbstractChannel, None]:
+        """
+        Create and yield a channel for the specified queue with retry mechanism.
+        This context manager ensures the channel is properly managed.
+        """
+        try:
+            # Create a new channel with retry
+            channel = await retry_with_backoff(
+                fn=lambda: self._establish_channel(queue_name),
+                retry_config=self.config.consumer_retry_config,
+                retry_exceptions=(
+                    aio_pika.exceptions.AMQPConnectionError,
+                    aio_pika.exceptions.AMQPChannelError,
+                    ConnectionError,
+                ),
+            )
+
+            # Save in the channels dict for tracking
+            self.channels[queue_name] = channel
+            logger.debug(f"Created new channel for queue {queue_name}")
+
+            try:
+                yield channel
+            finally:
+                # Don't close the channel here as it might be used later
+                # It will be closed during shutdown
+                pass
+        except aio_pika.exceptions.AMQPError as e:
+            logger.error(
+                f"Error creating channel for queue {queue_name} after retries: {e}"
+            )
+            raise
+
+    async def _establish_connection(self) -> aio_pika.abc.AbstractConnection:
+        """
+        Creates a new RabbitMQ connection with retry logic.
+        """
+        try:
+            logger.info("Establishing connection to RabbitMQ")
+            connection = await aio_pika.connect(self.config.url)
+            logger.info("Connected to RabbitMQ successfully")
+            return connection
+        except Exception as e:
+            logger.error(f"Failed to connect to RabbitMQ: {e}")
+            raise
 
-        main_ex = await RabbitmqUtils.declare_main_exchange(
-            channel=channel,
-            exchange_name=self.config.exchange,
-            passive=passive_declare,
+    @asynccontextmanager
+    async def connect(self) -> AsyncGenerator[aio_pika.abc.AbstractConnection, None]:
+        """
+        Create and manage the main connection to RabbitMQ with automatic retry.
+        """
+        if self.connection is not None and not self.connection.is_closed:
+            logger.debug("Connection already exists, reusing existing connection")
+            try:
+                yield self.connection
+            finally:
+                # The existing connection will be handled by close_channels_and_connection
+                pass
+            return
+
+        try:
+            # Create a new connection with retry
+            self.connection = await retry_with_backoff(
+                self._establish_connection,
+                retry_config=self.config.connection_retry_config,
+                retry_exceptions=(
+                    aio_pika.exceptions.AMQPConnectionError,
+                    ConnectionError,
+                    OSError,
+                    TimeoutError,
+                ),
+            )
+
+            try:
+                yield self.connection
+            finally:
+                # Don't close the connection here; it will be closed in close_channels_and_connection
+                pass
+        except Exception as e:
+            logger.error(
+                f"Failed to establish connection to RabbitMQ after retries: {e}"
+            )
+            if self.connection:
+                try:
+                    await self.connection.close()
+                except Exception as close_error:
+                    logger.error(
+                        f"Error closing connection after connect failure: {close_error}"
+                    )
+                self.connection = None
+            raise
+
+    @asynccontextmanager
+    async def get_channel_ctx(
+        self, queue_name: str
+    ) -> AsyncGenerator[aio_pika.abc.AbstractChannel, None]:
+        """
+        Get a channel for a specific queue as a context manager.
+        This is safer than using get_channel directly as it ensures proper error handling.
+        """
+        channel = await self.get_channel(queue_name)
+        if channel is None:
+            if self.connection and not self.connection.is_closed:
+                # Try to create a new channel
+                async with self.create_channel(queue_name) as new_channel:
+                    yield new_channel
+            else:
+                raise RuntimeError(
+                    f"Cannot get channel for queue {queue_name}: no connection available"
+                )
+        else:
+            try:
+                yield channel
+            finally:
+                # We don't close the channel here as it's managed by the consumer
+                pass
+
+
+def create_message_bus(
+    broker_url: str,
+    broker_backend: MessageBrokerBackend,
+    scheduled_actions: SCHEDULED_ACTION_DATA_SET,
+    message_handler_set: MESSAGE_HANDLER_DATA_SET,
+    uow_context_provider: UnitOfWorkContextProvider,
+) -> MessageBusConsumer:
+
+    parsed_url = urlparse(broker_url)
+
+    if parsed_url.scheme == "amqp" or parsed_url.scheme == "amqps":
+        assert parsed_url.query, "Query string must be set for AMQP URLs"
+
+        query_params: dict[str, list[str]] = parse_qs(parsed_url.query)
+
+        assert "exchange" in query_params, "Exchange must be set in the query string"
+        assert (
+            len(query_params["exchange"]) == 1
+        ), "Exchange must be set in the query string"
+        assert (
+            "prefetch_count" in query_params
+        ), "Prefetch count must be set in the query string"
+        assert (
+            len(query_params["prefetch_count"]) == 1
+        ), "Prefetch count must be set in the query string"
+        assert query_params["prefetch_count"][
+            0
+        ].isdigit(), "Prefetch count must be an integer in the query string"
+        assert query_params["exchange"][0], "Exchange must be set in the query string"
+        assert query_params["prefetch_count"][
+            0
+        ], "Prefetch count must be set in the query string"
+
+        exchange = query_params["exchange"][0]
+        prefetch_count = int(query_params["prefetch_count"][0])
+
+        # Parse optional retry configuration parameters
+        connection_retry_config = RetryConfig()
+        consumer_retry_config = RetryConfig(
+            max_retries=30, initial_delay=5, max_delay=60.0, backoff_factor=3.0
         )
 
-        dlx, dlq = await RabbitmqUtils.declare_dl_kit(channel=channel)
+        # Connection retry config parameters
+        if (
+            "connection_retry_max" in query_params
+            and query_params["connection_retry_max"][0].isdigit()
+        ):
+            connection_retry_config.max_retries = int(
+                query_params["connection_retry_max"][0]
+            )
 
-        for handler in self.message_handler_set:
+        if "connection_retry_delay" in query_params:
+            try:
+                connection_retry_config.initial_delay = float(
+                    query_params["connection_retry_delay"][0]
+                )
+            except ValueError:
+                pass
 
-            queue_name = f"{handler.message_type.MESSAGE_TOPIC}.{handler.callable.__module__}.{handler.callable.__qualname__}"
-            routing_key = f"{handler.message_type.MESSAGE_TOPIC}.#"
+        if "connection_retry_max_delay" in query_params:
+            try:
+                connection_retry_config.max_delay = float(
+                    query_params["connection_retry_max_delay"][0]
+                )
+            except ValueError:
+                pass
 
-            self.incoming_map[queue_name] = handler
+        if "connection_retry_backoff" in query_params:
+            try:
+                connection_retry_config.backoff_factor = float(
+                    query_params["connection_retry_backoff"][0]
+                )
+            except ValueError:
+                pass
+
+        # Consumer retry config parameters
+        if (
+            "consumer_retry_max" in query_params
+            and query_params["consumer_retry_max"][0].isdigit()
+        ):
+            consumer_retry_config.max_retries = int(
+                query_params["consumer_retry_max"][0]
+            )
+
+        if "consumer_retry_delay" in query_params:
+            try:
+                consumer_retry_config.initial_delay = float(
+                    query_params["consumer_retry_delay"][0]
+                )
+            except ValueError:
+                pass
+
+        if "consumer_retry_max_delay" in query_params:
+            try:
+                consumer_retry_config.max_delay = float(
+                    query_params["consumer_retry_max_delay"][0]
+                )
+            except ValueError:
+                pass
+
+        if "consumer_retry_backoff" in query_params:
+            try:
+                consumer_retry_config.backoff_factor = float(
+                    query_params["consumer_retry_backoff"][0]
+                )
+            except ValueError:
+                pass
+
+        config = AioPikaWorkerConfig(
+            url=broker_url,
+            exchange=exchange,
+            prefetch_count=prefetch_count,
+            connection_retry_config=connection_retry_config,
+            consumer_retry_config=consumer_retry_config,
+        )
+
+        return AioPikaMicroserviceConsumer(
+            config=config,
+            broker_backend=broker_backend,
+            message_handler_set=message_handler_set,
+            scheduled_actions=scheduled_actions,
+            uow_context_provider=uow_context_provider,
+        )
+
+    raise ValueError(
+        f"Unsupported broker URL scheme: {parsed_url.scheme}. Supported schemes are amqp and amqps"
+    )
 
-            queue: aio_pika.abc.AbstractQueue = await channel.declare_queue(
-                passive=passive_declare,
-                name=queue_name,
-                arguments={
-                    "x-dead-letter-exchange": dlx.name,
-                    "x-dead-letter-routing-key": dlq.name,
-                },
+
+class ScheduledMessageHandlerCallback:
+    def __init__(
+        self,
+        consumer: AioPikaMicroserviceConsumer,
+        queue_name: str,
+        routing_key: str,
+        scheduled_action: ScheduledActionData,
+    ):
+        self.consumer = consumer
+        self.queue_name = queue_name
+        self.routing_key = routing_key
+        self.scheduled_action = scheduled_action
+
+    async def __call__(
+        self, aio_pika_message: aio_pika.abc.AbstractIncomingMessage
+    ) -> None:
+
+        if self.consumer.shutdown_event.is_set():
+            logger.info(
+                f"Shutdown in progress. Requeuing scheduled message for {self.queue_name}"
             )
+            try:
+                # Use channel context for requeuing
+                async with self.consumer.get_channel_ctx(self.queue_name):
+                    await aio_pika_message.reject(requeue=True)
+            except RuntimeError:
+                logger.warning(
+                    f"Could not requeue scheduled message during shutdown - channel not available"
+                )
+            except Exception as e:
+                logger.error(
+                    f"Failed to requeue scheduled message during shutdown: {e}"
+                )
+            return
 
-            await queue.bind(exchange=main_ex, routing_key=routing_key)
+        async with self.consumer.lock:
+            task = asyncio.create_task(self.handle_message(aio_pika_message))
+            self.consumer.tasks.add(task)
+            task.add_done_callback(self.handle_message_consume_done)
 
-            await queue.consume(
-                callback=MessageHandlerCallback(
-                    consumer=self,
-                    queue_name=queue_name,
-                    routing_key=routing_key,
-                    message_handler=handler,
-                ),
-                no_ack=handler.spec.auto_ack,
+    def handle_message_consume_done(self, task: asyncio.Task[Any]) -> None:
+        self.consumer.tasks.discard(task)
+        if task.cancelled():
+            logger.warning(f"Scheduled task for {self.queue_name} was cancelled")
+            return
+
+        if (error := task.exception()) is not None:
+            logger.exception(
+                f"Error processing scheduled action {self.queue_name}", exc_info=error
             )
 
-            print(f"Consuming {queue_name}")
+    async def handle_message(
+        self, aio_pika_message: aio_pika.abc.AbstractIncomingMessage
+    ) -> None:
 
-        await self.shutdown_event.wait()
-        logger.info("Worker shutting down")
+        if self.consumer.shutdown_event.is_set():
+            logger.info(f"Shutdown event set. Requeuing message for {self.queue_name}")
+            try:
+                # Use channel context for requeuing
+                async with self.consumer.get_channel_ctx(self.queue_name):
+                    await aio_pika_message.reject(requeue=True)
+                return
+            except RuntimeError:
+                logger.warning(
+                    f"Could not requeue message during shutdown - channel not available"
+                )
+            except Exception as e:
+                logger.error(f"Failed to requeue message during shutdown: {e}")
+                return
+
+        sig = inspect.signature(self.scheduled_action.callable)
+        if len(sig.parameters) == 1:
+
+            task = asyncio.create_task(
+                self.run_with_context(
+                    self.scheduled_action,
+                    (ScheduleDispatchData(int(aio_pika_message.body.decode("utf-8"))),),
+                    {},
+                )
+            )
 
-        await self.wait_all_tasks_done()
+        elif len(sig.parameters) == 0:
+            task = asyncio.create_task(
+                self.run_with_context(
+                    self.scheduled_action,
+                    (),
+                    {},
+                )
+            )
+        else:
+            logger.warning(
+                "Scheduled action '%s' must have exactly one parameter of type ScheduleDispatchData or no parameters"
+                % self.queue_name
+            )
+            return
 
-        await channel.close()
-        await connection.close()
+        self.consumer.tasks.add(task)
+        task.add_done_callback(self.handle_message_consume_done)
 
-    async def wait_all_tasks_done(self) -> None:
-        async with self.lock:
-            await asyncio.gather(*self.tasks)
+        try:
+            await task
+        except Exception as e:
+
+            logger.exception(
+                f"Error processing scheduled action {self.queue_name}: {e}"
+            )
+
+    async def run_with_context(
+        self,
+        scheduled_action: ScheduledActionData,
+        args: tuple[Any, ...],
+        kwargs: dict[str, Any],
+    ) -> None:
+        async with self.consumer.uow_context_provider(
+            AppTransactionContext(
+                controller_member_reflect=scheduled_action.controller_member,
+                transaction_data=SchedulerTransactionData(
+                    scheduled_to=datetime.now(UTC),
+                    cron_expression=scheduled_action.spec.cron,
+                    triggered_at=datetime.now(UTC),
+                ),
+            )
+        ):
+
+            await scheduled_action.callable(*args, **kwargs)
 
 
 class MessageHandlerCallback:
@@ -163,11 +859,25 @@ class MessageHandlerCallback:
         self.queue_name = queue_name
         self.routing_key = routing_key
         self.message_handler = message_handler
+        self.retry_state: dict[str, dict[str, Any]] = {}
 
     async def message_consumer(
         self, aio_pika_message: aio_pika.abc.AbstractIncomingMessage
     ) -> None:
         if self.consumer.shutdown_event.is_set():
+            logger.info(
+                f"Shutdown in progress. Requeuing message for {self.queue_name}"
+            )
+            try:
+                # Use channel context for requeuing
+                async with self.consumer.get_channel_ctx(self.queue_name):
+                    await aio_pika_message.reject(requeue=True)
+            except RuntimeError:
+                logger.warning(
+                    f"Could not requeue message during shutdown - channel not available"
+                )
+            except Exception as e:
+                logger.error(f"Failed to requeue message during shutdown: {e}")
             return
 
         async with self.consumer.lock:
@@ -178,10 +888,13 @@ class MessageHandlerCallback:
     def handle_message_consume_done(self, task: asyncio.Task[Any]) -> None:
         self.consumer.tasks.discard(task)
         if task.cancelled():
+            logger.warning(f"Task for queue {self.queue_name} was cancelled")
             return
 
         if (error := task.exception()) is not None:
-            logger.exception("Error processing message", exc_info=error)
+            logger.exception(
+                f"Error processing message for queue {self.queue_name}", exc_info=error
+            )
 
     async def __call__(
         self, aio_pika_message: aio_pika.abc.AbstractIncomingMessage
@@ -192,13 +905,177 @@ class MessageHandlerCallback:
         self,
         aio_pika_message: aio_pika.abc.AbstractIncomingMessage,
         requeue: bool = False,
+        retry_count: int = 0,
+        exception: Optional[BaseException] = None,
     ) -> None:
-        if self.message_handler.spec.auto_ack is False:
-            await aio_pika_message.reject(requeue=requeue)
-        elif requeue:
-            logger.warning(
-                f"Message {aio_pika_message.message_id} ({self.queue_name}) cannot be requeued because auto_ack is enabled"
+        """
+        Handle rejecting a message, with support for retry with exponential backoff.
+
+        Args:
+            aio_pika_message: The message to reject
+            requeue: Whether to requeue the message directly (True) or handle with retry logic (False)
+            retry_count: The current retry count for this message
+            exception: The exception that caused the rejection, if any
+        """
+        message_id = aio_pika_message.message_id or str(uuid.uuid4())
+
+        # If auto_ack is enabled, we cannot retry the message through RabbitMQ reject mechanism
+        if self.message_handler.spec.auto_ack:
+            if requeue:
+                logger.warning(
+                    f"Message {message_id} ({self.queue_name}) cannot be requeued because auto_ack is enabled"
+                )
+            return
+
+        try:
+            # Check if we should retry with backoff
+            if (
+                not requeue
+                and self.message_handler.spec.requeue_on_exception
+                and exception is not None
+            ):
+                # Get retry config from consumer
+                retry_config = self.consumer.config.consumer_retry_config
+
+                # Check if we reached max retries
+                if retry_count >= retry_config.max_retries:
+                    logger.warning(
+                        f"Message {message_id} ({self.queue_name}) failed after {retry_count} retries, "
+                        f"dead-lettering: {str(exception)}"
+                    )
+                    # Dead-letter the message after max retries
+                    async with self.consumer.get_channel_ctx(self.queue_name):
+                        await aio_pika_message.reject(requeue=False)
+                    return
+
+                # Calculate delay for this retry attempt
+                delay = retry_config.initial_delay * (
+                    retry_config.backoff_factor**retry_count
+                )
+                if retry_config.jitter:
+                    jitter_amount = delay * 0.25
+                    delay = delay + random.uniform(-jitter_amount, jitter_amount)
+                    delay = max(
+                        delay, 0.1
+                    )  # Ensure delay doesn't go negative due to jitter
+
+                delay = min(delay, retry_config.max_delay)
+
+                logger.info(
+                    f"Message {message_id} ({self.queue_name}) failed with {str(exception)}, "
+                    f"retry {retry_count+1}/{retry_config.max_retries} scheduled in {delay:.2f}s"
+                )
+
+                # Store retry state for this message
+                self.retry_state[message_id] = {
+                    "retry_count": retry_count + 1,
+                    "last_exception": exception,
+                    "next_retry": time.time() + delay,
+                }
+
+                # Schedule retry after delay
+                asyncio.create_task(
+                    self._delayed_retry(
+                        aio_pika_message, delay, retry_count + 1, exception
+                    )
+                )
+
+                # Acknowledge the current message since we'll handle retry ourselves
+                async with self.consumer.get_channel_ctx(self.queue_name):
+                    await aio_pika_message.ack()
+                return
+
+            # Standard reject without retry or with immediate requeue
+            async with self.consumer.get_channel_ctx(self.queue_name):
+                await aio_pika_message.reject(requeue=requeue)
+                if requeue:
+                    logger.info(
+                        f"Message {message_id} ({self.queue_name}) requeued for immediate retry"
+                    )
+                else:
+                    logger.info(
+                        f"Message {message_id} ({self.queue_name}) rejected without requeue"
+                    )
+
+        except RuntimeError as e:
+            logger.error(
+                f"Error rejecting message {message_id} ({self.queue_name}): {e}"
+            )
+        except Exception as e:
+            logger.exception(
+                f"Unexpected error rejecting message {message_id} ({self.queue_name}): {e}"
+            )
+
+    async def _delayed_retry(
+        self,
+        aio_pika_message: aio_pika.abc.AbstractIncomingMessage,
+        delay: float,
+        retry_count: int,
+        exception: Optional[BaseException],
+    ) -> None:
+        """
+        Handle delayed retry of a message after exponential backoff delay.
+
+        Args:
+            aio_pika_message: The original message
+            delay: Delay in seconds before retry
+            retry_count: The current retry count (after increment)
+            exception: The exception that caused the failure
+        """
+        message_id = aio_pika_message.message_id or str(uuid.uuid4())
+
+        try:
+            # Wait for the backoff delay
+            await asyncio.sleep(delay)
+
+            # Get message body and properties for republishing
+            message_body = aio_pika_message.body
+            headers = (
+                aio_pika_message.headers.copy() if aio_pika_message.headers else {}
+            )
+
+            # Add retry information to headers
+            headers["x-retry-count"] = retry_count
+            if exception:
+                headers["x-last-error"] = str(exception)
+
+            # Clean up retry state
+            if message_id in self.retry_state:
+                del self.retry_state[message_id]
+
+            # Republish the message to the same queue
+            async with self.consumer.get_channel_ctx(self.queue_name) as channel:
+                exchange = await RabbitmqUtils.get_main_exchange(
+                    channel=channel,
+                    exchange_name=self.consumer.config.exchange,
+                )
+
+                await exchange.publish(
+                    aio_pika.Message(
+                        body=message_body,
+                        headers=headers,
+                        message_id=message_id,
+                        content_type=aio_pika_message.content_type,
+                        content_encoding=aio_pika_message.content_encoding,
+                        delivery_mode=aio_pika_message.delivery_mode,
+                    ),
+                    routing_key=self.routing_key,
+                )
+
+                logger.info(
+                    f"Message {message_id} ({self.queue_name}) republished for retry {retry_count}"
+                )
+
+        except Exception as e:
+            logger.exception(
+                f"Failed to execute delayed retry for message {message_id} ({self.queue_name}): {e}"
             )
+            # If we fail to republish, try to dead-letter the original message
+            try:
+                if message_id in self.retry_state:
+                    del self.retry_state[message_id]
+            except Exception:
+                pass
 
     async def handle_message(
         self, aio_pika_message: aio_pika.abc.AbstractIncomingMessage
@@ -211,15 +1088,9 @@ class MessageHandlerCallback:
             await self.handle_reject_message(aio_pika_message)
             return
 
-        handler_data = self.consumer.incoming_map.get(routing_key)
+        handler_data = self.message_handler
 
-        if handler_data is None:
-            logger.warning("No handler found for topic '%s'" % routing_key)
-            await self.handle_reject_message(aio_pika_message)
-
-            return
-
-        handler = handler_data.callable
+        handler = handler_data.instance_callable
 
         sig = inspect.signature(handler)
 
@@ -263,9 +1134,12 @@ class MessageHandlerCallback:
         assert incoming_message_spec is not None
 
         async with self.consumer.uow_context_provider(
-            MessageBusAppContext(
-                message=builded_message,
-                topic=routing_key,
+            AppTransactionContext(
+                controller_member_reflect=handler_data.controller_member,
+                transaction_data=MessageBusTransactionData(
+                    message=builded_message,
+                    topic=routing_key,
+                ),
             )
         ):
             ctx: AsyncContextManager[Any]
@@ -281,29 +1155,61 @@ class MessageHandlerCallback:
                         await handler(builded_message)
                     if not incoming_message_spec.auto_ack:
                         with suppress(aio_pika.MessageProcessError):
-                            await aio_pika_message.ack()
+                            # Use channel context for acknowledgement
+                            async with self.consumer.get_channel_ctx(self.queue_name):
+                                await aio_pika_message.ack()
                 except BaseException as base_exc:
+                    # Get message id for logging
+                    message_id = aio_pika_message.message_id or str(uuid.uuid4())
+
+                    # Extract retry count from headers if available
+                    headers = aio_pika_message.headers or {}
+                    retry_count = int(str(headers.get("x-retry-count", 0)))
+
+                    # Process exception handler if configured
                     if incoming_message_spec.exception_handler is not None:
                         try:
                             incoming_message_spec.exception_handler(base_exc)
                         except Exception as nested_exc:
                             logger.exception(
-                                f"Error processing exception handler: {base_exc} | {nested_exc}"
+                                f"Error processing exception handler for message {message_id}: {base_exc} | {nested_exc}"
                             )
                     else:
                         logger.exception(
-                            f"Error processing message on topic {routing_key}"
+                            f"Error processing message {message_id} on topic {routing_key}: {str(base_exc)}"
                         )
+
+                    # Handle rejection with retry logic
                     if incoming_message_spec.requeue_on_exception:
-                        await self.handle_reject_message(aio_pika_message, requeue=True)
+                        # Use our retry with backoff mechanism
+                        await self.handle_reject_message(
+                            aio_pika_message,
+                            requeue=False,  # Don't requeue directly, use our backoff mechanism
+                            retry_count=retry_count,
+                            exception=base_exc,
+                        )
                     else:
+                        # Message shouldn't be retried, reject it
                         await self.handle_reject_message(
-                            aio_pika_message, requeue=False
+                            aio_pika_message, requeue=False, exception=base_exc
                         )
                 else:
-                    logger.info(
-                        f"Message {aio_pika_message.message_id}#{self.queue_name} processed successfully"
-                    )
+                    # Message processed successfully, log and clean up any retry state
+                    message_id = aio_pika_message.message_id or str(uuid.uuid4())
+                    if message_id in self.retry_state:
+                        del self.retry_state[message_id]
+
+                    # Log success with retry information if applicable
+                    headers = aio_pika_message.headers or {}
+                    if "x-retry-count" in headers:
+                        retry_count = int(str(headers.get("x-retry-count", 0)))
+                        logger.info(
+                            f"Message {message_id}#{self.queue_name} processed successfully after {retry_count} retries"
+                        )
+                    else:
+                        logger.info(
+                            f"Message {message_id}#{self.queue_name} processed successfully"
+                        )
 
 
 @asynccontextmanager
@@ -312,9 +1218,18 @@ async def none_context() -> AsyncGenerator[None, None]:
 
 
 class MessageBusWorker:
-    def __init__(self, app: Microservice, config: AioPikaWorkerConfig) -> None:
+    def __init__(
+        self,
+        app: Microservice,
+        broker_url: str,
+        backend_url: str,
+        handler_names: set[str] | None = None,
+    ) -> None:
         self.app = app
-        self.config = config
+        self.backend_url = backend_url
+        self.broker_url = broker_url
+        self.handler_names = handler_names
+
         self.container = Container(app)
         self.lifecycle = AppLifecycle(app, self.container)
 
@@ -322,33 +1237,49 @@ class MessageBusWorker:
             app=app, container=self.container
         )
 
-        self._consumer: AioPikaMicroserviceConsumer | None = None
+        self._consumer: MessageBusConsumer | None = None
 
     @property
-    def consumer(self) -> AioPikaMicroserviceConsumer:
+    def consumer(self) -> MessageBusConsumer:
         if self._consumer is None:
             raise RuntimeError("Consumer not started")
         return self._consumer
 
-    async def start_async(self, passive_declare: bool) -> None:
+    async def start_async(self) -> None:
         all_message_handlers_set: MESSAGE_HANDLER_DATA_SET = set()
+        all_scheduled_actions_set: SCHEDULED_ACTION_DATA_SET = set()
         async with self.lifecycle():
-            for instance_type in self.app.controllers:
-                controller = MessageBusController.get_messagebus(instance_type)
+            for instance_class in self.app.controllers:
+                controller = MessageBusController.get_messagebus(instance_class)
 
                 if controller is None:
                     continue
 
-                instance: Any = self.container.get_by_type(instance_type)
+                instance: Any = self.container.get_by_type(instance_class)
 
                 factory = controller.get_messagebus_factory()
-                handlers, _ = factory(instance)
+                handlers, schedulers = factory(instance)
 
                 message_handler_data_map: dict[str, MessageHandlerData] = {}
-
+                all_scheduled_actions_set.update(schedulers)
                 for handler_data in handlers:
                     message_type = handler_data.spec.message_type
                     topic = message_type.MESSAGE_TOPIC
+
+                    # Filter handlers by name if specified
+                    if (
+                        self.handler_names is not None
+                        and handler_data.spec.name is not None
+                    ):
+                        if handler_data.spec.name not in self.handler_names:
+                            continue
+                    elif (
+                        self.handler_names is not None
+                        and handler_data.spec.name is None
+                    ):
+                        # Skip handlers without names when filtering is requested
+                        continue
+
                     if (
                         topic in message_handler_data_map
                         and message_type.MESSAGE_TYPE == "task"
@@ -361,30 +1292,79 @@ class MessageBusWorker:
                     message_handler_data_map[topic] = handler_data
                     all_message_handlers_set.add(handler_data)
 
-            consumer = self._consumer = AioPikaMicroserviceConsumer(
-                config=self.config,
+            broker_backend = get_message_broker_backend_from_url(url=self.backend_url)
+
+            consumer = self._consumer = create_message_bus(
+                broker_url=self.broker_url,
+                broker_backend=broker_backend,
+                scheduled_actions=all_scheduled_actions_set,
                 message_handler_set=all_message_handlers_set,
                 uow_context_provider=self.uow_context_provider,
             )
 
-            await consumer.consume(passive_declare=passive_declare)
+            await consumer.consume()
 
-    def start_sync(self, passive_declare: bool) -> None:
+    def start_sync(self) -> None:
 
         def on_shutdown(loop: asyncio.AbstractEventLoop) -> None:
-            logger.info("Shutting down")
-            self.consumer.shutdown_event.set()
+            logger.info("Shutting down - signal received")
+            # Schedule the shutdown to run in the event loop
+            asyncio.create_task(self._graceful_shutdown())
+            # wait until the shutdown is complete
 
         with asyncio.Runner(loop_factory=uvloop.new_event_loop) as runner:
-            runner.get_loop().add_signal_handler(
-                signal.SIGINT, on_shutdown, runner.get_loop()
-            )
-            runner.run(self.start_async(passive_declare=passive_declare))
+            loop = runner.get_loop()
+            loop.add_signal_handler(signal.SIGINT, on_shutdown, loop)
+            # Add graceful shutdown handler for SIGTERM as well
+            loop.add_signal_handler(signal.SIGTERM, on_shutdown, loop)
+            runner.run(self.start_async())
+
+    async def _graceful_shutdown(self) -> None:
+        """Handles graceful shutdown process"""
+        logger.info("Initiating graceful shutdown sequence")
+        # Use the comprehensive close method that handles shutdown, task waiting and connection cleanup
+
+        self.consumer.shutdown()
+        logger.info("Graceful shutdown completed")
 
 
 class AioPikaMessageBusController(BusMessageController):
     def __init__(self, aio_pika_message: aio_pika.abc.AbstractIncomingMessage):
         self.aio_pika_message = aio_pika_message
+        # We access consumer callback through context if available
+        self._callback: Optional[MessageHandlerCallback] = None
+
+    def _get_callback(self) -> MessageHandlerCallback:
+        """
+        Find the callback associated with this message.
+        This allows us to access the retry mechanisms.
+        """
+        if self._callback is None:
+            # Get the context from current frame's locals
+            frame = inspect.currentframe()
+            if frame is not None:
+                try:
+                    caller_frame = frame.f_back
+                    if caller_frame is not None:
+                        # Check for context with handler callback
+                        callback_ref = None
+                        # Look for handler_message call context
+                        while caller_frame is not None:
+                            if "self" in caller_frame.f_locals:
+                                self_obj = caller_frame.f_locals["self"]
+                                if isinstance(self_obj, MessageHandlerCallback):
+                                    callback_ref = self_obj
+                                    break
+                            caller_frame = caller_frame.f_back
+                        # Save callback reference if we found it
+                        self._callback = callback_ref
+                finally:
+                    del frame  # Avoid reference cycles
+
+            if self._callback is None:
+                raise RuntimeError("Could not find callback context for message retry")
+
+        return self._callback
 
     async def ack(self) -> None:
         await self.aio_pika_message.ack()
@@ -396,7 +1376,41 @@ class AioPikaMessageBusController(BusMessageController):
         await self.aio_pika_message.reject()
 
     async def retry(self) -> None:
-        await self.aio_pika_message.reject(requeue=True)
+        """
+        Retry the message immediately by rejecting with requeue flag.
+        This doesn't use the exponential backoff mechanism.
+        """
+        callback = self._get_callback()
+        await callback.handle_reject_message(self.aio_pika_message, requeue=True)
 
     async def retry_later(self, delay: int) -> None:
-        raise NotImplementedError("Not implemented")
+        """
+        Retry the message after a specified delay using the exponential backoff mechanism.
+
+        Args:
+            delay: Minimum delay in seconds before retrying
+        """
+        try:
+            callback = self._get_callback()
+
+            # Get current retry count from message headers
+            headers = self.aio_pika_message.headers or {}
+            retry_count = int(str(headers.get("x-retry-count", 0)))
+
+            # Handle retry with explicit delay
+            asyncio.create_task(
+                callback._delayed_retry(
+                    self.aio_pika_message,
+                    float(delay),
+                    retry_count + 1,
+                    None,  # No specific exception
+                )
+            )
+
+            # Acknowledge the current message since we'll republish
+            await self.aio_pika_message.ack()
+
+        except Exception as e:
+            logger.exception(f"Failed to schedule retry_later: {e}")
+            # Fall back to immediate retry
+            await self.aio_pika_message.reject(requeue=True)