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

jararaca/presentation/websocket/context.py

@@ -1,3 +1,7 @@
+# SPDX-FileCopyrightText: 2025 Lucas S
+#
+# SPDX-License-Identifier: GPL-3.0-or-later
+
 from contextlib import contextmanager, suppress
 from contextvars import ContextVar
 from typing import Generator, Protocol
@@ -16,12 +20,12 @@ class WebSocketConnectionManager(Protocol):
     async def remove_websocket(self, websocket: WebSocket) -> None: ...
 
 
-_ws_manage_ctx = ContextVar[WebSocketConnectionManager]("ws_manage_ctx")
+_ws_conn_manager_ctx = ContextVar[WebSocketConnectionManager]("ws_manage_ctx")
 
 
 def use_ws_manager() -> WebSocketConnectionManager:
     try:
-        return _ws_manage_ctx.get()
+        return _ws_conn_manager_ctx.get()
     except LookupError:
         raise RuntimeError("No WebSocketConnectionManager found")
 
@@ -30,9 +34,35 @@ def use_ws_manager() -> WebSocketConnectionManager:
 def provide_ws_manager(
     ws_manager: WebSocketConnectionManager,
 ) -> Generator[None, None, None]:
-    token = _ws_manage_ctx.set(ws_manager)
+    token = _ws_conn_manager_ctx.set(ws_manager)
+    try:
+        yield
+    finally:
+        with suppress(ValueError):
+            _ws_conn_manager_ctx.reset(token)
+
+
+class WebSocketMessageSender(Protocol):
+    async def send(self, rooms: list[str], message: WebSocketMessageBase) -> None: ...
+
+
+_ws_msg_sender_ctx = ContextVar[WebSocketMessageSender]("ws_msg_sender_ctx")
+
+
+def use_ws_message_sender() -> WebSocketMessageSender:
+    try:
+        return _ws_msg_sender_ctx.get()
+    except LookupError:
+        raise RuntimeError("No WebSocketMessageSender found")
+
+
+@contextmanager
+def provide_ws_message_sender(
+    ws_message_sender: WebSocketMessageSender,
+) -> Generator[None, None, None]:
+    token = _ws_msg_sender_ctx.set(ws_message_sender)
     try:
         yield
     finally:
         with suppress(ValueError):
-            _ws_manage_ctx.reset(token)
+            _ws_msg_sender_ctx.reset(token)

jararaca/presentation/websocket/decorators.py

@@ -1,6 +1,11 @@
-from typing import TypedDict, TypeVar, cast
+# SPDX-FileCopyrightText: 2025 Lucas S
+#
+# SPDX-License-Identifier: GPL-3.0-or-later
+
+from typing import TypedDict, TypeVar
 
 from jararaca.presentation.websocket.base_types import WebSocketMessageBase
+from jararaca.reflect.decorators import StackableDecorator
 
 DECORATED_CLASS = TypeVar("DECORATED_CLASS")
 
@@ -8,9 +13,8 @@ DECORATED_CLASS = TypeVar("DECORATED_CLASS")
 class WebSocketEndpointOptions(TypedDict): ...
 
 
-class WebSocketEndpoint:
+class WebSocketEndpoint(StackableDecorator):
 
-    WEBSOCKET_ENDPOINT_ATTR = "__websocket_endpoint__"
     ORDER_COUNTER = 0
 
     def __init__(self, path: str, options: WebSocketEndpointOptions = {}) -> None:
@@ -19,48 +23,11 @@ class WebSocketEndpoint:
         WebSocketEndpoint.ORDER_COUNTER += 1
         self.order = WebSocketEndpoint.ORDER_COUNTER
 
-    @staticmethod
-    def register(cls: DECORATED_CLASS, instance: "WebSocketEndpoint") -> None:
-        setattr(cls, WebSocketEndpoint.WEBSOCKET_ENDPOINT_ATTR, instance)
-
-    @staticmethod
-    def get(cls: DECORATED_CLASS) -> "WebSocketEndpoint | None":
-        if not hasattr(cls, WebSocketEndpoint.WEBSOCKET_ENDPOINT_ATTR):
-            return None
-
-        return cast(
-            WebSocketEndpoint, getattr(cls, WebSocketEndpoint.WEBSOCKET_ENDPOINT_ATTR)
-        )
-
-    def __call__(self, cls: DECORATED_CLASS) -> DECORATED_CLASS:
-        WebSocketEndpoint.register(cls, self)
-        return cls
-
 
 INHERITS_WS_MESSAGE = TypeVar("INHERITS_WS_MESSAGE", bound=WebSocketMessageBase)
 
 
-class RegisterWebSocketMessage:
-
-    REGISTER_WEBSOCKET_MESSAGE_ATTR = "__register_websocket_message__"
+class RegisterWebSocketMessage(StackableDecorator):
 
     def __init__(self, *message_types: type[INHERITS_WS_MESSAGE]) -> None:
         self.message_types = message_types
-
-    @staticmethod
-    def register(cls: DECORATED_CLASS, instance: "RegisterWebSocketMessage") -> None:
-        setattr(cls, RegisterWebSocketMessage.REGISTER_WEBSOCKET_MESSAGE_ATTR, instance)
-
-    @staticmethod
-    def get(cls: DECORATED_CLASS) -> "RegisterWebSocketMessage | None":
-        if not hasattr(cls, RegisterWebSocketMessage.REGISTER_WEBSOCKET_MESSAGE_ATTR):
-            return None
-
-        return cast(
-            RegisterWebSocketMessage,
-            getattr(cls, RegisterWebSocketMessage.REGISTER_WEBSOCKET_MESSAGE_ATTR),
-        )
-
-    def __call__(self, cls: DECORATED_CLASS) -> DECORATED_CLASS:
-        RegisterWebSocketMessage.register(cls, self)
-        return cls

jararaca/presentation/websocket/redis.py

@@ -1,4 +1,9 @@
+# SPDX-FileCopyrightText: 2025 Lucas S
+#
+# SPDX-License-Identifier: GPL-3.0-or-later
+
 import asyncio
+import logging
 from dataclasses import dataclass
 from typing import Any
 
@@ -10,6 +15,8 @@ from jararaca.presentation.websocket.websocket_interceptor import (
     WebSocketConnectionBackend,
 )
 
+logger = logging.getLogger(__name__)
+
 
 @dataclass
 class BroadcastMessage:
@@ -55,6 +62,8 @@ class RedisWebSocketConnectionBackend(WebSocketConnectionBackend):
         send_pubsub_channel: str,
         consume_broadcast_timeout: int = 1,
         consume_send_timeout: int = 1,
+        retry_delay: float = 5.0,
+        max_concurrent_tasks: int = 1000,
     ) -> None:
 
         self.redis = conn
@@ -63,89 +72,307 @@ 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
+        self.retry_delay = retry_delay
+        self.__shutdown_event: asyncio.Event | None = None
+
+        self.__send_func: SendFunc | None = None
+        self.__broadcast_func: BroadcastFunc | None = None
+
+    @property
+    def shutdown_event(self) -> asyncio.Event:
+        if self.__shutdown_event is None:
+            raise RuntimeError(
+                "Shutdown event is not set. Please configure the backend before using it."
+            )
+        return self.__shutdown_event
+
+    @property
+    def send_func(self) -> SendFunc:
+        if self.__send_func is None:
+            raise RuntimeError(
+                "Send function is not set. Please configure the backend before using it."
+            )
+        return self.__send_func
+
+    @property
+    def broadcast_func(self) -> BroadcastFunc:
+        if self.__broadcast_func is None:
+            raise RuntimeError(
+                "Broadcast function is not set. Please configure the backend before using it."
+            )
+        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(
+                "Failed to publish broadcast message to Redis: %s", 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(
+                "Failed to publish send message to Redis: %s", e, exc_info=True
+            )
+            raise
 
     def configure(
         self, broadcast: BroadcastFunc, send: SendFunc, shutdown_event: asyncio.Event
     ) -> None:
+        if self.__shutdown_event is not None:
+            raise RuntimeError("Backend is already configured.")
+        self.__shutdown_event = shutdown_event
+        self.__send_func = send
+        self.__broadcast_func = broadcast
+        self.setup_send_consumer()
+        self.setup_broadcast_consumer()
 
-        broadcast_task = asyncio.get_event_loop().create_task(
-            self.consume_broadcast(broadcast, shutdown_event)
-        )
+    def setup_send_consumer(self) -> None:
 
         send_task = asyncio.get_event_loop().create_task(
-            self.consume_send(send, shutdown_event)
+            self.consume_send(self.send_func, self.shutdown_event)
         )
 
-        self.tasks.add(broadcast_task)
-        self.tasks.add(send_task)
-
-    async def consume_broadcast(
-        self, broadcast: BroadcastFunc, shutdown_event: asyncio.Event
-    ) -> None:
+        # Use lock when modifying tasks set to prevent race conditions
+        async def add_task() -> None:
+            async with self.lock:
+                self.tasks.add(send_task)
 
-        async with self.redis.pubsub() as pubsub:
-            await pubsub.subscribe(self.broadcast_pubsub_channel)
+        asyncio.get_event_loop().create_task(add_task())
+        send_task.add_done_callback(self.handle_send_task_done)
 
-            while not shutdown_event.is_set():
-                message: dict[str, Any] | None = await pubsub.get_message(
-                    ignore_subscribe_messages=True,
-                    timeout=self.consume_broadcast_timeout,
-                )
+    def setup_broadcast_consumer(self) -> None:
 
-                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)
-                    )
+        broadcast_task = asyncio.get_event_loop().create_task(
+            self.consume_broadcast(self.broadcast_func, self.shutdown_event)
+        )
 
-                    self.tasks.add(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:
+            logger.exception(
+                "Broadcast task raised an exception:", exc_info=task.exception()
+            )
+        else:
+            logger.warning("Broadcast task somehow completed successfully.")
+
+        if not self.shutdown_event.is_set():
+            logger.warning(
+                "Broadcast task completed, but shutdown event is not set. This is unexpected."
+            )
+            # Add delay before retrying to avoid excessive CPU usage
+            asyncio.get_event_loop().create_task(
+                self._retry_broadcast_consumer_with_delay()
+            )
+
+    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:
+            logger.exception(
+                "Send task raised an exception:", exc_info=task.exception()
+            )
+        else:
+            logger.warning("Send task somehow completed successfully.")
+
+        if not self.shutdown_event.is_set():
+            logger.warning(
+                "Send task completed, but shutdown event is not set. This is unexpected."
+            )
+            # Add delay before retrying to avoid excessive CPU usage
+            asyncio.get_event_loop().create_task(self._retry_send_consumer_with_delay())
+
+    async def _retry_broadcast_consumer_with_delay(self) -> None:
+        """Retry setting up broadcast consumer after a delay to avoid excessive CPU usage."""
+        logger.warning(
+            "Waiting %s seconds before retrying broadcast consumer...", self.retry_delay
+        )
+        await asyncio.sleep(self.retry_delay)
 
-                    task.add_done_callback(self.tasks.discard)
+        if not self.shutdown_event.is_set():
+            logger.warning("Retrying broadcast consumer setup...")
+            self.setup_broadcast_consumer()
 
-    async def consume_send(self, send: SendFunc, shutdown_event: asyncio.Event) -> None:
-
-        async with self.redis.pubsub() as pubsub:
-            await pubsub.subscribe(self.send_pubsub_channel)
+    async def _retry_send_consumer_with_delay(self) -> None:
+        """Retry setting up send consumer after a delay to avoid excessive CPU usage."""
+        logger.warning(
+            "Waiting %s seconds before retrying send consumer...", self.retry_delay
+        )
+        await asyncio.sleep(self.retry_delay)
 
-            while not shutdown_event.is_set():
+        if not self.shutdown_event.is_set():
+            logger.warning("Retrying send consumer setup...")
+            self.setup_send_consumer()
 
-                message: dict[str, Any] | None = await pubsub.get_message(
-                    ignore_subscribe_messages=True, timeout=self.consume_send_timeout
+    async def consume_broadcast(
+        self, broadcast: BroadcastFunc, shutdown_event: asyncio.Event
+    ) -> None:
+        logger.debug("Starting broadcast consumer...")
+        try:
+            # Validate Redis connection before starting
+            try:
+                await self.redis.ping()
+                logger.debug("Redis connection validated for broadcast consumer")
+            except Exception as e:
+                logger.error("Redis connection validation failed: %s", e, exc_info=True)
+                raise
+
+            async with self.redis.pubsub() as pubsub:
+                await pubsub.subscribe(self.broadcast_pubsub_channel)
+                logger.debug(
+                    "Subscribed to broadcast channel: %s", self.broadcast_pubsub_channel
                 )
 
-                if message is None:
-                    continue
-
-                send_message = SendToRoomsMessage.decode(message["data"])
+                while not shutdown_event.is_set():
+                    message: dict[str, Any] | None = await pubsub.get_message(
+                        ignore_subscribe_messages=True,
+                        timeout=self.consume_broadcast_timeout,
+                    )
 
-                async with self.lock:
+                    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(
+                            "Error processing broadcast message: %s", e, exc_info=True
+                        )
+                        # Continue processing other messages
+                        continue
+        except Exception as e:
+            logger.error(
+                "Fatal error in broadcast consumer, will retry: %s", e, exc_info=True
+            )
+            raise
 
-                    task = asyncio.get_event_loop().create_task(
-                        send(send_message.rooms, send_message.message)
+    async def consume_send(self, send: SendFunc, shutdown_event: asyncio.Event) -> None:
+        logger.debug("Starting send consumer...")
+        try:
+            # Validate Redis connection before starting
+            try:
+                await self.redis.ping()
+                logger.debug("Redis connection validated for send consumer")
+            except Exception as e:
+                logger.error("Redis connection validation failed: %s", e, exc_info=True)
+                raise
+
+            async with self.redis.pubsub() as pubsub:
+                await pubsub.subscribe(self.send_pubsub_channel)
+                logger.debug("Subscribed to send channel: %s", 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(
+                            "Error processing send message: %s", e, exc_info=True
+                        )
+                        # Continue processing other messages
+                        continue
+        except Exception as e:
+            logger.error(
+                "Fatal error in send consumer, will retry: %s", e, exc_info=True
+            )
+            raise
 
     async def shutdown(self) -> None:
         async with self.lock:

jararaca/presentation/websocket/types.py

@@ -1,8 +1,12 @@
+# SPDX-FileCopyrightText: 2025 Lucas S
+#
+# SPDX-License-Identifier: GPL-3.0-or-later
+
 from jararaca.presentation.websocket.base_types import WebSocketMessageBase
-from jararaca.presentation.websocket.context import use_ws_manager
+from jararaca.presentation.websocket.context import use_ws_message_sender
 
 
 class WebSocketMessage(WebSocketMessageBase):
 
     async def send(self, *rooms: str) -> None:
-        await use_ws_manager().send(list(rooms), self)
+        await use_ws_message_sender().send(list(rooms), self)