jararaca 0.3.16__py3-none-any.whl → 0.3.18__py3-none-any.whl

jararaca/messagebus/decorators.py

@@ -69,7 +69,7 @@ class MessageHandler(Generic[INHERITS_MESSAGE_CO]):
 class MessageHandlerData:
     message_type: type[Any]
     spec: MessageHandler[Message]
-    instance_callable: Callable[[MessageOf[Any]], Awaitable[None]]
+    instance_callable: Callable[..., Awaitable[None]]
     controller_member: ControllerMemberReflect
 
 

jararaca/messagebus/worker.py

@@ -188,11 +188,8 @@ class AioPikaMicroserviceConsumer(MessageBusConsumer):
         # Connection resilience attributes
         self.connection_healthy = False
         self.connection_lock = asyncio.Lock()
-        self.reconnection_event = asyncio.Event()
-        self.reconnection_in_progress = False
         self.consumer_tags: dict[str, str] = {}  # Track consumer tags for cleanup
         self.health_check_task: asyncio.Task[Any] | None = None
-        self.reconnection_task: asyncio.Task[Any] | None = None
 
     async def _verify_infrastructure(self) -> bool:
         """
@@ -229,10 +226,6 @@ class AioPikaMicroserviceConsumer(MessageBusConsumer):
         routing_key = f"{handler.message_type.MESSAGE_TOPIC}.#"
 
         async def setup_consumer() -> None:
-            # Wait for connection to be healthy if reconnection is in progress
-            if self.reconnection_in_progress:
-                await self.reconnection_event.wait()
-
             # Create a channel using the context manager
             async with self.create_channel(queue_name) as channel:
                 queue = await RabbitmqUtils.get_queue(
@@ -289,10 +282,6 @@ class AioPikaMicroserviceConsumer(MessageBusConsumer):
         routing_key = queue_name
 
         async def setup_consumer() -> None:
-            # Wait for connection to be healthy if reconnection is in progress
-            if self.reconnection_in_progress:
-                await self.reconnection_event.wait()
-
             # Create a channel using the context manager
             async with self.create_channel(queue_name) as channel:
                 queue = await RabbitmqUtils.get_queue(
@@ -341,106 +330,107 @@ class AioPikaMicroserviceConsumer(MessageBusConsumer):
         Main consume method that sets up all message handlers and scheduled actions with retry mechanisms.
         """
         # Establish initial connection
-        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()
-            )
+        try:
+            async with self.connect() as connection:
+                self.connection_healthy = True
 
-            # Verify infrastructure with retry
-            infra_check_success = await retry_with_backoff(
-                self._verify_infrastructure,
-                retry_config=self.config.connection_retry_config,
-                retry_exceptions=(Exception,),
-            )
+                # Start connection health monitoring
+                self.health_check_task = asyncio.create_task(
+                    self._monitor_connection_health()
+                )
 
-            if not infra_check_success:
-                logger.critical(
-                    "Failed to verify RabbitMQ infrastructure. Shutting down."
+                # Verify infrastructure with retry
+                infra_check_success = await retry_with_backoff(
+                    self._verify_infrastructure,
+                    retry_config=self.config.connection_retry_config,
+                    retry_exceptions=(Exception,),
                 )
-                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
+                if not infra_check_success:
+                    logger.critical(
+                        "Failed to verify RabbitMQ infrastructure. Shutting down."
+                    )
+                    self.shutdown_event.set()
+                    return
 
-            tasks: set[asyncio.Task[tuple[str, str, bool]]] = set()
+                async def wait_for(
+                    type: str, name: str, coroutine: Awaitable[bool]
+                ) -> tuple[str, str, bool]:
+                    return type, name, await coroutine
 
-            # 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: set[asyncio.Task[tuple[str, str, bool]]] = set()
 
-                tasks.add(
-                    task := asyncio.create_task(
-                        wait_for(
-                            "message_handler",
-                            queue_name,
-                            self._setup_message_handler_consumer(handler),
-                        )
-                    )
-                )
+                # 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
 
-            # 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),
+                    tasks.add(
+                        task := asyncio.create_task(
+                            wait_for(
+                                "message_handler",
+                                queue_name,
+                                self._setup_message_handler_consumer(handler),
+                            )
                         )
                     )
-                )
 
-            async def handle_task_results() -> None:
-                for task in asyncio.as_completed(tasks):
-                    type, name, success = await task
-                    if success:
-                        logger.info(f"Successfully set up {type} consumer for {name}")
-                    else:
-                        logger.warning(
-                            f"Failed to set up {type} consumer for {name}, will not process messages from this queue"
+                # 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),
+                            )
                         )
+                    )
 
-            handle_task_results_task = asyncio.create_task(handle_task_results())
-
-            # Wait for shutdown signal
-            await self.shutdown_event.wait()
-            logger.info("Shutdown event received, stopping consumers")
+                async def handle_task_results() -> None:
+                    for task in asyncio.as_completed(tasks):
+                        type, name, success = await task
+                        if success:
+                            logger.info(
+                                f"Successfully set up {type} consumer for {name}"
+                            )
+                        else:
+                            logger.warning(
+                                f"Failed to set up {type} consumer for {name}, will not process messages from this queue"
+                            )
 
-            # 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 = asyncio.create_task(handle_task_results())
 
-            # Cancel reconnection task if running
-            if self.reconnection_task:
-                self.reconnection_task.cancel()
-                with suppress(asyncio.CancelledError):
-                    await self.reconnection_task
+                # Wait for shutdown signal
+                await self.shutdown_event.wait()
+                logger.info("Shutdown event received, stopping consumers")
 
-            handle_task_results_task.cancel()
-            with suppress(asyncio.CancelledError):
-                await handle_task_results_task
-            for task in tasks:
-                if not task.done():
-                    task.cancel()
+                # Cancel health monitoring
+                if self.health_check_task:
+                    self.health_check_task.cancel()
                     with suppress(asyncio.CancelledError):
-                        await task
-            logger.info("Worker shutting down")
+                        await self.health_check_task
 
-            # Wait for all tasks to complete
-            await self.wait_all_tasks_done()
-
-            # Close all channels and the connection
-            await self.close_channels_and_connection()
+                handle_task_results_task.cancel()
+                with suppress(asyncio.CancelledError):
+                    await handle_task_results_task
+                for task in tasks:
+                    if not task.done():
+                        task.cancel()
+                        with suppress(asyncio.CancelledError):
+                            await task
+                logger.info("Worker shutting down")
+
+                # Wait for all tasks to complete
+                await self.wait_all_tasks_done()
+
+                # Close all channels and the connection
+                await self.close_channels_and_connection()
+        except Exception as e:
+            logger.critical(f"Failed to establish initial connection to RabbitMQ: {e}")
+            # Re-raise the exception so it can be caught by the caller
+            raise
 
     async def wait_all_tasks_done(self) -> None:
         if not self.tasks:
@@ -478,12 +468,6 @@ class AioPikaMicroserviceConsumer(MessageBusConsumer):
             with suppress(asyncio.CancelledError):
                 await self.health_check_task
 
-        # Cancel reconnection task if running
-        if self.reconnection_task:
-            self.reconnection_task.cancel()
-            with suppress(asyncio.CancelledError):
-                await self.reconnection_task
-
         await self.wait_all_tasks_done()
         await self.close_channels_and_connection()
 
@@ -492,16 +476,6 @@ class AioPikaMicroserviceConsumer(MessageBusConsumer):
         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 reconnection is in progress, wait for it to complete
-        if self.reconnection_in_progress:
-            try:
-                await asyncio.wait_for(self.reconnection_event.wait(), timeout=30.0)
-            except asyncio.TimeoutError:
-                logger.warning(
-                    f"Timeout waiting for reconnection when getting channel for {queue_name}"
-                )
-                return None
-
         if queue_name not in self.channels:
             logger.warning(f"No channel found for queue {queue_name}")
             return None
@@ -530,17 +504,17 @@ class AioPikaMicroserviceConsumer(MessageBusConsumer):
                         logger.error(
                             f"Failed to recreate channel for {queue_name}: {e}"
                         )
-                        # Trigger reconnection if channel creation fails
+                        # Trigger shutdown if channel creation fails
                         self._trigger_reconnection()
                         return None
                 else:
-                    # Connection is not healthy, trigger reconnection
+                    # Connection is not healthy, trigger shutdown
                     self._trigger_reconnection()
                     return None
             return channel
         except Exception as e:
             logger.error(f"Error accessing channel for queue {queue_name}: {e}")
-            # Trigger reconnection on any channel access error
+            # Trigger shutdown on any channel access error
             self._trigger_reconnection()
             return None
 
@@ -691,33 +665,14 @@ class AioPikaMicroserviceConsumer(MessageBusConsumer):
                         yield new_channel
                         return
                 else:
-                    # Connection is not healthy, wait for reconnection
-                    if self.reconnection_in_progress:
-                        try:
-                            await asyncio.wait_for(
-                                self.reconnection_event.wait(), timeout=30.0
-                            )
-                            # Retry after reconnection
-                            continue
-                        except asyncio.TimeoutError:
-                            logger.warning(
-                                f"Timeout waiting for reconnection for queue {queue_name}"
-                            )
-
-                    # Still no connection, trigger reconnection
-                    if not self.reconnection_in_progress:
-                        self._trigger_reconnection()
-
-                    if attempt < max_retries - 1:
-                        logger.info(
-                            f"Retrying channel access for {queue_name} in {retry_delay}s"
-                        )
-                        await asyncio.sleep(retry_delay)
-                        retry_delay *= 2
-                    else:
-                        raise RuntimeError(
-                            f"Cannot get channel for queue {queue_name}: no connection available after {max_retries} attempts"
-                        )
+                    # Connection is not healthy, trigger shutdown
+                    logger.error(
+                        f"Connection not healthy while getting channel for {queue_name}, triggering shutdown"
+                    )
+                    self._trigger_reconnection()
+                    raise RuntimeError(
+                        f"Cannot get channel for queue {queue_name}: connection is not healthy"
+                    )
 
             except Exception as e:
                 if attempt < max_retries - 1:
@@ -734,7 +689,7 @@ class AioPikaMicroserviceConsumer(MessageBusConsumer):
 
     async def _monitor_connection_health(self) -> None:
         """
-        Monitor connection health and trigger reconnection if needed.
+        Monitor connection health and trigger shutdown if connection is lost.
         This runs as a background task.
         """
         while not self.shutdown_event.is_set():
@@ -746,11 +701,11 @@ class AioPikaMicroserviceConsumer(MessageBusConsumer):
 
                 # Check connection health
                 if not await self._is_connection_healthy():
-                    logger.warning(
-                        "Connection health check failed, triggering reconnection"
+                    logger.error(
+                        "Connection health check failed, initiating worker shutdown"
                     )
-                    if not self.reconnection_in_progress:
-                        self._trigger_reconnection()
+                    self.shutdown()
+                    break
 
             except asyncio.CancelledError:
                 logger.info("Connection health monitoring cancelled")
@@ -778,74 +733,12 @@ class AioPikaMicroserviceConsumer(MessageBusConsumer):
 
     def _trigger_reconnection(self) -> None:
         """
-        Trigger reconnection process.
+        Trigger worker shutdown due to connection loss.
         """
-        if not self.reconnection_in_progress and not self.shutdown_event.is_set():
-            self.reconnection_in_progress = True
+        if not self.shutdown_event.is_set():
+            logger.error("Connection lost, initiating worker shutdown")
             self.connection_healthy = False
-            self.reconnection_event.clear()
-
-            # Start reconnection task
-            self.reconnection_task = asyncio.create_task(self._handle_reconnection())
-            self.reconnection_task.add_done_callback(self._on_reconnection_done)
-
-    def _on_reconnection_done(self, task: asyncio.Task[Any]) -> None:
-        """
-        Handle completion of reconnection task.
-        """
-        self.reconnection_in_progress = False
-        if task.exception():
-            logger.error(f"Reconnection task failed: {task.exception()}")
-        else:
-            logger.info("Reconnection completed successfully")
-
-    async def _handle_reconnection(self) -> None:
-        """
-        Handle the reconnection process with exponential backoff.
-        """
-        logger.info("Starting reconnection process")
-
-        # Close existing connection and channels
-        await self._cleanup_connection()
-
-        reconnection_config = self.config.reconnection_backoff_config
-        attempt = 0
-
-        while not self.shutdown_event.is_set():
-            try:
-                attempt += 1
-                logger.info(f"Reconnection attempt {attempt}")
-
-                # Establish new connection
-                self.connection = await self._establish_connection()
-                self.connection_healthy = True
-
-                # Re-establish all consumers
-                await self._reestablish_consumers()
-
-                logger.info("Reconnection successful")
-                self.reconnection_event.set()
-                return
-
-            except Exception as e:
-                logger.error(f"Reconnection attempt {attempt} failed: {e}")
-
-                if self.shutdown_event.is_set():
-                    break
-
-                # Calculate backoff delay
-                delay = reconnection_config.initial_delay * (
-                    reconnection_config.backoff_factor ** (attempt - 1)
-                )
-                if reconnection_config.jitter:
-                    jitter_amount = delay * 0.25
-                    delay = delay + random.uniform(-jitter_amount, jitter_amount)
-                    delay = max(delay, 0.1)
-
-                delay = min(delay, reconnection_config.max_delay)
-
-                logger.info(f"Retrying reconnection in {delay:.2f} seconds")
-                await asyncio.sleep(delay)
+            self.shutdown()
 
     async def _cleanup_connection(self) -> None:
         """
@@ -889,32 +782,6 @@ class AioPikaMicroserviceConsumer(MessageBusConsumer):
         self.connection = None
         self.connection_healthy = False
 
-    async def _reestablish_consumers(self) -> None:
-        """
-        Re-establish all consumers after reconnection.
-        """
-        logger.info("Re-establishing consumers after reconnection")
-
-        # Re-establish 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__}"
-            try:
-                await self._setup_message_handler_consumer(handler)
-                logger.info(f"Re-established consumer for {queue_name}")
-            except Exception as e:
-                logger.error(f"Failed to re-establish consumer for {queue_name}: {e}")
-
-        # Re-establish scheduled actions
-        for scheduled_action in self.scheduled_actions:
-            queue_name = f"{scheduled_action.callable.__module__}.{scheduled_action.callable.__qualname__}"
-            try:
-                await self._setup_scheduled_action_consumer(scheduled_action)
-                logger.info(f"Re-established scheduler consumer for {queue_name}")
-            except Exception as e:
-                logger.error(
-                    f"Failed to re-establish scheduler consumer for {queue_name}: {e}"
-                )
-
 
 def create_message_bus(
     broker_url: str,
@@ -1798,7 +1665,14 @@ class MessageBusWorker:
             loop.add_signal_handler(signal.SIGINT, on_shutdown, loop)
             # Add graceful shutdown handler for SIGTERM as well
             loop.add_signal_handler(signal.SIGTERM, on_shutdown, loop)
-            runner.run(self.start_async())
+            try:
+                runner.run(self.start_async())
+            except Exception as e:
+                logger.critical(f"Worker failed to start due to connection error: {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"""

jararaca/presentation/websocket/redis.py

@@ -59,6 +59,7 @@ class RedisWebSocketConnectionBackend(WebSocketConnectionBackend):
         consume_broadcast_timeout: int = 1,
         consume_send_timeout: int = 1,
         retry_delay: float = 5.0,
+        max_concurrent_tasks: int = 1000,
     ) -> None:
 
         self.redis = conn
@@ -67,6 +68,8 @@ class RedisWebSocketConnectionBackend(WebSocketConnectionBackend):
 
         self.lock = asyncio.Lock()
         self.tasks: set[asyncio.Task[Any]] = set()
+        self.max_concurrent_tasks = max_concurrent_tasks
+        self.task_semaphore = asyncio.Semaphore(max_concurrent_tasks)
 
         self.consume_broadcast_timeout = consume_broadcast_timeout
         self.consume_send_timeout = consume_send_timeout
@@ -101,16 +104,26 @@ class RedisWebSocketConnectionBackend(WebSocketConnectionBackend):
         return self.__broadcast_func
 
     async def broadcast(self, message: bytes) -> None:
-        await self.redis.publish(
-            self.broadcast_pubsub_channel,
-            BroadcastMessage.from_message(message).encode(),
-        )
+        try:
+            await self.redis.publish(
+                self.broadcast_pubsub_channel,
+                BroadcastMessage.from_message(message).encode(),
+            )
+        except Exception as e:
+            logger.error(
+                f"Failed to publish broadcast message to Redis: {e}", exc_info=True
+            )
+            raise
 
     async def send(self, rooms: list[str], message: bytes) -> None:
-        await self.redis.publish(
-            self.send_pubsub_channel,
-            SendToRoomsMessage.from_message(rooms, message).encode(),
-        )
+        try:
+            await self.redis.publish(
+                self.send_pubsub_channel,
+                SendToRoomsMessage.from_message(rooms, message).encode(),
+            )
+        except Exception as e:
+            logger.error(f"Failed to publish send message to Redis: {e}", exc_info=True)
+            raise
 
     def configure(
         self, broadcast: BroadcastFunc, send: SendFunc, shutdown_event: asyncio.Event
@@ -129,7 +142,12 @@ class RedisWebSocketConnectionBackend(WebSocketConnectionBackend):
             self.consume_send(self.send_func, self.shutdown_event)
         )
 
-        self.tasks.add(send_task)
+        # Use lock when modifying tasks set to prevent race conditions
+        async def add_task() -> None:
+            async with self.lock:
+                self.tasks.add(send_task)
+
+        asyncio.get_event_loop().create_task(add_task())
         send_task.add_done_callback(self.handle_send_task_done)
 
     def setup_broadcast_consumer(self) -> None:
@@ -138,11 +156,23 @@ class RedisWebSocketConnectionBackend(WebSocketConnectionBackend):
             self.consume_broadcast(self.broadcast_func, self.shutdown_event)
         )
 
-        self.tasks.add(broadcast_task)
+        # Use lock when modifying tasks set to prevent race conditions
+        async def add_task() -> None:
+            async with self.lock:
+                self.tasks.add(broadcast_task)
+
+        asyncio.get_event_loop().create_task(add_task())
 
         broadcast_task.add_done_callback(self.handle_broadcast_task_done)
 
     def handle_broadcast_task_done(self, task: asyncio.Task[Any]) -> None:
+        # Remove task from set safely with lock
+        async def remove_task() -> None:
+            async with self.lock:
+                self.tasks.discard(task)
+
+        asyncio.get_event_loop().create_task(remove_task())
+
         if task.cancelled():
             logger.warning("Broadcast task was cancelled.")
         elif task.exception() is not None:
@@ -162,6 +192,13 @@ class RedisWebSocketConnectionBackend(WebSocketConnectionBackend):
             )
 
     def handle_send_task_done(self, task: asyncio.Task[Any]) -> None:
+        # Remove task from set safely with lock
+        async def remove_task() -> None:
+            async with self.lock:
+                self.tasks.discard(task)
+
+        asyncio.get_event_loop().create_task(remove_task())
+
         if task.cancelled():
             logger.warning("Send task was cancelled.")
         elif task.exception() is not None:
@@ -204,54 +241,132 @@ class RedisWebSocketConnectionBackend(WebSocketConnectionBackend):
         self, broadcast: BroadcastFunc, shutdown_event: asyncio.Event
     ) -> None:
         logger.info("Starting broadcast consumer...")
-        async with self.redis.pubsub() as pubsub:
-            await pubsub.subscribe(self.broadcast_pubsub_channel)
-
-            while not shutdown_event.is_set():
-                message: dict[str, Any] | None = await pubsub.get_message(
-                    ignore_subscribe_messages=True,
-                    timeout=self.consume_broadcast_timeout,
+        try:
+            # Validate Redis connection before starting
+            try:
+                await self.redis.ping()
+                logger.info("Redis connection validated for broadcast consumer")
+            except Exception as e:
+                logger.error(f"Redis connection validation failed: {e}", exc_info=True)
+                raise
+
+            async with self.redis.pubsub() as pubsub:
+                await pubsub.subscribe(self.broadcast_pubsub_channel)
+                logger.info(
+                    f"Subscribed to broadcast channel: {self.broadcast_pubsub_channel}"
                 )
 
-                if message is None:
-                    continue
-
-                broadcast_message = BroadcastMessage.decode(message["data"])
-
-                async with self.lock:
-                    task = asyncio.get_event_loop().create_task(
-                        broadcast(message=broadcast_message.message)
+                while not shutdown_event.is_set():
+                    message: dict[str, Any] | None = await pubsub.get_message(
+                        ignore_subscribe_messages=True,
+                        timeout=self.consume_broadcast_timeout,
                     )
 
-                    self.tasks.add(task)
-
-                    task.add_done_callback(self.tasks.discard)
+                    if message is None:
+                        continue
+
+                    broadcast_message = BroadcastMessage.decode(message["data"])
+
+                    # Use semaphore for backpressure control
+                    acquired = False
+                    try:
+                        await self.task_semaphore.acquire()
+                        acquired = True
+
+                        async def broadcast_with_cleanup(msg: bytes) -> None:
+                            try:
+                                await broadcast(message=msg)
+                            finally:
+                                self.task_semaphore.release()
+
+                        async with self.lock:
+                            task = asyncio.get_event_loop().create_task(
+                                broadcast_with_cleanup(broadcast_message.message)
+                            )
+
+                            self.tasks.add(task)
+
+                            task.add_done_callback(self.tasks.discard)
+                    except Exception as e:
+                        # Release semaphore if we acquired it but failed to create task
+                        if acquired:
+                            self.task_semaphore.release()
+                        logger.error(
+                            f"Error processing broadcast message: {e}", exc_info=True
+                        )
+                        # Continue processing other messages
+                        continue
+        except Exception as e:
+            logger.error(
+                f"Fatal error in broadcast consumer, will retry: {e}", exc_info=True
+            )
+            raise
 
     async def consume_send(self, send: SendFunc, shutdown_event: asyncio.Event) -> None:
         logger.info("Starting send consumer...")
-        async with self.redis.pubsub() as pubsub:
-            await pubsub.subscribe(self.send_pubsub_channel)
-
-            while not shutdown_event.is_set():
-
-                message: dict[str, Any] | None = await pubsub.get_message(
-                    ignore_subscribe_messages=True, timeout=self.consume_send_timeout
-                )
-
-                if message is None:
-                    continue
-
-                send_message = SendToRoomsMessage.decode(message["data"])
-
-                async with self.lock:
-
-                    task = asyncio.get_event_loop().create_task(
-                        send(send_message.rooms, send_message.message)
+        try:
+            # Validate Redis connection before starting
+            try:
+                await self.redis.ping()
+                logger.info("Redis connection validated for send consumer")
+            except Exception as e:
+                logger.error(f"Redis connection validation failed: {e}", exc_info=True)
+                raise
+
+            async with self.redis.pubsub() as pubsub:
+                await pubsub.subscribe(self.send_pubsub_channel)
+                logger.info(f"Subscribed to send channel: {self.send_pubsub_channel}")
+
+                while not shutdown_event.is_set():
+                    message: dict[str, Any] | None = await pubsub.get_message(
+                        ignore_subscribe_messages=True,
+                        timeout=self.consume_send_timeout,
                     )
 
-                    self.tasks.add(task)
-
-                    task.add_done_callback(self.tasks.discard)
+                    if message is None:
+                        continue
+
+                    send_message = SendToRoomsMessage.decode(message["data"])
+
+                    # Use semaphore for backpressure control
+                    acquired = False
+                    try:
+                        await self.task_semaphore.acquire()
+                        acquired = True
+
+                        async def send_with_cleanup(
+                            rooms: list[str], msg: bytes
+                        ) -> None:
+                            try:
+                                await send(rooms, msg)
+                            finally:
+                                self.task_semaphore.release()
+
+                        async with self.lock:
+
+                            task = asyncio.get_event_loop().create_task(
+                                send_with_cleanup(
+                                    send_message.rooms, send_message.message
+                                )
+                            )
+
+                            self.tasks.add(task)
+
+                            task.add_done_callback(self.tasks.discard)
+                    except Exception as e:
+                        # Release semaphore if we acquired it but failed to create task
+                        if acquired:
+                            self.task_semaphore.release()
+                        logger.error(
+                            f"Error processing send message: {e}", exc_info=True
+                        )
+                        # Continue processing other messages
+                        continue
+        except Exception as e:
+            logger.error(
+                f"Fatal error in send consumer, will retry: {e}", exc_info=True
+            )
+            raise
 
     async def shutdown(self) -> None:
         async with self.lock:

jararaca/presentation/websocket/websocket_interceptor.py

@@ -85,13 +85,24 @@ class WebSocketConnectionManagerImpl(WebSocketConnectionManager):
         await self.backend.broadcast(message)
 
     async def _broadcast_from_backend(self, message: bytes) -> None:
-        for websocket in self.all_websockets:
+        # Create a copy of the websockets set to avoid modification during iteration
+        async with self.lock:
+            websockets_to_send = list(self.all_websockets)
+
+        disconnected_websockets: list[WebSocket] = []
+
+        for websocket in websockets_to_send:
             try:
                 if websocket.client_state == WebSocketState.CONNECTED:
                     await websocket.send_bytes(message)
             except WebSocketDisconnect:
-                async with self.lock:  # TODO: check if this can cause concurrency slowdown issues
-                    self.all_websockets.remove(websocket)
+                disconnected_websockets.append(websocket)
+
+        # Clean up disconnected websockets in a single lock acquisition
+        if disconnected_websockets:
+            async with self.lock:
+                for websocket in disconnected_websockets:
+                    self.all_websockets.discard(websocket)
 
     async def send(self, rooms: list[str], message: WebSocketMessageBase) -> None:
 
@@ -103,16 +114,28 @@ class WebSocketConnectionManagerImpl(WebSocketConnectionManager):
         )
 
     async def _send_from_backend(self, rooms: list[str], message: bytes) -> None:
+        # Create a copy of room memberships to avoid modification during iteration
+        async with self.lock:
+            room_websockets: dict[str, list[WebSocket]] = {
+                room: list(self.rooms.get(room, set())) for room in rooms
+            }
+
+        disconnected_by_room: dict[str, list[WebSocket]] = {room: [] for room in rooms}
+
+        for room, websockets in room_websockets.items():
+            for websocket in websockets:
+                try:
+                    if websocket.client_state == WebSocketState.CONNECTED:
+                        await websocket.send_bytes(message)
+                except WebSocketDisconnect:
+                    disconnected_by_room[room].append(websocket)
+
+        # Clean up disconnected websockets in a single lock acquisition
         async with self.lock:
-            for room in rooms:
-                for websocket in self.rooms.get(room, set()):
-                    try:
-                        if websocket.client_state == WebSocketState.CONNECTED:
-                            await websocket.send_bytes(message)
-                    except WebSocketDisconnect:
-                        async with self.lock:
-                            if websocket in self.rooms[room]:
-                                self.rooms[room].remove(websocket)
+            for room, disconnected_websockets in disconnected_by_room.items():
+                if room in self.rooms:
+                    for websocket in disconnected_websockets:
+                        self.rooms[room].discard(websocket)
 
     async def join(self, rooms: list[str], websocket: WebSocket) -> None:
         for room in rooms:

{jararaca-0.3.16.dist-info → jararaca-0.3.18.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: jararaca
-Version: 0.3.16
+Version: 0.3.18
 Summary: A simple and fast API framework for Python
 Home-page: https://github.com/LuscasLeo/jararaca
 Author: Lucas S

{jararaca-0.3.16.dist-info → jararaca-0.3.18.dist-info}/RECORD

@@ -1,6 +1,6 @@
 LICENSE,sha256=OXLcl0T2SZ8Pmy2_dmlvKuetivmyPd5m1q-Gyd-zaYY,35149
 README.md,sha256=YmCngjU8llW0l7L3tuXkkfr8qH7V9aBMgfp2jEzeiKg,3517
-pyproject.toml,sha256=Qn_u7qKJbAetHjXlx0gwVX_gnib76NkIQd6zQMYfkro,2832
+pyproject.toml,sha256=qmLCDQEsY1Me9LPJoJM7-ZQM5qJZKE_uzyOazfgOP4I,2832
 jararaca/__init__.py,sha256=IMnvfDoyNWTGVittF_wq2Uxtv_BY_wLN5Om6C3vUsCw,22302
 jararaca/__main__.py,sha256=-O3vsB5lHdqNFjUtoELDF81IYFtR-DSiiFMzRaiSsv4,67
 jararaca/broker_backend/__init__.py,sha256=GzEIuHR1xzgCJD4FE3harNjoaYzxHMHoEL0_clUaC-k,3528
@@ -17,13 +17,13 @@ jararaca/lifecycle.py,sha256=qKlzLQQioS8QkxNJ_FC_5WbmT77cNbc_S7OcQeOoHkI,1895
 jararaca/messagebus/__init__.py,sha256=5jAqPqdcEMYBfQyfZDWPnplYdrfMyJLMcacf3qLyUhk,56
 jararaca/messagebus/bus_message_controller.py,sha256=Xd_qwnX5jUvgBTCarHR36fvtol9lPTsYp2IIGKyQQaE,1487
 jararaca/messagebus/consumers/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-jararaca/messagebus/decorators.py,sha256=P5z0BBL4hJfgCJHKbsDhSZmlRwzV4PlCPKvEZ4sguXM,6013
+jararaca/messagebus/decorators.py,sha256=71ZvZ5slKruAsTaUHzr2f-D8yaUFl9Xxh9chNLaAs9E,6000
 jararaca/messagebus/interceptors/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 jararaca/messagebus/interceptors/aiopika_publisher_interceptor.py,sha256=_DEHwIH9LYsA26Hu1mo9oHzLZuATgjilU9E3o-ecDjs,6520
 jararaca/messagebus/interceptors/publisher_interceptor.py,sha256=ojy1bRhqMgrkQljcGGS8cd8-8pUjL8ZHjIUkdmaAnNM,1325
 jararaca/messagebus/message.py,sha256=U6cyd2XknX8mtm0333slz5fanky2PFLWCmokAO56vvU,819
 jararaca/messagebus/publisher.py,sha256=JTkxdKbvxvDWT8nK8PVEyyX061vYYbKQMxRHXrZtcEY,2173
-jararaca/messagebus/worker.py,sha256=TdWdKHFUc3AAsqhLZI08Y3uc5iqs0GtDHmEhyqYWObI,73968
+jararaca/messagebus/worker.py,sha256=Feh7b3z68sj7a94z0lwMowLgU41pFF2PGgefhdH8rw8,68926
 jararaca/microservice.py,sha256=OYCw5C4797X_tVnM_9sEz8BdjbICPHSVsCixsA_FwE4,11419
 jararaca/observability/decorators.py,sha256=MOIr2PttPYYvRwEdfQZEwD5RxKHOTv8UEy9n1YQVoKw,2281
 jararaca/observability/interceptor.py,sha256=U4ZLM0f8j6Q7gMUKKnA85bnvD-Qa0ii79Qa_X8KsXAQ,1498
@@ -47,9 +47,9 @@ jararaca/presentation/websocket/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm
 jararaca/presentation/websocket/base_types.py,sha256=AvUeeZ1TFhSiRMcYqZU1HaQNqSrcgTkC5R0ArP5dGmA,146
 jararaca/presentation/websocket/context.py,sha256=A6K5W3kqo9Hgeh1m6JiI7Cdz5SfbXcaICSVX7u1ARZo,1903
 jararaca/presentation/websocket/decorators.py,sha256=ZNd5aoA9UkyfHOt1C8D2Ffy2gQUNDEsusVnQuTgExgs,2157
-jararaca/presentation/websocket/redis.py,sha256=XG_kfr-msgEdfIZkD5JB_GH5lWoEHXwTzpQNVOONvfc,8985
+jararaca/presentation/websocket/redis.py,sha256=1vykr3mcdSDGpSu1rbb4vGnUZNZEvjRfXlIR7TiSho8,13931
 jararaca/presentation/websocket/types.py,sha256=M8snAMSdaQlKrwEM2qOgF2qrefo5Meio_oOw620Joc8,308
-jararaca/presentation/websocket/websocket_interceptor.py,sha256=JWn_G8Q2WO0-1kmN7-Gv0HkIM6nZ_yjCdGRuXUS8F7A,9191
+jararaca/presentation/websocket/websocket_interceptor.py,sha256=c5q8sBi82jXidK4m9KJo-OXmwb-nKsW-dK1DfRqJnlc,10124
 jararaca/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 jararaca/reflect/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 jararaca/reflect/controller_inspect.py,sha256=UtV4pRIOqCoK4ogBTXQE0dyopEQ5LDFhwm-1iJvrkJc,2326
@@ -74,8 +74,8 @@ jararaca/tools/typescript/interface_parser.py,sha256=yOSuOXKOeG0soGFo0fKiZIabu4Y
 jararaca/utils/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 jararaca/utils/rabbitmq_utils.py,sha256=ytdAFUyv-OBkaVnxezuJaJoLrmN7giZgtKeet_IsMBs,10918
 jararaca/utils/retry.py,sha256=DzPX_fXUvTqej6BQ8Mt2dvLo9nNlTBm7Kx2pFZ26P2Q,4668
-jararaca-0.3.16.dist-info/LICENSE,sha256=OXLcl0T2SZ8Pmy2_dmlvKuetivmyPd5m1q-Gyd-zaYY,35149
-jararaca-0.3.16.dist-info/METADATA,sha256=mIIZae49Gn54udR4m7qk3ZwTjx66iSQpgIp9Abm_Ymo,5149
-jararaca-0.3.16.dist-info/WHEEL,sha256=sP946D7jFCHeNz5Iq4fL4Lu-PrWrFsgfLXbbkciIZwg,88
-jararaca-0.3.16.dist-info/entry_points.txt,sha256=WIh3aIvz8LwUJZIDfs4EeH3VoFyCGEk7cWJurW38q0I,45
-jararaca-0.3.16.dist-info/RECORD,,
+jararaca-0.3.18.dist-info/LICENSE,sha256=OXLcl0T2SZ8Pmy2_dmlvKuetivmyPd5m1q-Gyd-zaYY,35149
+jararaca-0.3.18.dist-info/METADATA,sha256=5BT5cExvZh792lstVc-YrVbWM51tGXEG1EE6VD-EH4k,5149
+jararaca-0.3.18.dist-info/WHEEL,sha256=sP946D7jFCHeNz5Iq4fL4Lu-PrWrFsgfLXbbkciIZwg,88
+jararaca-0.3.18.dist-info/entry_points.txt,sha256=WIh3aIvz8LwUJZIDfs4EeH3VoFyCGEk7cWJurW38q0I,45
+jararaca-0.3.18.dist-info/RECORD,,

pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "jararaca"
-version = "0.3.16"
+version = "0.3.18"
 description = "A simple and fast API framework for Python"
 authors = ["Lucas S <me@luscasleo.dev>"]
 readme = "README.md"