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

jararaca/presentation/server.py

@@ -1,7 +1,17 @@
+# SPDX-FileCopyrightText: 2025 Lucas S
+#
+# SPDX-License-Identifier: GPL-3.0-or-later
+
+import asyncio
+import logging
+import os
+import signal
+import threading
 from contextlib import asynccontextmanager
+from signal import SIGINT, SIGTERM
 from typing import Any, AsyncGenerator
 
-from fastapi import Depends, FastAPI, Request, WebSocket
+from fastapi import Depends, FastAPI, HTTPException, Request, Response, WebSocket
 from starlette.types import ASGIApp
 
 from jararaca.core.uow import UnitOfWorkContextProvider
@@ -10,12 +20,18 @@ from jararaca.lifecycle import AppLifecycle
 from jararaca.microservice import (
     AppTransactionContext,
     HttpTransactionData,
+    ShutdownState,
     WebSocketTransactionData,
+    provide_shutdown_state,
+    providing_app_type,
 )
 from jararaca.presentation.decorators import RestController
+from jararaca.presentation.exceptions import PresentationException
 from jararaca.presentation.http_microservice import HttpMicroservice
 from jararaca.reflect.controller_inspect import ControllerMemberReflect
 
+logger = logging.getLogger(__name__)
+
 
 class HttpAppLifecycle:
 
@@ -31,58 +47,83 @@ class HttpAppLifecycle:
 
     @asynccontextmanager
     async def __call__(self, api: FastAPI) -> AsyncGenerator[None, None]:
-        async with self.lifecycle():
+        with providing_app_type("http"):
+            async with self.lifecycle():
 
-            # websocket_interceptors = [
-            #     interceptor
-            #     for interceptor in self.lifecycle.initialized_interceptors
-            #     if isinstance(interceptor, WebSocketInterceptor)
-            # ]
+                for controller_t in self.lifecycle.app.controllers:
+                    controller = RestController.get_last(controller_t)
 
-            # for interceptor in websocket_interceptors:
-            #     router = interceptor.get_ws_router(
-            #         self.lifecycle.app, self.lifecycle.container, self.uow_provider
-            #     )
+                    if controller is None:
+                        continue
 
-            #     api.include_router(router)
+                    instance: Any = self.lifecycle.container.get_by_type(controller_t)
 
-            for controller_t in self.lifecycle.app.controllers:
-                controller = RestController.get_controller(controller_t)
+                    router = controller.get_router_factory()(self.lifecycle, instance)
 
-                if controller is None:
-                    continue
+                    api.include_router(router)
 
-                instance: Any = self.lifecycle.container.get_by_type(controller_t)
+                    for middleware in self.http_app.middlewares:
+                        middleware_instance = self.lifecycle.container.get_by_type(
+                            middleware
+                        )
+                        api.router.dependencies.append(
+                            Depends(middleware_instance.intercept)
+                        )
 
-                # dependencies: list[DependsCls] = []
-                # for middleware in controller.middlewares:
-                #     middleware_instance = self.lifecycle.container.get_by_type(
-                #         middleware
-                #     )
-                #     dependencies.append(Depends(middleware_instance.intercept))
+                yield
 
-                router = controller.get_router_factory()(self.lifecycle, instance)
 
-                api.include_router(router)
+class HttpShutdownState(ShutdownState):
+    def __init__(self) -> None:
+        self._requested = False
+        self.old_signal_handlers = {
+            SIGINT: signal.getsignal(SIGINT),
+            SIGTERM: signal.getsignal(SIGTERM),
+        }
+        self.thread_lock = threading.Lock()
+        self.aevent = asyncio.Event()
 
-                for middleware in self.http_app.middlewares:
-                    middleware_instance = self.lifecycle.container.get_by_type(
-                        middleware
-                    )
-                    api.router.dependencies.append(
-                        Depends(middleware_instance.intercept)
-                    )
+    def request_shutdown(self) -> None:
+        if not self._requested:
+            self._requested = True
+            os.kill(os.getpid(), SIGINT)
+
+    def is_shutdown_requested(self) -> bool:
+        return self._requested
+
+    def handle_signal(self, signum: int, frame: Any) -> None:
+        logger.warning(f"Received signal {signum}, initiating shutdown...")
+        self.aevent.set()
+        if self._requested:
+            logger.warning("Shutdown already requested, ignoring signal.")
+            return
+        logger.warning("Requesting shutdown...")
+        self._requested = True
 
-            yield
+        # remove the signal handler to prevent recursion
+        for sig in (SIGINT, SIGTERM):
+            if self.old_signal_handlers[sig] is not None:
+                signal.signal(sig, self.old_signal_handlers[sig])
+
+        signal.raise_signal(signum)
+
+    async def wait_for_shutdown(self) -> None:
+        await self.aevent.wait()
+
+    def setup_signal_handlers(self) -> None:
+        signal.signal(SIGINT, self.handle_signal)
+        signal.signal(SIGTERM, self.handle_signal)
 
 
 class HttpUowContextProviderDependency:
 
     def __init__(self, uow_provider: UnitOfWorkContextProvider) -> None:
         self.uow_provider = uow_provider
+        self.shutdown_state = HttpShutdownState()
+        self.shutdown_state.setup_signal_handlers()
 
     async def __call__(
-        self, websocket: WebSocket = None, request: Request = None  # type: ignore
+        self, websocket: WebSocket = None, request: Request = None, response: Response = None  # type: ignore
     ) -> AsyncGenerator[None, None]:
         if request:
             endpoint = request.scope["endpoint"]
@@ -101,17 +142,28 @@ class HttpUowContextProviderDependency:
             "ControllerMemberReflect, but got: {}".format(type(member))
         )
 
-        async with self.uow_provider(
-            AppTransactionContext(
-                controller_member_reflect=member,
-                transaction_data=(
-                    HttpTransactionData(request=request)
-                    if request
-                    else WebSocketTransactionData(websocket=websocket)
-                ),
-            )
-        ):
-            yield
+        with provide_shutdown_state(self.shutdown_state):
+            async with self.uow_provider(
+                AppTransactionContext(
+                    controller_member_reflect=member,
+                    transaction_data=(
+                        HttpTransactionData(request=request, response=response)
+                        if request
+                        else WebSocketTransactionData(websocket=websocket)
+                    ),
+                )
+            ):
+                try:
+                    yield
+                except HTTPException:
+                    raise
+                except Exception as e:
+                    raise PresentationException(
+                        original_exception=e,
+                        request=request,
+                        response=response,
+                        websocket=websocket,
+                    )
 
 
 def create_http_server(

jararaca/presentation/websocket/__init__.py

@@ -0,0 +1,3 @@
+# SPDX-FileCopyrightText: 2025 Lucas S
+#
+# SPDX-License-Identifier: GPL-3.0-or-later

jararaca/presentation/websocket/base_types.py

@@ -1,3 +1,7 @@
+# SPDX-FileCopyrightText: 2025 Lucas S
+#
+# SPDX-License-Identifier: GPL-3.0-or-later
+
 from typing import ClassVar
 
 from pydantic import BaseModel

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

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: