jararaca 0.3.11a16__py3-none-any.whl → 0.4.0a19__py3-none-any.whl

jararaca/scheduler/beat_worker.py

@@ -1,11 +1,15 @@
+# SPDX-FileCopyrightText: 2025 Lucas S
+#
+# SPDX-License-Identifier: GPL-3.0-or-later
+
 import asyncio
 import contextlib
 import logging
 import signal
 import time
 from abc import ABC, abstractmethod
+from dataclasses import dataclass, field
 from datetime import UTC, datetime
-from types import FrameType
 from typing import Any
 from urllib.parse import parse_qs
 
@@ -14,16 +18,24 @@ import croniter
 import urllib3
 import urllib3.util
 import uvloop
-from aio_pika import connect_robust
-from aio_pika.abc import AbstractChannel, AbstractRobustConnection
+from aio_pika import connect
+from aio_pika.abc import AbstractChannel, AbstractConnection
+from aio_pika.exceptions import (
+    AMQPChannelError,
+    AMQPConnectionError,
+    AMQPError,
+    ChannelClosed,
+    ConnectionClosed,
+)
 from aio_pika.pool import Pool
+from aiormq.exceptions import ChannelInvalidStateError
 
 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
-from jararaca.microservice import Microservice
+from jararaca.microservice import Microservice, providing_app_type
 from jararaca.scheduler.decorators import (
     ScheduledAction,
     ScheduledActionData,
@@ -31,6 +43,7 @@ from jararaca.scheduler.decorators import (
 )
 from jararaca.scheduler.types import DelayedMessageData
 from jararaca.utils.rabbitmq_utils import RabbitmqUtils
+from jararaca.utils.retry import RetryPolicy, retry_with_backoff
 
 logger = logging.getLogger(__name__)
 
@@ -101,17 +114,26 @@ class _MessageBrokerDispatcher(ABC):
 
 class _RabbitMQBrokerDispatcher(_MessageBrokerDispatcher):
 
-    def __init__(self, url: str) -> None:
+    def __init__(
+        self,
+        url: str,
+        config: "BeatWorkerConfig | None" = None,
+        shutdown_event: asyncio.Event | None = None,
+    ) -> None:
         self.url = url
+        self.config = config or BeatWorkerConfig()
+        self.connection_healthy = False
+        self.shutdown_event = shutdown_event or asyncio.Event()
+        self.health_check_task: asyncio.Task[Any] | None = None
 
-        self.conn_pool: "Pool[AbstractRobustConnection]" = Pool(
+        self.conn_pool: "Pool[AbstractConnection]" = Pool(
             self._create_connection,
-            max_size=10,
+            max_size=self.config.max_pool_size,
         )
 
         self.channel_pool: "Pool[AbstractChannel]" = Pool(
             self._create_channel,
-            max_size=10,
+            max_size=self.config.max_pool_size,
         )
 
         splitted = urllib3.util.parse_url(url)
@@ -130,88 +152,290 @@ class _RabbitMQBrokerDispatcher(_MessageBrokerDispatcher):
 
         self.exchange = str(query_params["exchange"][0])
 
-    async def _create_connection(self) -> AbstractRobustConnection:
+    async def _create_connection(self) -> AbstractConnection:
         """
-        Create a connection to the RabbitMQ server.
-        This is used to send messages to the RabbitMQ server.
+        Create a connection to the RabbitMQ server with retry logic.
         """
-        connection = await connect_robust(self.url)
-        return connection
+
+        async def _establish_connection() -> AbstractConnection:
+            logger.debug("Establishing connection to RabbitMQ")
+            connection = await connect(
+                self.url,
+                heartbeat=self.config.connection_heartbeat_interval,
+            )
+            logger.debug("Connected to RabbitMQ successfully")
+            return connection
+
+        return await retry_with_backoff(
+            _establish_connection,
+            retry_policy=self.config.connection_retry_config,
+            retry_exceptions=(
+                AMQPConnectionError,
+                ConnectionError,
+                OSError,
+                TimeoutError,
+            ),
+        )
 
     async def _create_channel(self) -> AbstractChannel:
         """
-        Create a channel to the RabbitMQ server.
-        This is used to send messages to the RabbitMQ server.
+        Create a channel to the RabbitMQ server with retry logic.
         """
-        async with self.conn_pool.acquire() as connection:
-            channel = await connection.channel()
-            return channel
+
+        async def _establish_channel() -> AbstractChannel:
+            async with self.conn_pool.acquire() as connection:
+                channel = await connection.channel()
+                return channel
+
+        return await retry_with_backoff(
+            _establish_channel,
+            retry_policy=self.config.connection_retry_config,
+            retry_exceptions=(
+                AMQPConnectionError,
+                AMQPChannelError,
+                ChannelClosed,
+                ConnectionError,
+            ),
+        )
 
     async def dispatch_scheduled_action(self, action_id: str, timestamp: int) -> None:
         """
-        Dispatch a message to the RabbitMQ server.
-        This is used to send a message to the RabbitMQ server
-        to trigger the scheduled action.
+        Dispatch a message to the RabbitMQ server with retry logic.
         """
+        if not self.connection_healthy:
+            await self._wait_for_connection()
 
-        logger.info(f"Dispatching message to {action_id} at {timestamp}")
-        async with self.channel_pool.acquire() as channel:
-            exchange = await RabbitmqUtils.get_main_exchange(channel, self.exchange)
+        async def _dispatch() -> None:
+            logger.debug("Dispatching message to %s at %s", action_id, timestamp)
+            async with self.channel_pool.acquire() as channel:
+                exchange = await RabbitmqUtils.get_main_exchange(channel, self.exchange)
 
-            await exchange.publish(
-                aio_pika.Message(body=str(timestamp).encode()),
-                routing_key=action_id,
+                await exchange.publish(
+                    aio_pika.Message(body=str(timestamp).encode()),
+                    routing_key=action_id,
+                )
+                logger.debug("Dispatched message to %s at %s", action_id, timestamp)
+
+        try:
+            await retry_with_backoff(
+                _dispatch,
+                retry_policy=self.config.dispatch_retry_config,
+                retry_exceptions=(
+                    AMQPConnectionError,
+                    AMQPChannelError,
+                    ChannelClosed,
+                    ConnectionClosed,
+                    AMQPError,
+                ),
+            )
+
+        except ChannelInvalidStateError as e:
+            logger.error(
+                "Channel invalid state error when dispatching to %s: %s", action_id, e
             )
-            logger.info(f"Dispatched message to {action_id} at {timestamp}")
+            # Trigger shutdown if dispatch fails
+            self.shutdown_event.set()
+            raise
+
+        except Exception as e:
+            logger.error(
+                "Failed to dispatch message to %s after retries: %s", action_id, e
+            )
+            # Trigger shutdown if dispatch fails
+            self.shutdown_event.set()
+            raise
 
     async def dispatch_delayed_message(
         self, delayed_message: DelayedMessageData
     ) -> None:
         """
-        Dispatch a delayed message to the RabbitMQ server.
-        This is used to send a message to the RabbitMQ server
-        to trigger the scheduled action.
+        Dispatch a delayed message to the RabbitMQ server with retry logic.
         """
-        async with self.channel_pool.acquire() as channel:
+        if not self.connection_healthy:
+            await self._wait_for_connection()
+
+        async def _dispatch() -> None:
+            async with self.channel_pool.acquire() as channel:
+                exchange = await RabbitmqUtils.get_main_exchange(channel, self.exchange)
+                await exchange.publish(
+                    aio_pika.Message(
+                        body=delayed_message.payload,
+                    ),
+                    routing_key=f"{delayed_message.message_topic}.#",
+                )
 
-            exchange = await RabbitmqUtils.get_main_exchange(channel, self.exchange)
-            await exchange.publish(
-                aio_pika.Message(
-                    body=delayed_message.payload,
+        try:
+            await retry_with_backoff(
+                _dispatch,
+                retry_policy=self.config.dispatch_retry_config,
+                retry_exceptions=(
+                    AMQPConnectionError,
+                    AMQPChannelError,
+                    ChannelClosed,
+                    ConnectionClosed,
+                    AMQPError,
                 ),
-                routing_key=f"{delayed_message.message_topic}.",
             )
+        except Exception as e:
+            logger.error("Failed to dispatch delayed message after retries: %s", e)
+            # Trigger shutdown if dispatch fails
+            self.shutdown_event.set()
+            raise
 
     async def initialize(self, scheduled_actions: list[ScheduledActionData]) -> None:
         """
-        Initialize the RabbitMQ server.
-        This is used to create the exchange and queues for the scheduled actions.
+        Initialize the RabbitMQ server with retry logic.
         """
 
-        async with self.channel_pool.acquire() as channel:
-            await RabbitmqUtils.get_main_exchange(channel, self.exchange)
+        async def _initialize() -> None:
+            async with self.channel_pool.acquire() as channel:
+                await RabbitmqUtils.get_main_exchange(channel, self.exchange)
 
-            for sched_act_data in scheduled_actions:
-                queue_name = ScheduledAction.get_function_id(sched_act_data.callable)
+                for sched_act_data in scheduled_actions:
+                    queue_name = ScheduledAction.get_function_id(
+                        sched_act_data.callable
+                    )
+
+                    # Try to get existing queue
+                    await RabbitmqUtils.get_scheduled_action_queue(
+                        channel=channel,
+                        queue_name=queue_name,
+                    )
 
-                # Try to get existing queue
-                await RabbitmqUtils.get_scheduled_action_queue(
-                    channel=channel,
-                    queue_name=queue_name,
+        try:
+            logger.debug("Initializing RabbitMQ connection...")
+            await retry_with_backoff(
+                _initialize,
+                retry_policy=self.config.connection_retry_config,
+                retry_exceptions=(
+                    AMQPConnectionError,
+                    AMQPChannelError,
+                    ChannelClosed,
+                    ConnectionClosed,
+                    AMQPError,
+                ),
+            )
+
+            # Verify connection is actually healthy after initialization
+            if not await self._is_connection_healthy():
+                logger.warning(
+                    "Connection health check failed after initialization, retrying..."
                 )
+                # Wait a bit and try again
+                await asyncio.sleep(2.0)
+                if not await self._is_connection_healthy():
+                    raise ConnectionError("Connection not healthy after initialization")
+
+            self.connection_healthy = True
+            logger.debug("RabbitMQ connection initialized successfully")
+
+            # Start health monitoring
+            self.health_check_task = asyncio.create_task(
+                self._monitor_connection_health()
+            )
+
+        except Exception as e:
+            logger.error("Failed to initialize RabbitMQ after retries: %s", e)
+            raise
 
     async def dispose(self) -> None:
-        await self.channel_pool.close()
-        await self.conn_pool.close()
+        """Clean up resources"""
+        logger.debug("Disposing RabbitMQ broker dispatcher")
+        self.shutdown_event.set()
+
+        # Cancel health monitoring
+        if self.health_check_task:
+            self.health_check_task.cancel()
+            try:
+                await self.health_check_task
+            except asyncio.CancelledError:
+                pass
+
+        # Clean up pools
+        await self._cleanup_pools()
+
+    async def _monitor_connection_health(self) -> None:
+        """Monitor connection health and trigger shutdown if needed"""
+        while not self.shutdown_event.is_set():
+            try:
+                await asyncio.sleep(self.config.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, triggering shutdown")
+                    self.shutdown_event.set()
+                    break
 
-def _get_message_broker_dispatcher_from_url(url: str) -> _MessageBrokerDispatcher:
+            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:
+            # Try to acquire a connection from the pool
+            async with self.conn_pool.acquire() as connection:
+                if connection.is_closed:
+                    return False
+
+                # Try to create a channel to test connection
+                channel = await connection.channel()
+                await channel.close()
+                return True
+
+        except Exception as e:
+            logger.debug("Connection health check failed: %s", e)
+            return False
+
+    async def _cleanup_pools(self) -> None:
+        """Clean up existing connection pools"""
+        try:
+            if hasattr(self, "channel_pool"):
+                await self.channel_pool.close()
+        except Exception as e:
+            logger.warning("Error closing channel pool: %s", e)
+
+        try:
+            if hasattr(self, "conn_pool"):
+                await self.conn_pool.close()
+        except Exception as e:
+            logger.warning("Error closing connection pool: %s", e)
+
+    async def _wait_for_connection(self) -> None:
+        """Wait for connection to be healthy"""
+        max_wait = 30.0  # Maximum wait time
+        wait_time = 0.0
+
+        while not self.connection_healthy and wait_time < max_wait:
+            if self.shutdown_event.is_set():
+                raise ConnectionError("Shutdown requested while waiting for connection")
+
+            await asyncio.sleep(0.5)
+            wait_time += 0.5
+
+        if not self.connection_healthy:
+            raise ConnectionError("Connection not healthy after maximum wait time")
+
+
+def _get_message_broker_dispatcher_from_url(
+    url: str,
+    config: "BeatWorkerConfig | None" = None,
+    shutdown_event: asyncio.Event | None = None,
+) -> _MessageBrokerDispatcher:
     """
     Factory function to create a message broker instance from a URL.
     Currently, only RabbitMQ is supported.
     """
     if url.startswith("amqp://") or url.startswith("amqps://"):
-        return _RabbitMQBrokerDispatcher(url=url)
+        return _RabbitMQBrokerDispatcher(
+            url=url, config=config, shutdown_event=shutdown_event
+        )
     else:
         raise ValueError(f"Unsupported message broker URL: {url}")
 
@@ -219,6 +443,39 @@ def _get_message_broker_dispatcher_from_url(url: str) -> _MessageBrokerDispatche
 # endregion
 
 
+@dataclass
+class BeatWorkerConfig:
+    """Configuration for beat worker connection resilience"""
+
+    connection_retry_config: RetryPolicy = field(
+        default_factory=lambda: RetryPolicy(
+            max_retries=10,
+            initial_delay=2.0,
+            max_delay=60.0,
+            backoff_factor=2.0,
+            jitter=True,
+        )
+    )
+    dispatch_retry_config: RetryPolicy = field(
+        default_factory=lambda: RetryPolicy(
+            max_retries=3,
+            initial_delay=1.0,
+            max_delay=10.0,
+            backoff_factor=2.0,
+            jitter=True,
+        )
+    )
+    connection_heartbeat_interval: float = 30.0
+    health_check_interval: float = 15.0
+
+    # Connection establishment timeouts
+    connection_wait_timeout: float = 300.0  # 5 minutes to wait for initial connection
+
+    # Pool configuration
+    max_pool_size: int = 10
+    pool_recycle_time: float = 3600.0  # 1 hour
+
+
 class BeatWorker:
 
     def __init__(
@@ -228,11 +485,16 @@ class BeatWorker:
         broker_url: str,
         backend_url: str,
         scheduled_action_names: set[str] | None = None,
+        config: "BeatWorkerConfig | None" = None,
+        shutdown_event: asyncio.Event | None = None,
     ) -> None:
+        self.shutdown_event = shutdown_event or asyncio.Event()
+
         self.app = app
+        self.config = config or BeatWorkerConfig()
 
         self.broker: _MessageBrokerDispatcher = _get_message_broker_dispatcher_from_url(
-            broker_url
+            broker_url, self.config, shutdown_event=self.shutdown_event
         )
         self.backend: MessageBrokerBackend = get_message_broker_backend_from_url(
             backend_url
@@ -243,41 +505,77 @@ class BeatWorker:
         self.container = Container(self.app)
         self.uow_provider = UnitOfWorkContextProvider(app, self.container)
 
-        self.shutdown_event = asyncio.Event()
-
         self.lifecycle = AppLifecycle(app, self.container)
 
     def run(self) -> None:
 
-        def on_signal_received(signal: int, frame_type: FrameType | None) -> None:
-            logger.info("Received shutdown signal")
-            self.shutdown_event.set()
-
-        signal.signal(signal.SIGINT, on_signal_received)
+        def on_shutdown(loop: asyncio.AbstractEventLoop) -> None:
+            logger.debug("Shutting down - signal received")
+            # Schedule the shutdown to run in the event loop
+            asyncio.create_task(self._graceful_shutdown())
 
         with asyncio.Runner(loop_factory=uvloop.new_event_loop) as runner:
-            runner.run(self.start_scheduler())
+            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_scheduler())
+            except Exception as e:
+                logger.critical(
+                    "Scheduler 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 start_scheduler(self) -> None:
         """
         Declares the scheduled actions and starts the scheduler.
         This is the main entry point for the scheduler.
         """
-        async with self.lifecycle():
+        with providing_app_type("beat"):
+            async with self.lifecycle():
 
-            scheduled_actions = _extract_scheduled_actions(
-                self.app, self.container, self.scheduler_names
-            )
+                scheduled_actions = _extract_scheduled_actions(
+                    self.app, self.container, self.scheduler_names
+                )
+
+                # Initialize and wait for connection to be established
+                logger.debug("Initializing broker connection...")
+                await self.broker.initialize(scheduled_actions)
 
-            await self.broker.initialize(scheduled_actions)
+                # Wait for connection to be healthy before starting scheduler
+                logger.debug("Waiting for connection to be established...")
+                await self._wait_for_broker_connection()
 
-            await self.run_scheduled_actions(scheduled_actions)
+                logger.debug("Connection established, starting scheduler...")
+                await self.run_scheduled_actions(scheduled_actions)
 
     async def run_scheduled_actions(
         self, scheduled_actions: list[ScheduledActionData]
     ) -> None:
 
+        logger.debug("Starting scheduled actions processing loop")
+
+        # Ensure we have a healthy connection before starting the main loop
+        if (
+            isinstance(self.broker, _RabbitMQBrokerDispatcher)
+            and not self.broker.connection_healthy
+        ):
+            logger.error("Connection not healthy at start of processing loop. Exiting.")
+            return
+
         while not self.shutdown_event.is_set():
+            # Check connection health before processing scheduled actions
+            if (
+                isinstance(self.broker, _RabbitMQBrokerDispatcher)
+                and not self.broker.connection_healthy
+            ):
+                logger.error("Broker connection is not healthy. Exiting.")
+                break
+
             now = int(time.time())
             for sched_act_data in scheduled_actions:
                 func = sched_act_data.callable
@@ -285,58 +583,176 @@ class BeatWorker:
                 if self.shutdown_event.is_set():
                     break
 
-                async with self.backend.lock():
+                try:
+                    async with self.backend.lock():
 
-                    last_dispatch_time: int | None = (
-                        await self.backend.get_last_dispatch_time(
-                            ScheduledAction.get_function_id(func)
+                        last_dispatch_time: int | None = (
+                            await self.backend.get_last_dispatch_time(
+                                ScheduledAction.get_function_id(func)
+                            )
                         )
-                    )
 
-                    if last_dispatch_time is not None:
-                        cron = croniter.croniter(
-                            scheduled_action.cron, last_dispatch_time
-                        )
-                        next_run: datetime = cron.get_next(datetime).replace(tzinfo=UTC)
-                        if next_run > datetime.now(UTC):
-                            logger.info(
-                                f"Skipping {func.__module__}.{func.__qualname__} until {next_run}"
+                        if last_dispatch_time is not None:
+                            cron = croniter.croniter(
+                                scheduled_action.cron, last_dispatch_time
                             )
-                            continue
+                            next_run: datetime = cron.get_next(datetime).replace(
+                                tzinfo=UTC
+                            )
+                            if next_run > datetime.now(UTC):
+                                logger.debug(
+                                    "Skipping %s.%s until %s",
+                                    func.__module__,
+                                    func.__qualname__,
+                                    next_run,
+                                )
+                                continue
+
+                        if not scheduled_action.allow_overlap:
+                            if (
+                                await self.backend.get_in_execution_count(
+                                    ScheduledAction.get_function_id(func)
+                                )
+                                > 0
+                            ):
+                                continue
+
+                        try:
+                            start_time = time.perf_counter()
+                            await self.broker.dispatch_scheduled_action(
+                                ScheduledAction.get_function_id(func),
+                                now,
+                            )
+                            elapsed_time = time.perf_counter() - start_time
 
-                    if not scheduled_action.allow_overlap:
-                        if (
-                            await self.backend.get_in_execution_count(
-                                ScheduledAction.get_function_id(func)
+                            await self.backend.set_last_dispatch_time(
+                                ScheduledAction.get_function_id(func), now
                             )
-                            > 0
-                        ):
+
+                            logger.debug(
+                                "Scheduled %s.%s at %s in %.4fs",
+                                func.__module__,
+                                func.__qualname__,
+                                now,
+                                elapsed_time,
+                            )
+                        except ChannelInvalidStateError as e:
+                            logger.error(
+                                "Channel invalid state error when dispatching %s.%s: %s",
+                                func.__module__,
+                                func.__qualname__,
+                                e,
+                            )
+                            # Trigger shutdown if dispatch fails
+                            self.shutdown_event.set()
+                            raise
+                        except Exception as e:
+                            logger.error(
+                                "Failed to dispatch scheduled action %s.%s: %s",
+                                func.__module__,
+                                func.__qualname__,
+                                e,
+                            )
+                            # Continue with other scheduled actions even if one fails
                             continue
 
-                    await self.broker.dispatch_scheduled_action(
-                        ScheduledAction.get_function_id(func),
-                        now,
+                except Exception as e:
+                    logger.error(
+                        "Error processing scheduled action %s.%s: %s",
+                        func.__module__,
+                        func.__qualname__,
+                        e,
                     )
+                    # Continue with other scheduled actions even if one fails
+                    continue
+
+            # Handle delayed messages
+            try:
+                delayed_messages = await self.backend.dequeue_next_delayed_messages(now)
+                for delayed_message_data in delayed_messages:
+                    try:
+                        start_time = time.perf_counter()
+                        await self.broker.dispatch_delayed_message(delayed_message_data)
+                        elapsed_time = time.perf_counter() - start_time
+                        logger.debug(
+                            "Dispatched delayed message for topic %s in %.4fs",
+                            delayed_message_data.message_topic,
+                            elapsed_time,
+                        )
+                    except Exception as e:
+                        logger.error("Failed to dispatch delayed message: %s", e)
+                        # Continue with other delayed messages even if one fails
+                        continue
+            except Exception as e:
+                logger.error("Error processing delayed messages: %s", e)
 
-                    await self.backend.set_last_dispatch_time(
-                        ScheduledAction.get_function_id(func), now
-                    )
+            with contextlib.suppress(asyncio.TimeoutError):
+                await asyncio.wait_for(self.shutdown_event.wait(), self.interval)
 
-                    logger.info(
-                        f"Scheduled {func.__module__}.{func.__qualname__} at {now}"
-                    )
+        logger.debug("Scheduler stopped")
 
-            for (
-                delayed_message_data
-            ) in await self.backend.dequeue_next_delayed_messages(now):
-                await self.broker.dispatch_delayed_message(delayed_message_data)
+        try:
+            await self.backend.dispose()
+        except Exception as e:
+            logger.error("Error disposing backend: %s", e)
 
-            with contextlib.suppress(asyncio.TimeoutError):
-                await asyncio.wait_for(self.shutdown_event.wait(), self.interval)
+        try:
+            await self.broker.dispose()
+        except Exception as e:
+            logger.error("Error disposing broker: %s", e)
 
-            # await self.shutdown_event.wait(self.interval)
+    async def _graceful_shutdown(self) -> None:
+        """Handles graceful shutdown process"""
+        logger.debug("Initiating graceful shutdown sequence")
+        self.shutdown_event.set()
+        logger.debug("Graceful shutdown completed")
 
-        logger.info("Scheduler stopped")
+    async def _wait_for_broker_connection(self) -> None:
+        """
+        Wait for the broker connection to be established and healthy.
+        This ensures the scheduler doesn't start until RabbitMQ is ready.
+        """
+        max_wait_time = self.config.connection_wait_timeout
+        check_interval = 2.0  # Check every 2 seconds
+        elapsed_time = 0.0
+
+        logger.debug(
+            "Waiting for broker connection to be established (timeout: %ss)...",
+            max_wait_time,
+        )
+
+        while elapsed_time < max_wait_time:
+            if self.shutdown_event.is_set():
+                raise ConnectionError(
+                    "Shutdown requested while waiting for broker connection"
+                )
+
+            # Check if broker connection is healthy
+            if (
+                isinstance(self.broker, _RabbitMQBrokerDispatcher)
+                and self.broker.connection_healthy
+            ):
+                logger.debug("Broker connection is healthy")
+                return
+
+            # If broker doesn't have health status, try a simple health check
+            if not hasattr(self.broker, "connection_healthy"):
+                try:
+                    # For non-RabbitMQ brokers, assume connection is ready after initialization
+                    logger.debug("Broker connection assumed to be ready")
+                    return
+                except Exception as e:
+                    logger.debug("Broker connection check failed: %s", e)
+
+            if elapsed_time % 10.0 == 0.0:  # Log every 10 seconds
+                logger.warning(
+                    "Still waiting for broker connection... (%.1fs elapsed)",
+                    elapsed_time,
+                )
 
-        await self.backend.dispose()
-        await self.broker.dispose()
+            await asyncio.sleep(check_interval)
+            elapsed_time += check_interval
+
+        raise ConnectionError(
+            f"Broker connection not established after {max_wait_time} seconds"
+        )