services/gr-processor/src/main.py - develop - laritsiuriumov/aef-manager-python-services

services/gr-processor/src/main.py
297 строк9 KB
Цюрюмов Лари Валерьевич [B]
[AEF-8279] Align gr-processor with docs: gr-core naming, ai_hub_id key, INCOMPLETE_ALERT stats, camelCase contract
07 июл 2026, 19:27
07 июл 2026, 19:279f2d501
import asyncio
import logging
from collections.abc import Awaitable, Callable
from contextlib import asynccontextmanager

import uvicorn
from aiokafka import AIOKafkaConsumer, AIOKafkaProducer
from fastapi import FastAPI, Request, status
from fastapi.responses import JSONResponse
from starlette.middleware.cors import CORSMiddleware

from aggregator import Aggregator
from config import Settings
from logging_utils import configure_logging
from runtime import RuntimeState

settings = Settings()  # type: ignore

configure_logging(
    service_name="gr-processor",
    log_path=settings.log_path,
    level=settings.basic_logger_level,
)
logging.getLogger("aiokafka").setLevel(settings.aiokafka_logger_level)
logger = logging.getLogger(__name__)
PROBE_PATHS = frozenset({"/health", "/ready", "/readiness", "/live", "/liveness"})


def _describe_error(exc: Exception) -> str:
    return exc.__class__.__name__


class ProbeAccessFilter(logging.Filter):
    def filter(self, record: logging.LogRecord) -> bool:
        args = record.args
        if not isinstance(args, (tuple, list)) or len(args) < 5:
            return True

        path = str(args[2]).split("?", 1)[0]
        status_code = str(args[4])
        return not (path in PROBE_PATHS and status_code == "200")


def _configure_uvicorn_access_logging() -> None:
    access_logger = logging.getLogger("uvicorn.access")
    if any(isinstance(filter_, ProbeAccessFilter) for filter_ in access_logger.filters):
        return
    access_logger.addFilter(ProbeAccessFilter())


def _build_runtime_state() -> RuntimeState:
    return RuntimeState(
        service_name="gr-processor",
        dependencies={
            "kafka_consumer": False,
            "kafka_producer": False,
        },
        core_workers={"aggregator": False},
    )


async def _run_with_retries(
    action: Callable[[], Awaitable[None]],
    *,
    attempts: int,
    delay: int,
    operation: str,
) -> None:
    last_error: Exception | None = None

    for attempt in range(1, attempts + 1):
        try:
            await action()
            logger.info("%s succeeded on attempt %d/%d", operation, attempt, attempts)
            return
        except Exception as exc:
            last_error = exc
            if attempt == attempts:
                break

            retry_delay = delay * attempt
            logger.warning(
                "%s failed on attempt %d/%d due to %s. Retrying in %s seconds...",
                operation,
                attempt,
                attempts,
                _describe_error(exc),
                retry_delay,
            )
            await asyncio.sleep(retry_delay)

    assert last_error is not None
    logger.error(
        "%s failed after %d attempts due to %s",
        operation,
        attempts,
        _describe_error(last_error),
    )
    raise RuntimeError(f"{operation} failed after {attempts} attempts") from last_error


async def _stop_kafka_clients(
    consumer: AIOKafkaConsumer, producer: AIOKafkaProducer
) -> None:
    try:
        await producer.stop()
    except Exception:
        logger.exception("Failed to stop Kafka producer cleanly")

    try:
        await consumer.stop()
    except Exception:
        logger.exception("Failed to stop Kafka consumer cleanly")


async def _connect_kafka(
    consumer: AIOKafkaConsumer,
    producer: AIOKafkaProducer,
    runtime_state: RuntimeState,
) -> None:
    async def connect() -> None:
        await consumer.start()
        await producer.start()

    await _run_with_retries(
        connect,
        attempts=settings.kafka_connection_max_retries,
        delay=settings.kafka_connection_retry_delay,
        operation="Kafka startup",
    )
    runtime_state.mark_dependency("kafka_consumer", True)
    runtime_state.mark_dependency("kafka_producer", True)


async def _run_aggregator_worker(runtime_state: RuntimeState) -> None:
    consumer = AIOKafkaConsumer(
        settings.consumer_kafka_topic,
        bootstrap_servers=settings.consumer_kafka_host_uri,  # type: ignore
        group_id=settings.consumer_group,
    )
    producer = AIOKafkaProducer(bootstrap_servers=settings.publisher_kafka_host_uri)  # type: ignore

    try:
        await _connect_kafka(consumer, producer, runtime_state)
        runtime_state.mark_core_worker("aggregator", True)
        aggregator = Aggregator(
            settings,
            consumer=consumer,
            producer=producer,
            runtime_state=runtime_state,
        )
        await aggregator.run()
    except asyncio.CancelledError:
        raise
    except Exception as exc:
        runtime_state.set_fatal_error(
            f"aggregator worker failed: {_describe_error(exc)}"
        )
        logger.exception("Aggregator worker failed")
        raise
    finally:
        runtime_state.mark_core_worker("aggregator", False)
        runtime_state.mark_dependency("kafka_consumer", False)
        runtime_state.mark_dependency("kafka_producer", False)
        await _stop_kafka_clients(consumer, producer)
        logger.info("Kafka connections closed")


async def _wait_for_worker_startup(
    task: asyncio.Task[None], runtime_state: RuntimeState
) -> None:
    attempts = settings.kafka_connection_max_retries
    retry_delay = settings.kafka_connection_retry_delay
    total_retry_wait = retry_delay * attempts * (attempts - 1) / 2
    timeout_seconds = total_retry_wait + 5
    deadline = asyncio.get_running_loop().time() + max(timeout_seconds, 5)

    while True:
        if runtime_state.is_ready and runtime_state.is_live:
            return

        if runtime_state.fatal_error is not None:
            raise RuntimeError(runtime_state.fatal_error)

        if task.done():
            error = task.exception()
            if error is not None:
                raise RuntimeError("Aggregator worker stopped during startup") from error
            raise RuntimeError("Aggregator worker exited during startup")

        if asyncio.get_running_loop().time() > deadline:
            raise TimeoutError("Timed out waiting for aggregator worker startup")

        await asyncio.sleep(0.05)


def _probe_response(
    runtime_state: RuntimeState,
    *,
    payload_status: str,
    status_code: int,
) -> JSONResponse:
    return JSONResponse(
        status_code=status_code,
        content=runtime_state.payload(payload_status),
    )


@asynccontextmanager
async def lifespan(app: FastAPI):
    task: asyncio.Task[None] | None = None
    runtime_state: RuntimeState = app.state.runtime_state

    if app.state.start_worker:
        task = asyncio.create_task(_run_aggregator_worker(runtime_state))
        app.state.worker_task = task
        await _wait_for_worker_startup(task, runtime_state)

    try:
        yield
    finally:
        if task is not None:
            task.cancel()
            try:
                await task
            except asyncio.CancelledError:
                pass


def create_app(*, start_worker: bool = True) -> FastAPI:
    app = FastAPI(title="GR Processor", lifespan=lifespan)
    app.state.runtime_state = _build_runtime_state()
    app.state.start_worker = start_worker
    app.state.worker_task = None

    @app.get("/health")
    async def health(request: Request) -> JSONResponse:
        runtime_state: RuntimeState = request.app.state.runtime_state
        return _probe_response(
            runtime_state,
            payload_status="ok",
            status_code=status.HTTP_200_OK,
        )

    @app.get("/readiness")
    async def readiness(request: Request) -> JSONResponse:
        runtime_state: RuntimeState = request.app.state.runtime_state
        return _probe_response(
            runtime_state,
            payload_status="ready" if runtime_state.is_ready else "not_ready",
            status_code=(
                status.HTTP_200_OK
                if runtime_state.is_ready
                else status.HTTP_503_SERVICE_UNAVAILABLE
            ),
        )

    @app.get("/liveness")
    async def liveness(request: Request) -> JSONResponse:
        runtime_state: RuntimeState = request.app.state.runtime_state
        return _probe_response(
            runtime_state,
            payload_status="alive" if runtime_state.is_live else "not_live",
            status_code=(
                status.HTTP_200_OK
                if runtime_state.is_live
                else status.HTTP_503_SERVICE_UNAVAILABLE
            ),
        )

    app.add_middleware(
        CORSMiddleware,
        allow_origins=["*"],
        allow_credentials=True,
        allow_methods=["*"],
        allow_headers=["*"],
    )
    return app


app = create_app()
_configure_uvicorn_access_logging()


def main() -> None:
    uvicorn.run(
        "main:app",
        host="0.0.0.0",
        port=8080,
        log_level=settings.uvicorn_logger_level.lower(),
        reload=settings.debug,
        log_config=None,
    )


if __name__ == "__main__":
    main()