diff --git a/src/backend/base/langflow/services/background_execution/metrics_collector.py b/src/backend/base/langflow/services/background_execution/metrics_collector.py index 8e41589224..f660dfd2d3 100644 --- a/src/backend/base/langflow/services/background_execution/metrics_collector.py +++ b/src/backend/base/langflow/services/background_execution/metrics_collector.py @@ -16,11 +16,16 @@ the lease window (``age <= lease_window``), stale once older. from __future__ import annotations +import asyncio +import contextlib from datetime import datetime, timezone +from lfx.log.logger import logger from sqlmodel import col, func, select +from langflow.services.background_execution.metrics import current_backend from langflow.services.database.models.jobs.model import Job, JobStatus +from langflow.services.deps import get_telemetry_service, session_scope # Non-terminal statuses: a job in one of these is still occupying the system. # QUEUED/IN_PROGRESS are the only non-terminal states; COMPLETED / FAILED / @@ -97,3 +102,74 @@ async def alive_worker_count(session, now: datetime, lease_window: float) -> int if age <= lease_window: alive_owners.add(owner) return len(alive_owners) + + +class BackgroundMetricsCollector: + """Periodically pushes DB-derived bg-execution gauges to the OTel registry. + + The collector owns the clock: each tick computes one aware-UTC ``now``, + runs the three query functions against a short-lived session, and writes the + gauges via ``get_telemetry_service().ot.update_gauge``. It is the only writer + of these gauges; an ObservableGauge reports the last value set per label-set, + so each tick zero-fills the canonical non-terminal status set to make a status + dropping to 0 overwrite a stale prior value. + + Best-effort: a failing tick logs a warning and returns without raising so the + loop keeps running — observability must never crash the service. + + ``lease_window`` defaults to ``45.0`` to match the runtime + ``background_lease_ttl_s`` setting that ``sweep_orphans`` / ``requeue_lost`` + are actually called with, so "alive workers" agrees with what the watchdog + reconciles. Task 8 passes ``settings.background_lease_ttl_s`` explicitly. + """ + + def __init__(self, *, interval: float, lease_window: float = 45.0): + self.interval = interval + self.lease_window = lease_window + self._stopped = False + self._task: asyncio.Task | None = None + + async def collect_once(self, session) -> None: + """Run the three queries and push the gauges. Never raises.""" + try: + now = datetime.now(timezone.utc) + counts = await count_nonterminal_jobs(session) + oldest = await oldest_queued_seconds(session, now) + alive = await alive_worker_count(session, now, self.lease_window) + + backend = current_backend() + ot = get_telemetry_service().ot + + # Zero-fill the canonical non-terminal set so a status that drops to 0 + # overwrites the gauge's stale prior value (last-value-wins semantics). + for status in (JobStatus.QUEUED, JobStatus.IN_PROGRESS): + ot.update_gauge( + "langflow_bg_jobs", + counts.get(status.value, 0), + {"status": status.value, "backend": backend}, + ) + ot.update_gauge("langflow_bg_oldest_queued_seconds", oldest, {"backend": backend}) + ot.update_gauge("langflow_bg_alive_workers", alive, {"backend": backend}) + except Exception as exc: # noqa: BLE001 - observability must never crash the loop + logger.warning(f"bg metrics collection tick skipped: {exc}") + + async def run(self) -> None: + """Loop: open a short-lived session per tick, collect, sleep the interval.""" + while not self._stopped: + async with session_scope() as session: + await self.collect_once(session) + await asyncio.sleep(self.interval) + + def start(self) -> None: + """Spawn the collector loop task.""" + self._stopped = False + self._task = asyncio.create_task(self.run()) + + async def stop(self) -> None: + """Signal stop and cancel/await the loop task.""" + self._stopped = True + if self._task is not None: + self._task.cancel() + with contextlib.suppress(asyncio.CancelledError): + await self._task + self._task = None diff --git a/src/backend/tests/unit/services/background_execution/test_metrics_collector.py b/src/backend/tests/unit/services/background_execution/test_metrics_collector.py index ce963f5f54..e0b4ed538d 100644 --- a/src/backend/tests/unit/services/background_execution/test_metrics_collector.py +++ b/src/backend/tests/unit/services/background_execution/test_metrics_collector.py @@ -17,13 +17,15 @@ from datetime import datetime, timedelta, timezone from uuid import uuid4 import pytest +from langflow.services.background_execution.metrics import current_backend from langflow.services.background_execution.metrics_collector import ( + BackgroundMetricsCollector, alive_worker_count, count_nonterminal_jobs, oldest_queued_seconds, ) from langflow.services.database.models.jobs.model import JobStatus -from langflow.services.deps import session_scope +from langflow.services.deps import get_telemetry_service, session_scope from langflow.services.jobs.service import JobService pytestmark = pytest.mark.usefixtures("client") @@ -146,3 +148,103 @@ async def test_alive_worker_count_distinct_owner(): alive = await alive_worker_count(session, now, lease_window) assert alive == 1 + + +def _gauge_value(metric_name: str, labels: dict[str, str]) -> float: + """Read a gauge value straight off the real OTel ObservableGaugeWrapper. + + The wrapper stores values keyed by ``tuple(sorted(labels.items()))`` — assert + the SPECIFIC label-set we set this tick so a prior test's other label-sets in + the process-wide singleton cannot interfere. + """ + gauge = get_telemetry_service().ot._metrics[metric_name] + return gauge._values[tuple(sorted(labels.items()))] + + +async def test_collect_once_sets_gauges(): + """One tick over a seeded mix sets each gauge to the DB-derived value.""" + service = JobService() + backend = current_backend() + + queued_a = uuid4() + queued_b = uuid4() + in_progress = uuid4() + completed = uuid4() + + await service.create_job(job_id=queued_a, flow_id=uuid4(), user_id=uuid4()) + await service.create_job(job_id=queued_b, flow_id=uuid4(), user_id=uuid4()) + + await service.create_job(job_id=in_progress, flow_id=uuid4(), user_id=uuid4()) + await service.update_job_status(in_progress, JobStatus.IN_PROGRESS) + await service.heartbeat(in_progress, owner="worker-A") + + await service.create_job(job_id=completed, flow_id=uuid4(), user_id=uuid4()) + await service.update_job_status(completed, JobStatus.COMPLETED, finished_timestamp=True) + + collector = BackgroundMetricsCollector(interval=15.0) + async with session_scope() as session: + await collector.collect_once(session) + + assert _gauge_value("langflow_bg_jobs", {"status": "queued", "backend": backend}) == 2 + assert _gauge_value("langflow_bg_jobs", {"status": "in_progress", "backend": backend}) == 1 + assert _gauge_value("langflow_bg_oldest_queued_seconds", {"backend": backend}) > 0 + assert _gauge_value("langflow_bg_alive_workers", {"backend": backend}) == 1 + + +async def test_collect_once_zero_fills_dropped_status(): + """A status dropping to 0 overwrites the stale prior value (zero-fill).""" + service = JobService() + backend = current_backend() + + queued_a = uuid4() + queued_b = uuid4() + await service.create_job(job_id=queued_a, flow_id=uuid4(), user_id=uuid4()) + await service.create_job(job_id=queued_b, flow_id=uuid4(), user_id=uuid4()) + + collector = BackgroundMetricsCollector(interval=15.0) + async with session_scope() as session: + await collector.collect_once(session) + + assert _gauge_value("langflow_bg_jobs", {"status": "queued", "backend": backend}) == 2 + + # Drain the queue: both queued jobs go terminal, so queued must fall to 0. + await service.update_job_status(queued_a, JobStatus.COMPLETED, finished_timestamp=True) + await service.update_job_status(queued_b, JobStatus.COMPLETED, finished_timestamp=True) + + async with session_scope() as session: + await collector.collect_once(session) + + assert _gauge_value("langflow_bg_jobs", {"status": "queued", "backend": backend}) == 0 + + +async def test_run_stop_lifecycle(): + """start() spawns the loop, a tick runs against the real DB, stop() cancels cleanly. + + The process-wide OTel singleton already holds the ``queued`` label-set from + prior tests, so we cannot key the wait on the key merely existing. Instead we + flip the value to a sentinel, let a real tick overwrite it with this DB's + count (>= the one queued job we seeded), then assert stop() ends the loop. + """ + import asyncio + + service = JobService() + backend = current_backend() + + queued = uuid4() + await service.create_job(job_id=queued, flow_id=uuid4(), user_id=uuid4()) + + sentinel = -1.0 + gauge = get_telemetry_service().ot._metrics["langflow_bg_jobs"] + key = tuple(sorted({"status": "queued", "backend": backend}.items())) + gauge._values[key] = sentinel + + collector = BackgroundMetricsCollector(interval=0.01) + collector.start() + for _ in range(200): + if gauge._values.get(key, sentinel) != sentinel: + break + await asyncio.sleep(0.01) + + await collector.stop() + assert collector._task is None + assert _gauge_value("langflow_bg_jobs", {"status": "queued", "backend": backend}) >= 1