jararaca 0.3.17__py3-none-any.whl → 0.3.19__py3-none-any.whl

jararaca/__init__.py

@@ -17,6 +17,8 @@ if TYPE_CHECKING:
         AppTransactionContext,
         use_app_type,
     )
+    from jararaca.observability.decorators import TracedClass, TracedFunc, traced_class
+    from jararaca.observability.hooks import spawn_trace
     from jararaca.observability.interceptor import ObservabilityInterceptor
     from jararaca.observability.providers.otel import OtelObservabilityProvider
     from jararaca.persistence.sort_filter import (
@@ -222,6 +224,10 @@ if TYPE_CHECKING:
         "HttpPut",
         "HttpDelete",
         "ObservabilityInterceptor",
+        "TracedFunc",
+        "TracedClass",
+        "traced_class",
+        "spawn_trace",
         "QueryInjector",
         "HttpMicroservice",
         "use_current_container",
@@ -392,6 +398,10 @@ _dynamic_imports: "dict[str, tuple[str, str, str | None]]" = {
     "HttpPut": (__SPEC_PARENT__, "rpc.http.decorators", "Put"),
     "HttpDelete": (__SPEC_PARENT__, "rpc.http.decorators", "Delete"),
     "ObservabilityInterceptor": (__SPEC_PARENT__, "observability.interceptor", None),
+    "TracedFunc": (__SPEC_PARENT__, "observability.decorators", None),
+    "TracedClass": (__SPEC_PARENT__, "observability.decorators", None),
+    "traced_class": (__SPEC_PARENT__, "observability.decorators", None),
+    "spawn_trace": (__SPEC_PARENT__, "observability.hooks", None),
     "QueryInjector": (__SPEC_PARENT__, "persistence.utilities", None),
     "HttpMicroservice": (__SPEC_PARENT__, "presentation.http_microservice", None),
     "use_current_container": (__SPEC_PARENT__, "microservice", None),

jararaca/core/uow.py

@@ -33,7 +33,6 @@ class UnitOfWorkContextProvider:
         self.container = container
         self.container_interceptor = ContainerInterceptor(container)
 
-    # TODO: Guarantee that the context is closed whenever an exception is raised
     # TODO: Guarantee a unit of work workflow for the whole request, including all the interceptors
 
     def factory_app_interceptors(self) -> Sequence[AppInterceptor]:
@@ -65,7 +64,36 @@ class UnitOfWorkContextProvider:
             for ctx in ctxs:
                 await ctx.__aenter__()
 
-            yield None
-
-            for ctx in reversed(ctxs):
-                await ctx.__aexit__(None, None, None)
+            exc_type = None
+            exc_value = None
+            exc_traceback = None
+
+            try:
+                yield None
+            except BaseException as e:
+                exc_type = type(e)
+                exc_value = e
+                exc_traceback = e.__traceback__
+                raise
+            finally:
+                # Exit interceptors in reverse order, propagating exception info
+                for ctx in reversed(ctxs):
+                    try:
+                        suppressed = await ctx.__aexit__(
+                            exc_type, exc_value, exc_traceback
+                        )
+                        # If an interceptor returns True, it suppresses the exception
+                        if suppressed and exc_type is not None:
+                            exc_type = None
+                            exc_value = None
+                            exc_traceback = None
+                    except BaseException as exit_exc:
+                        # If an interceptor raises an exception during cleanup,
+                        # replace the original exception with the new one
+                        exc_type = type(exit_exc)
+                        exc_value = exit_exc
+                        exc_traceback = exit_exc.__traceback__
+
+                # Re-raise the exception if it wasn't suppressed
+                if exc_type is not None and exc_value is not None:
+                    raise exc_value.with_traceback(exc_traceback)

jararaca/messagebus/implicit_headers.py

@@ -0,0 +1,45 @@
+import datetime
+import decimal
+import typing
+from contextlib import contextmanager
+from contextvars import ContextVar
+from typing import Any, Dict, Generator
+
+FieldArray = list["FieldValue"]
+"""A data structure for holding an array of field values."""
+
+FieldTable = typing.Dict[str, "FieldValue"]
+FieldValue = (
+    bool
+    | bytes
+    | bytearray
+    | decimal.Decimal
+    | FieldArray
+    | FieldTable
+    | float
+    | int
+    | None
+    | str
+    | datetime.datetime
+)
+
+ImplicitHeaders = Dict[str, FieldValue]
+
+implicit_headers_ctx = ContextVar[ImplicitHeaders | None](
+    "implicit_headers_ctx", default=None
+)
+
+
+def use_implicit_headers() -> ImplicitHeaders | None:
+    return implicit_headers_ctx.get()
+
+
+@contextmanager
+def provide_implicit_headers(
+    implicit_headers: ImplicitHeaders,
+) -> Generator[None, Any, None]:
+    token = implicit_headers_ctx.set(implicit_headers)
+    try:
+        yield
+    finally:
+        implicit_headers_ctx.reset(token)

jararaca/messagebus/interceptors/aiopika_publisher_interceptor.py

@@ -9,6 +9,7 @@ from aio_pika.abc import AbstractConnection
 from pydantic import BaseModel
 
 from jararaca.broker_backend import MessageBrokerBackend
+from jararaca.messagebus import implicit_headers
 from jararaca.messagebus.interceptors.publisher_interceptor import (
     MessageBusConnectionFactory,
 )
@@ -42,8 +43,12 @@ class AIOPikaMessagePublisher(MessagePublisher):
             logging.warning(f"Exchange {self.exchange_name} not found")
             return
         routing_key = f"{topic}."
+
+        implicit_headers_data = implicit_headers.use_implicit_headers()
         await exchange.publish(
-            aio_pika.Message(body=message.model_dump_json().encode()),
+            aio_pika.Message(
+                body=message.model_dump_json().encode(), headers=implicit_headers_data
+            ),
             routing_key=routing_key,
         )
 

jararaca/messagebus/worker.py

@@ -50,6 +50,7 @@ from jararaca.messagebus.decorators import (
     MessageHandlerData,
     ScheduleDispatchData,
 )
+from jararaca.messagebus.implicit_headers import provide_implicit_headers
 from jararaca.messagebus.message import Message, MessageOf
 from jararaca.microservice import (
     AppTransactionContext,
@@ -1130,6 +1131,8 @@ class ScheduledMessageHandlerCallback:
                 AppTransactionContext(
                     controller_member_reflect=scheduled_action.controller_member,
                     transaction_data=SchedulerTransactionData(
+                        task_name=scheduled_action.spec.name
+                        or scheduled_action.callable.__qualname__,
                         scheduled_to=datetime.now(UTC),
                         cron_expression=scheduled_action.spec.cron,
                         triggered_at=datetime.now(UTC),
@@ -1469,7 +1472,9 @@ class MessageHandlerCallback:
         incoming_message_spec = MessageHandler.get_message_incoming(handler)
         assert incoming_message_spec is not None
 
-        with provide_shutdown_state(self.consumer.shutdown_state):
+        with provide_implicit_headers(aio_pika_message.headers), provide_shutdown_state(
+            self.consumer.shutdown_state
+        ):
             async with self.consumer.uow_context_provider(
                 AppTransactionContext(
                     controller_member_reflect=handler_data.controller_member,

jararaca/microservice.py

@@ -19,7 +19,7 @@ from typing import (
     runtime_checkable,
 )
 
-from fastapi import Request, WebSocket
+from fastapi import Request, Response, WebSocket
 
 from jararaca.core.providers import ProviderSpec, T, Token
 from jararaca.messagebus import MessageOf
@@ -34,6 +34,7 @@ if TYPE_CHECKING:
 
 @dataclass
 class SchedulerTransactionData:
+    task_name: str
     triggered_at: datetime
     scheduled_to: datetime
     cron_expression: str
@@ -43,6 +44,7 @@ class SchedulerTransactionData:
 @dataclass
 class HttpTransactionData:
     request: Request
+    response: Response
     context_type: Literal["http"] = "http"
 
 

jararaca/observability/decorators.py

@@ -1,3 +1,4 @@
+import inspect
 from contextlib import contextmanager, suppress
 from contextvars import ContextVar
 from functools import wraps
@@ -8,21 +9,35 @@ from typing import (
     Callable,
     ContextManager,
     Generator,
-    ParamSpec,
+    Mapping,
     Protocol,
+    Sequence,
     TypeVar,
+    Union,
 )
 
 from jararaca.microservice import AppTransactionContext
 
-P = ParamSpec("P")
-R = TypeVar("R")
+F = TypeVar("F", bound=Callable[..., Awaitable[Any]])
+
+AttributeValue = Union[
+    str,
+    bool,
+    int,
+    float,
+    Sequence[str],
+    Sequence[bool],
+    Sequence[int],
+    Sequence[float],
+]
+
+AttributeMap = Mapping[str, AttributeValue]
 
 
 class TracingContextProvider(Protocol):
 
     def __call__(
-        self, trace_name: str, context_attributes: dict[str, str]
+        self, trace_name: str, context_attributes: AttributeMap | None
     ) -> ContextManager[Any]: ...
 
 
@@ -76,22 +91,123 @@ class TracedFunc:
 
     def __call__(
         self,
-        decorated: Callable[P, Awaitable[R]],
-    ) -> Callable[P, Awaitable[R]]:
+        decorated: F,
+    ) -> F:
 
         @wraps(decorated)
         async def wrapper(
-            *args: P.args,
-            **kwargs: P.kwargs,
-        ) -> R:
+            *args: Any,
+            **kwargs: Any,
+        ) -> Any:
 
             if ctx_provider := get_tracing_ctx_provider():
                 with ctx_provider(
                     self.trace_name,
-                    self.trace_mapper(**kwargs),
+                    self.trace_mapper(*args, **kwargs),
                 ):
                     return await decorated(*args, **kwargs)
 
             return await decorated(*args, **kwargs)
 
-        return wrapper
+        return wrapper  # type: ignore[return-value]
+
+
+C = TypeVar("C", bound=type)
+
+
+class TracedClass:
+    """
+    Class decorator that automatically applies tracing to all async methods in a class.
+
+    Usage:
+        @TracedClass()
+        class MyService:
+            async def method1(self) -> str:
+                return "hello"
+
+            async def method2(self, x: int) -> int:
+                return x * 2
+
+            def sync_method(self) -> str:  # Not traced
+                return "sync"
+    """
+
+    def __init__(
+        self,
+        trace_name_prefix: str | None = None,
+        trace_mapper: Callable[..., dict[str, str]] = default_trace_mapper,
+        include_private: bool = False,
+        exclude_methods: set[str] | None = None,
+    ):
+        """
+        Initialize the TracedClass decorator.
+
+        Args:
+            trace_name_prefix: Prefix for trace names. If None, uses class name.
+            trace_mapper: Function to map method arguments to trace attributes.
+            include_private: Whether to trace private methods (starting with _).
+            exclude_methods: Set of method names to exclude from tracing.
+        """
+        self.trace_name_prefix = trace_name_prefix
+        self.trace_mapper = trace_mapper
+        self.include_private = include_private
+        self.exclude_methods = exclude_methods or set()
+
+    def __call__(self, cls: C) -> C:
+        """Apply tracing to all async methods in the class."""
+
+        # Use class name as prefix if not provided
+        trace_prefix = self.trace_name_prefix or cls.__name__
+
+        # Get all methods in the class
+        for name, method in inspect.getmembers(cls, predicate=inspect.isfunction):
+            # Skip if method should be excluded
+            if name in self.exclude_methods:
+                continue
+
+            # Skip private methods unless explicitly included
+            if name.startswith("_") and not self.include_private:
+                continue
+
+            # Only trace async methods
+            if inspect.iscoroutinefunction(method):
+                trace_name = f"{trace_prefix}.{name}"
+                traced_method = TracedFunc(trace_name, self.trace_mapper)(method)
+                setattr(cls, name, traced_method)
+
+        return cls
+
+
+def traced_class(
+    trace_name_prefix: str | None = None,
+    trace_mapper: Callable[..., dict[str, str]] = default_trace_mapper,
+    include_private: bool = False,
+    exclude_methods: set[str] | None = None,
+) -> Callable[[C], C]:
+    """
+    Functional interface for TracedClass decorator.
+
+    Usage:
+        @traced_class(trace_name_prefix="MyService")
+        class MyService:
+            async def method1(self) -> str:
+                return "hello"
+    """
+    return TracedClass(
+        trace_name_prefix=trace_name_prefix,
+        trace_mapper=trace_mapper,
+        include_private=include_private,
+        exclude_methods=exclude_methods,
+    )
+
+
+__all__ = [
+    "TracingContextProvider",
+    "TracingContextProviderFactory",
+    "provide_tracing_ctx_provider",
+    "get_tracing_ctx_provider",
+    "default_trace_mapper",
+    "TracedFunc",
+    "TracedClass",
+    "traced_class",
+]

jararaca/observability/hooks.py

@@ -0,0 +1,20 @@
+from contextlib import contextmanager
+from typing import Any, Generator
+
+from jararaca.observability.decorators import AttributeMap, get_tracing_ctx_provider
+
+
+@contextmanager
+def spawn_trace(
+    name: str,
+    attributes: AttributeMap | None = None,
+) -> Generator[None, Any, None]:
+
+    if trace_context_provider := get_tracing_ctx_provider():
+        with trace_context_provider(trace_name=name, context_attributes=attributes):
+            yield
+    else:
+        yield
+
+
+spawn_trace

jararaca/observability/providers/otel.py

@@ -1,6 +1,6 @@
 import logging
 from contextlib import asynccontextmanager, contextmanager
-from typing import AsyncGenerator, Generator, Protocol
+from typing import Any, AsyncGenerator, Generator, Protocol
 
 from opentelemetry import metrics, trace
 from opentelemetry._logs import set_logger_provider
@@ -23,11 +23,21 @@ from opentelemetry.sdk.trace import TracerProvider
 from opentelemetry.sdk.trace.export import BatchSpanProcessor
 from opentelemetry.trace.propagation.tracecontext import TraceContextTextMapPropagator
 
-from jararaca.microservice import AppTransactionContext, Container, Microservice
+from jararaca.messagebus.implicit_headers import (
+    ImplicitHeaders,
+    provide_implicit_headers,
+    use_implicit_headers,
+)
+from jararaca.microservice import (
+    AppTransactionContext,
+    Container,
+    Microservice,
+    use_app_transaction_context,
+)
 from jararaca.observability.decorators import (
+    AttributeMap,
     TracingContextProvider,
     TracingContextProviderFactory,
-    get_tracing_ctx_provider,
 )
 from jararaca.observability.interceptor import ObservabilityProvider
 
@@ -43,7 +53,7 @@ class OtelTracingContextProvider(TracingContextProvider):
     def __call__(
         self,
         trace_name: str,
-        context_attributes: dict[str, str],
+        context_attributes: AttributeMap | None,
     ) -> Generator[None, None, None]:
 
         with tracer.start_as_current_span(trace_name, attributes=context_attributes):
@@ -63,7 +73,7 @@ class OtelTracingContextProviderFactory(TracingContextProviderFactory):
     ) -> AsyncGenerator[None, None]:
 
         title: str = "Unmapped App Context Execution"
-        headers = {}
+        headers: dict[str, Any] = {}
         tx_data = app_tx_ctx.transaction_data
         if tx_data.context_type == "http":
 
@@ -72,6 +82,19 @@ class OtelTracingContextProviderFactory(TracingContextProviderFactory):
 
         elif tx_data.context_type == "message_bus":
             title = f"Message Bus {tx_data.topic}"
+            headers = use_implicit_headers() or {}
+
+        elif tx_data.context_type == "websocket":
+            headers = dict(tx_data.websocket.headers)
+            title = f"WebSocket {tx_data.websocket.url}"
+
+        elif tx_data.context_type == "scheduler":
+            title = f"Scheduler Task {tx_data.task_name}"
+            headers = {
+                "scheduled_to": tx_data.scheduled_to.isoformat(),
+                "cron_expression": tx_data.cron_expression,
+                "triggered_at": tx_data.triggered_at.isoformat(),
+            }
 
         carrier = {
             key: value
@@ -90,8 +113,23 @@ class OtelTracingContextProviderFactory(TracingContextProviderFactory):
 
         ctx2 = W3CBaggagePropagator().extract(b2, context=ctx)
 
-        with tracer.start_as_current_span(name=title, context=ctx2):
-            yield
+        with tracer.start_as_current_span(
+            name=title,
+            context=ctx2,
+            attributes={
+                "app.context_type": tx_data.context_type,
+            },
+        ) as root_span:
+            cx = root_span.get_span_context()
+            if app_tx_ctx.transaction_data.context_type == "http":
+                app_tx_ctx.transaction_data.response.headers["traceparent"] = hex(
+                    cx.trace_id
+                )[2:].rjust(32, "0")
+            tracing_headers: ImplicitHeaders = {}
+            TraceContextTextMapPropagator().inject(tracing_headers)
+            W3CBaggagePropagator().inject(tracing_headers)
+            with provide_implicit_headers(tracing_headers):
+                yield
 
 
 class LoggerHandlerCallback(Protocol):
@@ -99,6 +137,21 @@ class LoggerHandlerCallback(Protocol):
     def __call__(self, logger_handler: logging.Handler) -> None: ...
 
 
+class CustomLoggingHandler(LoggingHandler):
+
+    def _translate(self, record: logging.LogRecord) -> dict[str, Any]:
+        try:
+            ctx = use_app_transaction_context()
+            return {
+                **super()._translate(record),
+                "attributes": {
+                    "context_type": ctx.transaction_data.context_type,
+                },
+            }
+        except LookupError:
+            return super()._translate(record)
+
+
 class OtelObservabilityProvider(ObservabilityProvider):
 
     def __init__(
@@ -143,11 +196,11 @@ class OtelObservabilityProvider(ObservabilityProvider):
             BatchLogRecordProcessor(self.logs_exporter)
         )
 
-        logging_handler = LoggingHandler(
+        logging_handler = CustomLoggingHandler(
             level=logging.DEBUG, logger_provider=logger_provider
         )
 
-        logging_handler.addFilter(lambda _: get_tracing_ctx_provider() is not None)
+        # logging_handler.addFilter(lambda _: get_tracing_ctx_provider() is not None)
 
         self.logging_handler_callback(logging_handler)
 

jararaca/presentation/server.py

@@ -1,3 +1,4 @@
+import logging
 import os
 import signal
 import threading
@@ -5,7 +6,7 @@ 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
@@ -23,6 +24,8 @@ from jararaca.presentation.decorators import RestController
 from jararaca.presentation.http_microservice import HttpMicroservice
 from jararaca.reflect.controller_inspect import ControllerMemberReflect
 
+logger = logging.getLogger(__name__)
+
 
 class HttpAppLifecycle:
 
@@ -109,7 +112,7 @@ class HttpUowContextProviderDependency:
         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"]
@@ -133,13 +136,34 @@ class HttpUowContextProviderDependency:
                 AppTransactionContext(
                     controller_member_reflect=member,
                     transaction_data=(
-                        HttpTransactionData(request=request)
+                        HttpTransactionData(request=request, response=response)
                         if request
                         else WebSocketTransactionData(websocket=websocket)
                     ),
                 )
             ):
-                yield
+                try:
+                    yield
+                except HTTPException:
+                    raise
+                except Exception as e:
+                    logger.exception("Unhandled exception in request handling.")
+                    raise HTTPException(
+                        status_code=500,
+                        detail={
+                            "message": "Internal server error occurred.",
+                            "x-traceparentid": (
+                                response.headers.get("traceparent")
+                                if response
+                                else None
+                            ),
+                        },
+                        headers=(
+                            {k: str(v) for k, v in response.headers.items()}
+                            if response
+                            else {}
+                        ),
+                    ) from e
 
 
 def create_http_server(

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.17.dist-info → jararaca-0.3.19.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: jararaca
-Version: 0.3.17
+Version: 0.3.19
 Summary: A simple and fast API framework for Python
 Home-page: https://github.com/LuscasLeo/jararaca
 Author: Lucas S
@@ -18,11 +18,11 @@ Requires-Dist: croniter (>=3.0.3,<4.0.0)
 Requires-Dist: fastapi (>=0.113.0,<0.114.0)
 Requires-Dist: frozendict (>=2.4.6,<3.0.0)
 Requires-Dist: mako (>=1.3.5,<2.0.0)
-Requires-Dist: opentelemetry-api (>=1.27.0,<2.0.0) ; extra == "opentelemetry"
-Requires-Dist: opentelemetry-distro (>=0.49b2,<0.50) ; extra == "opentelemetry"
+Requires-Dist: opentelemetry-api (>=1.38.0,<2.0.0) ; extra == "opentelemetry"
+Requires-Dist: opentelemetry-distro (>=0.59b0,<0.60) ; extra == "opentelemetry"
 Requires-Dist: opentelemetry-exporter-otlp (>=1.27.0,<2.0.0) ; extra == "opentelemetry"
 Requires-Dist: opentelemetry-exporter-otlp-proto-http (>=1.27.0,<2.0.0) ; extra == "opentelemetry"
-Requires-Dist: opentelemetry-sdk (>=1.27.0,<2.0.0) ; extra == "opentelemetry"
+Requires-Dist: opentelemetry-sdk (>=1.38.0,<2.0.0) ; extra == "opentelemetry"
 Requires-Dist: redis (>=5.0.8,<6.0.0)
 Requires-Dist: sqlalchemy (>=2.0.34,<3.0.0)
 Requires-Dist: types-croniter (>=3.0.3.20240731,<4.0.0.0)

{jararaca-0.3.17.dist-info → jararaca-0.3.19.dist-info}/RECORD

@@ -1,7 +1,7 @@
 LICENSE,sha256=OXLcl0T2SZ8Pmy2_dmlvKuetivmyPd5m1q-Gyd-zaYY,35149
 README.md,sha256=YmCngjU8llW0l7L3tuXkkfr8qH7V9aBMgfp2jEzeiKg,3517
-pyproject.toml,sha256=KJbNO4NulXPNhOpQ6kq0t3H8B6PzFpTOuULyhL_w8zc,2832
-jararaca/__init__.py,sha256=IMnvfDoyNWTGVittF_wq2Uxtv_BY_wLN5Om6C3vUsCw,22302
+pyproject.toml,sha256=JdgxP3SwJNa20HtUWhMdpWDnNTFyRb-sCzLStMMDOwM,2739
+jararaca/__init__.py,sha256=niZQiN-Whyw2ExcwQbB8uBDPy90pK36pPwa7dYkecuU,22822
 jararaca/__main__.py,sha256=-O3vsB5lHdqNFjUtoELDF81IYFtR-DSiiFMzRaiSsv4,67
 jararaca/broker_backend/__init__.py,sha256=GzEIuHR1xzgCJD4FE3harNjoaYzxHMHoEL0_clUaC-k,3528
 jararaca/broker_backend/mapper.py,sha256=vTsi7sWpNvlga1PWPFg0rCJ5joJ0cdzykkIc2Tuvenc,696
@@ -10,7 +10,7 @@ jararaca/cli.py,sha256=n3fTOVSNFNICmbZrLCJNGctpbqLWp39xTSkqlIB6Rds,32005
 jararaca/common/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 jararaca/core/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 jararaca/core/providers.py,sha256=wktH84FK7c1s2wNq-fudf1uMfi3CQBR0neU2czJ_L0U,434
-jararaca/core/uow.py,sha256=jKlhU_YsssiN0hC1s056JcNerNagXPcXdZJ5kJYG68c,2252
+jararaca/core/uow.py,sha256=J6shc1TwM7uyRDspekY2EdA8slVDPitvHKg4Bunb48E,3533
 jararaca/di.py,sha256=h3IsXdYZjJj8PJfnEDn0ZAwdd4EBfh8jU-wWO8ko_t4,76
 jararaca/files/entity.py.mako,sha256=tjQ-1udAMvVqgRokhsrR4uy3P_OVnbk3XZ8X69ixWhE,3098
 jararaca/lifecycle.py,sha256=qKlzLQQioS8QkxNJ_FC_5WbmT77cNbc_S7OcQeOoHkI,1895
@@ -18,17 +18,19 @@ jararaca/messagebus/__init__.py,sha256=5jAqPqdcEMYBfQyfZDWPnplYdrfMyJLMcacf3qLyU
 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=71ZvZ5slKruAsTaUHzr2f-D8yaUFl9Xxh9chNLaAs9E,6000
+jararaca/messagebus/implicit_headers.py,sha256=maMvSbWwMklVPW0hX6jZ8gAoKsYXB7saYod85jfoSq4,965
 jararaca/messagebus/interceptors/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-jararaca/messagebus/interceptors/aiopika_publisher_interceptor.py,sha256=_DEHwIH9LYsA26Hu1mo9oHzLZuATgjilU9E3o-ecDjs,6520
+jararaca/messagebus/interceptors/aiopika_publisher_interceptor.py,sha256=vs8kCZYslNwiQTy6wFAzuX_N6Oq9nx35RV0MBWwdM-c,6703
 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=Feh7b3z68sj7a94z0lwMowLgU41pFF2PGgefhdH8rw8,68926
-jararaca/microservice.py,sha256=OYCw5C4797X_tVnM_9sEz8BdjbICPHSVsCixsA_FwE4,11419
-jararaca/observability/decorators.py,sha256=MOIr2PttPYYvRwEdfQZEwD5RxKHOTv8UEy9n1YQVoKw,2281
+jararaca/messagebus/worker.py,sha256=vni08oRRB-Vz_qXOOoxRDnBTuO8a_rZ6Zue0qx3UUes,69202
+jararaca/microservice.py,sha256=DW4RVeqgrx4J-dAg17sbzTn_sLXuvV1UOQFde2fpqds,11471
+jararaca/observability/decorators.py,sha256=vKAPGLyKz092rd1K4k3xtI2FrRafN69EEuRZKEhaLSk,5627
+jararaca/observability/hooks.py,sha256=VRYg-vlNlFxu9FqOgxyS9VIY6zDcCcgOuZ0RhTUtjwY,488
 jararaca/observability/interceptor.py,sha256=U4ZLM0f8j6Q7gMUKKnA85bnvD-Qa0ii79Qa_X8KsXAQ,1498
 jararaca/observability/providers/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-jararaca/observability/providers/otel.py,sha256=8N1F32W43t7c8cwmtTh6Yz9b7HyfGFSRVrkMLf3NDXc,6157
+jararaca/observability/providers/otel.py,sha256=eYjejx7SW6FZ9-GTNiAxiTiae9RxQpiUSHiYN1rRZnI,8019
 jararaca/persistence/base.py,sha256=xnGUbsLNz3gO-9iJt-Sn5NY13Yc9-misP8wLwQuGGoM,1024
 jararaca/persistence/exports.py,sha256=Ghx4yoFaB4QVTb9WxrFYgmcSATXMNvrOvT8ybPNKXCA,62
 jararaca/persistence/interceptors/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
@@ -42,14 +44,14 @@ jararaca/presentation/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3h
 jararaca/presentation/decorators.py,sha256=dNcK1xYWhdHVBeWDa8dhqh4z8yQcPkJZTQkfgLXz6RI,11655
 jararaca/presentation/hooks.py,sha256=WBbU5DG3-MAm2Ro2YraQyYG_HENfizYfyShL2ktHi6k,1980
 jararaca/presentation/http_microservice.py,sha256=g771JosV6jTY3hQtG-HkLOo-T0e-r3b3rp1ddt99Qf0,533
-jararaca/presentation/server.py,sha256=YzTcG39xEmEstrtB6nmVUSf2cMP8cSguppKdz_jA_tE,5537
+jararaca/presentation/server.py,sha256=Kz7H_RuB68b5bphjNlS-7cseH_d-xZMa8V-ETFp_jzI,6563
 jararaca/presentation/websocket/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 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 +76,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.17.dist-info/LICENSE,sha256=OXLcl0T2SZ8Pmy2_dmlvKuetivmyPd5m1q-Gyd-zaYY,35149
-jararaca-0.3.17.dist-info/METADATA,sha256=PvHosK9SliL8gr08GUz7H9TVNStY1DMWehVN825-LL4,5149
-jararaca-0.3.17.dist-info/WHEEL,sha256=sP946D7jFCHeNz5Iq4fL4Lu-PrWrFsgfLXbbkciIZwg,88
-jararaca-0.3.17.dist-info/entry_points.txt,sha256=WIh3aIvz8LwUJZIDfs4EeH3VoFyCGEk7cWJurW38q0I,45
-jararaca-0.3.17.dist-info/RECORD,,
+jararaca-0.3.19.dist-info/LICENSE,sha256=OXLcl0T2SZ8Pmy2_dmlvKuetivmyPd5m1q-Gyd-zaYY,35149
+jararaca-0.3.19.dist-info/METADATA,sha256=Snc9WSwcSrBb-PyTB9v4BWurlSLPhuTaEHjKMZGdUck,5149
+jararaca-0.3.19.dist-info/WHEEL,sha256=sP946D7jFCHeNz5Iq4fL4Lu-PrWrFsgfLXbbkciIZwg,88
+jararaca-0.3.19.dist-info/entry_points.txt,sha256=WIh3aIvz8LwUJZIDfs4EeH3VoFyCGEk7cWJurW38q0I,45
+jararaca-0.3.19.dist-info/RECORD,,

pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "jararaca"
-version = "0.3.17"
+version = "0.3.19"
 description = "A simple and fast API framework for Python"
 authors = ["Lucas S <me@luscasleo.dev>"]
 readme = "README.md"
@@ -20,9 +20,9 @@ croniter = "^3.0.3"
 redis = "^5.0.8"
 websockets = "^13.0.1"
 opentelemetry-exporter-otlp-proto-http = { version = "^1.27.0", optional = true }
-opentelemetry-api = { version = "^1.27.0", optional = true }
-opentelemetry-sdk = { version = "^1.27.0", optional = true }
-opentelemetry-distro = { version = "^0.49b2", optional = true }
+opentelemetry-api = "^1.38.0"
+opentelemetry-sdk = "^1.38.0"
+opentelemetry-distro = "^0.59b0"
 opentelemetry-exporter-otlp = { version = "^1.27.0", optional = true }
 types-croniter = "^3.0.3.20240731"
 types-redis = "^4.6.0.20240903"