jararaca 0.2.37a12__py3-none-any.whl → 0.4.0a5__py3-none-any.whl

jararaca/messagebus/worker.py

@@ -1,16 +1,44 @@
+# SPDX-FileCopyrightText: 2025 Lucas S
+#
+# SPDX-License-Identifier: GPL-3.0-or-later
+
 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 (
+    AMQPChannelError,
+    AMQPConnectionError,
+    AMQPError,
+    ChannelClosed,
+    ChannelNotFoundEntity,
+    ConnectionClosed,
+)
 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,12 +48,26 @@ 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.types import Message, MessageOf
-from jararaca.microservice import MessageBusAppContext, Microservice
+from jararaca.messagebus.implicit_headers import provide_implicit_headers
+from jararaca.messagebus.message import Message, MessageOf
+from jararaca.microservice import (
+    AppTransactionContext,
+    MessageBusTransactionData,
+    Microservice,
+    SchedulerTransactionData,
+    ShutdownState,
+    provide_shutdown_state,
+    providing_app_type,
+)
+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__)
 
@@ -33,9 +75,27 @@ logger = logging.getLogger(__name__)
 @dataclass
 class AioPikaWorkerConfig:
     url: str
-    queue: 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,
+        )
+    )
+    # Connection health monitoring settings
+    connection_heartbeat_interval: float = 30.0  # seconds
+    connection_health_check_interval: float = 10.0  # seconds
 
 
 class AioPikaMessage(MessageOf[Message]):
@@ -76,79 +136,1047 @@ 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 _WorkerShutdownState(ShutdownState):
+    def __init__(self, shutdown_event: asyncio.Event):
+        self.shutdown_event = shutdown_event
+
+    def request_shutdown(self) -> None:
+        self.shutdown_event.set()
+
+    def is_shutdown_requested(self) -> bool:
+        return self.shutdown_event.is_set()
+
+    async def wait_for_shutdown(self) -> None:
+        await self.shutdown_event.wait()
+
+
+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.shutdown_state = _WorkerShutdownState(self.shutdown_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] = {}
+
+        # Connection resilience attributes
+        self.connection_healthy = False
+        self.connection_lock = asyncio.Lock()
+        self.consumer_tags: dict[str, str] = {}  # Track consumer tags for cleanup
+        self.health_check_task: asyncio.Task[Any] | None = None
+
+    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: aio_pika.abc.AbstractQueue = await RabbitmqUtils.get_queue(
+                    channel=channel, queue_name=queue_name
+                )
+
+                # Configure consumer and get the consumer tag
+                consumer_tag = await queue.consume(
+                    callback=MessageHandlerCallback(
+                        consumer=self,
+                        queue_name=queue_name,
+                        routing_key=routing_key,
+                        message_handler=handler,
+                    ),
+                    # no_ack=handler.spec.auto_ack,
+                )
+
+                # Store consumer tag for cleanup
+                self.consumer_tags[queue_name] = consumer_tag
+
+                logger.info(
+                    "Consuming message handler %s on dedicated channel", queue_name
+                )
+
+                await self.shutdown_event.wait()
+
+                logger.warning(
+                    "Shutdown event received, stopping consumer for %s", queue_name
+                )
+                await queue.cancel(consumer_tag)
+
+        try:
+            # Setup with retry
+            await retry_with_backoff(
+                setup_consumer,
+                retry_config=self.config.consumer_retry_config,
+                retry_exceptions=(
+                    ChannelNotFoundEntity,
+                    ChannelClosed,
+                    AMQPError,
+                    AMQPConnectionError,
+                    AMQPChannelError,
+                    ConnectionClosed,
+                ),
+            )
+            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 and get the consumer tag
+                consumer_tag = await queue.consume(
+                    callback=ScheduledMessageHandlerCallback(
+                        consumer=self,
+                        queue_name=queue_name,
+                        routing_key=routing_key,
+                        scheduled_action=scheduled_action,
+                    ),
+                    no_ack=True,
+                )
+
+                # Store consumer tag for cleanup
+                self.consumer_tags[queue_name] = consumer_tag
+
+                logger.debug("Consuming scheduler %s on dedicated channel", queue_name)
+
+                await self.shutdown_event.wait()
+
+                logger.warning(
+                    "Shutdown event received, stopping consumer for %s", queue_name
+                )
+                await queue.cancel(consumer_tag)
+
+        try:
+            # Setup with retry
+            await retry_with_backoff(
+                setup_consumer,
+                retry_config=self.config.consumer_retry_config,
+                retry_exceptions=(
+                    ChannelNotFoundEntity,
+                    ChannelClosed,
+                    AMQPError,
+                    AMQPConnectionError,
+                    AMQPChannelError,
+                    ConnectionClosed,
+                ),
+            )
+            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.
+        """
+        # Establish initial connection
+        try:
+            async with self.connect() as connection:
+                self.connection_healthy = True
+
+                # Start connection health monitoring
+                self.health_check_task = asyncio.create_task(
+                    self._monitor_connection_health(), name="ConnectionHealthMonitor"
+                )
+
+                # Verify infrastructure with retry
+                infra_check_success = await retry_with_backoff(
+                    self._verify_infrastructure,
+                    retry_config=self.config.connection_retry_config,
+                    retry_exceptions=(Exception,),
+                )
+
+                if not infra_check_success:
+                    logger.critical(
+                        "Failed to verify RabbitMQ infrastructure. Shutting down."
+                    )
+                    self.shutdown_event.set()
+                    return
+
+                async def wait_for(
+                    type: str, name: str, coroutine: Awaitable[bool]
+                ) -> tuple[str, str, bool]:
+                    return type, name, await coroutine
+
+                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),
+                            ),
+                            name=f"MessageHandler-{queue_name}-setup-consumer",
+                        )
+                    )
 
-        connection = await aio_pika.connect(self.config.url)
+                # 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),
+                            ),
+                            name=f"ScheduledAction-{queue_name}-setup-consumer",
+                        )
+                    )
 
-        channel = await connection.channel()
+                async def handle_task_results() -> None:
+                    for task in asyncio.as_completed(tasks):
+                        type, name, success = await task
+                        if success:
+                            logger.debug(
+                                "Successfully set up %s consumer for %s", type, name
+                            )
+                        else:
+                            logger.warning(
+                                "Failed to set up %s consumer for %s, will not process messages from this queue",
+                                type,
+                                name,
+                            )
+
+                handle_task_results_task = asyncio.create_task(
+                    handle_task_results(), name="HandleSetupTaskResults"
+                )
+
+                # Wait for shutdown signal
+                await self.shutdown_event.wait()
+                logger.debug("Shutdown event received, stopping consumers")
+
+                await self.cancel_queue_consumers()
+
+                # Cancel health monitoring
+                if self.health_check_task:
+                    self.health_check_task.cancel()
+                    with suppress(asyncio.CancelledError):
+                        await self.health_check_task
+
+                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.debug("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()
+
+        except Exception as e:
+            logger.critical("Failed to establish initial connection to RabbitMQ: %s", e)
+            # Re-raise the exception so it can be caught by the caller
+            raise
+
+    async def cancel_queue_consumers(self) -> None:
+        """
+        Cancel all active queue consumers.
+        """
+        logger.debug("Cancelling all active queue consumers...")
+        for queue_name, channel in self.channels.items():
+            try:
+                if not channel.is_closed:
+                    # Cancel consumer if we have its tag
+                    if queue_name in self.consumer_tags:
+                        try:
+                            queue = await channel.get_queue(queue_name, ensure=False)
+                            if queue:
+                                await queue.cancel(self.consumer_tags[queue_name])
+                        except Exception as cancel_error:
+                            logger.warning(
+                                "Error cancelling consumer for %s: %s",
+                                queue_name,
+                                cancel_error,
+                            )
+                        del self.consumer_tags[queue_name]
+            except Exception as e:
+                logger.warning("Error cancelling consumer for %s: %s", queue_name, e)
+
+    async def wait_all_tasks_done(self) -> None:
+        if not self.tasks:
+            return
+
+        logger.warning(
+            "Waiting for (%s) in-flight tasks to complete: %s",
+            len(self.tasks),
+            ", ".join((task.get_name()) for task in self.tasks),
+        )
+        # 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)
+        pending_tasks = [task for task in self.tasks if not task.done()]
+        while len(pending_tasks) > 0:
+            if not pending_tasks:
+                break
+            await asyncio.wait(pending_tasks, return_when=asyncio.FIRST_COMPLETED)
+
+            pending_tasks = [task for task in pending_tasks if not task.done()]
+            if len(pending_tasks) > 0:
+                logger.warning(
+                    "Waiting for (%s) in-flight tasks to complete: %s",
+                    len(pending_tasks),
+                    ", ".join((task.get_name()) for task in pending_tasks),
+                )
+            else:
+                logger.warning("All in-flight tasks have completed.")
+            # Log any exceptions that occurred
+            # for result in results:
+            #     if isinstance(result, Exception):
+            #         logger.error("Task raised an exception during shutdown: %s", result)
+
+    async def close_channels_and_connection(self) -> None:
+        """Close all channels and then the connection"""
+        logger.warning("Closing channels and connection...")
+        await self._cleanup_connection()
+
+    def shutdown(self) -> None:
+        """Signal for shutdown"""
+        logger.warning("Initiating graceful shutdown")
+        self.shutdown_event.set()
+
+    async def close(self) -> None:
+        """Implement MessageBusConsumer.close for cleanup"""
+        logger.warning("Closing consumer...")
+        self.shutdown()
+
+        # Cancel health monitoring
+        if self.health_check_task:
+            self.health_check_task.cancel()
+            with suppress(asyncio.CancelledError):
+                await self.health_check_task
+
+        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("No channel found for queue %s", queue_name)
+            return None
+
+        try:
+            channel = self.channels[queue_name]
+            if channel.is_closed:
+                logger.warning("Channel for queue %s is closed", queue_name)
+                # Remove the closed channel
+                del self.channels[queue_name]
+
+                # Attempt to recreate the channel if connection is healthy
+                if (
+                    self.connection
+                    and not self.connection.is_closed
+                    and self.connection_healthy
+                ):
+                    try:
+                        logger.debug("Creating new channel for %s", 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]
+                    except Exception as e:
+                        logger.error(
+                            "Failed to recreate channel for %s: %s", queue_name, e
+                        )
+                        # Trigger shutdown if channel creation fails
+                        self._trigger_shutdown()
+                        return None
+                else:
+                    # Connection is not healthy, trigger shutdown
+                    self._trigger_shutdown()
+                    return None
+            return channel
+        except Exception as e:
+            logger.error("Error accessing channel for queue %s: %s", queue_name, e)
+            # Trigger shutdown on any channel access error
+            self._trigger_shutdown()
+            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(
+                "Cannot create channel for %s: connection is not available", queue_name
+            )
+            raise RuntimeError("Connection is not available")
 
+        logger.debug("Creating channel for queue %s", queue_name)
+        channel = await self.connection.channel()
         await channel.set_qos(prefetch_count=self.config.prefetch_count)
+        logger.debug("Created channel for queue %s", 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=(
+                    AMQPConnectionError,
+                    AMQPChannelError,
+                    ConnectionError,
+                ),
+            )
+
+            # Save in the channels dict for tracking
+            self.channels[queue_name] = channel
+            logger.debug("Created new channel for queue %s", 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(
+                "Error creating channel for queue %s after retries: %s", queue_name, e
+            )
+            raise
+
+    async def _establish_connection(self) -> aio_pika.abc.AbstractConnection:
+        """
+        Creates a new RabbitMQ connection with retry logic.
+        """
+        try:
+            logger.debug("Establishing connection to RabbitMQ")
+            connection = await aio_pika.connect(
+                self.config.url,
+                heartbeat=self.config.connection_heartbeat_interval,
+            )
+            logger.debug("Connected to RabbitMQ successfully")
+            return connection
+        except Exception as e:
+            logger.error("Failed to connect to RabbitMQ: %s", e)
+            raise
+
+    @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=(
+                    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(
+                "Failed to establish connection to RabbitMQ after retries: %s", e
+            )
+            if self.connection:
+                try:
+                    await self.connection.close()
+                except Exception as close_error:
+                    logger.error(
+                        "Error closing connection after connect failure: %s",
+                        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.
+        """
+        max_retries = 3
+        retry_delay = 1.0
+
+        for attempt in range(max_retries):
+            try:
+                channel = await self.get_channel(queue_name)
+                if channel is not None:
+                    try:
+                        yield channel
+                        return
+                    finally:
+                        # We don't close the channel here as it's managed by the consumer
+                        pass
+
+                # No channel available, check connection state
+                if (
+                    self.connection
+                    and not self.connection.is_closed
+                    and self.connection_healthy
+                ):
+                    # Try to create a new channel
+                    async with self.create_channel(queue_name) as new_channel:
+                        yield new_channel
+                        return
+                else:
+                    # Connection is not healthy, trigger shutdown
+                    logger.error(
+                        "Connection not healthy while getting channel for %s, triggering shutdown",
+                        queue_name,
+                    )
+                    self._trigger_shutdown()
+                    raise RuntimeError(
+                        f"Cannot get channel for queue {queue_name}: connection is not healthy"
+                    )
+
+            except Exception as e:
+                if attempt < max_retries - 1:
+                    logger.warning(
+                        "Error getting channel for %s, retrying: %s", queue_name, e
+                    )
+                    await asyncio.sleep(retry_delay)
+                    retry_delay *= 2
+                else:
+                    logger.error(
+                        "Failed to get channel for %s after %s attempts: %s",
+                        queue_name,
+                        max_retries,
+                        e,
+                    )
+                    raise
+
+    async def _monitor_connection_health(self) -> None:
+        """
+        Monitor connection health and trigger shutdown if connection is lost.
+        This runs as a background task.
+        """
+        while not self.shutdown_event.is_set():
+            try:
+                await asyncio.sleep(self.config.connection_health_check_interval)
+
+                if self.shutdown_event.is_set():
+                    break
+
+                # Check connection health
+                if not await self._is_connection_healthy():
+                    logger.error(
+                        "Connection health check failed, initiating worker shutdown"
+                    )
+                    self.shutdown()
+                    break
+
+            except asyncio.CancelledError:
+                logger.debug("Connection health monitoring cancelled")
+                break
+            except Exception as e:
+                logger.error("Error in connection health monitoring: %s", e)
+                await asyncio.sleep(5)  # Wait before retrying
+
+    async def _is_connection_healthy(self) -> bool:
+        """
+        Check if the connection is healthy.
+        """
+        try:
+            if self.connection is None or self.connection.is_closed:
+                return False
+
+            # Try to create a temporary channel to test connection
+            async with self.connection.channel() as test_channel:
+                # If we can create a channel, connection is healthy
+                return True
+
+        except Exception as e:
+            logger.debug("Connection health check failed: %s", e)
+            return False
+
+    def _trigger_shutdown(self) -> None:
+        """
+        Trigger worker shutdown due to connection loss.
+        """
+        if not self.shutdown_event.is_set():
+            logger.error("Connection lost, initiating worker shutdown")
+            self.connection_healthy = False
+            self.shutdown()
+
+    async def _cleanup_connection(self) -> None:
+        """
+        Clean up existing connection and channels.
+        """
+        # Cancel existing consumers
+        for queue_name, channel in self.channels.items():
+            try:
+                if not channel.is_closed:
+                    # Cancel consumer if we have its tag
+                    if queue_name in self.consumer_tags:
+                        try:
+                            queue = await channel.get_queue(queue_name, ensure=False)
+                            if queue:
+                                await queue.cancel(self.consumer_tags[queue_name])
+                        except Exception as cancel_error:
+                            logger.warning(
+                                "Error cancelling consumer for %s: %s",
+                                queue_name,
+                                cancel_error,
+                            )
+                        del self.consumer_tags[queue_name]
+            except Exception as e:
+                logger.warning("Error cancelling consumer for %s: %s", queue_name, e)
+
+        # Close channels
+        for queue_name, channel in self.channels.items():
+            try:
+                if not channel.is_closed:
+                    await channel.close()
+            except Exception as e:
+                logger.warning("Error closing channel for %s: %s", queue_name, e)
+
+        self.channels.clear()
+
+        # Close connection
+        if self.connection and not self.connection.is_closed:
+            try:
+                await self.connection.close()
+            except Exception as e:
+                logger.warning("Error closing connection: %s", e)
+
+        self.connection = None
+        self.connection_healthy = False
+
+
+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
+        )
 
-        main_x = await channel.declare_exchange(self.config.exchange, type="topic")
+        # Parse heartbeat and health check intervals
+        connection_heartbeat_interval = 30.0
+        connection_health_check_interval = 10.0
 
-        dlx = await channel.declare_exchange("dlx", type="direct")
+        # 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]
+            )
+
+        if "connection_retry_delay" in query_params:
+            try:
+                connection_retry_config.initial_delay = float(
+                    query_params["connection_retry_delay"][0]
+                )
+            except ValueError:
+                pass
+
+        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
+
+        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
+
+        # Heartbeat and health check intervals
+        if "connection_heartbeat_interval" in query_params:
+            try:
+                connection_heartbeat_interval = float(
+                    query_params["connection_heartbeat_interval"][0]
+                )
+            except ValueError:
+                pass
+
+        if "connection_health_check_interval" in query_params:
+            try:
+                connection_health_check_interval = float(
+                    query_params["connection_health_check_interval"][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,
+            connection_heartbeat_interval=connection_heartbeat_interval,
+            connection_health_check_interval=connection_health_check_interval,
+        )
+
+        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"
+    )
+
+
+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
 
-        dlq = await channel.declare_queue("dlq")
+    async def __call__(
+        self, aio_pika_message: aio_pika.abc.AbstractIncomingMessage
+    ) -> None:
 
-        await dlq.bind(dlx, routing_key="dlq")
+        if self.consumer.shutdown_event.is_set():
+            logger.debug(
+                "Shutdown in progress. Requeuing scheduled message for %s",
+                self.queue_name,
+            )
+            try:
+                # Use channel context for requeuing
+                await aio_pika_message.reject(requeue=True)
+            except RuntimeError:
+                logger.warning(
+                    "Could not requeue scheduled message during shutdown - channel not available"
+                )
+            except Exception as e:
+                logger.error(
+                    "Failed to requeue scheduled message during shutdown: %s", e
+                )
+            return
 
-        for handler in self.message_handler_set:
+        # Check if connection is healthy before processing
+        if not self.consumer.connection_healthy:
+            logger.warning(
+                "Connection not healthy, requeuing scheduled message for %s",
+                self.queue_name,
+            )
+            try:
+                if not self.consumer.connection_healthy:
+                    # Still not healthy, requeue the message
+                    async with self.consumer.get_channel_ctx(self.queue_name):
+                        await aio_pika_message.reject(requeue=True)
+                    return
+            except Exception as e:
+                logger.error(
+                    "Failed to requeue scheduled message due to connection issues: %s",
+                    e,
+                )
+                return
 
-            queue_name = f"{self.config.exchange}.{handler.message_type.MESSAGE_TOPIC}.{handler.callable.__module__}.{handler.callable.__qualname__}"
-            routing_key = (
-                f"{self.config.exchange}.{handler.message_type.MESSAGE_TOPIC}.#"
+        async with self.consumer.lock:
+            task = asyncio.create_task(
+                self.handle_message(aio_pika_message),
+                name=f"ScheduledAction-{self.queue_name}-handle-message-{aio_pika_message.message_id}",
             )
+            self.consumer.tasks.add(task)
+            task.add_done_callback(self.handle_message_consume_done)
 
-            self.incoming_map[queue_name] = handler
+    def handle_message_consume_done(self, task: asyncio.Task[Any]) -> None:
+        self.consumer.tasks.discard(task)
+        if task.cancelled():
+            logger.warning("Scheduled task for %s was cancelled", self.queue_name)
+            return
 
-            queue = await channel.declare_queue(
-                queue_name,
-                arguments={
-                    "x-dead-letter-exchange": "dlx",
-                    "x-dead-letter-routing-key": "dlq",
-                },
+        if (error := task.exception()) is not None:
+            logger.exception(
+                "Error processing scheduled action %s", self.queue_name, exc_info=error
             )
 
-            await queue.bind(exchange=self.config.exchange, routing_key=routing_key)
+    async def handle_message(
+        self, aio_pika_message: aio_pika.abc.AbstractIncomingMessage
+    ) -> None:
 
-            await queue.consume(
-                callback=MessageHandlerCallback(
-                    consumer=self,
-                    queue_name=queue_name,
-                    routing_key=routing_key,
-                    message_handler=handler,
+        if self.consumer.shutdown_event.is_set():
+            logger.debug(
+                "Shutdown event set. Requeuing message for %s", 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(
+                    "Could not requeue message during shutdown - channel not available"
+                )
+            except Exception as e:
+                logger.error("Failed to requeue message during shutdown: %s", e)
+                return
+
+        # Check connection health before processing
+        if not self.consumer.connection_healthy:
+            logger.warning(
+                "Connection not healthy, requeuing scheduled message for %s",
+                self.queue_name,
+            )
+            try:
+                async with self.consumer.get_channel_ctx(self.queue_name):
+                    await aio_pika_message.reject(requeue=True)
+                return
+            except Exception as e:
+                logger.error(
+                    "Failed to requeue scheduled message due to connection issues: %s",
+                    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"))),),
+                    {},
                 ),
-                no_ack=handler.spec.auto_ack,
+                name=f"ScheduledAction-{self.queue_name}-handle-message-{aio_pika_message.message_id}",
             )
 
-            print(f"Consuming {queue_name}")
+        elif len(sig.parameters) == 0:
+            task = asyncio.create_task(
+                self.run_with_context(
+                    self.scheduled_action,
+                    (),
+                    {},
+                ),
+                name=f"ScheduledAction-{self.queue_name}-handle-message-{aio_pika_message.message_id}",
+            )
+        else:
+            logger.warning(
+                "Scheduled action '%s' must have exactly one parameter of type ScheduleDispatchData or no parameters"
+                % self.queue_name
+            )
+            return
 
-        await self.shutdown_event.wait()
-        logger.info("Worker shutting down")
+        self.consumer.tasks.add(task)
+        task.add_done_callback(self.handle_message_consume_done)
 
-        await self.wait_all_tasks_done()
+        try:
+            await task
+        except Exception as e:
+
+            logger.exception(
+                "Error processing scheduled action %s: %s", self.queue_name, e
+            )
 
-        await channel.close()
-        await connection.close()
+    async def run_with_context(
+        self,
+        scheduled_action: ScheduledActionData,
+        args: tuple[Any, ...],
+        kwargs: dict[str, Any],
+    ) -> None:
 
-    async def wait_all_tasks_done(self) -> None:
-        async with self.lock:
-            await asyncio.gather(*self.tasks)
+        with provide_shutdown_state(self.consumer.shutdown_state):
+            async with self.consumer.uow_context_provider(
+                AppTransactionContext(
+                    controller_member_reflect=scheduled_action.controller_member,
+                    transaction_data=SchedulerTransactionData(
+                        task_name=scheduled_action.spec.name
+                        or scheduled_action.callable.__qualname__,
+                        scheduled_to=datetime.now(UTC),
+                        cron_expression=scheduled_action.spec.cron,
+                        triggered_at=datetime.now(UTC),
+                    ),
+                )
+            ):
+
+                await scheduled_action.callable(*args, **kwargs)
 
 
 class MessageHandlerCallback:
@@ -164,25 +1192,62 @@ 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.debug(
+                "Shutdown in progress. Requeuing message for %s", 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(
+                    "Could not requeue message during shutdown - channel not available"
+                )
+            except Exception as e:
+                logger.error("Failed to requeue message during shutdown: %s", e)
             return
 
+        # Check if connection is healthy before processing
+        if not self.consumer.connection_healthy:
+            logger.warning(
+                "Connection not healthy, requeuing message for %s", self.queue_name
+            )
+            try:
+                if not self.consumer.connection_healthy:
+                    # Still not healthy, requeue the message
+                    async with self.consumer.get_channel_ctx(self.queue_name):
+                        await aio_pika_message.reject(requeue=True)
+                    return
+            except Exception as e:
+                logger.error(
+                    "Failed to requeue message due to connection issues: %s", e
+                )
+                return
+
         async with self.consumer.lock:
-            task = asyncio.create_task(self.handle_message(aio_pika_message))
+            task = asyncio.create_task(
+                self.handle_message(aio_pika_message),
+                name=f"MessageHandler-{self.queue_name}-handle-message-{aio_pika_message.message_id}",
+            )
             self.consumer.tasks.add(task)
             task.add_done_callback(self.handle_message_consume_done)
 
     def handle_message_consume_done(self, task: asyncio.Task[Any]) -> None:
         self.consumer.tasks.discard(task)
         if task.cancelled():
+            logger.warning("Task for queue %s was cancelled", self.queue_name)
             return
 
         if (error := task.exception()) is not None:
-            logger.exception("Error processing message", exc_info=error)
+            logger.exception(
+                "Error processing message for queue %s", self.queue_name, exc_info=error
+            )
 
     async def __call__(
         self, aio_pika_message: aio_pika.abc.AbstractIncomingMessage
@@ -193,14 +1258,242 @@ 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())
+
+        try:
+            # Check if we should retry with backoff
+            if (
+                not requeue
+                and self.message_handler.spec.nack_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(
+                        "Message %s (%s) failed after %s retries, dead-lettering: %s",
+                        message_id,
+                        self.queue_name,
+                        retry_count,
+                        str(exception),
+                    )
+                    # Dead-letter the message after max retries
+                    try:
+
+                        await aio_pika_message.reject(requeue=False)
+                    except Exception as e:
+                        logger.error(
+                            "Failed to dead-letter message %s: %s", message_id, e
+                        )
+                    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.warning(
+                    "Message %s (%s) failed with %s, retry %s/%s scheduled in %.2fs",
+                    message_id,
+                    self.queue_name,
+                    str(exception),
+                    retry_count + 1,
+                    retry_config.max_retries,
+                    delay,
+                )
+
+                # 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
+                    ),
+                    name=f"MessageHandler-{self.queue_name}-delayed-retry-{message_id}",
+                )
+
+                # Acknowledge the current message since we'll handle retry ourselves
+                try:
+                    await aio_pika_message.ack()
+                except Exception as e:
+                    logger.error(
+                        "Failed to acknowledge message %s for retry: %s", message_id, e
+                    )
+                return
+
+            # Standard reject without retry or with immediate requeue
+            try:
+                async with self.consumer.get_channel_ctx(self.queue_name):
+                    await aio_pika_message.reject(requeue=requeue)
+                    if requeue:
+                        logger.warning(
+                            "Message %s (%s) requeued for immediate retry",
+                            message_id,
+                            self.queue_name,
+                        )
+                    else:
+                        logger.warning(
+                            "Message %s (%s) rejected without requeue",
+                            message_id,
+                            self.queue_name,
+                        )
+            except Exception as e:
+                logger.error("Failed to reject message %s: %s", message_id, e)
+
+        except Exception as e:
+            logger.exception(
+                "Unexpected error in handle_reject_message for %s (%s): %s",
+                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 retrying
+            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 self._wait_delay_or_shutdown(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 with retry logic
+            max_attempts = 3
+            for attempt in range(max_attempts):
+                try:
+                    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.warning(
+                            "Message %s (%s) republished for retry %s",
+                            message_id,
+                            self.queue_name,
+                            retry_count,
+                        )
+                        return
+
+                except Exception as e:
+                    if attempt < max_attempts - 1:
+                        logger.warning(
+                            "Failed to republish message %s (attempt %s): %s",
+                            message_id,
+                            attempt + 1,
+                            e,
+                        )
+                        await asyncio.sleep(1.0 * (attempt + 1))  # Exponential backoff
+                    else:
+                        logger.error(
+                            "Failed to republish message %s after %s attempts: %s",
+                            message_id,
+                            max_attempts,
+                            e,
+                        )
+                        raise
+
+        except Exception as e:
+            logger.exception(
+                "Failed to execute delayed retry for message %s (%s): %s",
+                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 _wait_delay_or_shutdown(self, delay: float) -> None:
+        """
+        Wait for the specified delay or exit early if shutdown is initiated.
+
+        Args:
+            delay: Delay in seconds to wait
+        """
+
+        wait_cor = asyncio.create_task(asyncio.sleep(delay), name="delayed-retry-wait")
+        wait_shutdown_cor = asyncio.create_task(
+            self.consumer.shutdown_event.wait(), name="delayed-retry-shutdown-wait"
+        )
+
+        await asyncio.wait(
+            [wait_cor, wait_shutdown_cor],
+            return_when=asyncio.FIRST_COMPLETED,
+        )
+
     async def handle_message(
         self, aio_pika_message: aio_pika.abc.AbstractIncomingMessage
     ) -> None:
@@ -212,15 +1505,9 @@ class MessageHandlerCallback:
             await self.handle_reject_message(aio_pika_message)
             return
 
-        handler_data = self.consumer.incoming_map.get(routing_key)
-
-        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_data = self.message_handler
 
-        handler = handler_data.callable
+        handler = handler_data.instance_callable
 
         sig = inspect.signature(handler)
 
@@ -260,51 +1547,123 @@ class MessageHandlerCallback:
 
         builded_message = AioPikaMessage(aio_pika_message, message_type)
 
-        incoming_message_spec = MessageHandler.get_message_incoming(handler)
+        incoming_message_spec = MessageHandler.get_last(handler)
         assert incoming_message_spec is not None
 
-        async with self.consumer.uow_context_provider(
-            MessageBusAppContext(
-                message=builded_message,
-                topic=routing_key,
-            )
+        with provide_implicit_headers(aio_pika_message.headers), provide_shutdown_state(
+            self.consumer.shutdown_state
         ):
-            ctx: AsyncContextManager[Any]
-            if incoming_message_spec.timeout is not None:
-                ctx = asyncio.timeout(incoming_message_spec.timeout)
-            else:
-                ctx = none_context()
-            async with ctx:
-                try:
-                    with provide_bus_message_controller(
-                        AioPikaMessageBusController(aio_pika_message)
-                    ):
-                        await handler(builded_message)
-                    if not incoming_message_spec.auto_ack:
-                        with suppress(aio_pika.MessageProcessError):
-                            await aio_pika_message.ack()
-                except BaseException as base_exc:
-                    if incoming_message_spec.exception_handler is not None:
-                        try:
-                            incoming_message_spec.exception_handler(base_exc)
-                        except Exception as nested_exc:
+            async with self.consumer.uow_context_provider(
+                AppTransactionContext(
+                    controller_member_reflect=handler_data.controller_member,
+                    transaction_data=MessageBusTransactionData(
+                        message_type=message_type,
+                        message=builded_message,
+                        topic=routing_key,
+                    ),
+                )
+            ):
+                maybe_timeout_ctx: AsyncContextManager[Any]
+                if incoming_message_spec.timeout is not None:
+                    maybe_timeout_ctx = asyncio.timeout(incoming_message_spec.timeout)
+                else:
+                    maybe_timeout_ctx = none_context()
+
+                start_time = time.perf_counter()
+                async with maybe_timeout_ctx:
+                    try:
+                        with provide_bus_message_controller(
+                            AioPikaMessageBusController(aio_pika_message)
+                        ):
+                            await handler(builded_message)
+                            with suppress(aio_pika.MessageProcessError):
+                                # Use channel context for acknowledgement with retry
+                                try:
+                                    await aio_pika_message.ack()
+                                except Exception as ack_error:
+                                    logger.warning(
+                                        "Failed to acknowledge message %s: %s",
+                                        aio_pika_message.message_id or "unknown",
+                                        ack_error,
+                                    )
+                        successfully = True
+                    except BaseException as base_exc:
+                        successfully = False
+                        # Get message id for logging
+                        message_id = aio_pika_message.message_id or "unknown"
+
+                        # 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(
+                                    "Error processing exception handler for message %s: %s | %s",
+                                    message_id,
+                                    base_exc,
+                                    nested_exc,
+                                )
+                        else:
                             logger.exception(
-                                f"Error processing exception handler: {base_exc} | {nested_exc}"
+                                "Error processing message %s on topic %s: %s",
+                                message_id,
+                                routing_key,
+                                str(base_exc),
                             )
-                    else:
-                        logger.exception(
-                            f"Error processing message on topic {routing_key}"
+
+                        # Handle rejection with retry logic
+                        if incoming_message_spec.nack_on_exception:
+                            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, exception=base_exc
+                            )
+
+                    elapsed_time = time.perf_counter() - start_time
+                    # 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 {}
+                    traceparent = headers.get("traceparent")
+                    trace_info = (
+                        f" [traceparent={str(traceparent)}]" if traceparent else ""
+                    )
+
+                    if "x-retry-count" in headers:
+                        retry_count = int(str(headers.get("x-retry-count", 0)))
+                        logger.debug(
+                            "Message %s#%s processed "
+                            + ("successfully" if successfully else "with errors")
+                            + " after %s retries in %.4fs%s",
+                            message_id,
+                            self.queue_name,
+                            retry_count,
+                            elapsed_time,
+                            trace_info,
                         )
-                    if incoming_message_spec.requeue_on_exception:
-                        await self.handle_reject_message(aio_pika_message, requeue=True)
                     else:
-                        await self.handle_reject_message(
-                            aio_pika_message, requeue=False
+                        logger.debug(
+                            "Message %s#%s processed "
+                            + ("successfully" if successfully else "with errors")
+                            + " in %.4fs%s",
+                            message_id,
+                            self.queue_name,
+                            elapsed_time,
+                            trace_info,
                         )
-                else:
-                    logger.info(
-                        f"Message {aio_pika_message.message_id}#{self.queue_name} processed successfully"
-                    )
 
 
 @asynccontextmanager
@@ -313,9 +1672,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)
 
@@ -323,69 +1691,146 @@ 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) -> None:
         all_message_handlers_set: MESSAGE_HANDLER_DATA_SET = set()
-        async with self.lifecycle():
-            for instance_type in self.app.controllers:
-                controller = MessageBusController.get_messagebus(instance_type)
-
-                if controller is None:
-                    continue
+        all_scheduled_actions_set: SCHEDULED_ACTION_DATA_SET = set()
+        with providing_app_type("worker"):
+            async with self.lifecycle():
+                for instance_class in self.app.controllers:
+                    controller = MessageBusController.get_last(instance_class)
 
-                instance: Any = self.container.get_by_type(instance_type)
-
-                factory = controller.get_messagebus_factory()
-                handlers = factory(instance)
+                    if controller is None:
+                        continue
 
-                message_handler_data_map: dict[str, MessageHandlerData] = {}
+                    instance: Any = self.container.get_by_type(instance_class)
+
+                    factory = controller.get_messagebus_factory()
+                    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"
+                        ):
+                            logger.warning(
+                                "Task handler for topic '%s' already registered. Skipping"
+                                % topic
+                            )
+                            continue
+                        message_handler_data_map[topic] = handler_data
+                        all_message_handlers_set.add(handler_data)
 
-                for handler_data in handlers:
-                    message_type = handler_data.spec.message_type
-                    topic = message_type.MESSAGE_TOPIC
-                    if (
-                        topic in message_handler_data_map
-                        and message_type.MESSAGE_TYPE == "task"
-                    ):
-                        logger.warning(
-                            "Task handler for topic '%s' already registered. Skipping"
-                            % topic
-                        )
-                        continue
-                    message_handler_data_map[topic] = handler_data
-                    all_message_handlers_set.add(handler_data)
+                broker_backend = get_message_broker_backend_from_url(
+                    url=self.backend_url
+                )
 
-            consumer = self._consumer = AioPikaMicroserviceConsumer(
-                config=self.config,
-                message_handler_set=all_message_handlers_set,
-                uow_context_provider=self.uow_context_provider,
-            )
+                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()
+                await consumer.consume()
 
     def start_sync(self) -> None:
 
         def on_shutdown(loop: asyncio.AbstractEventLoop) -> None:
-            logger.info("Shutting down")
-            self.consumer.shutdown_event.set()
+            logger.warning("Shutting down - signal received")
+            # Schedule the shutdown to run in the event loop
+            asyncio.create_task(
+                self._graceful_shutdown(), name="Worker-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())
+            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)
+            try:
+                runner.run(self.start_async())
+            except Exception as e:
+                logger.critical("Worker failed to start due to connection error: %s", e)
+                # Exit with error code 1 to indicate startup failure
+                import sys
+
+                sys.exit(1)
+
+    async def _graceful_shutdown(self) -> None:
+        """Handles graceful shutdown process"""
+        logger.warning("Initiating graceful shutdown sequence")
+        # Use the comprehensive close method that handles shutdown, task waiting and connection cleanup
+
+        self.consumer.shutdown()
+        logger.warning("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()
@@ -397,7 +1842,42 @@ 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
+                ),
+                name=f"MessageHandler-{callback.queue_name}-delayed-retry-{self.aio_pika_message.message_id or 'unknown'}-{int(time.time())}",
+            )
+
+            # Acknowledge the current message since we'll republish
+            await self.aio_pika_message.ack()
+
+        except Exception as e:
+            logger.exception("Failed to schedule retry_later: %s", e)
+            # Fall back to immediate retry
+            await self.aio_pika_message.reject(requeue=True)