diff --git a/docs/index.rst b/docs/index.rst index b94289a..c5d6de5 100644 --- a/docs/index.rst +++ b/docs/index.rst @@ -16,7 +16,7 @@ A native distributed task queue for Python. .. code-block:: shell uv add kuu - # extras: msgspec, nats, prometheus, redis, dashboard + # extras: msgspec, nats, prometheus, redis, dashboard, di .. toctree:: :maxdepth: 2 @@ -24,6 +24,7 @@ A native distributed task queue for Python. quickstart usage + integrations config apidocs/index @@ -48,7 +49,8 @@ At a Glance uv run kuu start See :doc:`quickstart` for setup, :doc:`usage` for the full API tour, -:doc:`config` for tunables, or :doc:`apidocs/index` for the reference. +:doc:`integrations` for OpenTelemetry/Prometheus/Dishka/structlog, :doc:`config` +for tunables, or :doc:`apidocs/index` for the reference. Indices ------- diff --git a/docs/integrations.rst b/docs/integrations.rst new file mode 100644 index 0000000..01c8664 --- /dev/null +++ b/docs/integrations.rst @@ -0,0 +1,213 @@ +Integrations +============ + +Integrations under ``kuu.contrib`` are optional. Each one has its own install +extra: + +.. code-block:: shell + + uv add "kuu[otel]" # OpenTelemetry traces, metrics, logs + uv add "kuu[prometheus]" # Prometheus metrics + uv add "kuu[di]" # Dishka dependency injection + uv add "kuu[structlog]" # structlog task logging + +OpenTelemetry +------------- + +``kuu.contrib.otel`` adds traces, metrics, and log correlation. +``KuuOTELInstrumentor`` sets up all three: + +.. code-block:: python + + from kuu.contrib.otel import KuuOTELInstrumentor + + # Auto-configure the SDK from OTEL_EXPORTER_OTLP_ENDPOINT: + KuuOTELInstrumentor(app=app).instrument() + + # Or reuse SDK providers you configured yourself: + KuuOTELInstrumentor(app=app).instrument(setup_sdk=False) + +``instrument()`` inserts ``OtelTracingMiddleware`` at the front of the chain (a +producer span on enqueue, a consumer span on process, with W3C trace-context +propagated through message headers), registers ``OtelMetrics`` on +``app.events``, and routes stdlib ``logging`` to the OTel log provider. + +Pass your own exporters (for example gRPC): + +.. code-block:: python + + from opentelemetry.exporter.otlp.proto.grpc.trace_exporter import OTLPSpanExporter + + KuuOTELInstrumentor(app=app, span_exporter=OTLPSpanExporter(...)).instrument() + +Emitted metrics: ``kuu.task.enqueued``, ``kuu.task.processed`` (by ``status``), +``kuu.task.duration``, ``kuu.task.in_flight``, ``kuu.task.retried``. Call +``shutdown_telemetry()`` on shutdown to flush exporters. ``OtelTracingMiddleware``, +``OtelMetrics``, and ``OtelLoggingBridge`` can also be used on their own. + +Prometheus +---------- + +There are two parts, both using ``app.events``: a worker-side emitter and +client-side middleware. + +Worker metrics run in the control plane. With ``[metrics] enable = true`` the +orchestrator runs an aggregator HTTP server and each worker writes to a shared +multiprocess directory: + +.. code-block:: toml + + [metrics] + enable = true + host = "0.0.0.0" + port = 9191 + +Scrape ``http://host:9191/metrics``. + +Client-side middleware records enqueue metrics from every ``task.q(...)`` call: + +.. code-block:: python + + from kuu.contrib.prometheus import ClientMetrics + + app = Kuu(broker=..., middleware=[ClientMetrics()]) + +This emits ``client_enqueued_total``, ``client_enqueue_errors_total`` and +``client_enqueue_duration_seconds`` on the same registry the worker side uses. +Expose them from your own ASGI app: + +.. code-block:: python + + from fastapi import FastAPI + from kuu.contrib.prometheus import asgi_app + + api = FastAPI() + api.mount("/metrics", asgi_app()) + +Dishka (Dependency Injection) +----------------------------- + +``kuu.contrib.dishka`` resolves task dependencies from a `Dishka +`_ container. Declare them on the task signature +and each run gets them from a fresh ``REQUEST`` scope. + +Mark injected parameters with ``FromDishka()`` as the default value. The +dependency type comes from the annotation, so the parameter reads like a normal +typed argument. Apply ``@inject`` below ``@app.task`` and register the container +with ``setup_dishka``: + +.. code-block:: python + + from dishka import Provider, Scope, make_async_container, provide + from kuu import Kuu + from kuu.contrib.dishka import FromDishka, inject, setup_dishka + + + class QueryService: ... + class WriteService: ... + + + class AppProvider(Provider): + scope = Scope.REQUEST + read = provide(QueryService) + write = provide(WriteService) + + + app = Kuu(broker=...) + setup_dishka(app, make_async_container(AppProvider())) + + + @app.task + @inject + async def get_lot( + lot_id: str, + read: QueryService = FromDishka(), + write: WriteService = FromDishka(), + fetch: bool = True, + ) -> str: + ... + + # kuu sees only the business signature `(lot_id, fetch)`: + await get_lot.q("lot-1", fetch=False) + +``@inject`` removes the ``FromDishka()`` parameters from the public signature. +They are not serialised into the message or expected from the caller; the +container provides them when the worker runs the task. Without ``@inject`` the +worker rebinds the payload against the full signature and fails with +``Missing argument 'read'``. ``setup_dishka`` inserts ``KuuDishkaMiddleware`` at +the front of the chain, which opens a ``REQUEST`` scope around every task and +closes it afterwards. + +Sync (blocking) tasks +~~~~~~~~~~~~~~~~~~~~~~ + +``blocking=True`` tasks run in a worker thread and need a **sync** container. +The choice is strict: async tasks resolve from ``container`` and blocking tasks +from ``sync_container``. Provide both if you mix them: + +.. code-block:: python + + from dishka import make_container, make_async_container + + setup_dishka( + app, + make_async_container(AppProvider()), + sync_container=make_container(AppProvider()), + ) + + @app.task(blocking=True) + @inject + def crunch(lot_id: str, read: QueryService = FromDishka()) -> int: + ... + +Using your own ContextVar +~~~~~~~~~~~~~~~~~~~~~~~~~~~ + +If you already keep the container in a ``ContextVar`` (a request middleware, a +shared web app), point ``inject`` at it with ``from_contextvar`` and skip the +kuu middleware. You manage the scope yourself: + +.. code-block:: python + + from kuu.contrib.dishka import from_contextvar, inject + + from myapp.ioc import container_var # ContextVar[AsyncContainer] + + @app.task + @inject(container_getter=from_contextvar(container_var)) + async def handler(read: QueryService = FromDishka()) -> None: + ... + +To let kuu manage the scope but write the request container into your own var, +pass ``context_var=`` to the middleware and read the same var in ``inject``: + +.. code-block:: python + + from kuu.contrib.dishka import KuuDishkaMiddleware, from_contextvar + + app.middleware.insert(0, KuuDishkaMiddleware(container, context_var=container_var)) + +The subscript form ``read: FromDishka[QueryService]`` also works, so existing +Dishka signatures keep running. + +Structlog +--------- + +``kuu.contrib.structlog.StructlogMiddleware`` logs task start, success, and +failure through a ``structlog`` logger, during the ``process`` phase only. +Configure structlog yourself and pass a bound logger: + +.. code-block:: python + + import structlog + from kuu import Kuu + from kuu.contrib.structlog import StructlogMiddleware + + structlog.configure(...) + app = Kuu(broker=..., middleware=[StructlogMiddleware(structlog.get_logger())]) + +It emits ``task.start``, ``task.ok``, and ``task.fail`` events with +``task_name``, ``sched_id``, ``queue``, ``attempt``, ``duration`` and +``exc_type`` keys. It implements :class:`kuu.middleware.TaskLogSink`; implement +that protocol and drive it with :func:`kuu.middleware.run_process_task_logging` +to log through any other frontend. diff --git a/docs/usage.rst b/docs/usage.rst index 8296104..246c0b5 100644 --- a/docs/usage.rst +++ b/docs/usage.rst @@ -318,6 +318,14 @@ Custom Middleware ``ctx.phase`` is ``"enqueue"`` for the publish path and ``"process"`` for the worker path. Most middleware filters on phase. +Dependency Injection +-------------------- + +``kuu.contrib.dishka`` resolves task dependencies from a `Dishka +`_ container per execution. Declare them as +``FromDishka()`` defaults and apply ``@inject`` below ``@app.task``. See +:doc:`integrations` for setup, sync containers, and custom ``ContextVar`` wiring. + Scheduler --------- @@ -406,6 +414,7 @@ Jobs could be registered with a mapping, instead of being registered one-by-one. So, instead of: .. code-block:: python + async def generate_notifications( kind: Literal["regular", "unique"], ) -> None: @@ -424,6 +433,7 @@ So, instead of: One could write: .. code-block:: python + @broker.sched({ at(time(18)): Payload(kwargs={"kind":"regular"}), on_day(Thu) & at(time(18, 30)): Payload(kwargs={"kind":"unique"}), @@ -432,7 +442,7 @@ One could write: kind: Literal["regular", "unique"], ) -> None: # ... -``` + Serializers ----------- @@ -598,50 +608,13 @@ custom tooling: envelope_to_bytes, envelope_from_bytes, ) -Prometheus Metrics ------------------- - -Two integration points: a worker-side emitter and client-side middleware. -Both wire onto ``app.events``. - -Worker Metrics -~~~~~~~~~~~~~~ - -When ``[metrics] enable = true``, the orchestrator runs an aggregator HTTP -server on ``metrics.port`` and tells each worker subprocess to write to a -shared multiprocess directory. - -.. code-block:: toml - - [metrics] - enable = true - host = "0.0.0.0" - port = 9191 - -Hit ``http://host:9191/metrics``. - -Client-Side -~~~~~~~~~~~ - -.. code-block:: python - - from kuu.prometheus import ClientMetrics - - app = Kuu(broker=..., middleware=[ClientMetrics()]) - -This records ``client_enqueued_total``, ``client_enqueue_errors_total`` and -``client_enqueue_duration_seconds`` from every ``task.q(...)`` call, on the -same registry the worker side uses. - -To expose them from your own ASGI app: - -.. code-block:: python - - from fastapi import FastAPI - from kuu.prometheus import asgi_app +Metrics +------- - api = FastAPI() - api.mount("/metrics", asgi_app()) +Worker-side Prometheus metrics are driven by the control plane via +``[metrics] enable = true``; client-side enqueue metrics and OpenTelemetry +traces/metrics/logs are available under ``kuu.contrib``. See +:doc:`integrations`. Dashboard --------- diff --git a/kuu/_queue_drain.py b/kuu/_queue_drain.py index 574590b..8fd2dc9 100644 --- a/kuu/_queue_drain.py +++ b/kuu/_queue_drain.py @@ -3,18 +3,16 @@ import logging from collections.abc import Callable from queue import Empty -from typing import Protocol, TypeVar +from typing import Protocol import anyio -T = TypeVar("T") - -class SyncGetQueue(Protocol[T]): +class SyncGetQueue[T](Protocol): def get_nowait(self) -> T: ... -async def drain_sync_queue( +async def drain_sync_queue[T]( q: SyncGetQueue[T], handler: Callable[[T], None], *, @@ -34,7 +32,7 @@ async def drain_sync_queue( did_work = False for _ in range(batch_size): try: - item = q.get_nowait() + item: T = q.get_nowait() did_work = True except Empty: break diff --git a/kuu/contrib/__init__.py b/kuu/contrib/__init__.py new file mode 100644 index 0000000..75eeb18 --- /dev/null +++ b/kuu/contrib/__init__.py @@ -0,0 +1 @@ +"""Optional, dependency-bearing integrations for Kuu (OpenTelemetry, Prometheus, ...).""" diff --git a/kuu/contrib/dishka/__init__.py b/kuu/contrib/dishka/__init__.py new file mode 100644 index 0000000..d00eb4f --- /dev/null +++ b/kuu/contrib/dishka/__init__.py @@ -0,0 +1,191 @@ +from __future__ import annotations + +__all__ = ( + "CONTAINER_STATE_KEY", + "AnyContainer", + "FromDishka", + "KuuDishkaMiddleware", + "current_container", + "from_contextvar", + "inject", + "setup_dishka", +) + +import functools +import inspect +from collections.abc import Callable +from contextvars import ContextVar +from inspect import Parameter +from typing import Any, overload + +from dishka import AsyncContainer, Container +from dishka.entities.component import DEFAULT_COMPONENT, Component +from dishka.entities.key import DependencyKey +from dishka.integrations.base import default_parse_dependency, wrap_injection + +from kuu import Kuu +from kuu.context import Context +from kuu.middleware.base import Next + +type AnyContainer = AsyncContainer | Container +type ContainerGetter = Callable[..., AnyContainer] + + +class _FromDishkaMarker: + __slots__ = ("component",) + + def __init__(self, component: Component = DEFAULT_COMPONENT) -> None: + self.component = component + + def __repr__(self) -> str: # pragma: no cover - cosmetic + return f"FromDishka(component={self.component!r})" + + +def FromDishka(component: Component = DEFAULT_COMPONENT) -> Any: + return _FromDishkaMarker(component) + + +def _parse_dependency(parameter: Parameter, hint: Any) -> DependencyKey | None: + default = parameter.default + if isinstance(default, _FromDishkaMarker): + return DependencyKey(hint, default.component) + return default_parse_dependency(parameter, hint) + + +CONTAINER_STATE_KEY = "dishka_container" + +# kuu's built-in REQUEST container slot, populated by KuuDishkaMiddleware. +# If you already keep the container in your *own* ContextVar, you don't have to +# use this one — see `from_contextvar` and `inject(container_getter=...)`. +_request_container: ContextVar[AnyContainer | None] = ContextVar( + "kuu_dishka_request_container", + default=None, +) + + +def current_container() -> AnyContainer: + container = _request_container.get() + if container is None: + raise LookupError( + "No Dishka REQUEST container in the current context. " + "Add KuuDishkaMiddleware to the app's middleware chain " + "(e.g. via setup_dishka(app, container))." + ) + return container + + +def from_contextvar(var: ContextVar[Any]) -> ContainerGetter: + def _getter(*_: Any, **__: Any) -> AnyContainer: + container = var.get() + if container is None: + raise LookupError( + f"ContextVar {var.name!r} holds no Dishka container; " + "it must be set before the task runs." + ) + return container + + return _getter + + +_container_getter: ContainerGetter = from_contextvar(_request_container) + + +@overload +def inject[**P, R](func: Callable[P, R], /) -> Callable[P, R]: ... + + +@overload +def inject[**P, R]( + *, + container_getter: ContainerGetter | None = ..., +) -> Callable[[Callable[P, R]], Callable[P, R]]: ... + + +def inject[**P, R]( + func: Callable[P, R] | None = None, + *, + container_getter: ContainerGetter | None = None, +) -> Callable[P, R] | Callable[[Callable[P, R]], Callable[P, R]]: + if func is None: + return functools.partial(inject, container_getter=container_getter) + # pyrefly: ignore [no-matching-overload] + return wrap_injection( + func=func, + container_getter=container_getter or _container_getter, + is_async=inspect.iscoroutinefunction(func), + remove_depends=True, + parse_dependency=_parse_dependency, + ) + + +class KuuDishkaMiddleware: + def __init__( + self, + container: AsyncContainer | None = None, + *, + sync_container: Container | None = None, + context_var: ContextVar[Any] | None = None, + ) -> None: + if container is None and sync_container is None: + raise ValueError( + "KuuDishkaMiddleware needs at least one of `container` " + "(AsyncContainer) or `sync_container` (Container)." + ) + self._async_app = container + self._sync_app = sync_container + self._var = context_var if context_var is not None else _request_container + + async def __call__(self, ctx: Context, call_next: Next) -> Any: + if ctx.phase != "process": + return await call_next() + + # A `blocking=True` task is `inject`-ed with a *sync* getter, so it must + # resolve from a sync `Container`; an async task from an `AsyncContainer`. + # The two are not interchangeable, so select strictly by `blocking`. + blocking = bool(ctx.task and ctx.task.blocking) + if blocking: + if self._sync_app is None: + raise LookupError( + f"task {getattr(ctx.task, 'task_name', '?')!r} needs a sync " + "Dishka container (it is blocking=True), but KuuDishkaMiddleware " + "was created without `sync_container=`." + ) + return await self._with_sync(ctx, call_next) + if self._async_app is None: + raise LookupError( + f"task {getattr(ctx.task, 'task_name', '?')!r} needs an async " + "Dishka container, but KuuDishkaMiddleware was created without " + "`container=`." + ) + return await self._with_async(ctx, call_next) + + async def _with_async(self, ctx: Context, call_next: Next) -> Any: + assert self._async_app is not None + async with self._async_app() as request: + token = self._var.set(request) + ctx.state[CONTAINER_STATE_KEY] = request + try: + return await call_next() + finally: + self._var.reset(token) + + async def _with_sync(self, ctx: Context, call_next: Next) -> Any: + assert self._sync_app is not None + with self._sync_app() as request: + token = self._var.set(request) + ctx.state[CONTAINER_STATE_KEY] = request + try: + return await call_next() + finally: + self._var.reset(token) + + +def setup_dishka( + app: Kuu, + container: AsyncContainer | None = None, + *, + sync_container: Container | None = None, +) -> KuuDishkaMiddleware: + mw = KuuDishkaMiddleware(container, sync_container=sync_container) + app.middleware.insert(0, mw) + return mw diff --git a/kuu/contrib/otel/__init__.py b/kuu/contrib/otel/__init__.py new file mode 100644 index 0000000..f9cfc55 --- /dev/null +++ b/kuu/contrib/otel/__init__.py @@ -0,0 +1,12 @@ +from ._instrumentor import KuuOTELInstrumentor, shutdown_telemetry +from ._logging import OtelLoggingBridge +from ._metrics import OtelMetrics +from ._traces import OtelTracingMiddleware + +__all__ = ( + "KuuOTELInstrumentor", + "shutdown_telemetry", + "OtelLoggingBridge", + "OtelMetrics", + "OtelTracingMiddleware", +) diff --git a/kuu/contrib/otel/_instrumentor.py b/kuu/contrib/otel/_instrumentor.py new file mode 100644 index 0000000..f2b8703 --- /dev/null +++ b/kuu/contrib/otel/_instrumentor.py @@ -0,0 +1,185 @@ +from __future__ import annotations + +import logging +import os +import typing + +from opentelemetry import metrics, trace +from opentelemetry.sdk.metrics import MeterProvider +from opentelemetry.sdk.metrics.export import PeriodicExportingMetricReader +from opentelemetry.sdk.resources import Resource +from opentelemetry.sdk.trace import TracerProvider +from opentelemetry.sdk.trace.export import BatchSpanProcessor + +from ._logging import OtelLoggingBridge +from ._metrics import OtelMetrics +from ._traces import OtelTracingMiddleware + +if typing.TYPE_CHECKING: + from opentelemetry.sdk._logs._internal.export import LogRecordExporter + from opentelemetry.sdk.metrics.export import MetricExporter + from opentelemetry.sdk.trace.export import SpanExporter + + from kuu.app import Kuu + +log = logging.getLogger("kuu.otel.instrumentor") + + +def _auto_setup_sdk( + endpoint: str | None = None, + *, + span_exporter: SpanExporter | None = None, + metric_exporter: MetricExporter | None = None, +) -> Resource | None: + endpoint = endpoint or os.getenv("OTEL_EXPORTER_OTLP_ENDPOINT") + + if span_exporter is None and metric_exporter is None and not endpoint: + log.debug("event=otel.no_endpoint") + return None + + service_name = os.getenv("OTEL_SERVICE_NAME", "kuu") + resource = Resource.create({"service.name": service_name}) + log.info("event=otel.auto_setup service=%s", service_name) + + if span_exporter is None: + from opentelemetry.exporter.otlp.proto.http.trace_exporter import OTLPSpanExporter + + span_exporter = OTLPSpanExporter(endpoint=endpoint) + + tp = TracerProvider(resource=resource) + tp.add_span_processor(BatchSpanProcessor(span_exporter)) + trace.set_tracer_provider(tp) + + if metric_exporter is None: + from opentelemetry.exporter.otlp.proto.http.metric_exporter import OTLPMetricExporter + + metric_exporter = OTLPMetricExporter(endpoint=endpoint) + + reader = PeriodicExportingMetricReader(metric_exporter) + mp = MeterProvider(resource=resource, metric_readers=[reader]) + metrics.set_meter_provider(mp) + + return resource + + +class KuuOTELInstrumentor: + """ + Wires OTEL instrumentation (traces + metrics + logs) into a Kuu app. + + Basic usage (uses globally configured OTel providers: bring your own + SDK setup):: + + KuuOTELInstrumentor(app=app).instrument(setup_sdk=False) + + Auto-setup with OTLP HTTP (reads ``OTEL_EXPORTER_OTLP_ENDPOINT``):: + + KuuOTELInstrumentor(app=app).instrument() + + Bring your own exporters (e.g. gRPC):: + + from opentelemetry.exporter.otlp.proto.grpc.trace_exporter import OTLPSpanExporter + from opentelemetry.exporter.otlp.proto.grpc.metric_exporter import OTLPMetricExporter + from opentelemetry.exporter.otlp.proto.grpc._log_exporter import OTLPLogExporter + + KuuOTELInstrumentor( + app=app, + span_exporter=OTLPSpanExporter(...), + metric_exporter=OTLPMetricExporter(...), + log_exporter=OTLPLogExporter(...), + ).instrument() + + ``instrument()`` will: + 1. Configure OTel SDK providers when exporters are supplied or + ``OTEL_EXPORTER_OTLP_ENDPOINT`` is set (skipped with + ``setup_sdk=False``). + 2. Insert ``OtelTracingMiddleware`` at the front of the middleware + chain. + 3. Wire ``OtelMetrics`` onto app.events signals. + 4. Bridge stdlib ``logging`` to the log provider when a log exporter + or ``OTEL_EXPORTER_OTLP_ENDPOINT`` is available. + """ + + def __init__( + self, + *, + app: Kuu, + tracer_name: str = "kuu", + meter_name: str = "kuu", + propagate: bool = True, + log_level: int = logging.INFO, + logger_name: str = "kuu", + span_exporter: SpanExporter | None = None, + metric_exporter: MetricExporter | None = None, + log_exporter: LogRecordExporter | None = None, + ): + self._app = app + self._tracer_name = tracer_name + self._meter_name = meter_name + self._propagate = propagate + self._log_level = log_level + self._logger_name = logger_name + self._span_exporter: SpanExporter | None = span_exporter + self._metric_exporter = metric_exporter + self._log_exporter = log_exporter + + self._otel_middleware: OtelTracingMiddleware | None = None + self._otel_metrics: OtelMetrics | None = None + self._logging_bridge: OtelLoggingBridge | None = None + + def instrument(self, *, setup_sdk: bool = True) -> None: + resource: Resource | None = None + if setup_sdk: + resource = _auto_setup_sdk( + span_exporter=self._span_exporter, + metric_exporter=self._metric_exporter, + ) + + self._otel_middleware = OtelTracingMiddleware( + tracer_name=self._tracer_name, + propagate_ctx=self._propagate, + ) + self._app.middleware.insert(0, self._otel_middleware) + + self._otel_metrics = OtelMetrics( + app=self._app, + meter_name=self._meter_name, + ) + + self._logging_bridge = OtelLoggingBridge( + level=self._log_level, + logger_name=self._logger_name, + ) + self._logging_bridge.setup(resource=resource, log_exporter=self._log_exporter) + + def uninstrument(self) -> None: + if self._otel_middleware is not None: + try: + self._app.middleware.remove(self._otel_middleware) + except ValueError: + pass + self._otel_middleware = None + + if self._otel_metrics is not None: + self._otel_metrics.disconnect() + self._otel_metrics = None + + if self._logging_bridge is not None: + self._logging_bridge.shutdown() + self._logging_bridge = None + + +def shutdown_telemetry() -> None: + from contextlib import suppress + + with suppress(Exception): + prov = trace.get_tracer_provider() + if hasattr(prov, "shutdown") and callable(prov.shutdown): + prov.shutdown() + + with suppress(Exception): + prov = metrics.get_meter_provider() + if hasattr(prov, "shutdown") and callable(prov.shutdown): + prov.shutdown() + + +__all__ = ("shutdown_telemetry", "KuuOTELInstrumentor") diff --git a/kuu/contrib/otel/_logging.py b/kuu/contrib/otel/_logging.py new file mode 100644 index 0000000..7e127a7 --- /dev/null +++ b/kuu/contrib/otel/_logging.py @@ -0,0 +1,90 @@ +from __future__ import annotations + +import logging +import os +import typing + +from opentelemetry.instrumentation.logging.handler import LoggingHandler +from opentelemetry.sdk._logs import LoggerProvider +from opentelemetry.sdk._logs._internal.export import BatchLogRecordProcessor +from opentelemetry.sdk.resources import Resource + +if typing.TYPE_CHECKING: + from opentelemetry.sdk._logs._internal.export import LogRecordExporter + +log = logging.getLogger("kuu.otel") + + +class OtelLoggingBridge: + """Bridges Python's ``logging`` to an OpenTelemetry ``LoggerProvider``. + + Attaches an OTel ``LoggingHandler`` to the named stdlib logger so that + any log records emitted through it are forwarded to the provider. Does + **not** configure any particular logging frontend : use ``logging``, + ``structlog``, or anything else that ultimately produces + ``logging.LogRecord`` instances. + + Provider resolution order (first wins): + 1. ``logger_provider`` passed to :meth:`setup`. + 2. ``log_exporter`` passed to :meth:`setup` : a new ``LoggerProvider`` + is created around it. + 3. Falls back to OTLP HTTP (``OTLPLogExporter``) only when + ``OTEL_EXPORTER_OTLP_ENDPOINT`` is set and neither of the above + was supplied. + """ + + def __init__( + self, + *, + level: int = logging.INFO, + logger_name: str = "kuu", + ): + self._level = level + self._logger_name = logger_name + self._provider: LoggerProvider | None = None + self._handler: LoggingHandler | None = None + self._owns_provider: bool = False + + def setup( + self, + resource: Resource | None = None, + *, + logger_provider: LoggerProvider | None = None, + log_exporter: LogRecordExporter | None = None, + ) -> logging.Logger: + if logger_provider is not None: + self._provider = logger_provider + self._owns_provider = False + else: + exporter = log_exporter + if exporter is None and os.getenv("OTEL_EXPORTER_OTLP_ENDPOINT"): + from opentelemetry.exporter.otlp.proto.http._log_exporter import OTLPLogExporter + + exporter = OTLPLogExporter() + + if exporter is not None: + res = resource or Resource.create({}) + self._provider = LoggerProvider(resource=res) + self._provider.add_log_record_processor(BatchLogRecordProcessor(exporter)) + self._owns_provider = True + + if self._provider is None: + log.debug("event=otel.no_provider") + return logging.getLogger(self._logger_name) + + self._handler = LoggingHandler(level=self._level, logger_provider=self._provider) + logger = logging.getLogger(self._logger_name) + logger.addHandler(self._handler) + log.info("event=otel.bridge_attached logger=%s", self._logger_name) + return logger + + def shutdown(self) -> None: + if self._handler is not None: + logging.getLogger(self._logger_name).removeHandler(self._handler) + self._handler = None + if self._owns_provider and self._provider is not None: + self._provider.shutdown() + self._provider = None + + +__all__ = ("OtelLoggingBridge",) diff --git a/kuu/contrib/otel/_metrics.py b/kuu/contrib/otel/_metrics.py new file mode 100644 index 0000000..62a67a4 --- /dev/null +++ b/kuu/contrib/otel/_metrics.py @@ -0,0 +1,114 @@ +from __future__ import annotations + +import typing +from typing import Any + +from opentelemetry import metrics + +if typing.TYPE_CHECKING: + from kuu.app import Kuu + from kuu.message import Message + + +class OtelMetrics: + """ + Wires OTEL metrics onto Kuu task lifecycle events. + + Instruments created (all under ``meter_name`` scope): + - ``kuu.task.enqueued`` (Counter) : tasks enqueued + - ``kuu.task.processed`` (Counter) : tasks processed, with ``status`` label + - ``kuu.task.duration`` (Histogram, seconds) : task execution duration + - ``kuu.task.in_flight`` (UpDownCounter) : tasks currently executing + - ``kuu.task.retried`` (Counter) : retry attempts + """ + + def __init__( + self, + *, + app: Kuu, + meter_name: str = "kuu", + ): + self._app = app + self._meter = metrics.get_meter(meter_name) + + self._enqueued = self._meter.create_counter( + "kuu.task.enqueued", + "tasks", + "Tasks enqueued", + ) + self._processed = self._meter.create_counter( + "kuu.task.processed", + "tasks", + "Tasks processed by outcome", + ) + self._duration = self._meter.create_histogram( + "kuu.task.duration", + "s", + "Task execution duration", + ) + self._in_flight = self._meter.create_up_down_counter( + "kuu.task.in_flight", + "tasks", + "Tasks currently executing", + ) + self._retried = self._meter.create_counter( + "kuu.task.retried", + "tasks", + "Task retry attempts", + ) + + self._inflight_ids = set() + + ev = app.events + ev.task_enqueued.connect(self._on_enqueued) + ev.task_started.connect(self._on_started) + ev.task_succeeded.connect(self._on_succeeded) + ev.task_failed.connect(self._on_failed) + ev.task_retried.connect(self._on_retried) + ev.task_dead.connect(self._on_dead) + + def _attrs(self, msg: Message, **extra: str) -> dict[str, Any]: + base: dict[str, Any] = {"task.name": msg.task, "queue": msg.queue} + base.update(extra) + return base + + def _settle(self, msg: Message) -> None: + if msg.id in self._inflight_ids: + self._inflight_ids.discard(msg.id) + self._in_flight.add(-1, self._attrs(msg)) + + def _on_enqueued(self, msg: Message) -> None: + self._enqueued.add(1, self._attrs(msg)) + + def _on_started(self, msg: Message) -> None: + self._inflight_ids.add(msg.id) + self._in_flight.add(1, self._attrs(msg)) + + def _on_succeeded(self, msg: Message, elapsed: float) -> None: + self._duration.record(elapsed, self._attrs(msg)) + self._processed.add(1, self._attrs(msg, status="ok")) + self._settle(msg) + + def _on_failed(self, msg: Message, exc: BaseException) -> None: + self._processed.add(1, self._attrs(msg, status="error")) + self._settle(msg) + + def _on_retried(self, msg: Message, delay: float) -> None: + self._retried.add(1, self._attrs(msg)) + self._settle(msg) + + def _on_dead(self, msg: Message) -> None: + self._processed.add(1, self._attrs(msg, status="dead")) + self._settle(msg) + + def disconnect(self) -> None: + ev = self._app.events + ev.task_enqueued.disconnect(self._on_enqueued) + ev.task_started.disconnect(self._on_started) + ev.task_succeeded.disconnect(self._on_succeeded) + ev.task_failed.disconnect(self._on_failed) + ev.task_retried.disconnect(self._on_retried) + ev.task_dead.disconnect(self._on_dead) + + +__all__ = ("OtelMetrics",) diff --git a/kuu/contrib/otel/_traces.py b/kuu/contrib/otel/_traces.py new file mode 100644 index 0000000..c0acb20 --- /dev/null +++ b/kuu/contrib/otel/_traces.py @@ -0,0 +1,118 @@ +from __future__ import annotations + +import typing +from typing import Any + +from opentelemetry import context as otel_context +from opentelemetry import propagate, trace +from opentelemetry.trace import SpanKind, Status, StatusCode, Tracer + +if typing.TYPE_CHECKING: + from kuu.context import Context + from kuu.middleware.base import Next + + +def _span_name(ctx: Context) -> str: + task_name = ctx.task.task_name if ctx.task else ctx.message.task + verb = "publish" if ctx.phase == "enqueue" else "process" + return f"{task_name} {verb}" + + +def _set_messaging_attributes(span: Any, ctx: Context) -> None: + msg = ctx.message + span.set_attribute("messaging.system", "kuu") + span.set_attribute("messaging.operation", "publish" if ctx.phase == "enqueue" else "process") + span.set_attribute("messaging.destination", msg.queue) + span.set_attribute("messaging.destination.name", msg.queue) + span.set_attribute("messaging.message.id", str(msg.id)) + span.set_attribute("kuu.task.name", msg.task) + span.set_attribute("kuu.task.attempt", msg.attempt) + span.set_attribute("kuu.task.max_attempts", msg.max_attempts) + + +class OtelTracingMiddleware: + """ + Middleware that creates OTEL spans around enqueue and process phases. + + Enqueue phase (PRODUCER): + 1. Start span ``{task_name} publish`` with PRODUCER kind. + 2. Set messaging attributes. + 3. Inject W3C TraceContext into ``ctx.message.headers``. + 4. Execute the middleware chain. + 5. End the span (set ERROR status + record exception on failure). + + Process phase (CONSUMER): + 1. Extract W3C TraceContext from ``ctx.message.headers``. + 2. Attach extracted context. + 3. Start span ``{task_name} process`` with CONSUMER kind. + 4. Execute the middleware chain. + 5. Set OK / ERROR status. + 6. Detach context and end span. + """ + + def __init__( + self, + *, + tracer_name: str = "kuu", + propagate_ctx: bool = True, + ): + self._tracer_name = tracer_name + self._propagate = propagate_ctx + self._tracer: Tracer | None = None + + @property + def tracer(self) -> Tracer: + if self._tracer is None: + self._tracer = trace.get_tracer(self._tracer_name) + return self._tracer + + async def __call__(self, ctx: Context, call_next: Next) -> Any: + if ctx.phase == "enqueue": + return await self._enqueue(ctx, call_next) + if ctx.phase == "process": + return await self._process(ctx, call_next) + return await call_next() + + async def _enqueue(self, ctx: Context, call_next: Next) -> Any: + span = self.tracer.start_span( + _span_name(ctx), + kind=SpanKind.PRODUCER, + ) + _set_messaging_attributes(span, ctx) + + with trace.use_span(span, end_on_exit=True): + if self._propagate: + propagate.inject(ctx.message.headers) + + try: + return await call_next() + except Exception as exc: + span.set_status(Status(StatusCode.ERROR)) + span.record_exception(exc) + raise + + async def _process(self, ctx: Context, call_next: Next) -> Any: + extracted = propagate.extract(carrier=ctx.message.headers) if self._propagate else None + token = otel_context.attach(extracted) if extracted else None + + span = self.tracer.start_span( + _span_name(ctx), + kind=SpanKind.CONSUMER, + ) + _set_messaging_attributes(span, ctx) + + with trace.use_span(span, end_on_exit=True): + try: + result = await call_next() + span.set_status(Status(StatusCode.OK)) + return result + except Exception as exc: + span.set_status(Status(StatusCode.ERROR)) + span.record_exception(exc) + raise + finally: + if token is not None: + otel_context.detach(token) + + +__all__ = ("OtelTracingMiddleware",) diff --git a/kuu/contrib/prometheus/__init__.py b/kuu/contrib/prometheus/__init__.py new file mode 100644 index 0000000..8481eed --- /dev/null +++ b/kuu/contrib/prometheus/__init__.py @@ -0,0 +1,311 @@ +from __future__ import annotations + +import os +import time +from threading import Thread +from typing import TYPE_CHECKING, Annotated, Any +from uuid import UUID +from wsgiref.simple_server import WSGIServer + +from prometheus_client import REGISTRY, CollectorRegistry, Counter, Gauge, Histogram +from typing_extensions import Doc + +if TYPE_CHECKING: + from kuu.app import Kuu + from kuu.context import Context + from kuu.message import Message + from kuu.middleware.base import Next + + +__all__ = ( + "WorkerMetrics", + "ClientMetrics", + "serve", + "asgi_app", + "mark_worker_dead", +) + + +_DEFAULT_DURATION_BUCKETS = ( + 0.005, + 0.01, + 0.025, + 0.05, + 0.1, + 0.25, + 0.5, + 1.0, + 2.5, + 5.0, + 10.0, + 30.0, + 60.0, +) +_DEFAULT_RETRY_DELAY_BUCKETS = (0.0, 0.1, 0.5, 1.0, 5.0, 15.0, 60.0, 300.0, 900.0) +_DEFAULT_ENQUEUE_BUCKETS = (0.001, 0.005, 0.01, 0.025, 0.05, 0.1, 0.25, 1.0, 5.0) + + +class WorkerMetrics: + """ + Worker-side, event-driven Prometheus metrics. Wires onto `app.events`. + + Instantiate inside each worker subprocess after the app/registry is loaded. + Exposes per-task counters/histograms and an in-flight gauge. In multiprocess + mode (PROMETHEUS_MULTIPROC_DIR set) the gauge sums across living workers. + + Metric names ({ns} = namespace, default "qq"): + - {ns}_tasks_received_total{task,queue} + - {ns}_tasks_started_total{task,queue} + - {ns}_tasks_succeeded_total{task,queue} + - {ns}_tasks_retried_total{task,queue} + - {ns}_tasks_failed_total{task,queue,exc} + - {ns}_tasks_dead_total{task,queue} + - {ns}_task_duration_seconds{task,queue} (from task_succeeded) + - {ns}_task_retry_delay_seconds{task,queue} (from task_retried) + - {ns}_tasks_in_flight{task,queue} gauge (livesum across procs) + + `in_flight` is decremented on the first terminal event per message sched_id + (succeeded/retried/failed/dead) to handle the Fail->Dead chain and the + standalone Reject->Dead path without double-counting. + """ + + def __init__( + self, + app: Kuu, + namespace: str = "kuu", + registry: CollectorRegistry = REGISTRY, + duration_buckets: tuple[float, ...] = _DEFAULT_DURATION_BUCKETS, + retry_delay_buckets: tuple[float, ...] = _DEFAULT_RETRY_DELAY_BUCKETS, + ): + self._inflight_ids: set[UUID] = set() + + self.received = Counter( + f"{namespace}_tasks_received_total", + "Tasks received by worker", + ("task", "queue"), + registry=registry, + ) + self.started = Counter( + f"{namespace}_tasks_started_total", + "Tasks started processing", + ("task", "queue"), + registry=registry, + ) + self.succeeded = Counter( + f"{namespace}_tasks_succeeded_total", + "Tasks succeeded", + ("task", "queue"), + registry=registry, + ) + self.retried = Counter( + f"{namespace}_tasks_retried_total", + "Tasks scheduled for retry", + ("task", "queue"), + registry=registry, + ) + self.failed = Counter( + f"{namespace}_tasks_failed_total", + "Tasks raised an exception", + ("task", "queue", "exc"), + registry=registry, + ) + self.dead = Counter( + f"{namespace}_tasks_dead_total", + "Tasks reached terminal state without success", + ("task", "queue"), + registry=registry, + ) + self.duration = Histogram( + f"{namespace}_task_duration_seconds", + "Task processing duration", + ("task", "queue"), + buckets=duration_buckets, + registry=registry, + ) + self.retry_delay = Histogram( + f"{namespace}_task_retry_delay_seconds", + "Scheduled retry delay", + ("task", "queue"), + buckets=retry_delay_buckets, + registry=registry, + ) + self.in_flight = Gauge( + f"{namespace}_tasks_in_flight", + "Tasks currently being processed", + ("task", "queue"), + multiprocess_mode="livesum", + registry=registry, + ) + + ev = app.events + ev.task_received.connect(self._on_received) + ev.task_started.connect(self._on_started) + ev.task_succeeded.connect(self._on_succeeded) + ev.task_retried.connect(self._on_retried) + ev.task_failed.connect(self._on_failed) + ev.task_dead.connect(self._on_dead) + + def _settle(self, msg: Message) -> None: + if msg.id in self._inflight_ids: + self._inflight_ids.discard(msg.id) + self.in_flight.labels(msg.task, msg.queue).dec() + + def _on_received(self, msg: Message) -> None: + self.received.labels(msg.task, msg.queue).inc() + + def _on_started(self, msg: Message) -> None: + self.started.labels(msg.task, msg.queue).inc() + self._inflight_ids.add(msg.id) + self.in_flight.labels(msg.task, msg.queue).inc() + + def _on_succeeded(self, msg: Message, elapsed: float) -> None: + self.duration.labels(msg.task, msg.queue).observe(elapsed) + self.succeeded.labels(msg.task, msg.queue).inc() + self._settle(msg) + + def _on_retried(self, msg: Message, delay: float) -> None: + self.retry_delay.labels(msg.task, msg.queue).observe(delay) + self.retried.labels(msg.task, msg.queue).inc() + self._settle(msg) + + def _on_failed(self, msg: Message, exc: BaseException) -> None: + self.failed.labels(msg.task, msg.queue, type(exc).__name__).inc() + self._settle(msg) + + def _on_dead(self, msg: Message) -> None: + self.dead.labels(msg.task, msg.queue).inc() + self._settle(msg) + + +class ClientMetrics: + """ + Client-side metrics, implemented as a Middleware. Add to `Kuu(middleware=[...])`. + + Acts only on the `enqueue` phase. Records: + - {ns}_client_enqueued_total{task,queue} + - {ns}_client_enqueue_errors_total{task,queue,exc} + - {ns}_client_enqueue_duration_seconds{task,queue} (broker handoff latency) + """ + + def __init__( + self, + namespace: str = "kuu", + registry: CollectorRegistry = REGISTRY, + duration_buckets: tuple[float, ...] = _DEFAULT_ENQUEUE_BUCKETS, + ): + self.enqueued = Counter( + f"{namespace}_client_enqueued_total", + "Tasks enqueued by client", + ("task", "queue"), + registry=registry, + ) + self.errors = Counter( + f"{namespace}_client_enqueue_errors_total", + "Enqueue attempts that raised", + ("task", "queue", "exc"), + registry=registry, + ) + self.duration = Histogram( + f"{namespace}_client_enqueue_duration_seconds", + "Time spent enqueueing to broker", + ("task", "queue"), + buckets=duration_buckets, + registry=registry, + ) + + async def __call__(self, ctx: Context, call_next: Next) -> Any: + if ctx.phase != "enqueue": + return await call_next() + + msg = ctx.message + t0 = time.perf_counter() + try: + result = await call_next() + except BaseException as exc: + self.errors.labels(msg.task, msg.queue, type(exc).__name__).inc() + raise + else: + self.enqueued.labels(msg.task, msg.queue).inc() + return result + finally: + self.duration.labels(msg.task, msg.queue).observe(time.perf_counter() - t0) + + +def serve( + host: Annotated[str, Doc("Host to bind the /metrics HTTP server to.")] = "0.0.0.0", + port: Annotated[int, Doc("Port to bind the /metrics HTTP server to.")] = 9100, + multiprocess_dir: Annotated[ + str | None, + Doc( + "Directory for prometheus_client multiprocess shared state. When set, aggregates metrics across all worker subprocesses that wrote to this directory." + ), + ] = None, +) -> tuple[WSGIServer, Thread]: + """ + Start a `/metrics` HTTP endpoint. Returns (server, thread) from prometheus_client. + + If `multiprocess_dir` is given (or `PROMETHEUS_MULTIPROC_DIR` env is set), + aggregates metrics from on-disk shared state across all worker subprocesses + that wrote to that dir. Otherwise serves the in-process REGISTRY directly. + + Call this in the orchestrator (parent) process when running multiple workers, + or in-process for a single-worker setup. + """ + from prometheus_client import start_http_server + + mp_dir = multiprocess_dir or os.environ.get("PROMETHEUS_MULTIPROC_DIR") + if mp_dir: + from prometheus_client.multiprocess import MultiProcessCollector + + os.environ["PROMETHEUS_MULTIPROC_DIR"] = mp_dir + registry = CollectorRegistry() + MultiProcessCollector(registry, path=mp_dir) + return start_http_server(port, addr=host, registry=registry) + return start_http_server(port, addr=host) + + +def asgi_app(multiprocess_dir: str | None = None) -> Any: + """ + Build an ASGI app exposing `/metrics`. Mount it into your own ASGI + application (FastAPI, Starlette, etc.). + + If `multiprocess_dir` is given (or `PROMETHEUS_MULTIPROC_DIR` env is set), + aggregates metrics from on-disk shared state across worker subprocesses. + Otherwise serves the in-process REGISTRY directly. + + Examples: + # Client side, FastAPI: + from fastapi import FastAPI + from kuu.contrib.prometheus import ClientMetrics, asgi_app + + app = FastAPI() + kuu_app = Kuu(broker=..., middleware=[ClientMetrics()]) + app.mount("/metrics", asgi_app()) + + # Worker side, exposing aggregated multiproc metrics from a sidecar + # ASGI server (e.g. uvicorn) instead of the orchestrator's WSGI server: + app = Starlette(routes=[Mount("/metrics", app=asgi_app(multiprocess_dir="/tmp/kuu"))]) + """ + from prometheus_client import make_asgi_app + + mp_dir = multiprocess_dir or os.environ.get("PROMETHEUS_MULTIPROC_DIR") + if mp_dir: + from prometheus_client.multiprocess import MultiProcessCollector + + os.environ["PROMETHEUS_MULTIPROC_DIR"] = mp_dir + registry = CollectorRegistry() + MultiProcessCollector(registry, path=mp_dir) + return make_asgi_app(registry=registry) + return make_asgi_app() + + +def mark_worker_dead(pid: int) -> None: + """ + Tell prometheus_client a worker subprocess is gone, so its files stop + contributing to `livesum` gauges. No-op outside multiprocess mode. + """ + if not os.environ.get("PROMETHEUS_MULTIPROC_DIR"): + return + from prometheus_client.multiprocess import mark_process_dead + + mark_process_dead(pid) diff --git a/kuu/contrib/structlog/__init__.py b/kuu/contrib/structlog/__init__.py new file mode 100644 index 0000000..54d57c2 --- /dev/null +++ b/kuu/contrib/structlog/__init__.py @@ -0,0 +1,59 @@ +from __future__ import annotations + +from typing import Any + +import structlog + +from kuu.context import Context +from kuu.middleware.base import Next +from kuu.middleware.task_log import TaskLogSink, run_process_task_logging + +__all__ = ("StructlogMiddleware",) + + +class _StructlogTaskLogSink(TaskLogSink): + def __init__(self, logger: structlog.BoundLogger) -> None: + self._log = logger + + def task_start(self, *, task_name: str, sched_id: str, queue: str, attempt: int) -> None: + self._log.info( + "task.start", + task_name=task_name, + sched_id=sched_id, + queue=queue, + attempt=attempt, + ) + + def task_ok(self, *, task_name: str, sched_id: str, duration: float) -> None: + self._log.info( + "task.ok", + task_name=task_name, + sched_id=sched_id, + duration=duration, + ) + + def task_fail(self, *, task_name: str, sched_id: str, duration: float, exc_type: str) -> None: + self._log.warning( + "task.fail", + task_name=task_name, + sched_id=sched_id, + duration=duration, + exc_type=exc_type, + ) + + +class StructlogMiddleware: + """Log task start, success, and failure via a structlog logger. + + Only active during the ``process`` phase. + + - `logger`: a ``structlog.BoundLogger`` instance (required). + The caller is responsible for configuring structlog globally + (``structlog.configure(...)``) before passing the logger. + """ + + def __init__(self, logger: structlog.BoundLogger): + self._sink = _StructlogTaskLogSink(logger) + + async def __call__(self, ctx: Context, call_next: Next) -> Any: + return await run_process_task_logging(ctx, call_next, self._sink) diff --git a/kuu/middleware/__init__.py b/kuu/middleware/__init__.py index 2eea622..b00d675 100644 --- a/kuu/middleware/__init__.py +++ b/kuu/middleware/__init__.py @@ -1,3 +1,4 @@ +from .task_log import TaskLogSink, run_process_task_logging from .base import Middleware, run_chain from .logging import LoggingMiddleware from .retry import RetryMiddleware @@ -9,4 +10,6 @@ "RetryMiddleware", "TimeoutMiddleware", "LoggingMiddleware", + "TaskLogSink", + "run_process_task_logging", ] diff --git a/kuu/middleware/logging.py b/kuu/middleware/logging.py index 9e5a4ac..5b7e83d 100644 --- a/kuu/middleware/logging.py +++ b/kuu/middleware/logging.py @@ -5,7 +5,7 @@ from ..context import Context from .base import Next -from ._task_log import TaskLogSink, run_process_task_logging +from .task_log import TaskLogSink, run_process_task_logging class _StdlibTaskLogSink(TaskLogSink): diff --git a/kuu/middleware/structlog.py b/kuu/middleware/structlog.py index aaddb89..fd7e46a 100644 --- a/kuu/middleware/structlog.py +++ b/kuu/middleware/structlog.py @@ -1,57 +1,14 @@ from __future__ import annotations -from typing import Any +import warnings -import structlog +from kuu.contrib.structlog import StructlogMiddleware -from ..context import Context -from .base import Next -from ._task_log import TaskLogSink, run_process_task_logging +__all__ = ("StructlogMiddleware",) - -class _StructlogTaskLogSink(TaskLogSink): - def __init__(self, logger: structlog.BoundLogger) -> None: - self._log = logger - - def task_start(self, *, task_name: str, sched_id: str, queue: str, attempt: int) -> None: - self._log.info( - "task.start", - task_name=task_name, - sched_id=sched_id, - queue=queue, - attempt=attempt, - ) - - def task_ok(self, *, task_name: str, sched_id: str, duration: float) -> None: - self._log.info( - "task.ok", - task_name=task_name, - sched_id=sched_id, - duration=duration, - ) - - def task_fail(self, *, task_name: str, sched_id: str, duration: float, exc_type: str) -> None: - self._log.warning( - "task.fail", - task_name=task_name, - sched_id=sched_id, - duration=duration, - exc_type=exc_type, - ) - - -class StructlogMiddleware: - """Log task start, success, and failure via a structlog logger. - - Only active during the ``process`` phase. - - - `logger`: a ``structlog.BoundLogger`` instance (required). - The caller is responsible for configuring structlog globally - (``structlog.configure(...)``) before passing the logger. - """ - - def __init__(self, logger: structlog.BoundLogger): - self._sink = _StructlogTaskLogSink(logger) - - async def __call__(self, ctx: Context, call_next: Next) -> Any: - return await run_process_task_logging(ctx, call_next, self._sink) +warnings.warn( + "`kuu.middleware.structlog` has moved to `kuu.contrib.structlog`; import from " + "there instead. This module is scheduled to be completely removed in 0.3.0.", + DeprecationWarning, + stacklevel=2, +) diff --git a/kuu/middleware/_task_log.py b/kuu/middleware/task_log.py similarity index 58% rename from kuu/middleware/_task_log.py rename to kuu/middleware/task_log.py index 003cbd8..13dd641 100644 --- a/kuu/middleware/_task_log.py +++ b/kuu/middleware/task_log.py @@ -8,11 +8,19 @@ class TaskLogSink(Protocol): + """Receiver for task lifecycle log events. + + Implement this protocol to back a logging middleware with any frontend + (stdlib :mod:`logging`, ``structlog``, a metrics emitter, ...). Drive it + with :func:`run_process_task_logging` from a middleware ``__call__``. + + ``LoggingMiddleware`` and ``kuu.contrib.structlog.StructlogMiddleware`` are + the bundled implementations. + """ + def task_start(self, *, task_name: str, sched_id: str, queue: str, attempt: int) -> None: ... - def task_ok( - self, *, task_name: str, sched_id: str, duration: float - ) -> None: ... + def task_ok(self, *, task_name: str, sched_id: str, duration: float) -> None: ... def task_fail( self, *, task_name: str, sched_id: str, duration: float, exc_type: str @@ -24,6 +32,12 @@ async def run_process_task_logging( call_next: Next, sink: TaskLogSink, ) -> Any: + """Drive a :class:`TaskLogSink` around the ``process`` phase of a task. + + Emits ``task_start`` before the task runs, then ``task_ok`` on success or + ``task_fail`` on exception (re-raising). The ``enqueue`` phase is passed + through untouched. + """ if ctx.phase != "process": return await call_next() msg = ctx.message diff --git a/kuu/orchestrator/_serve.py b/kuu/orchestrator/_serve.py index 204d0db..51968ed 100644 --- a/kuu/orchestrator/_serve.py +++ b/kuu/orchestrator/_serve.py @@ -23,7 +23,7 @@ def __init__(self) -> None: async def start(self, cfg: MetricsSettings, *, log_event: str) -> None: if not cfg.enable: return - from kuu.prometheus import serve + from kuu.contrib.prometheus import serve self._metrics_dir = await anyio.mkdtemp(prefix="kuu-prom-") os.environ["PROMETHEUS_MULTIPROC_DIR"] = self._metrics_dir diff --git a/kuu/orchestrator/_worker.py b/kuu/orchestrator/_worker.py index 5cd522d..182b4bd 100644 --- a/kuu/orchestrator/_worker.py +++ b/kuu/orchestrator/_worker.py @@ -119,7 +119,7 @@ def _terminate_and_wait(self, processes: list[SpawnProcess]) -> None: p.join(timeout=5) if self._config.metrics.enable: - from kuu.prometheus import mark_worker_dead + from kuu.contrib.prometheus import mark_worker_dead for p in processes: if p.pid is not None: @@ -135,7 +135,7 @@ def _run_worker(config: Settings, events_queue: mp.Queue | None = None) -> None: import_tasks(config.task_modules, "", False) if config.metrics.enable: - from kuu.prometheus import WorkerMetrics + from kuu.contrib.prometheus import WorkerMetrics WorkerMetrics(app) diff --git a/kuu/otel.py b/kuu/otel.py index 8546b1a..9de4676 100644 --- a/kuu/otel.py +++ b/kuu/otel.py @@ -1,473 +1,26 @@ from __future__ import annotations -import logging -import os -import typing -from typing import Any +import warnings -from opentelemetry import context as otel_context -from opentelemetry import metrics, propagate, trace -from opentelemetry.instrumentation.logging.handler import LoggingHandler -from opentelemetry.sdk._logs import LoggerProvider -from opentelemetry.sdk._logs._internal.export import BatchLogRecordProcessor -from opentelemetry.sdk.metrics import MeterProvider -from opentelemetry.sdk.metrics.export import PeriodicExportingMetricReader -from opentelemetry.sdk.resources import Resource -from opentelemetry.sdk.trace import TracerProvider -from opentelemetry.sdk.trace.export import BatchSpanProcessor -from opentelemetry.trace import SpanKind, Status, StatusCode, Tracer +from kuu.contrib.otel import ( + KuuOTELInstrumentor, + OtelLoggingBridge, + OtelMetrics, + OtelTracingMiddleware, + shutdown_telemetry, +) -if typing.TYPE_CHECKING: - from opentelemetry.sdk._logs._internal.export import LogRecordExporter - from opentelemetry.sdk.metrics.export import MetricExporter - from opentelemetry.sdk.trace.export import SpanExporter - - from .app import Kuu - from .context import Context - from .message import Message - from .middleware.base import Next - - -__all__ = [ +__all__ = ( "KuuOTELInstrumentor", "OtelTracingMiddleware", "OtelMetrics", "OtelLoggingBridge", "shutdown_telemetry", -] - - -log = logging.getLogger("kuu.otel") - - -def _span_name(ctx: Context) -> str: - task_name = ctx.task.task_name if ctx.task else ctx.message.task - verb = "publish" if ctx.phase == "enqueue" else "process" - return f"{task_name} {verb}" - - -def _set_messaging_attributes(span: Any, ctx: Context) -> None: - msg = ctx.message - span.set_attribute("messaging.system", "kuu") - span.set_attribute("messaging.operation", "publish" if ctx.phase == "enqueue" else "process") - span.set_attribute("messaging.destination", msg.queue) - span.set_attribute("messaging.destination.name", msg.queue) - span.set_attribute("messaging.message.id", str(msg.id)) - span.set_attribute("kuu.task.name", msg.task) - span.set_attribute("kuu.task.attempt", msg.attempt) - span.set_attribute("kuu.task.max_attempts", msg.max_attempts) - - -class OtelTracingMiddleware: - """ - Middleware that creates OTEL spans around enqueue and process phases. - - Enqueue phase (PRODUCER): - 1. Start span ``{task_name} publish`` with PRODUCER kind. - 2. Set messaging attributes. - 3. Inject W3C TraceContext into ``ctx.message.headers``. - 4. Execute the middleware chain. - 5. End the span (set ERROR status + record exception on failure). - - Process phase (CONSUMER): - 1. Extract W3C TraceContext from ``ctx.message.headers``. - 2. Attach extracted context. - 3. Start span ``{task_name} process`` with CONSUMER kind. - 4. Execute the middleware chain. - 5. Set OK / ERROR status. - 6. Detach context and end span. - """ - - def __init__( - self, - *, - tracer_name: str = "kuu", - propagate_ctx: bool = True, - ): - self._tracer_name = tracer_name - self._propagate = propagate_ctx - self._tracer: Tracer | None = None - - @property - def tracer(self) -> Tracer: - if self._tracer is None: - self._tracer = trace.get_tracer(self._tracer_name) - return self._tracer - - async def __call__(self, ctx: Context, call_next: Next) -> Any: - if ctx.phase == "enqueue": - return await self._enqueue(ctx, call_next) - if ctx.phase == "process": - return await self._process(ctx, call_next) - return await call_next() - - async def _enqueue(self, ctx: Context, call_next: Next) -> Any: - span = self.tracer.start_span( - _span_name(ctx), - kind=SpanKind.PRODUCER, - ) - _set_messaging_attributes(span, ctx) - - with trace.use_span(span, end_on_exit=True): - if self._propagate: - propagate.inject(ctx.message.headers) - - try: - return await call_next() - except Exception as exc: - span.set_status(Status(StatusCode.ERROR)) - span.record_exception(exc) - raise - - async def _process(self, ctx: Context, call_next: Next) -> Any: - extracted = propagate.extract(carrier=ctx.message.headers) if self._propagate else None - token = otel_context.attach(extracted) if extracted else None - - span = self.tracer.start_span( - _span_name(ctx), - kind=SpanKind.CONSUMER, - ) - _set_messaging_attributes(span, ctx) - - with trace.use_span(span, end_on_exit=True): - try: - result = await call_next() - span.set_status(Status(StatusCode.OK)) - return result - except Exception as exc: - span.set_status(Status(StatusCode.ERROR)) - span.record_exception(exc) - raise - finally: - if token is not None: - otel_context.detach(token) - - -class OtelMetrics: - """ - Wires OTEL metrics onto Kuu task lifecycle events. - - Instruments created (all under ``meter_name`` scope): - - ``kuu.task.enqueued`` (Counter) : tasks enqueued - - ``kuu.task.processed`` (Counter) : tasks processed, with ``status`` label - - ``kuu.task.duration`` (Histogram, seconds) : task execution duration - - ``kuu.task.in_flight`` (UpDownCounter) : tasks currently executing - - ``kuu.task.retried`` (Counter) : retry attempts - """ - - def __init__( - self, - *, - app: Kuu, - meter_name: str = "kuu", - ): - self._app = app - self._meter = metrics.get_meter(meter_name) - - self._enqueued = self._meter.create_counter( - "kuu.task.enqueued", - "tasks", - "Tasks enqueued", - ) - self._processed = self._meter.create_counter( - "kuu.task.processed", - "tasks", - "Tasks processed by outcome", - ) - self._duration = self._meter.create_histogram( - "kuu.task.duration", - "s", - "Task execution duration", - ) - self._in_flight = self._meter.create_up_down_counter( - "kuu.task.in_flight", - "tasks", - "Tasks currently executing", - ) - self._retried = self._meter.create_counter( - "kuu.task.retried", - "tasks", - "Task retry attempts", - ) - - self._inflight_ids = set() - - ev = app.events - ev.task_enqueued.connect(self._on_enqueued) - ev.task_started.connect(self._on_started) - ev.task_succeeded.connect(self._on_succeeded) - ev.task_failed.connect(self._on_failed) - ev.task_retried.connect(self._on_retried) - ev.task_dead.connect(self._on_dead) - - def _attrs(self, msg: Message, **extra: str) -> dict[str, Any]: - base: dict[str, Any] = {"task.name": msg.task, "queue": msg.queue} - base.update(extra) - return base - - def _settle(self, msg: Message) -> None: - if msg.id in self._inflight_ids: - self._inflight_ids.discard(msg.id) - self._in_flight.add(-1, self._attrs(msg)) - - def _on_enqueued(self, msg: Message) -> None: - self._enqueued.add(1, self._attrs(msg)) - - def _on_started(self, msg: Message) -> None: - self._inflight_ids.add(msg.id) - self._in_flight.add(1, self._attrs(msg)) - - def _on_succeeded(self, msg: Message, elapsed: float) -> None: - self._duration.record(elapsed, self._attrs(msg)) - self._processed.add(1, self._attrs(msg, status="ok")) - self._settle(msg) - - def _on_failed(self, msg: Message, exc: BaseException) -> None: - self._processed.add(1, self._attrs(msg, status="error")) - self._settle(msg) - - def _on_retried(self, msg: Message, delay: float) -> None: - self._retried.add(1, self._attrs(msg)) - self._settle(msg) - - def _on_dead(self, msg: Message) -> None: - self._processed.add(1, self._attrs(msg, status="dead")) - self._settle(msg) - - def disconnect(self) -> None: - ev = self._app.events - ev.task_enqueued.disconnect(self._on_enqueued) - ev.task_started.disconnect(self._on_started) - ev.task_succeeded.disconnect(self._on_succeeded) - ev.task_failed.disconnect(self._on_failed) - ev.task_retried.disconnect(self._on_retried) - ev.task_dead.disconnect(self._on_dead) - - -class OtelLoggingBridge: - """Bridges Python's ``logging`` to an OpenTelemetry ``LoggerProvider``. - - Attaches an OTel ``LoggingHandler`` to the named stdlib logger so that - any log records emitted through it are forwarded to the provider. Does - **not** configure any particular logging frontend : use ``logging``, - ``structlog``, or anything else that ultimately produces - ``logging.LogRecord`` instances. - - Provider resolution order (first wins): - 1. ``logger_provider`` passed to :meth:`setup`. - 2. ``log_exporter`` passed to :meth:`setup` : a new ``LoggerProvider`` - is created around it. - 3. Falls back to OTLP HTTP (``OTLPLogExporter``) only when - ``OTEL_EXPORTER_OTLP_ENDPOINT`` is set and neither of the above - was supplied. - """ - - def __init__( - self, - *, - level: int = logging.INFO, - logger_name: str = "kuu", - ): - self._level = level - self._logger_name = logger_name - self._provider: LoggerProvider | None = None - self._handler: LoggingHandler | None = None - self._owns_provider: bool = False - - def setup( - self, - resource: Resource | None = None, - *, - logger_provider: LoggerProvider | None = None, - log_exporter: LogRecordExporter | None = None, - ) -> logging.Logger: - if logger_provider is not None: - self._provider = logger_provider - self._owns_provider = False - else: - exporter = log_exporter - if exporter is None and os.getenv("OTEL_EXPORTER_OTLP_ENDPOINT"): - from opentelemetry.exporter.otlp.proto.http._log_exporter import OTLPLogExporter - - exporter = OTLPLogExporter() - - if exporter is not None: - res = resource or Resource.create({}) - self._provider = LoggerProvider(resource=res) - self._provider.add_log_record_processor(BatchLogRecordProcessor(exporter)) - self._owns_provider = True - - if self._provider is None: - log.debug("event=otel.no_provider") - return logging.getLogger(self._logger_name) - - self._handler = LoggingHandler(level=self._level, logger_provider=self._provider) - logger = logging.getLogger(self._logger_name) - logger.addHandler(self._handler) - log.info("event=otel.bridge_attached logger=%s", self._logger_name) - return logger - - def shutdown(self) -> None: - if self._handler is not None: - logging.getLogger(self._logger_name).removeHandler(self._handler) - self._handler = None - if self._owns_provider and self._provider is not None: - self._provider.shutdown() - self._provider = None - - -def _auto_setup_sdk( - endpoint: str | None = None, - *, - span_exporter: SpanExporter | None = None, - metric_exporter: MetricExporter | None = None, -) -> Resource | None: - endpoint = endpoint or os.getenv("OTEL_EXPORTER_OTLP_ENDPOINT") - - if span_exporter is None and metric_exporter is None and not endpoint: - log.debug("event=otel.no_endpoint") - return None - - service_name = os.getenv("OTEL_SERVICE_NAME", "kuu") - resource = Resource.create({"service.name": service_name}) - log.info("event=otel.auto_setup service=%s", service_name) - - if span_exporter is None: - from opentelemetry.exporter.otlp.proto.http.trace_exporter import OTLPSpanExporter - - span_exporter = OTLPSpanExporter(endpoint=endpoint) - - tp = TracerProvider(resource=resource) - tp.add_span_processor(BatchSpanProcessor(span_exporter)) - trace.set_tracer_provider(tp) - - if metric_exporter is None: - from opentelemetry.exporter.otlp.proto.http.metric_exporter import OTLPMetricExporter - - metric_exporter = OTLPMetricExporter(endpoint=endpoint) - - reader = PeriodicExportingMetricReader(metric_exporter) - mp = MeterProvider(resource=resource, metric_readers=[reader]) - metrics.set_meter_provider(mp) - - return resource - - -class KuuOTELInstrumentor: - """ - Wires OTEL instrumentation (traces + metrics + logs) into a Kuu app. - - Basic usage (uses globally configured OTel providers: bring your own - SDK setup):: - - KuuOTELInstrumentor(app=app).instrument(setup_sdk=False) - - Auto-setup with OTLP HTTP (reads ``OTEL_EXPORTER_OTLP_ENDPOINT``):: - - KuuOTELInstrumentor(app=app).instrument() - - Bring your own exporters (e.g. gRPC):: - - from opentelemetry.exporter.otlp.proto.grpc.trace_exporter import OTLPSpanExporter - from opentelemetry.exporter.otlp.proto.grpc.metric_exporter import OTLPMetricExporter - from opentelemetry.exporter.otlp.proto.grpc._log_exporter import OTLPLogExporter - - KuuOTELInstrumentor( - app=app, - span_exporter=OTLPSpanExporter(...), - metric_exporter=OTLPMetricExporter(...), - log_exporter=OTLPLogExporter(...), - ).instrument() - - ``instrument()`` will: - 1. Configure OTel SDK providers when exporters are supplied or - ``OTEL_EXPORTER_OTLP_ENDPOINT`` is set (skipped with - ``setup_sdk=False``). - 2. Insert ``OtelTracingMiddleware`` at the front of the middleware - chain. - 3. Wire ``OtelMetrics`` onto app.events signals. - 4. Bridge stdlib ``logging`` to the log provider when a log exporter - or ``OTEL_EXPORTER_OTLP_ENDPOINT`` is available. - """ - - def __init__( - self, - *, - app: Kuu, - tracer_name: str = "kuu", - meter_name: str = "kuu", - propagate: bool = True, - log_level: int = logging.INFO, - logger_name: str = "kuu", - span_exporter: SpanExporter | None = None, - metric_exporter: MetricExporter | None = None, - log_exporter: LogRecordExporter | None = None, - ): - self._app = app - self._tracer_name = tracer_name - self._meter_name = meter_name - self._propagate = propagate - self._log_level = log_level - self._logger_name = logger_name - self._span_exporter = span_exporter - self._metric_exporter = metric_exporter - self._log_exporter = log_exporter - - self._otel_middleware: OtelTracingMiddleware | None = None - self._otel_metrics: OtelMetrics | None = None - self._logging_bridge: OtelLoggingBridge | None = None - - def instrument(self, *, setup_sdk: bool = True) -> None: - resource: Resource | None = None - if setup_sdk: - resource = _auto_setup_sdk( - span_exporter=self._span_exporter, - metric_exporter=self._metric_exporter, - ) - - self._otel_middleware = OtelTracingMiddleware( - tracer_name=self._tracer_name, - propagate_ctx=self._propagate, - ) - self._app.middleware.insert(0, self._otel_middleware) - - self._otel_metrics = OtelMetrics( - app=self._app, - meter_name=self._meter_name, - ) - - self._logging_bridge = OtelLoggingBridge( - level=self._log_level, - logger_name=self._logger_name, - ) - self._logging_bridge.setup(resource=resource, log_exporter=self._log_exporter) - - def uninstrument(self) -> None: - if self._otel_middleware is not None: - try: - self._app.middleware.remove(self._otel_middleware) - except ValueError: - pass - self._otel_middleware = None - - if self._otel_metrics is not None: - self._otel_metrics.disconnect() - self._otel_metrics = None - - if self._logging_bridge is not None: - self._logging_bridge.shutdown() - self._logging_bridge = None - - -def shutdown_telemetry() -> None: - from contextlib import suppress - - with suppress(Exception): - prov = trace.get_tracer_provider() - if hasattr(prov, "shutdown") and callable(prov.shutdown): - prov.shutdown() - - with suppress(Exception): - prov = metrics.get_meter_provider() - if hasattr(prov, "shutdown") and callable(prov.shutdown): - prov.shutdown() +) + +warnings.warn( + "`kuu.otel` has moved to `kuu.contrib.otel`; import from there instead. " + "This module is scheduled to be completely removed in 0.3.0.", + DeprecationWarning, + stacklevel=2, +) diff --git a/kuu/prometheus.py b/kuu/prometheus.py index b56c6c9..6ff63ff 100644 --- a/kuu/prometheus.py +++ b/kuu/prometheus.py @@ -1,302 +1,26 @@ from __future__ import annotations -import os -import time -from threading import Thread -from typing import TYPE_CHECKING, Annotated, Any -from uuid import UUID -from wsgiref.simple_server import WSGIServer - -from prometheus_client import REGISTRY, CollectorRegistry, Counter, Gauge, Histogram -from typing_extensions import Doc - -if TYPE_CHECKING: - from .app import Kuu - from .context import Context - from .message import Message - from .middleware.base import Next - - -_DEFAULT_DURATION_BUCKETS = ( - 0.005, - 0.01, - 0.025, - 0.05, - 0.1, - 0.25, - 0.5, - 1.0, - 2.5, - 5.0, - 10.0, - 30.0, - 60.0, +import warnings + +from kuu.contrib.prometheus import ( + ClientMetrics, + WorkerMetrics, + asgi_app, + mark_worker_dead, + serve, ) -_DEFAULT_RETRY_DELAY_BUCKETS = (0.0, 0.1, 0.5, 1.0, 5.0, 15.0, 60.0, 300.0, 900.0) -_DEFAULT_ENQUEUE_BUCKETS = (0.001, 0.005, 0.01, 0.025, 0.05, 0.1, 0.25, 1.0, 5.0) - - -class WorkerMetrics: - """ - Worker-side, event-driven Prometheus metrics. Wires onto `app.events`. - - Instantiate inside each worker subprocess after the app/registry is loaded. - Exposes per-task counters/histograms and an in-flight gauge. In multiprocess - mode (PROMETHEUS_MULTIPROC_DIR set) the gauge sums across living workers. - - Metric names ({ns} = namespace, default "qq"): - - {ns}_tasks_received_total{task,queue} - - {ns}_tasks_started_total{task,queue} - - {ns}_tasks_succeeded_total{task,queue} - - {ns}_tasks_retried_total{task,queue} - - {ns}_tasks_failed_total{task,queue,exc} - - {ns}_tasks_dead_total{task,queue} - - {ns}_task_duration_seconds{task,queue} (from task_succeeded) - - {ns}_task_retry_delay_seconds{task,queue} (from task_retried) - - {ns}_tasks_in_flight{task,queue} gauge (livesum across procs) - - `in_flight` is decremented on the first terminal event per message sched_id - (succeeded/retried/failed/dead) to handle the Fail->Dead chain and the - standalone Reject->Dead path without double-counting. - """ - - def __init__( - self, - app: Kuu, - namespace: str = "kuu", - registry: CollectorRegistry = REGISTRY, - duration_buckets: tuple[float, ...] = _DEFAULT_DURATION_BUCKETS, - retry_delay_buckets: tuple[float, ...] = _DEFAULT_RETRY_DELAY_BUCKETS, - ): - self._inflight_ids: set[UUID] = set() - - self.received = Counter( - f"{namespace}_tasks_received_total", - "Tasks received by worker", - ("task", "queue"), - registry=registry, - ) - self.started = Counter( - f"{namespace}_tasks_started_total", - "Tasks started processing", - ("task", "queue"), - registry=registry, - ) - self.succeeded = Counter( - f"{namespace}_tasks_succeeded_total", - "Tasks succeeded", - ("task", "queue"), - registry=registry, - ) - self.retried = Counter( - f"{namespace}_tasks_retried_total", - "Tasks scheduled for retry", - ("task", "queue"), - registry=registry, - ) - self.failed = Counter( - f"{namespace}_tasks_failed_total", - "Tasks raised an exception", - ("task", "queue", "exc"), - registry=registry, - ) - self.dead = Counter( - f"{namespace}_tasks_dead_total", - "Tasks reached terminal state without success", - ("task", "queue"), - registry=registry, - ) - self.duration = Histogram( - f"{namespace}_task_duration_seconds", - "Task processing duration", - ("task", "queue"), - buckets=duration_buckets, - registry=registry, - ) - self.retry_delay = Histogram( - f"{namespace}_task_retry_delay_seconds", - "Scheduled retry delay", - ("task", "queue"), - buckets=retry_delay_buckets, - registry=registry, - ) - self.in_flight = Gauge( - f"{namespace}_tasks_in_flight", - "Tasks currently being processed", - ("task", "queue"), - multiprocess_mode="livesum", - registry=registry, - ) - - ev = app.events - ev.task_received.connect(self._on_received) - ev.task_started.connect(self._on_started) - ev.task_succeeded.connect(self._on_succeeded) - ev.task_retried.connect(self._on_retried) - ev.task_failed.connect(self._on_failed) - ev.task_dead.connect(self._on_dead) - - def _settle(self, msg: Message) -> None: - if msg.id in self._inflight_ids: - self._inflight_ids.discard(msg.id) - self.in_flight.labels(msg.task, msg.queue).dec() - - def _on_received(self, msg: Message) -> None: - self.received.labels(msg.task, msg.queue).inc() - - def _on_started(self, msg: Message) -> None: - self.started.labels(msg.task, msg.queue).inc() - self._inflight_ids.add(msg.id) - self.in_flight.labels(msg.task, msg.queue).inc() - - def _on_succeeded(self, msg: Message, elapsed: float) -> None: - self.duration.labels(msg.task, msg.queue).observe(elapsed) - self.succeeded.labels(msg.task, msg.queue).inc() - self._settle(msg) - - def _on_retried(self, msg: Message, delay: float) -> None: - self.retry_delay.labels(msg.task, msg.queue).observe(delay) - self.retried.labels(msg.task, msg.queue).inc() - self._settle(msg) - - def _on_failed(self, msg: Message, exc: BaseException) -> None: - self.failed.labels(msg.task, msg.queue, type(exc).__name__).inc() - self._settle(msg) - - def _on_dead(self, msg: Message) -> None: - self.dead.labels(msg.task, msg.queue).inc() - self._settle(msg) - - -class ClientMetrics: - """ - Client-side metrics, implemented as a Middleware. Add to `Kuu(middleware=[...])`. - Acts only on the `enqueue` phase. Records: - - {ns}_client_enqueued_total{task,queue} - - {ns}_client_enqueue_errors_total{task,queue,exc} - - {ns}_client_enqueue_duration_seconds{task,queue} (broker handoff latency) - """ - - def __init__( - self, - namespace: str = "kuu", - registry: CollectorRegistry = REGISTRY, - duration_buckets: tuple[float, ...] = _DEFAULT_ENQUEUE_BUCKETS, - ): - self.enqueued = Counter( - f"{namespace}_client_enqueued_total", - "Tasks enqueued by client", - ("task", "queue"), - registry=registry, - ) - self.errors = Counter( - f"{namespace}_client_enqueue_errors_total", - "Enqueue attempts that raised", - ("task", "queue", "exc"), - registry=registry, - ) - self.duration = Histogram( - f"{namespace}_client_enqueue_duration_seconds", - "Time spent enqueueing to broker", - ("task", "queue"), - buckets=duration_buckets, - registry=registry, - ) - - async def __call__(self, ctx: Context, call_next: Next) -> Any: - if ctx.phase != "enqueue": - return await call_next() - - msg = ctx.message - t0 = time.perf_counter() - try: - result = await call_next() - except BaseException as exc: - self.errors.labels(msg.task, msg.queue, type(exc).__name__).inc() - raise - else: - self.enqueued.labels(msg.task, msg.queue).inc() - return result - finally: - self.duration.labels(msg.task, msg.queue).observe(time.perf_counter() - t0) - - -def serve( - host: Annotated[str, Doc("Host to bind the /metrics HTTP server to.")] = "0.0.0.0", - port: Annotated[int, Doc("Port to bind the /metrics HTTP server to.")] = 9100, - multiprocess_dir: Annotated[ - str | None, - Doc( - "Directory for prometheus_client multiprocess shared state. When set, aggregates metrics across all worker subprocesses that wrote to this directory." - ), - ] = None, -) -> tuple[WSGIServer, Thread]: - """ - Start a `/metrics` HTTP endpoint. Returns (server, thread) from prometheus_client. - - If `multiprocess_dir` is given (or `PROMETHEUS_MULTIPROC_DIR` env is set), - aggregates metrics from on-disk shared state across all worker subprocesses - that wrote to that dir. Otherwise serves the in-process REGISTRY directly. - - Call this in the orchestrator (parent) process when running multiple workers, - or in-process for a single-worker setup. - """ - from prometheus_client import start_http_server - - mp_dir = multiprocess_dir or os.environ.get("PROMETHEUS_MULTIPROC_DIR") - if mp_dir: - from prometheus_client.multiprocess import MultiProcessCollector - - os.environ["PROMETHEUS_MULTIPROC_DIR"] = mp_dir - registry = CollectorRegistry() - MultiProcessCollector(registry, path=mp_dir) - return start_http_server(port, addr=host, registry=registry) - return start_http_server(port, addr=host) - - -def asgi_app(multiprocess_dir: str | None = None) -> Any: - """ - Build an ASGI app exposing `/metrics`. Mount it into your own ASGI - application (FastAPI, Starlette, etc.). - - If `multiprocess_dir` is given (or `PROMETHEUS_MULTIPROC_DIR` env is set), - aggregates metrics from on-disk shared state across worker subprocesses. - Otherwise serves the in-process REGISTRY directly. - - Examples: - # Client side, FastAPI: - from fastapi import FastAPI - from kuu.prometheus import ClientMetrics, asgi_app - - app = FastAPI() - kuu_app = Kuu(broker=..., middleware=[ClientMetrics()]) - app.mount("/metrics", asgi_app()) - - # Worker side, exposing aggregated multiproc metrics from a sidecar - # ASGI server (e.g. uvicorn) instead of the orchestrator's WSGI server: - app = Starlette(routes=[Mount("/metrics", app=asgi_app(multiprocess_dir="/tmp/kuu"))]) - """ - from prometheus_client import make_asgi_app - - mp_dir = multiprocess_dir or os.environ.get("PROMETHEUS_MULTIPROC_DIR") - if mp_dir: - from prometheus_client.multiprocess import MultiProcessCollector - - os.environ["PROMETHEUS_MULTIPROC_DIR"] = mp_dir - registry = CollectorRegistry() - MultiProcessCollector(registry, path=mp_dir) - return make_asgi_app(registry=registry) - return make_asgi_app() - - -def mark_worker_dead(pid: int) -> None: - """ - Tell prometheus_client a worker subprocess is gone, so its files stop - contributing to `livesum` gauges. No-op outside multiprocess mode. - """ - if not os.environ.get("PROMETHEUS_MULTIPROC_DIR"): - return - from prometheus_client.multiprocess import mark_process_dead +__all__ = ( + "WorkerMetrics", + "ClientMetrics", + "serve", + "asgi_app", + "mark_worker_dead", +) - mark_process_dead(pid) +warnings.warn( + "`kuu.prometheus` has moved to `kuu.contrib.prometheus`; import from there " + "instead. This module is scheduled to be completely removed in 0.3.0.", + DeprecationWarning, + stacklevel=2, +) diff --git a/pyproject.toml b/pyproject.toml index 8962820..5885aaa 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -60,6 +60,12 @@ otel = [ "opentelemetry-instrumentation-logging>=0.60,<1", "structlog>=24,<27", ] +di = [ + "dishka>=1.10.1", +] +structlog = [ + "structlog>=24,<27", +] [dependency-groups] dev = [ @@ -76,6 +82,7 @@ dev = [ "opentelemetry-sdk>=1.41,<2", "structlog>=24,<26", ] + docs = [ "furo>=2024.8", "sphinx>=8.1,<9", @@ -83,9 +90,6 @@ docs = [ "sphinx-autodoc2>=0.5", "sphinx-copybutton>=0.5", ] -postgres = [ - "pypika>=0.51.1", -] [build-system] requires = ["uv_build>=0.11,<0.13"] diff --git a/tests/test_dishka.py b/tests/test_dishka.py new file mode 100644 index 0000000..4ae82f7 --- /dev/null +++ b/tests/test_dishka.py @@ -0,0 +1,422 @@ +from __future__ import annotations + +import inspect +import threading +from contextvars import ContextVar +from typing import Any + +import anyio +import pytest +from dishka import ( + Provider, + Scope, + make_async_container, + make_container, + provide, +) + +from kuu.app import Kuu +from kuu.brokers.memory import MemoryBroker +from kuu.config import Settings +from kuu.context import Context +from kuu.contrib.dishka import ( + CONTAINER_STATE_KEY, + FromDishka, + KuuDishkaMiddleware, + current_container, + from_contextvar, + inject, + setup_dishka, +) +from kuu.contrib.dishka import _request_container # type: ignore[attr-defined] +from kuu.message import Message, Payload +from kuu.middleware.base import Next, run_chain +from kuu.worker import Worker +from kuu._types import _coerce_payload + +pytestmark = pytest.mark.anyio + + +# --------------------------------------------------------------------------- # +# Services / providers +# --------------------------------------------------------------------------- # +class ReadService: + tag = "read" + + +class WriteService: + tag = "write" + + +class AppProvider(Provider): + scope = Scope.REQUEST + + read = provide(ReadService) + write = provide(WriteService) + + +class CountingProvider(Provider): + """Counts how many ReadService instances are created (scope freshness).""" + + scope = Scope.REQUEST + instances = 0 + + @provide + def read(self) -> ReadService: + CountingProvider.instances += 1 + return ReadService() + + +@pytest.fixture +def async_container(): + c = make_async_container(AppProvider()) + yield c + + +@pytest.fixture +def sync_container(): + c = make_container(AppProvider()) + yield c + c.close() + + +# --------------------------------------------------------------------------- # +# FromDishka marker + inject signature rewriting (pure unit) +# --------------------------------------------------------------------------- # +def test_fromdishka_returns_marker_with_default_component(): + from kuu.contrib.dishka import _FromDishkaMarker # type: ignore[attr-defined] + + marker = FromDishka() + assert isinstance(marker, _FromDishkaMarker) + assert marker.component == "" + + +def test_fromdishka_carries_custom_component(): + marker = FromDishka(component="secondary") + assert marker.component == "secondary" + + +def test_inject_strips_di_params_from_signature_async(): + @inject + async def handler( + lot_id: str, + read: ReadService = FromDishka(), + write: WriteService = FromDishka(), + fetch: bool = True, + ) -> str: + return f"{lot_id}:{read.tag}:{write.tag}:{fetch}" + + params = list(inspect.signature(handler).parameters) + assert params == ["lot_id", "fetch"] + # the cleaned signature must still bind a payload that omits the deps + inspect.signature(handler).bind("lot-1", fetch=False) + + +def test_inject_strips_di_params_from_signature_sync(): + @inject + def handler(lot_id: str, read: ReadService = FromDishka()) -> str: + return f"{lot_id}:{read.tag}" + + assert list(inspect.signature(handler).parameters) == ["lot_id"] + + +def test_inject_preserves_asyncness(): + @inject + async def a(read: ReadService = FromDishka()) -> None: ... + + @inject + def s(read: ReadService = FromDishka()) -> None: ... + + assert inspect.iscoroutinefunction(a) + assert not inspect.iscoroutinefunction(s) + + +def test_inject_is_noop_without_dependencies(): + async def plain(x: int) -> int: + return x + + wrapped = inject(plain) + assert list(inspect.signature(wrapped).parameters) == ["x"] + + +def test_inject_factory_form_returns_decorator(): + var: ContextVar[Any] = ContextVar("v", default=None) + deco = inject(container_getter=from_contextvar(var)) + assert callable(deco) + + @deco + async def handler(read: ReadService = FromDishka()) -> None: ... + + assert list(inspect.signature(handler).parameters) == [] + + +# --------------------------------------------------------------------------- # +# current_container / from_contextvar error paths +# --------------------------------------------------------------------------- # +def test_current_container_raises_without_container(): + with pytest.raises(LookupError, match="No Dishka REQUEST container"): + current_container() + + +def test_from_contextvar_raises_when_unset(): + var: ContextVar[Any] = ContextVar("empty", default=None) + getter = from_contextvar(var) + with pytest.raises(LookupError, match="holds no Dishka container"): + getter() + + +async def test_from_contextvar_returns_set_container(async_container): + var: ContextVar[Any] = ContextVar("set", default=None) + async with async_container() as request: + var.set(request) + assert from_contextvar(var)() is request + + +# --------------------------------------------------------------------------- # +# KuuDishkaMiddleware construction +# --------------------------------------------------------------------------- # +def test_middleware_requires_a_container(): + with pytest.raises(ValueError, match="at least one"): + KuuDishkaMiddleware() + + +def test_setup_dishka_inserts_middleware_first(async_container): + app = Kuu(broker=MemoryBroker()) + + async def existing(ctx: Context, call_next: Next) -> Any: # pre-existing middleware + return await call_next() + + app.middleware.append(existing) + mw = setup_dishka(app, async_container) + assert app.middleware[0] is mw + assert existing in app.middleware + + +# --------------------------------------------------------------------------- # +# Middleware behaviour via run_chain (faithful, fast) +# --------------------------------------------------------------------------- # +def _ctx(app: Kuu, task: Any, payload: Payload) -> Context: + msg = Message(task=task.task_name, queue=task.task_queue, payload=payload) + return Context(app=app, message=msg, phase="process", task=task) + + +async def _terminal_for(task: Any): + """Mirror Worker._handle_inner terminal for a single task.""" + + async def _terminal(c: Context) -> Any: + p = _coerce_payload(task.original_func, c.message.payload) + if task.blocking: + return await anyio.to_thread.run_sync( + lambda: task.original_func(*p.args, **p.kwargs), abandon_on_cancel=True + ) + r = task.original_func(*p.args, **p.kwargs) + return await r if inspect.isawaitable(r) else r + + return _terminal + + +async def test_async_injection_via_middleware(async_container): + app = Kuu(broker=MemoryBroker()) + + @app.task + @inject + async def handler( + lot_id: str, + read: ReadService = FromDishka(), + write: WriteService = FromDishka(), + fetch: bool = True, + ) -> str: + return f"{lot_id}:{read.tag}:{write.tag}:{fetch}" + + mw = KuuDishkaMiddleware(async_container) + ctx = _ctx(app, handler, Payload(args=("lot-1",), kwargs={"fetch": False})) + result = await run_chain(ctx, [mw], await _terminal_for(handler)) + + assert result == "lot-1:read:write:False" + # request container was exposed in ctx.state + assert CONTAINER_STATE_KEY in ctx.state + + +async def test_sync_injection_runs_off_thread(sync_container): + app = Kuu(broker=MemoryBroker()) + main_thread = threading.get_ident() + seen: dict[str, Any] = {} + + @app.task(blocking=True) + @inject + def handler(lot_id: str, read: ReadService = FromDishka()) -> str: + seen["thread"] = threading.get_ident() + return f"{lot_id}:{read.tag}" + + mw = KuuDishkaMiddleware(sync_container=sync_container) + ctx = _ctx(app, handler, Payload(args=("lot-2",))) + result = await run_chain(ctx, [mw], await _terminal_for(handler)) + + assert result == "lot-2:read" + assert seen["thread"] != main_thread + + +async def test_middleware_resets_contextvar_after_task(async_container): + app = Kuu(broker=MemoryBroker()) + + @app.task + @inject + async def handler(read: ReadService = FromDishka()) -> str: + return read.tag + + mw = KuuDishkaMiddleware(async_container) + ctx = _ctx(app, handler, Payload()) + await run_chain(ctx, [mw], await _terminal_for(handler)) + + # contextvar must be cleared once the task is done + assert _request_container.get() is None + + +async def test_middleware_uses_custom_contextvar(async_container): + app = Kuu(broker=MemoryBroker()) + my_var: ContextVar[Any] = ContextVar("my_ioc", default=None) + + @app.task + @inject(container_getter=from_contextvar(my_var)) + async def handler(read: ReadService = FromDishka()) -> str: + return read.tag + + mw = KuuDishkaMiddleware(async_container, context_var=my_var) + ctx = _ctx(app, handler, Payload()) + result = await run_chain(ctx, [mw], await _terminal_for(handler)) + + assert result == "read" + # built-in slot stays untouched; custom var is reset afterwards + assert _request_container.get() is None + assert my_var.get() is None + + +async def test_middleware_passes_through_enqueue_phase(async_container): + app = Kuu(broker=MemoryBroker()) + mw = KuuDishkaMiddleware(async_container) + msg = Message(task="t", queue="default", payload=Payload()) + ctx = Context(app=app, message=msg, phase="enqueue", task=None) + + called = False + + async def _next() -> str: + nonlocal called + called = True + return "ok" + + result = await mw(ctx, _next) + assert result == "ok" and called + assert CONTAINER_STATE_KEY not in ctx.state + + +async def test_middleware_raises_when_required_container_missing(async_container): + """A blocking task with only an async container available is misconfigured.""" + app = Kuu(broker=MemoryBroker()) + + @app.task(blocking=True) + @inject + def handler(read: ReadService = FromDishka()) -> str: + return read.tag + + mw = KuuDishkaMiddleware(async_container) # no sync_container + ctx = _ctx(app, handler, Payload()) + + async def _next() -> None: ... + + with pytest.raises(LookupError, match="needs a sync Dishka container"): + await mw(ctx, _next) + + +async def test_middleware_raises_when_async_container_missing(sync_container): + """An async task with only a sync container available is misconfigured.""" + app = Kuu(broker=MemoryBroker()) + + @app.task + @inject + async def handler(read: ReadService = FromDishka()) -> str: + return read.tag + + mw = KuuDishkaMiddleware(sync_container=sync_container) # no async container + ctx = _ctx(app, handler, Payload()) + + async def _next() -> None: ... + + with pytest.raises(LookupError, match="needs an async Dishka container"): + await mw(ctx, _next) + + +async def test_request_scope_is_fresh_per_task(async_container): + app = Kuu(broker=MemoryBroker()) + container = make_async_container(CountingProvider()) + CountingProvider.instances = 0 + + @app.task + @inject + async def handler(read: ReadService = FromDishka()) -> str: + return read.tag + + mw = KuuDishkaMiddleware(container) + for _ in range(3): + ctx = _ctx(app, handler, Payload()) + await run_chain(ctx, [mw], await _terminal_for(handler)) + + assert CountingProvider.instances == 3 + await container.close() + + +# --------------------------------------------------------------------------- # +# Full end-to-end through the real Worker (.q -> broker -> worker -> inject) +# --------------------------------------------------------------------------- # +async def _run_worker_until(app: Kuu, predicate, *, timeout: float = 3.0) -> None: + config = Settings(app="test:app", task_modules=["t"], queues=["default"], concurrency=4) + worker = Worker(config, app=app) + + async def _supervise(scope: anyio.CancelScope): + while not predicate(): + await anyio.sleep(0.02) + scope.cancel() + + with anyio.fail_after(timeout): + async with anyio.create_task_group() as tg: + tg.start_soon(_supervise, tg.cancel_scope) + tg.start_soon(worker.run) + + +async def test_end_to_end_async_task(async_container): + app = Kuu(broker=MemoryBroker()) + setup_dishka(app, async_container) + captured: dict[str, Any] = {} + + @app.task + @inject + async def get_lot( + lot_id: str, + read: ReadService = FromDishka(), + write: WriteService = FromDishka(), + fetch: bool = True, + ) -> str: + captured["result"] = f"{lot_id}:{read.tag}:{write.tag}:{fetch}" + return captured["result"] + + # .q omits read/write entirely + await get_lot.q("lot-1", fetch=False) + await _run_worker_until(app, lambda: "result" in captured) + + assert captured["result"] == "lot-1:read:write:False" + + +async def test_end_to_end_blocking_task(sync_container): + app = Kuu(broker=MemoryBroker()) + setup_dishka(app, None, sync_container=sync_container) + captured: dict[str, Any] = {} + + @app.task(blocking=True) + @inject + def get_lot(lot_id: str, read: ReadService = FromDishka()) -> str: + captured["result"] = f"{lot_id}:{read.tag}" + return captured["result"] + + await get_lot.q("lot-9") + await _run_worker_until(app, lambda: "result" in captured) + + assert captured["result"] == "lot-9:read"