diff --git a/backend/backend/internal_base_urls.py b/backend/backend/internal_base_urls.py index 0354a691ae..30d4c8df22 100644 --- a/backend/backend/internal_base_urls.py +++ b/backend/backend/internal_base_urls.py @@ -191,6 +191,13 @@ def test_middleware_debug(request): include("workflow_manager.workflow_execution_internal_urls"), name="workflow_execution_internal", ), + # Dashboard-metrics periodics (UN-3796) — called by the thin worker tasks the PG + # scheduler fires, replacing Beat + workerMetrics. + path( + "v1/dashboard-metrics/", + include("dashboard_metrics.internal_urls"), + name="dashboard_metrics_internal", + ), # Workflow management and pipeline APIs path( "v1/workflow-manager/", diff --git a/backend/dashboard_metrics/internal_urls.py b/backend/dashboard_metrics/internal_urls.py new file mode 100644 index 0000000000..bc9d698f75 --- /dev/null +++ b/backend/dashboard_metrics/internal_urls.py @@ -0,0 +1,29 @@ +"""Internal API URLs for the dashboard-metrics periodics (UN-3796). + +Called by the thin worker tasks that the PG scheduler fires, replacing Beat + +``workerMetrics``. Mirrors the shape of ``execution_log_internal_urls``. +""" + +from django.urls import path + +from . import internal_views + +app_name = "dashboard_metrics_internal" + +urlpatterns = [ + path( + "aggregate/", + internal_views.AggregateMetricsAPIView.as_view(), + name="aggregate_metrics", + ), + path( + "cleanup/hourly/", + internal_views.CleanupHourlyMetricsAPIView.as_view(), + name="cleanup_hourly_metrics", + ), + path( + "cleanup/daily/", + internal_views.CleanupDailyMetricsAPIView.as_view(), + name="cleanup_daily_metrics", + ), +] diff --git a/backend/dashboard_metrics/internal_views.py b/backend/dashboard_metrics/internal_views.py new file mode 100644 index 0000000000..f776633944 --- /dev/null +++ b/backend/dashboard_metrics/internal_views.py @@ -0,0 +1,115 @@ +"""Internal API for running the dashboard-metrics periodics (UN-3796). + +Beat + ``workerMetrics`` are the last two Celery deployments blocking the +"scale every Celery deployment to zero" gate. The PG scheduler replaces Beat, but the +three ``dashboard_metrics.*`` tasks are **Django** — ORM, cache, a Redis lock — and the +PG consumer bootstraps its tasks from ``workers/`` with no Django, so it cannot run them +directly. + +These endpoints are the execution half, following the pattern already used by +``ProcessLogHistoryAPIView`` / ``process_notification_buffer``: a thin worker-side task +POSTs here, and the backend runs the real function. The task bodies are plain functions +that happen to carry ``@shared_task``, so they are called **verbatim** — the aggregation +logic, its windows and its Redis lock are reused unchanged, not reimplemented. + +Auth is entirely ``InternalAPIAuthMiddleware``, which fires on any ``/internal/`` path. + +**Org context is deliberately absent.** The Celery path runs with no organization set — +which is why ``tasks.py`` uses ``_base_manager`` throughout — and the middleware +populates ``StateStore`` from ``X-Organization-ID``. The callers must not send that +header, and each view clears the slot defensively: ``StateStore`` is a thread-local and +gunicorn reuses threads, so a value left behind by an earlier request on the same thread +would silently scope these global aggregations to one tenant. +""" + +import contextlib +import logging +from typing import Any + +from rest_framework import status +from rest_framework.request import Request +from rest_framework.response import Response +from rest_framework.views import APIView +from utils.constants import Account +from utils.local_context import StateStore + +from dashboard_metrics.tasks import ( + aggregate_metrics_from_sources, + cleanup_daily_metrics, + cleanup_hourly_metrics, +) + +logger = logging.getLogger(__name__) + +# Mirrors the defaults the Beat rows carry in their kwargs +# (dashboard_metrics/migrations/0002_setup_periodic_tasks.py), so a caller that omits +# them gets the same retention the Celery path applies. +DEFAULT_HOURLY_RETENTION_DAYS = 30 +DEFAULT_DAILY_RETENTION_DAYS = 365 + + +def _clear_org_context() -> None: + """Drop any organization left in this thread's StateStore. See the module docstring. + + ``suppress`` because ``StateStore.clear`` raises when the slot was never set, which + is the normal case and not an error. + """ + with contextlib.suppress(Exception): + StateStore.clear(Account.ORGANIZATION_ID) + + +def _int_arg(request: Request, key: str, default: int) -> int: + """Read an optional positive integer from the request body.""" + raw = request.data.get(key, default) if isinstance(request.data, dict) else default + try: + value = int(raw) + except (TypeError, ValueError) as exc: + raise ValueError(f"{key} must be an integer, got {raw!r}") from exc + if value < 1: + raise ValueError(f"{key} must be >= 1, got {value}") + return value + + +class _MetricsTaskAPIView(APIView): + """Shared plumbing: clear org context, run, translate errors.""" + + def _run(self, fn, *args: Any, **kwargs: Any) -> Response: + _clear_org_context() + try: + return Response(fn(*args, **kwargs)) + except ValueError as exc: # bad request body + return Response({"error": str(exc)}, status=status.HTTP_400_BAD_REQUEST) + except Exception as exc: + logger.error("dashboard-metrics internal call failed: %s", exc, exc_info=True) + return Response( + {"error": str(exc)}, status=status.HTTP_500_INTERNAL_SERVER_ERROR + ) + + +class AggregateMetricsAPIView(_MetricsTaskAPIView): + """Run the metrics aggregation. + + Calls the Celery task body verbatim, Redis lock included — this endpoint exists + only because the PG consumer has no Django, not to change what the job does. + """ + + def post(self, request: Request) -> Response: + return self._run(aggregate_metrics_from_sources) + + +class CleanupHourlyMetricsAPIView(_MetricsTaskAPIView): + def post(self, request: Request) -> Response: + try: + days = _int_arg(request, "retention_days", DEFAULT_HOURLY_RETENTION_DAYS) + except ValueError as exc: + return Response({"error": str(exc)}, status=status.HTTP_400_BAD_REQUEST) + return self._run(cleanup_hourly_metrics, retention_days=days) + + +class CleanupDailyMetricsAPIView(_MetricsTaskAPIView): + def post(self, request: Request) -> Response: + try: + days = _int_arg(request, "retention_days", DEFAULT_DAILY_RETENTION_DAYS) + except ValueError as exc: + return Response({"error": str(exc)}, status=status.HTTP_400_BAD_REQUEST) + return self._run(cleanup_daily_metrics, retention_days=days) diff --git a/backend/dashboard_metrics/migrations/0004_pg_periodic_tasks.py b/backend/dashboard_metrics/migrations/0004_pg_periodic_tasks.py new file mode 100644 index 0000000000..463966a5b3 --- /dev/null +++ b/backend/dashboard_metrics/migrations/0004_pg_periodic_tasks.py @@ -0,0 +1,97 @@ +"""Declare the dashboard-metrics periodics for the PG scheduler (UN-3796). + +The PG twin of ``0002_setup_periodic_tasks``, which declares the same three schedules +for Celery Beat. **Deliberately in this app rather than in ``pg_queue``**: the failure +mode that matters is the two declarations drifting apart, and a reviewer editing a +schedule sees both only if they sit side by side. + +Why a data migration here, when the pipeline mirror is a management command: these are +three fixed rows known at build time, so a migration is the right tool — tiny, idempotent +(``update_or_create``), and it reaches every environment including on-prem with no +operator step. Pipeline schedules are bulk and per-environment, so they stay a chunked +command that runs outside the migrate transaction. + +Rows land **inert**: ``pg_owned=False`` and ``next_run_at=NULL``. Nothing fires from this +migration — the PG scheduler skips rows it does not own, and a NULL ``next_run_at`` means +"record a baseline next tick" rather than "overdue, fire now", so enabling the flag never +produces a burst of catch-up runs. + +``task_kwargs`` is stored **decoded**: Beat keeps ``kwargs`` as a JSON *string* +(``'{"retention_days": 30}'``) while ``PgPeriodicTask.task_kwargs`` is a JSONField, so the +dispatcher can build a payload without re-parsing per tick. +""" + +from django.db import migrations + +# Single source for both directions, and importable by the drift test. Mirrors the Beat +# declarations in 0002_setup_periodic_tasks one-for-one — same names (the mirror key is +# PeriodicTask.name), same queue, same kwargs, and cron strings equivalent to the +# IntervalSchedule/CrontabSchedule rows there. +PG_PERIODIC_TASKS = [ + { + "name": "dashboard_metrics_aggregate_from_sources", + "task_name": "dashboard_metrics.aggregate_from_sources", + "queue": "dashboard_metric_events", + "task_args": [], + "task_kwargs": {}, + # Beat: IntervalSchedule(every=15, period="minutes") + "cron_string": "*/15 * * * *", + }, + { + "name": "dashboard_metrics_cleanup_hourly", + "task_name": "dashboard_metrics.cleanup_hourly_data", + "queue": "dashboard_metric_events", + "task_args": [], + "task_kwargs": {"retention_days": 30}, + # Beat: CrontabSchedule(minute=0, hour=2, every day) UTC + "cron_string": "0 2 * * *", + }, + { + "name": "dashboard_metrics_cleanup_daily", + "task_name": "dashboard_metrics.cleanup_daily_data", + "queue": "dashboard_metric_events", + "task_args": [], + "task_kwargs": {"retention_days": 365}, + # Beat: CrontabSchedule(minute=0, hour=3, day_of_week=0 → Sunday) UTC + "cron_string": "0 3 * * 0", + }, +] + + +def create_pg_periodic_tasks(apps, schema_editor): + PgPeriodicTask = apps.get_model("pg_queue", "PgPeriodicTask") + for spec in PG_PERIODIC_TASKS: + PgPeriodicTask.objects.update_or_create( + name=spec["name"], + defaults={ + "task_name": spec["task_name"], + "queue": spec["queue"], + "task_args": spec["task_args"], + "task_kwargs": spec["task_kwargs"], + "cron_string": spec["cron_string"], + "org_id": "", + "enabled": True, + # Inert until the rollout flag decides otherwise; never fired by + # applying this migration. + "pg_owned": False, + }, + ) + + +def remove_pg_periodic_tasks(apps, schema_editor): + PgPeriodicTask = apps.get_model("pg_queue", "PgPeriodicTask") + PgPeriodicTask.objects.filter( + name__in=[spec["name"] for spec in PG_PERIODIC_TASKS] + ).delete() + + +class Migration(migrations.Migration): + dependencies = [ + ("dashboard_metrics", "0003_alter_eventmetricsdaily_organization_and_more"), + # The table this seeds. + ("pg_queue", "0003_pgperiodictask"), + ] + + operations = [ + migrations.RunPython(create_pg_periodic_tasks, remove_pg_periodic_tasks), + ] diff --git a/backend/dashboard_metrics/tests/test_pg_periodic_task_declarations.py b/backend/dashboard_metrics/tests/test_pg_periodic_task_declarations.py new file mode 100644 index 0000000000..85ea407899 --- /dev/null +++ b/backend/dashboard_metrics/tests/test_pg_periodic_task_declarations.py @@ -0,0 +1,113 @@ +"""Drift guard between the Beat and PG declarations of the metrics periodics (UN-3796). + +Two migrations declare the same three schedules — ``0002_setup_periodic_tasks`` for Celery +Beat and ``0004_pg_periodic_tasks`` for the PG scheduler. They are separate rows in +separate tables, so nothing stops someone editing one and forgetting the other. That is +the whole failure mode this file exists for: a schedule changed on Beat but not on PG means +the task silently runs on a different cadence the moment the flag flips. + +DB-free — both migration modules are imported and their declared specs compared directly, +so this runs in the unit tier rather than needing a migrated database. +""" + +from __future__ import annotations + +import importlib +import json + +import pytest + +_BEAT_MIGRATION = "dashboard_metrics.migrations.0002_setup_periodic_tasks" +_PG_MIGRATION = "dashboard_metrics.migrations.0004_pg_periodic_tasks" + +# Cron equivalent of each Beat schedule, asserted against what the Beat migration builds. +# Written out rather than derived: deriving it from the same code under test would make +# the comparison vacuous. +_EXPECTED_CRON = { + "dashboard_metrics_aggregate_from_sources": "*/15 * * * *", + "dashboard_metrics_cleanup_hourly": "0 2 * * *", + "dashboard_metrics_cleanup_daily": "0 3 * * 0", +} + + +@pytest.fixture(scope="module") +def pg_specs() -> dict[str, dict]: + mod = importlib.import_module(_PG_MIGRATION) + return {spec["name"]: spec for spec in mod.PG_PERIODIC_TASKS} + + +class _FakeQuerySet: + """Captures update_or_create calls from the Beat migration without a database.""" + + def __init__(self, sink: dict): + self._sink = sink + + def get_or_create(self, **kwargs): + # Schedule rows (Interval/Crontab) — return the kwargs so the PeriodicTask + # call can be inspected for which schedule it was given. + return kwargs, True + + def update_or_create(self, name=None, defaults=None, **_kw): + self._sink[name] = defaults or {} + return defaults, True + + def filter(self, *_a, **_k): + return self + + def delete(self): + return (0, {}) + + +@pytest.fixture(scope="module") +def beat_specs() -> dict[str, dict]: + """Run the Beat migration's forward function against fakes and capture what it declares.""" + mod = importlib.import_module(_BEAT_MIGRATION) + captured: dict[str, dict] = {} + + class _Apps: + def get_model(self, _app, model): + if model == "PeriodicTask": + return type("PT", (), {"objects": _FakeQuerySet(captured)}) + return type("S", (), {"objects": _FakeQuerySet({})}) + + mod.create_periodic_tasks(_Apps(), None) + return captured + + +class TestDeclarationsAgree: + def test_same_set_of_schedules(self, beat_specs, pg_specs): + # A schedule added to Beat but not PG stops firing the moment the flag flips; + # the reverse fires something Beat never knew about. + assert set(beat_specs) == set(pg_specs) + + @pytest.mark.parametrize("name", sorted(_EXPECTED_CRON)) + def test_task_path_and_queue_match(self, beat_specs, pg_specs, name): + assert pg_specs[name]["task_name"] == beat_specs[name]["task"] + assert pg_specs[name]["queue"] == beat_specs[name]["queue"] + + @pytest.mark.parametrize("name", sorted(_EXPECTED_CRON)) + def test_kwargs_match_once_decoded(self, beat_specs, pg_specs, name): + # Beat stores kwargs as a JSON *string*; PgPeriodicTask.task_kwargs is a + # JSONField. A mismatch here means the cleanup runs with the wrong retention. + beat_kwargs = json.loads(beat_specs[name].get("kwargs") or "{}") + assert pg_specs[name]["task_kwargs"] == beat_kwargs + + @pytest.mark.parametrize("name,cron", sorted(_EXPECTED_CRON.items())) + def test_cron_matches_the_beat_cadence(self, pg_specs, name, cron): + assert pg_specs[name]["cron_string"] == cron + + +class TestSeededInert: + """Applying the migration must not cause anything to fire.""" + + def test_no_spec_declares_itself_pg_owned(self, pg_specs): + # pg_owned is set to False in the migration's defaults, never from the spec — + # this pins that no spec can smuggle ownership in. + assert not any("pg_owned" in spec for spec in pg_specs.values()) + + def test_no_spec_presets_a_run_time(self, pg_specs): + # A non-NULL next_run_at in the past would read as "overdue" and fire a burst + # of catch-up runs the moment the flag is enabled. + for spec in pg_specs.values(): + assert "next_run_at" not in spec + assert "last_run_at" not in spec diff --git a/backend/entrypoint.sh b/backend/entrypoint.sh index 20441c58fb..9766b01add 100755 --- a/backend/entrypoint.sh +++ b/backend/entrypoint.sh @@ -27,6 +27,49 @@ done if [ "$migrate" = true ]; then echo "Migration initiated" .venv/bin/python manage.py migrate + + # Converge schedule ownership to whatever PG_SCHEDULER_ENABLED declares (UN-3796) + # — adopt when on, release back to Beat when off. NOT a migration: an ordinary + # management command, sequenced after migrate only because pg_periodic_schedule + # must exist first. + # + # READ THIS BEFORE ASSUMING A RESTART IS INERT. This is NOT the earlier + # `--mirror-only` invocation, and it does NOT leave Beat alone: with + # PG_SCHEDULER_ENABLED=true it flips pg_owned and DISABLES the Beat PeriodicTask + # row for every mirrored pipeline THE pg_queue_enabled FLAG SELECTS — ownership is + # additionally per-schedule Flipt-gated and fails closed, so with Flipt blind or the + # flag at 0% this moves nothing and Beat keeps every row. A rolling restart therefore can move schedules + # off Beat. (A comment here used to describe --mirror-only and state that + # "ownership hand-over stays an explicit operator action" — it survived the switch + # to converge_pg_scheduler and told an SRE auditing exactly this question the + # opposite of the truth.) + # + # Running it on every start is what makes the ROLLBACK real: flipping the env var + # back is enough, with no operator remembering a management command in every + # environment (on-prem included, where `manage.py` is not something a deploy can + # reach). Idempotent both ways in OUTCOME, but NOT free: with the gate on it + # rewrites pg_periodic_schedule and the Beat PeriodicTask row and bumps + # PeriodicTasks.update_changed() for every schedule, whether or not ownership + # changed — roughly 2N row writes and N Beat reloads per pod start. (A previous + # version of this line said "a restart that changes nothing writes nothing"; that + # was wrong, and contradicted converge_pg_scheduler's own docstring.) The release + # direction runs --mirror-only and skips the reconcile, so it costs nothing here. + # + # PG_SCHEDULER_ADOPT_PERIODICS additionally moves the dashboard_metrics.* rows; + # it is separate because adopting them needs workerPgMetrics deployed. + # + # Best-effort by design: a convergence failure must never stop the backend from + # starting. Note it is NOT all-or-nothing though — convergence commits per schedule + # as it goes, so a failure can leave a partial hand-over. The command's summary line + # names the failed count; re-run it or flip the env var back. + PERIODICS_FLAG="" + if [ "$(printf '%s' "${PG_SCHEDULER_ADOPT_PERIODICS:-}" | tr '[:upper:]' '[:lower:]')" = "true" ]; then + PERIODICS_FLAG="--periodics" + fi + echo "PG schedule ownership convergence initiated" + # shellcheck disable=SC2086 # intentional word-splitting: empty = flag absent + .venv/bin/python manage.py converge_pg_scheduler $PERIODICS_FLAG \ + || echo "WARNING: PG schedule convergence failed; continuing startup (schedules keep their current firer)" fi # Configure Gunicorn based on --dev flag. diff --git a/backend/pg_queue/management/commands/converge_pg_scheduler.py b/backend/pg_queue/management/commands/converge_pg_scheduler.py new file mode 100644 index 0000000000..a4341f5827 --- /dev/null +++ b/backend/pg_queue/management/commands/converge_pg_scheduler.py @@ -0,0 +1,115 @@ +"""Converge schedule ownership to the state ``PG_SCHEDULER_ENABLED`` declares. + +One idempotent entry point, **both directions**: + + PG_SCHEDULER_ENABLED=true → adopt (PG fires; Beat rows disabled) + PG_SCHEDULER_ENABLED=false → release (Beat fires; pg_owned cleared) + +Why one converging command rather than an adopt command and a release command: the +reverse direction is the one that matters under pressure, and a rollback nobody can +run is not a rollback. Splitting it left the reverse as a *procedure* — remember the +command, remember the flags, remember to run it in every environment. Here the env +var declares intent and the deploy makes the database match, so reverting is a values +change like any other, and re-running changes nothing once converged. + +**Safe to run unattended**, which is what lets ``entrypoint.sh`` call it on every +start: + +* Both directions are idempotent in OUTCOME — ``_set_ownership`` skips periodic rows + already in the target state. The pipeline path is idempotent but **not** free: it + rewrites ``pg_periodic_schedule`` and the Beat ``PeriodicTask`` row and bumps + ``PeriodicTasks.update_changed()`` for every schedule on every call, whether or not + ownership changed — so a backend start costs 2N row writes and N Beat reloads. Safe + to repeat, not a no-op; an earlier version of this line claimed the latter. +* Neither direction invents state. Beat's ``PeriodicTask`` rows are only ever + *disabled* and *re-enabled*, never created or deleted, and the value written on + release is the one recorded before adoption (``pg_periodic_task.enabled``, and the + pipeline's own ``active``) — so a schedule an operator had switched off stays off + through a full adopt→release cycle. +* Adoption is not unilateral: it happens only because someone set the env var. + +**What it cannot do, and you must:** converging to Beat restores *ownership*, not +*capacity*. Beat publishes to RabbitMQ, so a released schedule only fires again if +``workerSchedulerV2`` (and ``workerMetrics`` for the periodics) are running. Flip +those back in the SAME change that sets ``PG_SCHEDULER_ENABLED=false`` — otherwise +you have simply moved the outage. The release path logs a warning saying so. + +Periodics (``dashboard_metrics.*``) are opt-in via ``--periodics``: pipelines and +metrics are separate rollout decisions, and metrics have their own consumer +(``workerPgMetrics``) that has to be deployed before they can be adopted. +""" + +from typing import Any + +from django.core.management import call_command +from django.core.management.base import BaseCommand +from scheduler.ownership import pg_scheduler_enabled + + +class Command(BaseCommand): + help = ( + "Converge schedule ownership to what PG_SCHEDULER_ENABLED declares: adopt to " + "the PG scheduler when on, release back to Celery Beat when off. Idempotent " + "in both directions and safe to run on every deploy." + ) + + def add_arguments(self, parser: Any) -> None: + parser.add_argument( + "--dry-run", + action="store_true", + help="Report what would change without writing.", + ) + parser.add_argument( + "--periodics", + action="store_true", + help=( + "Also converge the non-pipeline periodics (dashboard_metrics.*). " + "Off by default: adopting them requires workerPgMetrics to be " + "deployed, so it is a separate rollout decision from pipelines." + ), + ) + + def handle(self, *args: Any, **options: Any) -> None: + dry_run = options["dry_run"] + periodics = options["periodics"] + target_pg = pg_scheduler_enabled() + + self.stdout.write( + self.style.MIGRATE_HEADING( + f"Converging schedule ownership → " + f"{'PG scheduler' if target_pg else 'Celery Beat'} " + f"(PG_SCHEDULER_ENABLED={'true' if target_pg else 'false'})" + ) + ) + + if target_pg: + # Mirrors are backfilled first: a pipeline with no mirror row cannot be + # owned, and reconcile reports it as still-on-Beat rather than adopting it. + call_command("reconcile_pg_schedules", dry_run=dry_run) + if periodics: + # An empty name list means "every mirrored row" (the flag is + # `nargs="*"`, and the command distinguishes absent from empty). + # mirror_pg_periodic_tasks always backfills before flipping + # ownership, so this one call covers both halves. + call_command("mirror_pg_periodic_tasks", adopt=[], dry_run=dry_run) + else: + # --mirror-only keeps the backfill (inert at any flag state) while + # --release-stale hands back anything still marked pg_owned. Together + # they are the pipeline rollback. + call_command( + "reconcile_pg_schedules", + mirror_only=True, + release_stale=True, + dry_run=dry_run, + ) + if periodics: + call_command("mirror_pg_periodic_tasks", release=[], dry_run=dry_run) + self.stdout.write( + self.style.WARNING( + "Released to Beat. Beat publishes over RabbitMQ — these schedules " + "fire again ONLY if workerSchedulerV2 (and workerMetrics, with " + "--periodics) are running. Check that before relying on this." + ) + ) + + self.stdout.write(self.style.SUCCESS("Convergence complete.")) diff --git a/backend/pg_queue/management/commands/mirror_pg_periodic_tasks.py b/backend/pg_queue/management/commands/mirror_pg_periodic_tasks.py new file mode 100644 index 0000000000..7073d98c3c --- /dev/null +++ b/backend/pg_queue/management/commands/mirror_pg_periodic_tasks.py @@ -0,0 +1,427 @@ +"""Mirror non-pipeline Beat periodics into ``pg_periodic_task``, and hand them over. + +The Beat-replacement half for everything that is not a scheduled pipeline +(UN-3796): ``dashboard_metrics.*``, log-history, audit, and anything an operator +has added. Pipeline schedules keep their own mirror and their own percentage ramp +(``reconcile_pg_schedules``); this command deliberately skips them. + +Two steps, deliberately separate: + +* **mirror** (default) — upsert a ``PgPeriodicTask`` row for every non-pipeline + ``PeriodicTask``. Purely additive and inert: rows land ``pg_owned=False``, so the + PG scheduler still fires nothing and Beat keeps firing everything. +* **adopt / release** (explicit flags) — the actual hand-over. ``--adopt`` flips + ``pg_owned=True`` **and** disables the matching Beat ``PeriodicTask``, in one + transaction. ``--release`` reverses it. Doing both halves atomically is the whole + point: a row that is ``pg_owned`` while Beat still has it enabled fires **twice**, + which for ``cleanup_*`` means two concurrent deletes and for ``aggregate_*`` means + double-counted metrics. + +Idempotent and safe to re-run. Mirroring changes no behaviour on its own; only +``--adopt`` does, and only for the rows it names. + +Unlike the pipeline ramp there is no percentage: these are a handful of global +singletons, and the acceptance gate (Celery scaled to zero) needs all of them on PG, +so the meaningful states are "all Beat" and "all PG" with a per-name escape hatch. +""" + +import json +import logging +from typing import Any, NamedTuple + +from django.core.management.base import BaseCommand, CommandError +from django.db import transaction +from django.utils import timezone +from django_celery_beat.models import IntervalSchedule, PeriodicTask, PeriodicTasks + +from pg_queue.models import PgPeriodicTask + +logger = logging.getLogger(__name__) + +# Rows per DB round trip. Matches the batch size the repo's other bounded loops use +# (e.g. workflow_v2 migration 0012's "delete in batches to avoid long-running +# transactions"). Overridable with --batch-size. +DEFAULT_BATCH_SIZE = 1000 + +# Task paths this command must NOT mirror. +# +# Pipeline triggers own a separate mirror (pg_periodic_schedule) and a separate +# percentage ramp; mirroring one here too would give it two owners. BOTH known task +# paths are listed even though ``SchedulerHelper._schedule_task_job`` only writes the +# first today: a real Beat table was found carrying a legacy ``execute_pipeline_task_v2`` +# row, and it is only excluded here by luck of being disabled. Matching on the task +# path rather than a name convention keeps that from depending on luck. +# +# ``celery.backend_cleanup`` is Celery's own built-in: it prunes the Celery RESULT +# BACKEND. It is meaningless on PG (nothing there registers it, and the backend it +# cleans stops existing once Celery is off), so it stays with Beat and simply retires +# alongside it. +_EXCLUDED_TASK_PATHS = frozenset( + { + "scheduler.tasks.execute_pipeline_task", + "execute_pipeline_task_v2", + "celery.backend_cleanup", + } +) + + +def cron_from_periodic_task(task: PeriodicTask) -> str: + """Best-effort 5-field cron for a Beat periodic, or "" if not expressible. + + Beat schedules a task by ``crontab``, ``interval``, ``solar`` or ``clocked``. + Only the first two are in use here, and only they map onto cron: + + * ``crontab`` — a direct field-for-field reconstruction. + * ``interval`` — expressed as a step cron ONLY where one exists exactly. + + Returns ``""`` for anything else — notably **second**-resolution intervals, + which have no cron expression at all. Coarsening one to a minute would silently + change how often it runs, so the caller skips those and says so rather than + guessing. + + **``*/N`` is not "every N".** A step cron restarts at each field boundary, so it + is faithful only when N divides the field's range exactly. Otherwise the last + step of one period runs into the first of the next and the task fires early — + silently, and only at the boundary, which is the hardest kind of drift to spot: + + * ``*/7`` minutes → :00 :07 … :56, then **:00** — a 4-minute gap, not 7. + * ``*/45`` minutes → :00 :45, then **:00** — fires roughly twice as often. + * ``0 */5`` hours → 0 5 10 15 20, then **0** — a 4-hour gap, not 5. + + So minutes need ``60 % every == 0`` and hours ``24 % every == 0``, not merely a + range check. + + **Days are worse and are refused beyond 1.** ``*/N`` on day-of-month restarts + every month, and months are 28-31 days, so the gap at the boundary varies by + month and even by year: ``0 0 */7 * *`` fires on the 1st, 8th, 15th, 22nd, 29th + and then the 1st again — 2 to 4 days later depending on the month. There is no + correct cron for "every N days" at N > 1, so only ``every == 1`` maps (to a + plain daily), and the rest fall through to the caller's skip-and-explain path. + + An earlier version range-checked (``every < 60`` / ``< 24`` / ``< 32``) while + the docstring claimed exactness, so an "every 45 minutes" periodic mirrored to + something that fires twice as often. Nothing in integration hit it (15 divides + 60), but Beat schedules are per-environment DB rows that exist in no source + file, so staging or production can carry one. + """ + if task.crontab is not None: + c = task.crontab + return f"{c.minute} {c.hour} {c.day_of_month} {c.month_of_year} {c.day_of_week}" + interval = task.interval + if interval is None: + return "" + every = interval.every + if every < 1: + return "" + if interval.period == IntervalSchedule.MINUTES and every < 60 and 60 % every == 0: + return f"*/{every} * * * *" + if interval.period == IntervalSchedule.HOURS and every < 24 and 24 % every == 0: + return f"0 */{every} * * *" + if interval.period == IntervalSchedule.DAYS and every == 1: + return "0 0 * * *" + # SECONDS, a step that does not divide its field, or anything else. + return "" + + +def _decode(raw: str | None, fallback: Any) -> Any: + """Beat stores args/kwargs as JSON text; decode once here so a malformed value + fails at mirror time (visible, fixable) instead of at fire time (a periodic + that silently stops running). + """ + if not raw: + return fallback + return json.loads(raw) + + +class MirrorPlan(NamedTuple): + """What to do with one Beat periodic: mirror it, or skip it and say why. + + Every decision this command makes lives here rather than inside the loop that + talks to the database, so the rules can be tested against plain stand-ins + instead of a live multi-tenant schema. The command becomes glue: plan, then + apply. + """ + + name: str + fields: dict[str, Any] | None # None => skip + skip_reason: str | None = None + + @property + def should_mirror(self) -> bool: + return self.fields is not None + + +def plan_mirror(task: Any) -> MirrorPlan: + """Decide whether one Beat periodic can be mirrored, and with what fields. + + Accepts anything exposing the ``PeriodicTask`` attributes used here, so the + rules are testable without the ORM. + """ + if task.task in _EXCLUDED_TASK_PATHS: + return MirrorPlan( + task.name, + None, + f"{task.task} is owned elsewhere (pipeline mirror or Celery-internal)", + ) + cron = cron_from_periodic_task(task) + if not cron: + return MirrorPlan( + task.name, + None, + "schedule has no cron equivalent (second-resolution interval, solar " + "or clocked) — it must stay on Beat or move to a different mechanism", + ) + try: + task_args = _decode(task.args, []) + task_kwargs = _decode(task.kwargs, {}) + except ValueError as exc: + return MirrorPlan(task.name, None, f"malformed args/kwargs JSON ({exc})") + return MirrorPlan( + task.name, + { + "task_name": task.task, + # Beat falls back to the default queue when unset; mirror the same + # fallback so the row targets where Beat would have. + "queue": task.queue or "celery", + "task_args": task_args, + "task_kwargs": task_kwargs, + "cron_string": cron, + "enabled": task.enabled, + }, + ) + + +class Command(BaseCommand): + help = ( + "Mirror non-pipeline Beat periodics into pg_periodic_task. Additive and " + "inert by default; --adopt hands rows over to PG (and disables them in " + "Beat) atomically, --release reverses it." + ) + + def add_arguments(self, parser: Any) -> None: + parser.add_argument( + "--dry-run", + action="store_true", + help="Report what would change without writing.", + ) + parser.add_argument( + "--adopt", + nargs="*", + metavar="NAME", + help=( + "Hand rows over to PG: set pg_owned=True and DISABLE the matching " + "Beat PeriodicTask, atomically. Pass names to adopt those only, or " + "no names to adopt every mirrored row." + ), + ) + parser.add_argument( + "--release", + nargs="*", + metavar="NAME", + help="Reverse of --adopt: pg_owned=False and re-enable the Beat task.", + ) + parser.add_argument( + "--batch-size", + type=int, + default=DEFAULT_BATCH_SIZE, + metavar="N", + help=( + f"Rows fetched per DB round trip (default {DEFAULT_BATCH_SIZE}). " + "Bounds memory on a large Beat table." + ), + ) + parser.add_argument( + "--limit", + type=int, + default=0, + metavar="N", + help=( + "Stop after mirroring N rows (0 = no limit). A safety valve for a " + "first cautious run, NOT a resume cursor: the scan always starts " + "from the beginning, so re-running with --limit re-visits the same " + "rows. Drop it to mirror everything." + ), + ) + + def handle(self, *args: Any, **options: Any) -> None: + if options["adopt"] is not None and options["release"] is not None: + raise CommandError("--adopt and --release are mutually exclusive") + if options["batch_size"] < 1: + raise CommandError("--batch-size must be >= 1") + if options["limit"] < 0: + raise CommandError("--limit must be >= 0") + dry_run = options["dry_run"] + mirrored, skipped = self._mirror( + dry_run, batch_size=options["batch_size"], limit=options["limit"] + ) + + moved = 0 + if options["adopt"] is not None: + moved = self._set_ownership( + options["adopt"], + to_pg=True, + dry_run=dry_run, + batch_size=options["batch_size"], + ) + elif options["release"] is not None: + moved = self._set_ownership( + options["release"], + to_pg=False, + dry_run=dry_run, + batch_size=options["batch_size"], + ) + + prefix = "[dry-run] " if dry_run else "" + summary = ( + f"{prefix}mirrored={mirrored} skipped={skipped} ownership_changed={moved}" + ) + if skipped: + # Not fatal, but it means Beat is still the only firer for those — which + # blocks scaling Beat to zero. Surface it where the operator looks. + self.stderr.write(self.style.WARNING(summary)) + else: + self.stdout.write(self.style.SUCCESS(summary)) + + def _mirror(self, dry_run: bool, *, batch_size: int, limit: int) -> tuple[int, int]: + """Apply :func:`plan_mirror` to every Beat periodic. Never changes ownership. + + Deliberately thin: every rule lives in ``plan_mirror``; this only reports and + writes. A skip is loud because it means Beat remains the sole firer for that + row, which is what blocks scaling Beat to zero. + + Scale: the excluded task paths are filtered **in SQL**, not in Python. There is + one ``PeriodicTask`` row per scheduled pipeline, so matching them in Python + would drag the entire pipeline population through memory only to discard it. + ``.iterator()`` then bounds what is held at once. ``--limit`` caps a single + run for a cautious first pass; it is not a resume cursor, since the upsert is + idempotent and the scan always restarts from the beginning. + """ + mirrored = skipped = 0 + # order_by(pk) makes the scan deterministic, so a --limit run and its + # follow-ups walk the table in a stable order rather than re-treading rows. + queryset = ( + PeriodicTask.objects.exclude(task__in=_EXCLUDED_TASK_PATHS) + .select_related("crontab", "interval") + .order_by("pk") + ) + for task in queryset.iterator(chunk_size=batch_size): + if limit and mirrored >= limit: + self.stdout.write( + f"--limit {limit} reached; {queryset.count() - mirrored} row(s) " + "not visited. Re-run WITHOUT --limit to mirror everything " + "(the scan restarts from the beginning either way)." + ) + break + plan = plan_mirror(task) + if not plan.should_mirror: + # Excluded paths never reach here (filtered in SQL above), so every + # skip is something an operator must resolve before Beat can go away. + skipped += 1 + self.stderr.write( + self.style.WARNING(f"skipping {plan.name!r}: {plan.skip_reason}") + ) + continue + self.stdout.write( + f"mirror {plan.name!r} task={plan.fields['task_name']} " + f"queue={plan.fields['queue']!r} cron={plan.fields['cron_string']!r} " + f"enabled={plan.fields['enabled']}" + ) + if not dry_run: + fields = dict(plan.fields) + # `enabled` is the ONE mirrored field that stops tracking Beat once a + # row is adopted: after --adopt Beat's copy is False by definition, so + # re-mirroring would copy that back and leave pg_owned=True with + # enabled=False — matching NEITHER firer (the PG tick selects + # `WHERE pg_owned AND enabled`, Beat's row is disabled). The periodic + # would stop silently, and the pre-migration value needed to release it + # would be gone. Cron/args/queue still track Beat, so a schedule edit + # made while PG owns the row is still picked up. + if PgPeriodicTask.objects.filter(name=plan.name, pg_owned=True).exists(): + fields.pop("enabled", None) + PgPeriodicTask.objects.update_or_create(name=plan.name, defaults=fields) + mirrored += 1 + return mirrored, skipped + + def _set_ownership( + self, names: list[str], *, to_pg: bool, dry_run: bool, batch_size: int + ) -> int: + """Flip pg_owned and the matching Beat task's enabled flag together. + + One transaction per row: the two halves must not be separable, since the + window between them is exactly the double-fire (or no-fire) window. + """ + rows = PgPeriodicTask.objects.all().order_by("pk") + if names: + rows = rows.filter(name__in=names) + missing = set(names) - set(rows.values_list("name", flat=True)) + if missing: + raise CommandError( + f"no mirror row for: {', '.join(sorted(missing))} " + f"(run without --adopt/--release first to mirror them)" + ) + changed = 0 + # Deliberately NOT limited by --limit: an ownership flip is the operator's + # explicit, named intent, and stopping half way would leave some rows on PG + # and some on Beat with no record of where the boundary fell. It is chunked + # only to bound memory. In practice this table holds a handful of rows. + for row in rows.iterator(chunk_size=batch_size): + if row.pg_owned == to_pg: + continue + # RESTORE, don't blanket-enable. On release this used to write + # `enabled=True` unconditionally, which RESURRECTS a periodic an operator + # had deliberately switched off in Beat before the migration — a rollback + # is supposed to restore the previous state, not invent a new one. + # `row.enabled` is Beat's own value, captured by the mirror (plan_mirror + # copies task.enabled), so a row that was off mirrors off, is adopted off, + # and is released off. On adopt the answer is always False: PG owns it. + beat_enabled = False if to_pg else row.enabled + verb = "adopt" if to_pg else "release" + self.stdout.write(f"{verb} {row.name!r} (beat enabled -> {beat_enabled})") + if not dry_run: + with transaction.atomic(): + row.pg_owned = to_pg + # Clear the baseline on release so a later re-adopt records a + # fresh next_run_at instead of firing immediately for a + # next_run_at that went stale while Beat owned the schedule. + if not to_pg: + row.next_run_at = None + row.save(update_fields=["pg_owned", "next_run_at", "updated_at"]) + beat_updates: dict = {"enabled": beat_enabled} + # Baseline Beat's clock on release, the mirror of clearing + # next_run_at above. DatabaseScheduler derives due-ness from + # PeriodicTask.last_run_at against the crontab, so a periodic that + # spent days PG-owned is overdue by every interval it missed and + # Beat replays them the moment `enabled` flips back. Seen on + # integration 2026-08-24: all three dashboard_metrics.* fired + # inside 30 ms of the release, alongside four pipelines. + # + # Release-only: on adopt Beat is being switched OFF, so its clock + # is irrelevant, and touching it would corrupt the value a later + # release needs to restore. + if not to_pg: + beat_updates["last_run_at"] = timezone.now() + matched = PeriodicTask.objects.filter(name=row.name).update( + **beat_updates + ) + if not matched and not to_pg: + # Release only. A bulk .update() returning 0 is indistinguishable + # from success, and on release it means a periodic PG has let go + # of that Beat cannot pick up — no firer at all, on the rollback + # path. On ADOPT a missing Beat row just means there was nothing + # to disable, which is benign; warning on both is how the real + # one gets ignored. + self.stderr.write( + f"WARNING: no Beat PeriodicTask named {row.name!r} to " + f"release to — it may now have no firer" + ) + # Bulk .update() bypasses django-celery-beat's post_save signal, + # so PeriodicTasks.last_update never bumps and DatabaseScheduler + # never reloads. Without this, --adopt would set pg_owned=True and + # flip the DB row to disabled while Beat kept firing it from its + # stale in-memory copy — a DOUBLE FIRE, precisely what doing both + # halves in one transaction exists to prevent. --release has the + # mirror failure: Beat would never resume. Same fix and same + # reason as scheduler/ownership.py:132 on the pipeline path. + PeriodicTasks.update_changed() + changed += 1 + return changed diff --git a/backend/pg_queue/management/commands/reconcile_pg_schedules.py b/backend/pg_queue/management/commands/reconcile_pg_schedules.py index deaf059ee0..f5e9cd6f44 100644 --- a/backend/pg_queue/management/commands/reconcile_pg_schedules.py +++ b/backend/pg_queue/management/commands/reconcile_pg_schedules.py @@ -18,11 +18,19 @@ from django.core.management.base import BaseCommand, CommandError from django_celery_beat.models import CrontabSchedule, PeriodicTask -from scheduler.ownership import reconcile_ownership_for, resolve_schedule_owner +from scheduler.ownership import ( + pg_scheduler_enabled, + reconcile_ownership_for, + resolve_schedule_owner, +) from scheduler.tasks import mirror_periodic_schedule_upsert from pg_queue.models import PgPeriodicSchedule +# Rows per DB round trip; mirrors mirror_pg_periodic_tasks.DEFAULT_BATCH_SIZE and the +# batch size used by the repo's other bounded loops (workflow_v2 migration 0012). +DEFAULT_BATCH_SIZE = 1000 + # Only the pipeline-trigger PeriodicTasks are scheduled pipelines (other periodic # tasks — metrics, audit — are not mirrored). _PIPELINE_TASK_PATH = "scheduler.tasks.execute_pipeline_task" @@ -51,17 +59,75 @@ def add_arguments(self, parser: Any) -> None: action="store_true", help="Report what would change without writing.", ) + parser.add_argument( + "--batch-size", + type=int, + default=DEFAULT_BATCH_SIZE, + metavar="N", + help=( + f"Rows fetched per DB round trip (default {DEFAULT_BATCH_SIZE}). " + "There is one row per scheduled pipeline, so this bounds memory on " + "a large installation." + ), + ) + parser.add_argument( + "--mirror-only", + action="store_true", + help=( + "Backfill missing mirror rows and skip the ownership reconcile. " + "Purely additive: touches only pg_periodic_schedule, never a Beat " + "PeriodicTask. This is the mode automation runs ONLY on the release path (converge_pg_scheduler with PG_SCHEDULER_ENABLED=false); the adopt path runs the full reconcile — see handle()." + ), + ) + parser.add_argument( + "--release-stale", + action="store_true", + help=( + "Release schedules marked pg_owned while PG_SCHEDULER_ENABLED is " + "off, handing them back to Beat. Safe for unattended automation at " + "any flag state: with the gate ON it is a no-op, and with it off it " + "only ever moves a schedule TO Beat — the fail-safe direction. " + "Composes with --mirror-only." + ), + ) def handle(self, *args: Any, **options: Any) -> None: dry_run = options["dry_run"] - backfilled = self._backfill_mirrors(dry_run) - reconciled, pg_owned, failed = self._reconcile_all(dry_run) + batch_size = options["batch_size"] + mirror_only = options["mirror_only"] + if batch_size < 1: + raise CommandError("--batch-size must be >= 1") + backfilled = self._backfill_mirrors(dry_run, batch_size) + + # --mirror-only exists because it is safe to run at ANY flag state: backfilling + # is inert, while the reconcile below — fail-closed when the rollout is off — + # flips ownership *and disables the matching Beat PeriodicTask* when it is on. + # + # It is NO LONGER what the deploy-time automation runs. entrypoint.sh invokes + # converge_pg_scheduler on every backend start, without --mirror-only, so + # ownership hand-over IS an unattended action now; the env var is the operator's + # consent, given once, rather than a command typed per environment. That is + # deliberate — it is what makes rollback a values change — but it means this + # path's caller is no longer the only way ownership moves. A previous version of + # this comment ended "ownership stays an operator action", which stopped being + # true when the entrypoint switched commands. + if mirror_only: + reconciled, pg_owned, failed = 0, 0, 0 + else: + reconciled, pg_owned, failed = self._reconcile_all(dry_run, batch_size) + + released = 0 + if options["release_stale"]: + released, release_failed = self._release_stale(dry_run, batch_size) + failed += release_failed prefix = "[dry-run] " if dry_run else "" summary = ( f"{prefix}backfilled={backfilled} reconciled={reconciled} " - f"pg_owned={pg_owned} failed={failed}" + f"pg_owned={pg_owned} released={released} failed={failed}" ) + if mirror_only: + summary = f"{summary} (mirror-only: ownership reconcile skipped)" if failed: # Surface failures where the operator looks (and to automation). self.stderr.write(self.style.ERROR(summary)) @@ -93,15 +159,30 @@ def _mirror_fields_from_args(self, pt: Any, pipeline_id: str) -> dict | None: "pipeline_name": task_args[6] if len(task_args) > 6 else "", } - def _backfill_mirrors(self, dry_run: bool) -> int: - """Create a mirror row for every pipeline-trigger PeriodicTask lacking one.""" - # Pre-fetch the already-mirrored ids in one query (avoid an EXISTS per row). + def _backfill_mirrors(self, dry_run: bool, batch_size: int) -> int: + """Create a mirror row for every pipeline-trigger PeriodicTask lacking one. + + Both reads are bounded: the already-mirrored ids stream in via ``iterator`` + rather than materialising the whole table as a Python set, and the + PeriodicTask scan is chunked. There is one row per scheduled pipeline, so on + a large installation the unbounded version held the entire pipeline + population in memory twice. + """ + # Still one query, still an id set (the membership test below needs it), but + # streamed and values-only — flat ids, never model instances. mirrored = { str(pk) - for pk in PgPeriodicSchedule.objects.values_list("pipeline_id", flat=True) + for pk in PgPeriodicSchedule.objects.values_list( + "pipeline_id", flat=True + ).iterator(chunk_size=batch_size) } backfilled = 0 - for pt in PeriodicTask.objects.filter(task=_PIPELINE_TASK_PATH): + periodic_tasks = ( + PeriodicTask.objects.filter(task=_PIPELINE_TASK_PATH) + .select_related("crontab") + .order_by("pk") + ) + for pt in periodic_tasks.iterator(chunk_size=batch_size): pipeline_id = pt.name # = str(pipeline.pk) if pipeline_id in mirrored: continue @@ -121,12 +202,60 @@ def _backfill_mirrors(self, dry_run: bool) -> int: backfilled += 1 return backfilled - def _reconcile_all(self, dry_run: bool) -> tuple[int, int, int]: + def _release_stale(self, dry_run: bool, batch_size: int) -> tuple[int, int]: + """Hand every stale pg_owned row back to Beat. Returns (released, failed). + + Scoped to ``pg_owned=True`` rows so a clean installation does no work at all, + and gated on the env switch being OFF: with the ramp ON, a pg_owned row is + legitimate and releasing it would silently undo the rollout. + + Unlike :meth:`_reconcile_all` this IS safe to run unattended, because its only + possible effect is moving a schedule to Beat — the same direction the system + already fails to. That is what lets the deploy run it; see entrypoint.sh. + """ + if pg_scheduler_enabled(): + self.stdout.write( + "--release-stale: PG_SCHEDULER_ENABLED is on; pg_owned rows are " + "legitimate here, nothing released." + ) + return 0, 0 + + released = failed = 0 + for row in ( + PgPeriodicSchedule.objects.filter(pg_owned=True) + .order_by("pk") + .iterator(chunk_size=batch_size) + ): + if dry_run: + released += 1 + self.stdout.write( + f"[dry-run] would release pipeline {row.pipeline_id} " + f"({row.pipeline_name or 'unnamed'}) back to Beat" + ) + continue + # Routes through the same transaction the gate-off repair path uses, so + # Beat's PeriodicTask is re-enabled and next_run_at cleared in step. + result = reconcile_ownership_for( + str(row.pipeline_id), row.organization_id, active=row.enabled + ) + if result is None: # transaction failed (already logged) + failed += 1 + continue + released += 1 + return released, failed + + def _reconcile_all(self, dry_run: bool, batch_size: int) -> tuple[int, int, int]: """Reconcile ownership for every mirror row against the current rollout. Returns (reconciled, pg_owned, failed). + + Chunked: this loads full model instances, one per scheduled pipeline, and + each iteration does a Flipt evaluation plus a write — so it is the longest + loop in the command and the one worth bounding. """ reconciled = pg_owned = failed = 0 - for row in PgPeriodicSchedule.objects.all(): + for row in PgPeriodicSchedule.objects.order_by("pk").iterator( + chunk_size=batch_size + ): if dry_run: # Preview only — read the would-be owner (no DB write) so an # operator can see how many a ramp change would hand to PG. diff --git a/backend/pg_queue/migrations/0002_pgqueuemessage_available_at.py b/backend/pg_queue/migrations/0002_pgqueuemessage_available_at.py new file mode 100644 index 0000000000..d3eed988d2 --- /dev/null +++ b/backend/pg_queue/migrations/0002_pgqueuemessage_available_at.py @@ -0,0 +1,67 @@ +"""Delayed visibility for pg_queue_message (UN-3843). + +Adds ``available_at`` + the ``scheduled`` state so a dispatch can defer delivery +(Celery ``countdown``/``eta`` parity). Additive and inert: existing rows and every +existing enqueue resolve to ``available_at = now()`` / ``state = 'ready'``, which is +exactly today's behaviour. + +**Why the hand-written ``SET DEFAULT now()`` step.** Django's ``AddField`` adds the +column with a one-off literal default and then DROPS that default, leaving the column +``NOT NULL`` with no DB default. The workers' enqueue is raw SQL with an explicit +column list that does not mention ``available_at`` +(``queue_backend/pg_queue/client.py``) — so without a persistent DB default, every +worker enqueue would fail with a not-null violation the moment this migration landed, +and would keep failing regardless of deploy order. The DB default also encodes the +right semantic on its own: a row that does not ask to be deferred is available now. + +Ordering is therefore safe in both directions (migrate-before-workers or +workers-before-migrate), which is the posture the rest of the PG rollout assumes. +""" + +import django.utils.timezone +from django.db import migrations, models + + +class Migration(migrations.Migration): + dependencies = [ + ("pg_queue", "0001_initial_squashed"), + ] + + operations = [ + # Widen the closed state enum BEFORE anything can write 'scheduled'. + migrations.RemoveConstraint( + model_name="pgqueuemessage", + name="pg_queue_message_state_valid", + ), + migrations.AddField( + model_name="pgqueuemessage", + name="available_at", + field=models.DateTimeField(default=django.utils.timezone.now), + ), + # Restore a PERSISTENT default so the workers' raw INSERT (which omits this + # column) keeps working. Django dropped the AddField default above; state_ + # operations is empty because this changes only the DB, not the model Django + # tracks — the model keeps its Python-level default and stays in sync. + migrations.RunSQL( + sql="ALTER TABLE pg_queue_message ALTER COLUMN available_at SET DEFAULT now()", + reverse_sql=( + "ALTER TABLE pg_queue_message ALTER COLUMN available_at DROP DEFAULT" + ), + state_operations=[], + ), + migrations.AddIndex( + model_name="pgqueuemessage", + index=models.Index( + models.F("available_at"), + condition=models.Q(("state", "scheduled")), + name="pg_queue_message_scheduled_idx", + ), + ), + migrations.AddConstraint( + model_name="pgqueuemessage", + constraint=models.CheckConstraint( + check=models.Q(("state__in", ["ready", "claimed", "scheduled"])), + name="pg_queue_message_state_valid", + ), + ), + ] diff --git a/backend/pg_queue/migrations/0003_pgperiodictask.py b/backend/pg_queue/migrations/0003_pgperiodictask.py new file mode 100644 index 0000000000..cbeef73c9f --- /dev/null +++ b/backend/pg_queue/migrations/0003_pgperiodictask.py @@ -0,0 +1,40 @@ +# Generated by Django 4.2.30 on 2026-08-05 08:32 + +import django.utils.timezone +from django.db import migrations, models + + +class Migration(migrations.Migration): + dependencies = [ + ("pg_queue", "0002_pgqueuemessage_available_at"), + ] + + operations = [ + migrations.CreateModel( + name="PgPeriodicTask", + fields=[ + ("name", models.TextField(primary_key=True, serialize=False)), + ("task_name", models.TextField()), + ("queue", models.TextField()), + ("task_args", models.JSONField(default=list)), + ("task_kwargs", models.JSONField(default=dict)), + ("cron_string", models.TextField()), + ("org_id", models.TextField(blank=True, default="")), + ("enabled", models.BooleanField(default=True)), + ("pg_owned", models.BooleanField(default=False)), + ("last_run_at", models.DateTimeField(blank=True, null=True)), + ("next_run_at", models.DateTimeField(blank=True, null=True)), + ("created_at", models.DateTimeField(default=django.utils.timezone.now)), + ("updated_at", models.DateTimeField(auto_now=True)), + ], + options={ + "db_table": "pg_periodic_task", + "indexes": [ + models.Index( + fields=["pg_owned", "enabled", "next_run_at"], + name="pg_periodic_task_due_idx", + ) + ], + }, + ), + ] diff --git a/backend/pg_queue/models.py b/backend/pg_queue/models.py index abd5bab447..500c6d9084 100644 --- a/backend/pg_queue/models.py +++ b/backend/pg_queue/models.py @@ -8,6 +8,7 @@ # shared with the workers' raw SQL (client.py / reaper.py). See QueueMessageState. _READY = QueueMessageState.READY.value _CLAIMED = QueueMessageState.CLAIMED.value +_SCHEDULED = QueueMessageState.SCHEDULED.value class PgQueueMessage(models.Model): @@ -71,6 +72,24 @@ class PgQueueMessage(models.Model): # aggressive autovacuum) — kept out of the migration so the schema stays # portable across Postgres providers. state = models.TextField(default=_READY) + # Delayed visibility (UN-3843) — Celery ``countdown``/``eta`` parity. Set to a + # FUTURE timestamp only together with ``state='scheduled'``; the pair is what + # defers the row. Default now() keeps every pre-existing dispatch byte-identical + # (immediate rows are enqueued 'ready' and this column is never read for them). + # + # Deliberately NOT a claim predicate. The claim's partial index holds only + # 'ready' rows, so a scheduled row is physically absent from it and a pending + # delay costs the hot path nothing. Adding `available_at <= now()` to the claim + # would instead park every not-yet-due row inside the claim index to be walked + # and discarded on every claim — reintroducing the O(in-flight) scan-past cost + # that the state machine above was introduced to remove. The reaper promotes + # 'scheduled' -> 'ready' when due (reaper.promote_due_scheduled), so delivery is + # "not before available_at", never early, at reaper-tick granularity. + # + # Like `vt`, this column is absent from the claim index, so it never burdens the + # hot path; unlike `vt` it is written exactly once (at enqueue) and then only + # read by the promotion sweep. + available_at = models.DateTimeField(default=timezone.now) class Meta: db_table = "pg_queue_message" @@ -82,7 +101,7 @@ class Meta: # Backstop no writer can bypass: state is a closed enum. Values sourced # from QueueMessageState (single source of truth; drift-tested). models.CheckConstraint( - check=models.Q(state__in=[_READY, _CLAIMED]), + check=models.Q(state__in=[_READY, _CLAIMED, _SCHEDULED]), name="pg_queue_message_state_valid", ), ] @@ -114,6 +133,19 @@ class Meta: condition=models.Q(state=_CLAIMED), name="pg_queue_message_claimed_idx", ), + # PROMOTION path (UN-3843) — partial index over ONLY deferred rows. The + # reaper's promotion sweep (`WHERE state='scheduled' AND + # available_at<=now()`) walks this in due order, so a large backlog of + # far-future rows costs one index seek, not a scan. Bounded by pending + # delayed messages, which is a small set by construction (a delay is a + # stagger or a retry backoff, not a queue). Keyed on `available_at` alone: + # the sweep is time-ordered and queue-agnostic — it promotes every due row + # in one statement rather than per queue. + models.Index( + F("available_at"), + condition=models.Q(state=_SCHEDULED), + name="pg_queue_message_scheduled_idx", + ), ] @@ -350,6 +382,84 @@ class Meta: ] +class PgPeriodicTask(models.Model): + """Generic (non-pipeline) periodic, mirrored from ``django_celery_beat`` (UN-3796). + + The Beat-replacement sibling of :class:`PgPeriodicSchedule`. That table fires + **pipeline** triggers: it is keyed on ``pipeline_id`` and its dispatcher rebuilds + the fixed ``execute_pipeline_task`` argument list onto one fixed queue. Every + other Beat periodic — ``dashboard_metrics.*``, log-history, audit, anything an + operator adds — is an arbitrary ``(task, queue, args, kwargs)`` on a cron, which + that shape cannot express. + + A **sibling table rather than a widened one**: ``PgPeriodicSchedule`` is live, + flag-gated and mid-ramp, and widening it would mean a surrogate PK plus rewriting + every existing row to carry columns only the other kind of schedule uses. The two + genuinely differ in payload and target queue; they share only the cron mechanics, + which live in the dispatcher, not the row. + + Keyed on ``name`` — ``PeriodicTask.name``, the natural mirror key and already + unique in Beat, so reconciliation is a plain upsert with no id mapping to keep. + + Ownership and the no-double-fire guarantee are identical to the sibling's and + depend on the same discipline: ``pg_owned`` defaults False, so the table is inert + until a row is explicitly handed over, and handing one over MUST disable the + matching Beat ``PeriodicTask`` in the same breath. A row flipped ``pg_owned=True`` + while Beat still has it enabled fires twice. + + Managed=True / generated migration, extension-free. + """ + + # = PeriodicTask.name. Beat's own uniqueness key, so the mirror needs no + # surrogate id and reconciliation is an upsert on the same identity Beat uses. + name = models.TextField(primary_key=True) + # Dotted task path (PeriodicTask.task), e.g. "dashboard_metrics.aggregate_from_sources". + task_name = models.TextField() + # Target queue (PeriodicTask.queue). Unlike the pipeline scheduler's single + # hardcoded queue, each periodic carries its own — dashboard_metrics.* go to + # dashboard_metric_events, log-history to celery_periodic_logs, and so on. + queue = models.TextField() + # PeriodicTask.args / .kwargs, already decoded from Beat's JSON strings. Stored + # decoded so the dispatcher builds a TaskPayload without re-parsing per tick, and + # so a malformed value fails loudly at reconcile time rather than at fire time. + task_args = models.JSONField(default=list) + task_kwargs = models.JSONField(default=dict) + # Cron for crontab-backed periodics. Beat also supports IntervalSchedule + # (every N seconds/minutes); those are mirrored as an equivalent cron where one + # exists and skipped with a warning where it does not — sub-minute intervals have + # no cron expression, and silently coarsening one to a minute would change + # behaviour. See reconcile_pg_schedules. + cron_string = models.TextField() + # Owning org, "" = none. Global periodics (dashboard_metrics.*) have no org, so + # this is empty for every row today and the dispatcher writes it straight through + # to the queue message's org_id. It exists now rather than later for one reason: + # PgPeriodicSchedule carries organization_id, so without this column a future + # unification of the two tables would be a schema redesign instead of a data + # migration. Same no-NULL-text convention as elsewhere. + org_id = models.TextField(blank=True, default="") + # Mirrors PeriodicTask.enabled. + enabled = models.BooleanField(default=True) + # Per-row rollout switch; see the class docstring. Defaults False → inert. + pg_owned = models.BooleanField(default=False) + # Owned by the scheduler tick. NULL next_run_at = "record a baseline next time, + # don't fire this cycle" — no burst when a row is handed over. + last_run_at = models.DateTimeField(null=True, blank=True) + next_run_at = models.DateTimeField(null=True, blank=True) + created_at = models.DateTimeField(default=timezone.now) + updated_at = models.DateTimeField(auto_now=True) + + class Meta: + db_table = "pg_periodic_task" + indexes = [ + # Drives the due scan, mirroring pg_periodic_schedule_due_idx: + # WHERE pg_owned AND enabled AND (next_run_at IS NULL OR <= now()). + models.Index( + fields=["pg_owned", "enabled", "next_run_at"], + name="pg_periodic_task_due_idx", + ), + ] + + class PgTaskResult(models.Model): """Request-reply result store for the executor RPC on PG. diff --git a/backend/pg_queue/producer.py b/backend/pg_queue/producer.py index 028fddfb6d..3fc69bd199 100644 --- a/backend/pg_queue/producer.py +++ b/backend/pg_queue/producer.py @@ -18,8 +18,11 @@ import json import logging +from datetime import UTC, datetime, timedelta from typing import Any +from django.utils import timezone + from pg_queue.models import PgQueueMessage from unstract.core.data_models import ( FAIRNESS_DEFAULT_PRIORITY, @@ -27,9 +30,13 @@ FAIRNESS_MIN_PRIORITY, ContinuationSpec, FairnessPayload, + QueueMessageState, TaskPayload, ) +_READY = QueueMessageState.READY.value +_SCHEDULED = QueueMessageState.SCHEDULED.value + logger = logging.getLogger(__name__) # Fairness L3 priority default, re-exported under the producer's name so callers @@ -61,6 +68,33 @@ def _json_safe(value: Any) -> Any: return json.loads(json.dumps(value, default=str, allow_nan=False)) +def _resolve_visibility( + countdown: float | None, eta: datetime | None +) -> tuple[datetime, str]: + """Map Celery-style ``countdown``/``eta`` onto ``(available_at, state)``. + + Returns ``(now, 'ready')`` for the immediate case — which is every pre-existing + call site, so their rows stay byte-identical to before this parameter existed. + A delay in the past or a non-positive countdown also resolves to immediate: + "deliver no earlier than T" is already satisfied when T has passed, and a + computed-zero stagger step must not cost a reaper tick. + + The caller has already rejected countdown+eta together. + """ + now = timezone.now() + if countdown is not None: + target = now + timedelta(seconds=countdown) + elif eta is not None: + # A naive datetime would compare-fail against the aware `now` under + # USE_TZ; treat it as UTC rather than raising on an otherwise valid call. + target = eta if timezone.is_aware(eta) else timezone.make_aware(eta, UTC) + else: + return now, _READY + if target <= now: + return now, _READY + return target, _SCHEDULED + + def enqueue_task( *, task_name: str, @@ -74,6 +108,8 @@ def enqueue_task( on_success: ContinuationSpec | None = None, on_error: ContinuationSpec | None = None, task_id: str | None = None, + countdown: float | None = None, + eta: datetime | None = None, ) -> int: """Enqueue a task onto the PG queue; returns the new ``msg_id``. @@ -92,7 +128,25 @@ def enqueue_task( ``on_error`` as the failed id (Celery ``link_error`` parity). Mutually exclusive with ``reply_key`` — passing both is rejected (the consumer checks ``reply_key`` first and would silently drop the callback). + + ``countdown`` (seconds from now) / ``eta`` (absolute time) defer delivery, + mirroring Celery's kwargs of the same names; they are mutually exclusive. A + deferred row is written ``state='scheduled'`` and is invisible to the claim + until the reaper promotes it once ``available_at`` passes — so the guarantee is + **not before** the requested time, at reaper-tick granularity (default 5s), not + exactly at it. That matches Celery, whose countdown is also approximate, and is + the right semantic for a stagger or a retry backoff: never early. A + non-positive ``countdown`` or a past ``eta`` is treated as "send now" rather + than an error, so a computed-zero stagger step (``i * delay`` with ``i == 0``) + stays on the immediate path instead of paying a reaper tick. + + **Requires a running reaper.** Nothing else promotes a scheduled row. The reaper + is already the mandatory singleton for crash redelivery, so this adds no new + deployment dependency — but a queue that uses delays inherits the reaper's + liveness alert (PG_QUEUE_CLAIM_STATE_RUNBOOK). """ + if countdown is not None and eta is not None: + raise ValueError("countdown and eta are mutually exclusive") if reply_key is not None and (on_success is not None or on_error is not None): raise ValueError( "reply_key (request-reply) and on_success/on_error (callback) are " @@ -104,6 +158,7 @@ def enqueue_task( f"[{FAIRNESS_MIN_PRIORITY}, {FAIRNESS_MAX_PRIORITY}]: {priority!r}" ) pg_queue = queue or DEFAULT_GENERAL_QUEUE + available_at, state = _resolve_visibility(countdown, eta) # Mirror the worker _enqueue_pg path: log the failure with breadcrumbs before # it propagates, so a DB/constraint/serialization error isn't mislabeled by # the caller's broad handler. Message construction (the _json_safe coercions) @@ -140,6 +195,8 @@ def enqueue_task( message=message, org_id=org_id or "", priority=priority, + available_at=available_at, + state=state, ) except Exception: logger.exception( diff --git a/backend/pg_queue/tests/test_converge_pg_scheduler.py b/backend/pg_queue/tests/test_converge_pg_scheduler.py new file mode 100644 index 0000000000..f0edae1fe6 --- /dev/null +++ b/backend/pg_queue/tests/test_converge_pg_scheduler.py @@ -0,0 +1,91 @@ +"""The converge command routes by PG_SCHEDULER_ENABLED — in BOTH directions. + +Thin glue, but the branch is load-bearing: it is what makes a rollback a values +change rather than a procedure someone has to remember under pressure. Getting the +direction wrong would either adopt on a rollback (the opposite of intent) or leave +the reverse a no-op — the exact failure that stranded a pipeline on integration and +made Beat and the PG scheduler both fire it. + +The sub-commands are stubbed: their own behaviour is covered by +test_reconcile_pg_schedules_command / test_mirror_pg_periodic_tasks. What is pinned +here is *which* is called, with *which* flags. +""" + +from unittest.mock import patch + +import pytest +from django.core.management import call_command + +_CMD = "pg_queue.management.commands.converge_pg_scheduler" + + +def _run(monkeypatch, *, gate: bool, periodics: bool = False, dry_run: bool = False): + """Run the command with the gate set, returning the sub-command calls made.""" + monkeypatch.setenv("PG_SCHEDULER_ENABLED", "true" if gate else "false") + with patch(f"{_CMD}.call_command") as sub: + args = ["converge_pg_scheduler"] + if periodics: + args.append("--periodics") + if dry_run: + args.append("--dry-run") + call_command(*args) + # {command_name: kwargs} — no sub-command is invoked twice in either direction. + return {c.args[0]: c.kwargs for c in sub.call_args_list} + + +class TestDirection: + def test_gate_on_adopts_pipelines(self, monkeypatch): + calls = _run(monkeypatch, gate=True) + assert "reconcile_pg_schedules" in calls + # NOT mirror-only: that is the inert mode, which would adopt nothing. + assert not calls["reconcile_pg_schedules"].get("mirror_only") + assert not calls["reconcile_pg_schedules"].get("release_stale") + + def test_gate_off_releases_pipelines_back_to_beat(self, monkeypatch): + calls = _run(monkeypatch, gate=False) + kwargs = calls["reconcile_pg_schedules"] + # Both halves: mirror stays inert-safe, release_stale is the actual rollback. + assert kwargs["mirror_only"] is True + assert kwargs["release_stale"] is True + + def test_an_unset_gate_converges_to_beat(self, monkeypatch): + """Absent must mean Beat, not "do nothing" — the default has to be the + safe direction, since that is what an environment that never opted in has. + """ + monkeypatch.delenv("PG_SCHEDULER_ENABLED", raising=False) + with patch(f"{_CMD}.call_command") as sub: + call_command("converge_pg_scheduler") + kwargs = {c.args[0]: c.kwargs for c in sub.call_args_list} + assert kwargs["reconcile_pg_schedules"]["release_stale"] is True + + +class TestPeriodicsAreOptIn: + def test_periodics_are_untouched_by_default(self, monkeypatch): + """Deferring metrics must be real. + + Adopting them requires workerPgMetrics to be deployed; doing it implicitly + would fire dashboard_metrics.* into a queue with no consumer. + """ + assert "mirror_pg_periodic_tasks" not in _run(monkeypatch, gate=True) + assert "mirror_pg_periodic_tasks" not in _run(monkeypatch, gate=False) + + def test_gate_on_with_periodics_adopts_every_mirrored_row(self, monkeypatch): + calls = _run(monkeypatch, gate=True, periodics=True) + # Empty list = "all rows" (nargs="*"); None would mean the flag was absent. + assert calls["mirror_pg_periodic_tasks"]["adopt"] == [] + + def test_gate_off_with_periodics_releases_them(self, monkeypatch): + calls = _run(monkeypatch, gate=False, periodics=True) + assert calls["mirror_pg_periodic_tasks"]["release"] == [] + assert "adopt" not in calls["mirror_pg_periodic_tasks"] + + +class TestDryRunReaches_Everything: + """A dry run that silently writes through one sub-command is worse than none.""" + + @pytest.mark.parametrize("gate", [True, False]) + def test_dry_run_propagates_to_every_sub_command(self, monkeypatch, gate): + calls = _run(monkeypatch, gate=gate, periodics=True, dry_run=True) + assert calls, "no sub-commands were invoked" + for name, kwargs in calls.items(): + assert kwargs.get("dry_run") is True, name diff --git a/backend/pg_queue/tests/test_mirror_pg_periodic_tasks.py b/backend/pg_queue/tests/test_mirror_pg_periodic_tasks.py new file mode 100644 index 0000000000..d5fe342907 --- /dev/null +++ b/backend/pg_queue/tests/test_mirror_pg_periodic_tasks.py @@ -0,0 +1,371 @@ +"""Unit tests for the non-pipeline Beat mirror (UN-3796). + +DB-free by construction: every decision lives in ``plan_mirror`` / +``cron_from_periodic_task``, which take anything shaped like a ``PeriodicTask``. +That matters here — no backend DB-bound test currently runs in this repo (the +multi-tenant schema isn't created for the test database, and pytest is not wired +into CI for backend at all), so a test written against the ORM would look like +coverage while never executing. + +What's worth pinning is the exclusions (what must NOT be mirrored, and why) and the +cron conversion (where a wrong answer silently changes how often a job runs). +""" + +import json +from types import SimpleNamespace +from unittest.mock import MagicMock, patch + +import pytest +from django.core.management import call_command + +from pg_queue.management.commands.mirror_pg_periodic_tasks import ( + cron_from_periodic_task, + plan_mirror, +) + +# django_celery_beat's IntervalSchedule period values are plain strings. +SECONDS, MINUTES, HOURS, DAYS = "seconds", "minutes", "hours", "days" + + +def _crontab(minute="0", hour="2", dom="*", moy="*", dow="*"): + return SimpleNamespace( + minute=minute, hour=hour, day_of_month=dom, month_of_year=moy, day_of_week=dow + ) + + +# Sentinel so an explicit `crontab=None` (a solar/clocked periodic) is +# distinguishable from "caller didn't say", which defaults to a plain crontab. +_UNSET = object() + + +def _task( + name="t", + task="some.task", + queue="q", + args="[]", + kwargs="{}", + enabled=True, + crontab=_UNSET, + interval=None, +): + if crontab is _UNSET: + crontab = None if interval else _crontab() + return SimpleNamespace( + name=name, + task=task, + queue=queue, + args=args, + kwargs=kwargs, + enabled=enabled, + crontab=crontab, + interval=interval, + ) + + +def _interval(every, period): + return SimpleNamespace(every=every, period=period) + + +class TestCronConversion: + def test_crontab_is_reconstructed_field_for_field(self): + assert cron_from_periodic_task(_task(crontab=_crontab("30", "4"))) == "30 4 * * *" + + @pytest.mark.parametrize( + "every,period,expected", + [ + (15, MINUTES, "*/15 * * * *"), # dashboard_metrics.aggregate_from_sources + (1, MINUTES, "*/1 * * * *"), + (30, MINUTES, "*/30 * * * *"), + (2, HOURS, "0 */2 * * *"), + (12, HOURS, "0 */12 * * *"), + # every==1 day is the ONLY expressible day interval; emitted as a plain + # daily rather than `*/1` so the stored cron says what it means. + (1, DAYS, "0 0 * * *"), + ], + ) + def test_interval_maps_to_an_exact_step_cron(self, every, period, expected): + assert cron_from_periodic_task(_task(interval=_interval(every, period))) == expected + + def test_second_resolution_interval_is_refused_not_approximated(self): + # The important negative: cron has no sub-minute field. Rounding a + # 30-second periodic up to a minute would silently halve how often it runs, + # so it must be refused. This is the real `workflow_log_history_v2` shape. + assert cron_from_periodic_task(_task(interval=_interval(30, SECONDS))) == "" + + @pytest.mark.parametrize( + "every,period", [(90, MINUTES), (36, HOURS), (40, DAYS), (0, MINUTES)] + ) + def test_step_too_large_for_one_field_is_refused(self, every, period): + # `*/90` in a 0-59 minute field does not mean "every 90 minutes". + assert cron_from_periodic_task(_task(interval=_interval(every, period))) == "" + + @pytest.mark.parametrize( + "every,period,fires_instead", + [ + (7, MINUTES, ":00 :07 … :56 then :00 — a 4-minute gap, not 7"), + (45, MINUTES, ":00 :45 then :00 — roughly twice as often"), + (8, MINUTES, ":00 :08 … :56 then :00 — a 4-minute gap"), + (25, MINUTES, ":00 :25 :50 then :00 — a 10-minute gap"), + (5, HOURS, "0 5 10 15 20 then 0 — a 4-hour gap, not 5"), + (7, HOURS, "0 7 14 21 then 0 — a 3-hour gap"), + (9, HOURS, "0 9 18 then 0 — a 6-hour gap"), + ], + ) + def test_an_interval_that_does_not_divide_its_field_is_refused( + self, every, period, fires_instead + ): + """`*/N` restarts at each field boundary, so it is only "every N" when N + divides the range. The old code range-checked (`every < 60` / `< 24`) while + the docstring claimed exactness, so these mirrored to a cron that fires at + the wrong rate — silently, and only at the boundary. + + Refusing sends them down plan_mirror's skip-and-explain path, which is what + already happens for second-resolution intervals: they stay on Beat, visibly, + instead of being adopted at a frequency nobody chose. + """ + assert cron_from_periodic_task(_task(interval=_interval(every, period))) == "", ( + f"every {every} {period} must be refused — `*/{every}` fires {fires_instead}" + ) + + @pytest.mark.parametrize("every", [2, 3, 7, 15, 31]) + def test_multi_day_intervals_are_refused_because_months_vary(self, every): + """`0 0 */N * *` restarts every month and months are 28-31 days, so the + boundary gap depends on the month and the year. `0 0 */7 * *` fires on the + 1st, 8th, 15th, 22nd, 29th, then the 1st again — 2 to 4 days later. There is + no correct cron for "every N days" at N > 1, so only every==1 is expressible. + """ + assert cron_from_periodic_task(_task(interval=_interval(every, DAYS))) == "" + + def test_no_schedule_at_all_is_refused(self): + # solar/clocked periodics have neither crontab nor interval. + assert cron_from_periodic_task(_task(crontab=None, interval=None)) == "" + + +class TestPlanMirror: + def test_mirrors_with_its_own_queue_and_decoded_kwargs(self): + plan = plan_mirror( + _task( + name="dashboard_metrics_cleanup_hourly", + task="dashboard_metrics.cleanup_hourly_data", + queue="dashboard_metric_events", + kwargs=json.dumps({"retention_days": 30}), + ) + ) + assert plan.should_mirror + # Each periodic carries its OWN target queue — the whole reason this can't + # reuse the pipeline scheduler, which enqueues onto one fixed queue. + assert plan.fields["queue"] == "dashboard_metric_events" + assert plan.fields["task_kwargs"] == {"retention_days": 30} + assert plan.fields["cron_string"] == "0 2 * * *" + + @pytest.mark.parametrize( + "task_path", + [ + "scheduler.tasks.execute_pipeline_task", + # A real Beat table was found carrying this legacy path; excluding by + # name convention alone would have let it through into a second owner. + "execute_pipeline_task_v2", + # Celery's own result-backend housekeeping — nothing registers it on PG, + # and what it cleans stops existing once Celery is off. + "celery.backend_cleanup", + ], + ) + def test_excluded_task_paths_are_never_mirrored(self, task_path): + plan = plan_mirror(_task(task=task_path)) + assert not plan.should_mirror + assert "owned elsewhere" in plan.skip_reason + + def test_uncronnable_schedule_is_skipped_with_a_reason(self): + plan = plan_mirror(_task(interval=_interval(30, SECONDS))) + assert not plan.should_mirror + assert "no cron equivalent" in plan.skip_reason + + def test_malformed_args_is_skipped_rather_than_mirrored_empty(self): + # Beat stores args as JSON text, but a hand-edited row can hold a Python + # repr (a real one was found). Defaulting to [] would fire the task with the + # wrong arguments — silently, forever. + plan = plan_mirror(_task(args="['not', 'json']")) + assert not plan.should_mirror + assert "malformed" in plan.skip_reason + + def test_missing_queue_falls_back_the_way_beat_does(self): + assert plan_mirror(_task(queue=None)).fields["queue"] == "celery" + + def test_empty_args_and_kwargs_decode_to_empty_containers(self): + fields = plan_mirror(_task(args="", kwargs=None)).fields + assert fields["task_args"] == [] and fields["task_kwargs"] == {} + + def test_disabled_state_is_carried_across(self): + # A row disabled in Beat must mirror as disabled, or adopting it would start + # running something an operator had deliberately turned off. + assert plan_mirror(_task(enabled=False)).fields["enabled"] is False + + +class TestBeatReloadSignal: + """A hand-over must tell Beat to reload, or the atomicity is worthless. + + `PeriodicTask.objects.filter(...).update(...)` is a bulk update: it bypasses + django-celery-beat's post_save signal, so `PeriodicTasks.last_update` never bumps + and `DatabaseScheduler` keeps running from its stale in-memory copy. The DB would + say "disabled" while Beat carried on firing — a double fire alongside the PG + scheduler, which is the exact failure doing both halves in one transaction exists + to prevent. `--release` fails the mirror way: Beat never resumes. + + Pinned here because the symptom is invisible in the DB — you only see it in + duplicated side effects. + """ + + _CMD = "pg_queue.management.commands.mirror_pg_periodic_tasks" + + def _run(self, flag): + # `enabled` mirrors Beat's value and is what --release restores; a real + # PgPeriodicTask always has it, so the stand-in must too. + row = SimpleNamespace( + name="h", pg_owned=(flag == "--release"), next_run_at=None, enabled=True + ) + with ( + patch(f"{self._CMD}.PgPeriodicTask") as Model, + patch(f"{self._CMD}.PeriodicTask") as Beat, + patch(f"{self._CMD}.PeriodicTasks") as BeatSignal, + patch(f"{self._CMD}.transaction.atomic"), + ): + qs = MagicMock() + qs.iterator.return_value = [row] + qs.values_list.return_value = ["h"] + qs.filter.return_value = qs + Model.objects.all.return_value.order_by.return_value = qs + # No Beat rows to mirror; we only exercise the ownership flip. + Beat.objects.exclude.return_value.select_related.return_value.order_by.return_value.iterator.return_value = [] + row.save = MagicMock() + call_command("mirror_pg_periodic_tasks", flag) + return Beat, BeatSignal + + @pytest.mark.parametrize("flag", ["--adopt", "--release"]) + def test_ownership_flip_bumps_the_beat_reload_signal(self, flag): + Beat, BeatSignal = self._run(flag) + Beat.objects.filter.return_value.update.assert_called_once() + BeatSignal.update_changed.assert_called_once() + + +class TestReleaseRestoresRatherThanResurrects: + """A rollback restores the previous state; it must not invent a new one. + + `--release` wrote `enabled=True` unconditionally, so a periodic an operator had + deliberately switched OFF in Beat before the migration came back **on** after a + rollback — silently re-arming a job someone had stopped on purpose. The + pre-migration value is already recorded: `plan_mirror` copies `task.enabled` into + the mirror row, so `row.enabled` is Beat's own value. + """ + + _CMD = "pg_queue.management.commands.mirror_pg_periodic_tasks" + + def _flip(self, flag: str, mirrored_enabled: bool = True): + row = SimpleNamespace( + name="h", + pg_owned=(flag == "--release"), + next_run_at=None, + enabled=mirrored_enabled, + ) + with ( + patch(f"{self._CMD}.PgPeriodicTask") as Model, + patch(f"{self._CMD}.PeriodicTask") as Beat, + patch(f"{self._CMD}.PeriodicTasks"), + patch(f"{self._CMD}.transaction.atomic"), + ): + qs = MagicMock() + qs.iterator.return_value = [row] + qs.values_list.return_value = ["h"] + qs.filter.return_value = qs + Model.objects.all.return_value.order_by.return_value = qs + Beat.objects.exclude.return_value.select_related.return_value.order_by.return_value.iterator.return_value = [] + row.save = MagicMock() + call_command("mirror_pg_periodic_tasks", flag) + return Beat.objects.filter.return_value.update.call_args.kwargs + + def _release(self, mirrored_enabled: bool): + return self._flip("--release", mirrored_enabled) + + def test_a_row_that_was_enabled_is_restored_enabled(self): + assert self._release(mirrored_enabled=True)["enabled"] is True + + def test_a_row_that_was_DISABLED_stays_disabled(self): + # The regression: blanket `enabled=True` re-armed a job someone stopped. + assert self._release(mirrored_enabled=False)["enabled"] is False + + +class TestReleaseBaselinesBeatsClock: + """Release must reset Beat's clock, or Beat replays every missed interval. + + DatabaseScheduler keeps no next_run_at — it derives due-ness from + ``PeriodicTask.last_run_at`` against the crontab. A periodic that spent days + PG-owned still carries the last_run_at from before the hand-over, so the moment + ``enabled`` flips back it is overdue by every interval it missed and Beat fires + them all at once. + + Observed on integration 2026-08-24: releasing the fleet fired all three + ``dashboard_metrics.*`` plus four pipelines within 30 ms of "Released to Beat". + This is the exact mirror of the next_run_at baseline on the adopt side + (``scheduler/ownership.py``, OSS 2088d6962) — that half was fixed, this one was + not, and a comment there wrongly claimed Beat does not catch up. + + These assertions previously read ``== {"enabled": ...}`` on the whole kwargs dict. + Changed deliberately, not loosened: the write now carries a second key, and the + per-key assertions below plus the two new cases pin strictly more than the exact + dict did. + """ + + _CMD = "pg_queue.management.commands.mirror_pg_periodic_tasks" + + def _flip(self, flag: str): + return TestReleaseRestoresRatherThanResurrects()._flip(flag) + + def test_release_stamps_last_run_at(self): + kwargs = self._flip("--release") + assert "last_run_at" in kwargs, ( + "release must baseline Beat's clock; without it DatabaseScheduler sees " + "every missed interval as overdue and replays them" + ) + assert kwargs["last_run_at"] is not None + + def test_adopt_does_NOT_touch_last_run_at(self): + """On adopt Beat is being switched OFF, so its clock is irrelevant — and + overwriting it would destroy the value a later release needs to restore.""" + assert "last_run_at" not in self._flip("--adopt") + + +class TestMirrorDoesNotClobberAnAdoptedRow: + """Re-mirroring after adoption must not strand the row with no firer. + + `enabled` is the one mirrored field that stops tracking Beat once PG owns the row: + after `--adopt`, Beat's copy is False *by definition*. Copying that back leaves + `pg_owned=True, enabled=False`, which matches NEITHER firer — the PG tick selects + `WHERE pg_owned AND enabled`, and Beat's row is disabled. The periodic stops + silently, and the value needed to release it correctly is gone. + """ + + _CMD = "pg_queue.management.commands.mirror_pg_periodic_tasks" + + def _mirror_with(self, already_adopted: bool): + beat_row = _task(name="h", enabled=False, crontab=_crontab()) + with ( + patch(f"{self._CMD}.PgPeriodicTask") as Model, + patch(f"{self._CMD}.PeriodicTask") as Beat, + patch(f"{self._CMD}.PeriodicTasks"), + ): + Beat.objects.exclude.return_value.select_related.return_value.order_by.return_value.iterator.return_value = [ + beat_row + ] + Model.objects.filter.return_value.exists.return_value = already_adopted + call_command("mirror_pg_periodic_tasks") + return Model.objects.update_or_create.call_args.kwargs["defaults"] + + def test_an_adopted_row_keeps_its_own_enabled(self): + assert "enabled" not in self._mirror_with(already_adopted=True) + + def test_a_beat_owned_row_still_tracks_beat(self): + # The normal path must keep mirroring enabled, or a pause in Beat is missed. + assert self._mirror_with(already_adopted=False)["enabled"] is False + + def test_cron_still_tracks_beat_even_when_adopted(self): + # Only `enabled` diverges — a schedule edit made while PG owns it must land. + assert "cron_string" in self._mirror_with(already_adopted=True) diff --git a/backend/pg_queue/tests/test_producer.py b/backend/pg_queue/tests/test_producer.py index c94d219f01..246855a959 100644 --- a/backend/pg_queue/tests/test_producer.py +++ b/backend/pg_queue/tests/test_producer.py @@ -10,8 +10,10 @@ from unittest.mock import MagicMock, patch import pytest +from django.utils import timezone as django_timezone from pg_queue import producer +from unstract.core.data_models import QueueMessageState _MODEL = "pg_queue.producer.PgQueueMessage" @@ -176,3 +178,90 @@ def test_continuation_specs_are_json_coerced(self): ) msg = model.objects.create.call_args.kwargs["message"] assert msg["on_success"]["kwargs"]["callback_kwargs"]["doc_id"] == str(uid) + + +class TestDelayedVisibility: + """UN-3843 — ``countdown``/``eta`` defer delivery via ``scheduled`` + ``available_at``. + + A deferred row is written ``state='scheduled'`` so it is absent from the claim's + partial index; the reaper promotes it when due. These pin the producer half of + that contract — that the right ``(available_at, state)`` pair reaches the row, and + that every non-deferred call still writes exactly what it wrote before this + parameter existed. + """ + + @staticmethod + def _create_kwargs(model): + return model.objects.create.call_args.kwargs + + def test_no_delay_is_unchanged_and_ready(self): + # The zero-regression case: every pre-existing call site lands here. + before = django_timezone.now() + with patch(_MODEL) as model: + model.objects.create.return_value = MagicMock(msg_id=1) + producer.enqueue_task(task_name="t", queue="q") + kw = self._create_kwargs(model) + assert kw["state"] == QueueMessageState.READY.value + assert before <= kw["available_at"] <= django_timezone.now() + + def test_countdown_defers_and_marks_scheduled(self): + with patch(_MODEL) as model: + model.objects.create.return_value = MagicMock(msg_id=1) + before = django_timezone.now() + producer.enqueue_task(task_name="t", queue="q", countdown=90) + kw = self._create_kwargs(model) + assert kw["state"] == QueueMessageState.SCHEDULED.value + # ~90s out; generous window so a slow CI box can't flake it. + delta = (kw["available_at"] - before).total_seconds() + assert 89 <= delta <= 95 + + def test_eta_defers_and_marks_scheduled(self): + eta = django_timezone.now() + datetime.timedelta(minutes=5) + with patch(_MODEL) as model: + model.objects.create.return_value = MagicMock(msg_id=1) + producer.enqueue_task(task_name="t", queue="q", eta=eta) + kw = self._create_kwargs(model) + assert kw["state"] == QueueMessageState.SCHEDULED.value + assert kw["available_at"] == eta + + def test_naive_eta_is_read_as_utc(self): + # USE_TZ makes `now()` aware; a naive eta would raise on comparison. Treat + # it as UTC rather than rejecting an otherwise valid call. + naive = ( + django_timezone.now() + datetime.timedelta(hours=1) + ).replace(tzinfo=None) + with patch(_MODEL) as model: + model.objects.create.return_value = MagicMock(msg_id=1) + producer.enqueue_task(task_name="t", queue="q", eta=naive) + kw = self._create_kwargs(model) + assert kw["state"] == QueueMessageState.SCHEDULED.value + assert kw["available_at"].tzinfo is not None + + @pytest.mark.parametrize("countdown", [0, -5]) + def test_non_positive_countdown_stays_on_the_immediate_path(self, countdown): + # A stagger computes `i * delay`; step 0 must not pay a reaper tick just to + # become claimable. + with patch(_MODEL) as model: + model.objects.create.return_value = MagicMock(msg_id=1) + producer.enqueue_task(task_name="t", queue="q", countdown=countdown) + assert self._create_kwargs(model)["state"] == QueueMessageState.READY.value + + def test_past_eta_stays_on_the_immediate_path(self): + past = django_timezone.now() - datetime.timedelta(minutes=1) + with patch(_MODEL) as model: + model.objects.create.return_value = MagicMock(msg_id=1) + producer.enqueue_task(task_name="t", queue="q", eta=past) + assert self._create_kwargs(model)["state"] == QueueMessageState.READY.value + + def test_countdown_and_eta_together_are_rejected(self): + with patch(_MODEL) as model: + model.objects.create.return_value = MagicMock(msg_id=1) + with pytest.raises(ValueError, match="mutually exclusive"): + producer.enqueue_task( + task_name="t", + queue="q", + countdown=10, + eta=django_timezone.now(), + ) + # Rejected before any row is written — no half-enqueued message. + model.objects.create.assert_not_called() diff --git a/backend/pg_queue/tests/test_reconcile_pg_schedules_command.py b/backend/pg_queue/tests/test_reconcile_pg_schedules_command.py index 081aff2cf6..04d8d07753 100644 --- a/backend/pg_queue/tests/test_reconcile_pg_schedules_command.py +++ b/backend/pg_queue/tests/test_reconcile_pg_schedules_command.py @@ -11,6 +11,10 @@ from django.core.management import call_command from django.core.management.base import CommandError +from pg_queue.management.commands.reconcile_pg_schedules import ( + DEFAULT_BATCH_SIZE, +) + _CMD = "pg_queue.management.commands.reconcile_pg_schedules" @@ -35,6 +39,32 @@ def _row(pid, org="org", enabled=True): return m +def _periodic_tasks(rows): + """Mock the PeriodicTask queryset chain the command now uses. + + The command reads `.filter(...).select_related(...).order_by(...).iterator(...)` + so the scan is chunked; these tests pin behaviour, not the chain, so the helper + absorbs the plumbing. + """ + qs = MagicMock() + qs.select_related.return_value.order_by.return_value.iterator.return_value = rows + return qs + + +def _sched_ids(ids): + """Mock `.values_list("pipeline_id", flat=True).iterator(...)`.""" + vl = MagicMock() + vl.iterator.return_value = ids + return vl + + +def _sched_rows(rows): + """Mock `.order_by("pk").iterator(...)`.""" + qs = MagicMock() + qs.iterator.return_value = rows + return qs + + class TestReconcileCommand: def test_backfills_only_unmirrored_and_reconciles(self): pt_new = _pt("pid-new", args='["wf", "org", "", "", "pid-new", false, "n"]') @@ -45,16 +75,35 @@ def test_backfills_only_unmirrored_and_reconciles(self): patch(f"{_CMD}.mirror_periodic_schedule_upsert") as upsert, patch(f"{_CMD}.reconcile_ownership_for", return_value=False) as reconcile, ): - PT.objects.filter.return_value = [pt_new, pt_exists] + PT.objects.filter.return_value = _periodic_tasks([pt_new, pt_exists]) # pid-exists already mirrored; pid-new not (one prefetch query). - Sched.objects.values_list.return_value = ["pid-exists"] - Sched.objects.all.return_value = [_row("pid-new"), _row("pid-exists")] + Sched.objects.values_list.return_value = _sched_ids(["pid-exists"]) + Sched.objects.order_by.return_value = _sched_rows([_row("pid-new"), _row("pid-exists")]) call_command("reconcile_pg_schedules") upsert.assert_called_once() # only the unmirrored one backfilled assert upsert.call_args.kwargs["pipeline_id"] == "pid-new" assert reconcile.call_count == 2 # both rows reconciled + def test_mirror_only_backfills_but_never_reconciles(self): + """--mirror-only is what deploy automation runs, so it must be safe at ANY + flag state. The reconcile it skips is the step that disables Beat rows once + the rollout is on — an unattended job must never make that change.""" + pt_new = _pt("pid-new", args='["wf", "org", "", "", "pid-new", false, "n"]') + with ( + patch(f"{_CMD}.PeriodicTask") as PT, + patch(f"{_CMD}.PgPeriodicSchedule") as Sched, + patch(f"{_CMD}.mirror_periodic_schedule_upsert") as upsert, + patch(f"{_CMD}.reconcile_ownership_for") as reconcile, + ): + PT.objects.filter.return_value = _periodic_tasks([pt_new]) + Sched.objects.values_list.return_value = _sched_ids([]) + Sched.objects.order_by.return_value = _sched_rows([_row("pid-new")]) + call_command("reconcile_pg_schedules", "--mirror-only") + + upsert.assert_called_once() # the backfill still happens + reconcile.assert_not_called() # the ownership flip does not + def test_malformed_args_skipped_not_fatal(self): bad = _pt("pid-bad", args="{ this is not json") good = _pt("pid-good", args='["wf", "org", "", "", "pid-good", false, "n"]') @@ -64,9 +113,9 @@ def test_malformed_args_skipped_not_fatal(self): patch(f"{_CMD}.mirror_periodic_schedule_upsert") as upsert, patch(f"{_CMD}.reconcile_ownership_for", return_value=False), ): - PT.objects.filter.return_value = [bad, good] - Sched.objects.values_list.return_value = [] - Sched.objects.all.return_value = [] + PT.objects.filter.return_value = _periodic_tasks([bad, good]) + Sched.objects.values_list.return_value = _sched_ids([]) + Sched.objects.order_by.return_value = _sched_rows([]) # Must not raise despite the bad row. call_command("reconcile_pg_schedules") @@ -82,9 +131,9 @@ def test_non_list_args_skipped(self): patch(f"{_CMD}.mirror_periodic_schedule_upsert") as upsert, patch(f"{_CMD}.reconcile_ownership_for", return_value=False), ): - PT.objects.filter.return_value = [weird] - Sched.objects.values_list.return_value = [] - Sched.objects.all.return_value = [] + PT.objects.filter.return_value = _periodic_tasks([weird]) + Sched.objects.values_list.return_value = _sched_ids([]) + Sched.objects.order_by.return_value = _sched_rows([]) call_command("reconcile_pg_schedules") upsert.assert_not_called() @@ -97,11 +146,11 @@ def test_dry_run_writes_nothing_but_previews_owner(self): patch(f"{_CMD}.reconcile_ownership_for") as reconcile, patch(f"{_CMD}.resolve_schedule_owner", return_value=True) as resolve, ): - PT.objects.filter.return_value = [ - _pt("pid-1", args='["wf", "org", "", "", "pid-1", false, "n"]') - ] - Sched.objects.values_list.return_value = [] - Sched.objects.all.return_value = [_row("pid-1")] + PT.objects.filter.return_value = _periodic_tasks( + [_pt("pid-1", args='["wf", "org", "", "", "pid-1", false, "n"]')] + ) + Sched.objects.values_list.return_value = _sched_ids([]) + Sched.objects.order_by.return_value = _sched_rows([_row("pid-1")]) call_command("reconcile_pg_schedules", "--dry-run") upsert.assert_not_called() # no backfill write @@ -115,8 +164,65 @@ def test_failure_raises_command_error(self): patch(f"{_CMD}.mirror_periodic_schedule_upsert"), patch(f"{_CMD}.reconcile_ownership_for", return_value=None), # failed ): - PT.objects.filter.return_value = [] - Sched.objects.values_list.return_value = [] - Sched.objects.all.return_value = [_row("pid-1")] + PT.objects.filter.return_value = _periodic_tasks([]) + Sched.objects.values_list.return_value = _sched_ids([]) + Sched.objects.order_by.return_value = _sched_rows([_row("pid-1")]) with pytest.raises(CommandError): call_command("reconcile_pg_schedules") + + +class TestChunking: + """Both loops must stay bounded — there is one row per scheduled pipeline, so an + unbounded scan holds the whole pipeline population in memory (twice, before this). + Pinned so a later refactor can't quietly drop `.iterator()`. + """ + + def _run(self, extra_args=()): + with ( + patch(f"{_CMD}.PeriodicTask") as PT, + patch(f"{_CMD}.PgPeriodicSchedule") as Sched, + patch(f"{_CMD}.mirror_periodic_schedule_upsert"), + patch(f"{_CMD}.reconcile_ownership_for", return_value=False), + ): + PT.objects.filter.return_value = _periodic_tasks([]) + Sched.objects.values_list.return_value = _sched_ids([]) + Sched.objects.order_by.return_value = _sched_rows([]) + call_command("reconcile_pg_schedules", *extra_args) + return PT, Sched + + def test_both_scans_are_chunked_with_the_default_batch_size(self): + PT, Sched = self._run() + pt_chain = PT.objects.filter.return_value.select_related.return_value + pt_chain.order_by.return_value.iterator.assert_called_once_with( + chunk_size=DEFAULT_BATCH_SIZE + ) + Sched.objects.order_by.return_value.iterator.assert_called_once_with( + chunk_size=DEFAULT_BATCH_SIZE + ) + + def test_mirrored_id_prefetch_streams_rather_than_materialising(self): + # This one used to build a set from the full table with no bound. + _, Sched = self._run() + Sched.objects.values_list.assert_called_once_with("pipeline_id", flat=True) + Sched.objects.values_list.return_value.iterator.assert_called_once_with( + chunk_size=DEFAULT_BATCH_SIZE + ) + + def test_batch_size_flag_is_honoured(self): + PT, Sched = self._run(("--batch-size", "7")) + pt_chain = PT.objects.filter.return_value.select_related.return_value + pt_chain.order_by.return_value.iterator.assert_called_once_with(chunk_size=7) + Sched.objects.order_by.return_value.iterator.assert_called_once_with(chunk_size=7) + + def test_scans_are_ordered_so_a_chunked_walk_is_deterministic(self): + # Without an ORDER BY, a chunked scan has no guaranteed row order between + # batches, so rows can be visited twice or not at all. + PT, Sched = self._run() + pt_chain = PT.objects.filter.return_value.select_related.return_value + assert pt_chain.order_by.call_args[0] == ("pk",) + assert Sched.objects.order_by.call_args[0] == ("pk",) + + @pytest.mark.parametrize("bad", ["0", "-1"]) + def test_non_positive_batch_size_is_rejected(self, bad): + with pytest.raises(CommandError, match="--batch-size"): + call_command("reconcile_pg_schedules", "--batch-size", bad) diff --git a/backend/scheduler/ownership.py b/backend/scheduler/ownership.py index 9c1e1f2878..c7c9c28173 100644 --- a/backend/scheduler/ownership.py +++ b/backend/scheduler/ownership.py @@ -11,10 +11,17 @@ *conditional* on this slice): a ``pg_owned`` row always has its Beat ``PeriodicTask`` disabled, so the two can't both fire. -Inert by default: ``resolve_schedule_owner`` fails closed to Beat -(``pg_owned=False``) until ops ramps the single ``pg_queue_enabled`` Flipt flag — so -reconciling on every schedule edit is a no-op (everything stays Beat-owned) until the -rollout starts. +**Two gates, and both must be on.** ``PG_SCHEDULER_ENABLED`` (env, default off) comes +first; the ``pg_queue_enabled`` Flipt flag then decides per schedule, failing closed to +Beat on a blind Flipt or any error. + +They are separate because Beat stays the **sole scheduler for the whole PG rollout**: +pipelines already reach PG without this module (``execute_pipeline_task_v2`` → +``complete_execution`` → ``execute_workflow_async`` → ``resolve_transport``), so the PG +scheduler exists only to retire the Beat *deployment* — a later, deliberate step. While +the env gate is off this module writes **nothing at all**, so Beat's tables are untouched +for the entire rollout and turning the flag off again is a pure Flipt flip with nothing +to restore. """ from __future__ import annotations @@ -37,6 +44,34 @@ # buckets the %-rollout on pipeline_id (each subsystem keys on its own entity). +# Second gate, ahead of the Flipt flag: hand-over to the PG scheduler happens only +# when this is explicitly on. Referenced as "pg_scheduler_enabled" by this module's +# callers' comments (scheduler/helper.py, reconcile_pg_schedules) since the ramp was +# designed, but never actually implemented — resolve_schedule_owner gated on +# pg_queue_enabled alone, so the two flips were welded together. +# +# They must be separable, because Beat stays the SOLE scheduler for the whole PG +# rollout: pipelines already reach PG without it (execute_pipeline_task_v2 → +# complete_execution → execute_workflow_async → resolve_transport), so the PG +# scheduler exists only to retire the Beat *deployment* — a later, deliberate step. +# +# Without this gate, turning pg_queue_enabled on would hand schedules to a PG +# scheduler that is not running: reconcile_ownership_for disables the Beat +# PeriodicTask, nothing polls the PG side, and the pipeline has no firer at all. It +# runs in the backend on every schedule save, so scaling Beat to zero does not avoid +# it and no ramp command is needed to trigger it — a user saving a schedule suffices. +_PG_SCHEDULER_ENABLED_ENV = "PG_SCHEDULER_ENABLED" + + +def pg_scheduler_enabled() -> bool: + """Whether schedule hand-over to the PG scheduler is switched on at all. + + Defaults **off**. Flip it only alongside deploying the PG scheduler worker, as + part of retiring Celery Beat. + """ + return os.environ.get(_PG_SCHEDULER_ENABLED_ENV, "false").strip().lower() == "true" + + def resolve_schedule_owner(pipeline_id: str, organization_id: str | None) -> bool: """True → the PG scheduler owns this schedule; False → Celery Beat does. @@ -44,7 +79,13 @@ def resolve_schedule_owner(pipeline_id: str, organization_id: str | None) -> boo flag, keyed on ``pipeline_id`` for a stable percentage bucket. **Fails closed to Beat** on a blind Flipt or any error — so a schedule never silently loses its firer. + + Requires ``PG_SCHEDULER_ENABLED`` **as well as** the flag: during the PG rollout + the flag is on while the PG scheduler is still dark, and a schedule handed to a + scheduler that is not running would simply stop firing. """ + if not pg_scheduler_enabled(): + return False if os.environ.get("FLIPT_SERVICE_AVAILABLE", "false").lower() != "true": logger.warning( "resolve_schedule_owner: FLIPT_SERVICE_AVAILABLE != true " @@ -68,8 +109,7 @@ def resolve_schedule_owner(pipeline_id: str, organization_id: str | None) -> boo # edit — warn with traceback rather than logger.exception so a persistently # down Flipt doesn't bury real errors as a per-edit Sentry exception. logger.warning( - "resolve_schedule_owner: Flipt check failed for pipeline %s; " - "leaving on Beat", + "resolve_schedule_owner: Flipt check failed for pipeline %s; leaving on Beat", pipeline_id, exc_info=True, ) @@ -77,6 +117,27 @@ def resolve_schedule_owner(pipeline_id: str, organization_id: str | None) -> boo return bool(owned) +def _has_stale_pg_ownership(pipeline_id: str) -> bool: + """True if this row claims PG ownership while the gate is off — i.e. it needs + releasing back to Beat. + + **Fails closed to today's behaviour**: on any DB error this returns False, so the + caller takes the historical `return False` path and writes nothing. A repair that + cannot confirm it is needed must not run. + """ + try: + return PgPeriodicSchedule.objects.filter( + pipeline_id=pipeline_id, pg_owned=True + ).exists() + except Exception: + logger.exception( + "reconcile_ownership_for: could not check stale PG ownership for " + "pipeline %s; leaving the schedule on its current firer", + pipeline_id, + ) + return False + + def reconcile_ownership_for( pipeline_id: str, organization_id: str | None, *, active: bool ) -> bool | None: @@ -94,16 +155,73 @@ def reconcile_ownership_for( Best-effort: a DB failure is logged and swallowed so it can never break the caller. Returns the resolved ``pg_owned`` on success, or **None** if the transaction failed (so the ramp command can tally + surface failures). + + **No-op while ``PG_SCHEDULER_ENABLED`` is off** — writing NEITHER table, *unless* + the row contradicts the gate (see below). Returning early rather than relying on + ``resolve_schedule_owner`` returning False matters: that path would still issue + ``PeriodicTask.update(enabled=active)`` and bump ``PeriodicTasks.update_changed()`` + on every schedule save, forcing a Beat reload to write back a value Beat already + had. Beat's tables stay untouched for the whole rollout, so flag-off is a pure + Flipt flip with nothing to restore. """ - pg_owned = resolve_schedule_owner(pipeline_id, organization_id) + if not pg_scheduler_enabled(): + # Gate off ⇒ Beat owns everything. ASSERT that instead of assuming it. + # + # The original blanket `return False` assumed pg_owned could never be True + # while the gate is off. An image predating this gate breaks that assumption: + # there, resolve_schedule_owner keys on the Flipt flag ALONE, so flipping + # pg_queue_enabled hands schedules to PG. Roll forward onto this gate and the + # row is stranded — writing nothing can never correct it, and NO code path + # (not even `reconcile_pg_schedules` without --mirror-only, which routes back + # through here) can clear it. Only hand-written SQL could, which is not a + # deploy procedure. Observed on integration 2026-08-12: Beat and the PG + # scheduler both fired one pipeline 2.4s apart; a single execution resulted + # only because the Celery consumers happened to be scaled to zero. + # + # This is reachable in production WITHOUT a mis-built image: deploy this + # branch, roll back to a pre-gate build while the flag is on, roll forward. + # Rollback is precisely when a stranded scheduler is least affordable. + # + # Cost in the normal case is one indexed existence check and still zero + # writes, so the "Beat's tables stay untouched" guarantee holds for every + # environment that never had a contradictory row. + if not _has_stale_pg_ownership(pipeline_id): + return False + logger.warning( + "reconcile_ownership_for: pipeline %s is pg_owned while " + "PG_SCHEDULER_ENABLED is off (stale hand-over from a pre-gate build); " + "releasing it back to Beat", + pipeline_id, + ) + pg_owned = False + else: + pg_owned = resolve_schedule_owner(pipeline_id, organization_id) try: with transaction.atomic(): + was_pg_owned = ( + PgPeriodicSchedule.objects.filter(pipeline_id=pipeline_id) + .values_list("pg_owned", flat=True) + .first() + ) # queryset .update() doesn't fire auto_now, so bump updated_at # explicitly (mirrors _mirror_periodic_schedule_set_enabled). updates: dict = {"pg_owned": pg_owned, "updated_at": timezone.now()} - if not pg_owned: - # Back on Beat → clear the PG next-run so a future re-hand-over - # baselines instead of firing immediately on a stale timestamp. + # Baseline on an ownership CHANGE, in either direction. NULL means + # "record a baseline next tick, don't fire this cycle" + # (pg_queue/models.py:366), and that is the only thing standing between a + # hand-over and an immediate unscheduled run: the tick selects + # `WHERE pg_owned AND enabled AND (next_run_at IS NULL OR <= now())`, so a + # next_run_at left over from an earlier PG-ownership period is already in + # the past and fires at once. Observed on integration 2026-08-14 — + # gallh_load_test carried next_run_at=2026-08-12 06:08 from a stale row and + # fired ~2s after being re-enabled, two days late, alongside the operator's + # manual run. + # + # Scoped to the TRANSITION, not every call: this runs on every pipeline + # save, and clearing unconditionally would re-baseline mid-cycle — a save at + # 12:07:59 against a 12:08 next_run_at would skip that fire entirely. + handing_over_to_pg = pg_owned and not was_pg_owned + if not pg_owned or handing_over_to_pg: updates["next_run_at"] = None # The mirror row exists from the dual-write / backfill; guard # anyway — a missing row means nothing to own yet. @@ -121,9 +239,51 @@ def reconcile_ownership_for( # Return False so the ramp count isn't inflated past what's live. return False # Beat owns it only when active AND not handed to PG. - PeriodicTask.objects.filter(name=pipeline_id).update( - enabled=active and not pg_owned - ) + beat_updates: dict = {"enabled": active and not pg_owned} + # Baseline Beat's clock on the way BACK, for the same reason next_run_at + # is baselined on the way out — and this half was missing. + # + # DatabaseScheduler keeps no next_run_at; it derives due-ness from + # PeriodicTask.last_run_at against the crontab. A schedule that spent days + # PG-owned carries a last_run_at from before the hand-over, so the instant + # `enabled` flips back every missed interval is overdue and Beat replays + # them at once. Observed on integration 2026-08-24: releasing 23 schedules + # fired 4 pipelines plus 3 periodics within 30 ms of "Released to Beat". + # + # An earlier comment here asserted the opposite — that DatabaseScheduler + # "recomputes due-ness from the crontab each tick, so re-enabling never + # produced a catch-up run". That was wrong: it recomputes from last_run_at, + # which is exactly what makes it catch up. + # + # Scoped to the RELEASE transition, mirroring the next_run_at rule above: + # stamping it on every call would push the clock forward on an ordinary + # pipeline save and silently skip a due fire. + if was_pg_owned and not pg_owned: + beat_updates["last_run_at"] = timezone.now() + matched = PeriodicTask.objects.filter(name=pipeline_id).update(**beat_updates) + if not matched: + # Branch the WHOLE message, not just the verb. The two directions mean + # opposite things, and an earlier version applied the alarming tail to + # both — which is how the direction that matters gets filtered out. + if pg_owned: + # Adopt: nothing to switch off. The mirror row was just written + # pg_owned=True and the PG tick selects WHERE pg_owned AND enabled, + # so PG *is* the firer. Benign, and this path runs on every pipeline + # save, so a warning here would be noise. + logger.info( + "reconcile_ownership_for: no Beat PeriodicTask named %s to " + "disable — PG owns it now, nothing needed switching off", + pipeline_id, + ) + else: + # Release: PG has let go and Beat has no row to take over, so the + # schedule has no firer at all — on the rollback path. Silent before + # this: a bulk .update() reports success by returning 0. + logger.warning( + "reconcile_ownership_for: no Beat PeriodicTask named %s to " + "release to — the schedule may now have no firer", + pipeline_id, + ) # Bulk .update() bypasses django-celery-beat's post_save signal, so # PeriodicTasks.last_update never bumps and DatabaseScheduler never # reloads — Beat would keep firing the schedule from its stale diff --git a/backend/scheduler/tasks.py b/backend/scheduler/tasks.py index 7db1cad534..755d878b20 100644 --- a/backend/scheduler/tasks.py +++ b/backend/scheduler/tasks.py @@ -34,47 +34,86 @@ # --------------------------------------------------------------------------- -def _retargeted_next_run_at(pipeline_id: str, cron_string: str) -> datetime | None: - """The ``next_run_at`` a cron EDIT should set for Beat parity, or ``None`` to - leave it as the PG scheduler's baseline owns it. - - Returns a recomputed next-match ONLY when ``cron_string`` differs from an - already-baselined mirror row (``next_run_at`` set). Returns ``None`` for a - brand-new row (no mirror yet) and a not-yet-baselined one (``next_run_at`` - NULL) so the scheduler's no-burst baseline records the first next-time. +# Sentinel: leave the ``next_run_at`` column alone entirely. Distinct from the +# value ``None``, which is a deliberate instruction to WRITE NULL — the PG tick +# reads NULL as "record a baseline next pass, don't fire this one". Conflating +# the two is what let a stale next_run_at survive a resume. +_LEAVE_NEXT_RUN_AT = object() + + +def _next_run_at_for_upsert( + pipeline_id: str, cron_string: str, enabled: bool +) -> datetime | None | object: + """What the mirror upsert should do with ``next_run_at``. Three outcomes: + + * :data:`_LEAVE_NEXT_RUN_AT` — don't touch the column. A brand-new row, a + not-yet-baselined one (``next_run_at`` NULL), or an unchanged cron on a + row that was already enabled. The PG tick's no-burst baseline owns it. + * ``None`` — write NULL, i.e. "baseline on the next tick, don't fire". A + RESUME: the mirror row was ``enabled=False`` and the incoming state is + ``True``. + * a ``datetime`` — a cron EDIT on an already-baselined row: retarget to the + new cron's next match for Beat parity (UN-3690), else the scheduler fires + once more at the stale old-cron time. + + **Why resume needs the explicit NULL, and why it is checked first.** While a + schedule is paused its ``next_run_at`` keeps drifting into the past, so the + instant ``enabled`` flips back the tick's + ``WHERE pg_owned AND enabled AND (next_run_at IS NULL OR <= now())`` matches + and the pipeline runs immediately — days late, on top of whatever the + operator triggered by hand. ``enable_task`` already baselines for exactly + this reason (see ``_mirror_periodic_schedule_set_enabled``), but the UI + resume does not take that path: it re-saves the pipeline, which reaches + ``SchedulerHelper._schedule_task_job`` → here and ``reconcile_ownership_for`` + — and that one deliberately baselines only on a Beat→PG hand-over + (``pg_owned and not was_pg_owned``), so an already-``pg_owned`` row keeps its + stale value and nothing clears it. Observed on integration 2026-08-12, -14 + and -20: three enables of ``gallh_load_test``, three spurious runs 2-3s + later, each settling back onto the cron afterwards. Exactly one extra run per + enable — bounded, but it costs a full LLM pass over the source. + + Resume is checked BEFORE the cron comparison because a resume that also + edited the cron must still baseline; ordering it the other way would let the + ``cron_string == existing`` early-out swallow the resume case, which is the + common one (a plain pause/resume does not change the cron). + + Scoped to the TRANSITION, not to every call — the same discipline + ``reconcile_ownership_for`` uses. This runs on every pipeline save, so + clearing unconditionally would re-baseline mid-cycle: a save at 12:07:59 + against a 12:08 ``next_run_at`` would skip that fire entirely. Fully guarded on purpose: this is an OPTIONAL enhancement over the mandatory ``cron_string``/``enabled`` mirror write, so a read failure, a ``None``/invalid cron (normally rejected upstream by ``PipelineSerializer.validate_cron_string`` - in ``pipeline_v2``), or a croniter error must degrade to ``None`` — never take - the base mirror write down with it. + in ``pipeline_v2``), or a croniter error must degrade to leaving the column + alone — never take the base mirror write down with it. """ try: existing = ( PgPeriodicSchedule.objects.filter(pipeline_id=pipeline_id) - .values("cron_string", "next_run_at") + .values("cron_string", "next_run_at", "enabled") .first() ) - if ( - existing is None - or existing["next_run_at"] is None - or existing["cron_string"] == cron_string - ): + if existing is None or existing["next_run_at"] is None: + return _LEAVE_NEXT_RUN_AT + if enabled and not existing["enabled"]: return None + if existing["cron_string"] == cron_string: + return _LEAVE_NEXT_RUN_AT return croniter(cron_string, timezone.now()).get_next(datetime) except Exception as exc: # Log the ACTUAL cause (bad read; cron=None → AttributeError; # CroniterBadCronError / CroniterBadDateError; a croniter API change) — a # generic "could not recompute" is a dead end when debugging a stale fire. logger.warning( - "pg_periodic_schedule: could not retarget next_run_at for pipeline %s " + "pg_periodic_schedule: could not resolve next_run_at for pipeline %s " "(cron %r): %s", pipeline_id, cron_string, exc, exc_info=True, ) - return None + return _LEAVE_NEXT_RUN_AT def mirror_periodic_schedule_upsert( @@ -94,18 +133,15 @@ def mirror_periodic_schedule_upsert( "cron_string": cron_string, "enabled": enabled, } - # Beat parity: a cron EDIT on an already-baselined row must retarget - # next_run_at to the new cron's next match, else the PG scheduler fires - # once more at the STALE old-cron time — Beat has no such staleness, it - # recomputes due-ness live from the crontab each tick (UN-3690). Left - # untouched (NULL) for a brand-new row — the scheduler's no-burst baseline - # is correct there (a new schedule fires at its next match) — and for a - # Beat→PG hand-over, where next_run_at is already NULL (never set while - # Beat-owned; ownership clears it only on rollback to Beat, see - # reconcile_ownership_for). Computed via a fully-guarded helper so it can - # never take down the mandatory cron_string/enabled mirror write below. - next_run_at = _retargeted_next_run_at(pipeline_id, cron_string) - if next_run_at is not None: + # next_run_at is owned by the PG tick, with two exceptions this path must + # honour: a cron EDIT retargets it (Beat parity, UN-3690) and a RESUME + # clears it to NULL so a value that went stale during the pause can't fire + # a catch-up run the moment the schedule is re-enabled. The helper decides + # which — including "leave it alone", which is NOT the same as writing NULL + # — and is fully guarded so it can never take down the mandatory + # cron_string/enabled mirror write below. + next_run_at = _next_run_at_for_upsert(pipeline_id, cron_string, enabled) + if next_run_at is not _LEAVE_NEXT_RUN_AT: defaults["next_run_at"] = next_run_at PgPeriodicSchedule.objects.update_or_create( pipeline_id=pipeline_id, @@ -123,8 +159,33 @@ def _mirror_periodic_schedule_set_enabled(pipeline_id: str, enabled: bool) -> No # Bump updated_at explicitly: queryset .update() does NOT trigger the # field's auto_now, so without this a pause/resume would change enabled # without advancing the "last changed" timestamp. + updates: dict = {"enabled": enabled, "updated_at": timezone.now()} + if enabled: + # Resume → baseline instead of firing a stale next_run_at. While the + # schedule was paused its next_run_at kept drifting into the past, so the + # PG tick's `next_run_at <= now()` matches on the very next pass and the + # pipeline runs IMMEDIATELY on resume — the longer the pause, the more + # certain. NULL means "record a baseline next tick, don't fire this cycle" + # (pg_queue/models.py:366), which resumes at the next cron match instead. + # + # Observed on integration 2026-08-14: gallh_load_test fired ~2s after + # being re-enabled, against a next_run_at two days old. + # + # CORRECTION (2026-08-24): this used to add "Beat parity, not a new rule: + # DatabaseScheduler ... recomputes due-ness from the crontab each tick, so + # re-enabling never produced a catch-up run there." That was WRONG. + # DatabaseScheduler recomputes from PeriodicTask.last_run_at, which is + # exactly what makes it catch up too — releasing 23 schedules back to Beat + # fired 4 pipelines and 3 periodics within 30 ms. Beat needs the same + # baseline on its own clock; see scheduler/ownership.py. So this is a + # shared hazard of both schedulers, not a PG-only regression. + # + # Safe to do unconditionally here: this helper is reached only from + # enable_task/disable_task (an explicit pause/resume), never from the + # per-save path — so it cannot re-baseline mid-cycle and skip a fire. + updates["next_run_at"] = None matched = PgPeriodicSchedule.objects.filter(pipeline_id=pipeline_id).update( - enabled=enabled, updated_at=timezone.now() + **updates ) if matched == 0: # No mirror row — e.g. a pipeline scheduled before this shipped, or diff --git a/backend/scheduler/tests/test_pg_periodic_schedule_mirror.py b/backend/scheduler/tests/test_pg_periodic_schedule_mirror.py index 9d094f9a9f..8b4439f7a8 100644 --- a/backend/scheduler/tests/test_pg_periodic_schedule_mirror.py +++ b/backend/scheduler/tests/test_pg_periodic_schedule_mirror.py @@ -83,40 +83,49 @@ def test_failure_is_swallowed(self): ) -class TestUpsertNextRunRecompute: - """UN-3690 — a cron EDIT must retarget ``next_run_at`` (Beat parity), so the - new time takes effect this cycle instead of the pipeline firing once more at - the stale old-cron time. Only for an already-baselined row: a fresh row and a - Beat→PG hand-over keep ``next_run_at`` NULL for the scheduler's no-burst - baseline. - """ +class _UpsertNextRunHelpers: + """Shared setup for the two ``next_run_at`` suites below. Helpers only — no + tests live here, so neither suite re-runs the other's cases.""" _OLD = datetime(2026, 1, 1, 9, 0, tzinfo=dt_timezone.utc) @staticmethod - def _mock_existing(sched, *, cron_string, next_run_at): + def _mock_existing(sched, *, cron_string, next_run_at, enabled=True): _stub_existing_row( sched, None if cron_string is None - else {"cron_string": cron_string, "next_run_at": next_run_at}, + else { + "cron_string": cron_string, + "next_run_at": next_run_at, + "enabled": enabled, + }, ) @staticmethod - def _upsert(cron): + def _upsert(cron, enabled=True): tasks.mirror_periodic_schedule_upsert( pipeline_id=_PIPELINE_ID, organization_id=_ORG, workflow_id=_WORKFLOW_ID, pipeline_name=_REAL_NAME, cron_string=cron, - enabled=True, + enabled=enabled, ) @staticmethod def _defaults(sched): return sched.objects.update_or_create.call_args.kwargs["defaults"] + +class TestUpsertNextRunRecompute(_UpsertNextRunHelpers): + """UN-3690 — a cron EDIT must retarget ``next_run_at`` (Beat parity), so the + new time takes effect this cycle instead of the pipeline firing once more at + the stale old-cron time. Only for an already-baselined row: a fresh row and a + Beat→PG hand-over keep ``next_run_at`` NULL for the scheduler's no-burst + baseline. + """ + def test_cron_change_on_baselined_row_retargets_next_run_at(self): now = datetime(2026, 1, 1, 8, 0, tzinfo=dt_timezone.utc) with ( @@ -172,6 +181,86 @@ def test_read_failure_degrades_to_plain_upsert(self): assert "next_run_at" not in self._defaults(sched) +class TestUpsertResumeBaselinesStaleNextRun(_UpsertNextRunHelpers): + """A UI resume must not fire a catch-up run. + + ``enable_task`` baselines ``next_run_at`` on resume, but the UI does not take + that path — it re-saves the pipeline, reaching + ``SchedulerHelper._schedule_task_job`` → ``mirror_periodic_schedule_upsert`` + (here) and ``reconcile_ownership_for``. The latter baselines only on a + Beat→PG hand-over, so a row that is ALREADY ``pg_owned`` keeps whatever + ``next_run_at`` it had when it was paused — days in the past — and the PG + tick's ``next_run_at <= now()`` fires it on the very next pass. + + Observed three times on integration (2026-08-12, -14, -20): each enable of + ``gallh_load_test`` produced a spurious run 2-3s later, then settled onto the + cron correctly. + """ + + def test_resume_clears_a_stale_next_run_at(self): + """The bug, directly: paused row + past next_run_at, re-enabled.""" + with patch("scheduler.tasks.PgPeriodicSchedule") as sched: + self._mock_existing( + sched, cron_string="8 * * * *", next_run_at=self._OLD, enabled=False + ) + self._upsert("8 * * * *", enabled=True) + defaults = self._defaults(sched) + # Present AND None: writing NULL is the baseline instruction. Merely + # omitting the key would leave the stale value in place — the bug. + assert "next_run_at" in defaults + assert defaults["next_run_at"] is None + + def test_resume_that_also_edits_the_cron_still_baselines(self): + """Resume is checked before the cron comparison, so a combined + resume+edit baselines rather than retargeting. Both are non-firing, but + this pins the ordering — reversing it would let the equal-cron early-out + swallow the common plain-resume case.""" + with patch("scheduler.tasks.PgPeriodicSchedule") as sched: + self._mock_existing( + sched, cron_string="0 9 * * *", next_run_at=self._OLD, enabled=False + ) + self._upsert("0 10 * * *", enabled=True) + assert self._defaults(sched)["next_run_at"] is None + + def test_a_save_on_an_already_enabled_row_is_untouched(self): + """The mid-cycle guard. This helper runs on EVERY pipeline save, so a + rename at 12:07:59 against a 12:08 next_run_at must not re-baseline and + skip that fire. Only the False→True transition clears.""" + with patch("scheduler.tasks.PgPeriodicSchedule") as sched: + self._mock_existing( + sched, cron_string="8 * * * *", next_run_at=self._OLD, enabled=True + ) + self._upsert("8 * * * *", enabled=True) + assert "next_run_at" not in self._defaults(sched) + + def test_pause_does_not_clear_next_run_at(self): + """Nothing fires while disabled, so there is nothing to guard against — + and clearing here would discard the value resume needs to detect.""" + with patch("scheduler.tasks.PgPeriodicSchedule") as sched: + self._mock_existing( + sched, cron_string="8 * * * *", next_run_at=self._OLD, enabled=True + ) + self._upsert("8 * * * *", enabled=False) + assert "next_run_at" not in self._defaults(sched) + + def test_resume_of_a_never_baselined_row_writes_nothing(self): + """next_run_at already NULL is already the baseline — no write needed, + and the read short-circuits before the enabled check.""" + with patch("scheduler.tasks.PgPeriodicSchedule") as sched: + self._mock_existing( + sched, cron_string="8 * * * *", next_run_at=None, enabled=False + ) + self._upsert("8 * * * *", enabled=True) + assert "next_run_at" not in self._defaults(sched) + + def test_the_sentinel_is_not_none(self): + """The whole fix rests on 'leave alone' being distinguishable from 'write + NULL'. If the sentinel were None-y, the resume baseline would silently + degrade back to the old skip-the-write behaviour and every test above + would still pass on the omission path.""" + assert tasks._LEAVE_NEXT_RUN_AT is not None + + class TestHelperWiringSourcesRealName: """The High contract: the mirror must store the user-facing pipeline name, NOT the synthetic ``"Pipeline job-"`` label that the PeriodicTask args @@ -336,3 +425,53 @@ def test_delete_mirror_failure_is_swallowed(self): pt.objects.get.return_value = MagicMock() sched.objects.filter.return_value.delete.side_effect = RuntimeError("db down") tasks.delete_periodic_task(_PIPELINE_ID) # must not raise + + +class TestResumeBaselinesInsteadOfCatchingUp: + """Re-enabling a paused schedule must resume at the next cron match, not fire now. + + While a schedule is paused its `next_run_at` keeps drifting into the past, so the + PG tick's `next_run_at <= now()` matches on the very next pass and the pipeline runs + IMMEDIATELY on resume — the longer the pause, the more certain. NULL is the guard: + "record a baseline next tick, don't fire this cycle" (pg_queue/models.py:366). + + Observed on integration 2026-08-14: gallh_load_test fired ~2s after being + re-enabled, against a next_run_at two days old, on top of the operator's manual run. + + CORRECTION (2026-08-24): this docstring used to claim "Beat parity, not a new rule — + DatabaseScheduler ... recomputes due-ness from the crontab each tick, so resume + never produced a catch-up run there." Wrong on the mechanism: it recomputes from + PeriodicTask.last_run_at, so Beat catches up as well. Both schedulers need their + clock baselined on a hand-over; the Beat half is pinned by + scheduler/tests/test_pg_schedule_ownership.py::TestBeatClockBaselineOnRelease. + """ + + def _run(self, fn): + with ( + patch("scheduler.tasks.PeriodicTask") as pt, + patch("scheduler.tasks.PeriodicTasks"), + patch("scheduler.tasks.PipelineProcessor"), + patch("scheduler.tasks.PgPeriodicSchedule") as sched, + patch("scheduler.tasks.transaction.atomic", return_value=nullcontext()), + ): + pt.objects.get.return_value = MagicMock() + ( + sched.objects.select_for_update.return_value.filter.return_value + .values_list.return_value.first.return_value + ) = False + sched.objects.filter.return_value.update.return_value = 1 + fn(_PIPELINE_ID) + return sched.objects.filter.return_value.update.call_args.kwargs + + def test_resume_clears_next_run_at(self): + kwargs = self._run(tasks.enable_task) + assert kwargs["enabled"] is True + assert kwargs["next_run_at"] is None + + def test_pause_leaves_next_run_at_alone(self): + """Nothing fires while disabled, so there is nothing to guard against — and + clearing here would discard the value resume needs to baseline against. + """ + kwargs = self._run(tasks.disable_task) + assert kwargs["enabled"] is False + assert "next_run_at" not in kwargs diff --git a/backend/scheduler/tests/test_pg_schedule_ownership.py b/backend/scheduler/tests/test_pg_schedule_ownership.py index bb4c1af48a..4f05e8d150 100644 --- a/backend/scheduler/tests/test_pg_schedule_ownership.py +++ b/backend/scheduler/tests/test_pg_schedule_ownership.py @@ -18,6 +18,16 @@ class TestResolveScheduleOwner: + @pytest.fixture(autouse=True) + def _scheduler_gate_on(self, monkeypatch): + """These pin the FINAL-phase behaviour, when hand-over is switched on. + + PG_SCHEDULER_ENABLED defaults off for the whole rollout (Beat stays the sole + scheduler), so without this every case below would short-circuit to Beat and + assert nothing about the logic it was written for. + """ + monkeypatch.setenv("PG_SCHEDULER_ENABLED", "true") + def test_flipt_unavailable_is_beat(self, monkeypatch): monkeypatch.setenv("FLIPT_SERVICE_AVAILABLE", "false") with patch("scheduler.ownership.check_feature_flag_status") as flag: @@ -52,6 +62,16 @@ def test_flipt_error_fails_closed_to_beat(self, monkeypatch): class TestReconcileOwnership: + @pytest.fixture(autouse=True) + def _scheduler_gate_on(self, monkeypatch): + """These pin the FINAL-phase behaviour, when hand-over is switched on. + + PG_SCHEDULER_ENABLED defaults off for the whole rollout (Beat stays the sole + scheduler), so without this every case below would short-circuit to Beat and + assert nothing about the logic it was written for. + """ + monkeypatch.setenv("PG_SCHEDULER_ENABLED", "true") + @pytest.fixture(autouse=True) def _mock_periodic_tasks(self): # reconcile now bumps PeriodicTasks.last_update after the bulk .update(); @@ -104,18 +124,26 @@ def test_not_pg_owned_enables_beat_and_clears_next_run(self): PT.objects.filter.return_value.update.call_args.kwargs["enabled"] is True ) - def test_pg_owned_does_not_clear_next_run(self): + def test_an_ALREADY_pg_owned_schedule_does_not_clear_next_run(self): + """Narrowed: this is the True→True case, not "adopt never clears". + + Hand-over now DOES baseline (see TestNextRunBaselineOnTransition) — a + next_run_at surviving adoption fires the pipeline immediately. What must not + happen is re-baselining a schedule that was already PG-owned: reconcile runs on + every pipeline save, so a save at 12:07:59 would push a 12:08 next_run_at out to + 13:08 and skip that fire. + + `was_pg_owned` is set explicitly — a bare MagicMock is truthy, so this would + otherwise pass for the wrong reason and keep passing if the scoping regressed. + """ sched, pt, resolve, txn = self._patches(owner=True) with sched as Sched, pt, resolve, txn: - Sched.objects.filter.return_value.update.return_value = 1 + qs = Sched.objects.filter.return_value + qs.values_list.return_value.first.return_value = True + qs.update.return_value = 1 ownership.reconcile_ownership_for(_PID, _ORG, active=True) - # An active PG-owned schedule must NOT have its next_run_at reset (that - # would re-baseline and skip a fire). - assert ( - "next_run_at" - not in Sched.objects.filter.return_value.update.call_args.kwargs - ) + assert "next_run_at" not in qs.update.call_args.kwargs def test_paused_pipeline_keeps_beat_disabled_even_if_not_pg_owned(self): sched, pt, resolve, txn = self._patches(owner=False) @@ -167,6 +195,172 @@ def test_beat_reload_not_signalled_when_no_mirror_row(self, _mock_periodic_tasks _mock_periodic_tasks.update_changed.assert_not_called() +class TestPgSchedulerGate: + """The gate that keeps Beat the sole scheduler for the whole PG rollout. + + Pipelines already reach PG without the PG scheduler (execute_pipeline_task_v2 → + complete_execution → execute_workflow_async → resolve_transport), so the PG + scheduler exists only to retire the Beat *deployment* — a later, deliberate step. + Handing a schedule over before then leaves it with NO firer: Beat's PeriodicTask + disabled, nothing polling the PG side. + """ + + def test_defaults_off(self, monkeypatch): + monkeypatch.delenv("PG_SCHEDULER_ENABLED", raising=False) + assert ownership.pg_scheduler_enabled() is False + + @pytest.mark.parametrize("value", ["false", "False", "0", "", "yes", "TRUE "]) + def test_only_an_exact_true_opts_in(self, monkeypatch, value): + monkeypatch.setenv("PG_SCHEDULER_ENABLED", value) + assert ownership.pg_scheduler_enabled() is (value.strip().lower() == "true") + + def test_owner_stays_beat_even_with_the_flag_on(self, monkeypatch): + """The rollout state: pg_queue_enabled on, PG scheduler still dark.""" + monkeypatch.delenv("PG_SCHEDULER_ENABLED", raising=False) + monkeypatch.setenv("FLIPT_SERVICE_AVAILABLE", "true") + with patch( + "scheduler.ownership.check_feature_flag_status", return_value=True + ) as flag: + assert ownership.resolve_schedule_owner(_PID, _ORG) is False + # Short-circuits before Flipt — no evaluation on every schedule save. + flag.assert_not_called() + + def test_reconcile_writes_NEITHER_beat_table(self, monkeypatch): + """The regression this gate exists to prevent. + + Returning early matters over merely resolving to Beat: that path would still + issue PeriodicTask.update(enabled=active) and bump PeriodicTasks.update_changed() + on every schedule save — writing back a value Beat already had and forcing a + reload, on a table we promised not to touch. + """ + monkeypatch.delenv("PG_SCHEDULER_ENABLED", raising=False) + monkeypatch.setenv("FLIPT_SERVICE_AVAILABLE", "true") + with ( + patch("scheduler.ownership.PeriodicTask") as PT, + patch("scheduler.ownership.PeriodicTasks") as PTs, + patch("scheduler.ownership.PgPeriodicSchedule") as Sched, + patch("scheduler.ownership.check_feature_flag_status", return_value=True), + ): + # A CLEAN row: not pg_owned, so there is nothing to repair and the + # write-free path must be taken. Must be set explicitly — a bare + # MagicMock's .exists() is truthy, which would look like a stale row. + Sched.objects.filter.return_value.exists.return_value = False + assert ownership.reconcile_ownership_for(_PID, _ORG, active=True) is False + + PT.objects.filter.assert_not_called() + PTs.update_changed.assert_not_called() + # pg_owned is not WRITTEN either — otherwise switching the PG scheduler on + # later would fire every already-owned schedule while Beat still fires it too. + # (.filter() is now called once, for the read-only staleness check.) + Sched.objects.filter.return_value.update.assert_not_called() + + +class TestStalePgOwnershipIsReleased: + """A row left ``pg_owned`` by a build predating the gate must be repairable. + + Reproduced on integration 2026-08-12: `snapshot.993`'s backend came from OSS + `main`, where ``resolve_schedule_owner`` keys on the ``pg_queue_enabled`` Flipt + flag ALONE. Turning the flag on handed one pipeline to PG. Rolling forward onto + the gate then STRANDED it: the gate-off path wrote nothing, so nothing — + including ``reconcile_pg_schedules``, which routes back through here — could + clear it. Beat and the PG scheduler both fired that pipeline 2.4s apart; only + one execution resulted because the Celery consumers happened to be at zero. + + Reachable in production without a mis-built image: deploy this branch, roll back + to a pre-gate build while the flag is on, roll forward. + """ + + def _mocks(self, stale: bool): + sched = patch("scheduler.ownership.PgPeriodicSchedule").start() + sched.objects.filter.return_value.exists.return_value = stale + sched.objects.filter.return_value.update.return_value = 1 + # These tests are DB-free, but the repair path (unlike the write-free one) + # enters the real transaction.atomic() — same no-op stand-in as above. + patch( + "scheduler.ownership.transaction.atomic", + return_value=contextlib.nullcontext(), + ).start() + return sched + + def teardown_method(self): + patch.stopall() + + def test_a_stale_row_is_handed_back_to_beat(self, monkeypatch): + monkeypatch.delenv("PG_SCHEDULER_ENABLED", raising=False) + sched = self._mocks(stale=True) + with ( + patch("scheduler.ownership.PeriodicTask") as PT, + patch("scheduler.ownership.PeriodicTasks") as PTs, + ): + assert ownership.reconcile_ownership_for(_PID, _ORG, active=True) is False + + # pg_owned cleared, and next_run_at reset so a later re-hand-over baselines + # instead of firing immediately on a stale timestamp. + updates = sched.objects.filter.return_value.update.call_args.kwargs + assert updates["pg_owned"] is False + assert updates["next_run_at"] is None + # Beat re-enabled in the same breath — releasing pg_owned without this would + # leave the schedule with NO firer at all. + # + # UN-3796 (2026-08-24): the write now also baselines Beat's clock, so this + # asserts per-key rather than on the whole kwargs dict. Not a loosening — the + # repair path is a RELEASE, and it is the one most likely to hand back a row + # whose last_run_at is deeply stale, which is exactly what makes Beat replay + # every missed interval at once. TestBeatClockBaselineOnRelease pins the rule. + beat = PT.objects.filter.return_value.update.call_args.kwargs + assert beat["enabled"] is True + assert "last_run_at" in beat + # ...and Beat told to reload, else it keeps its stale in-memory copy. + PTs.update_changed.assert_called_once() + + def test_a_paused_pipeline_is_not_resurrected_by_the_repair(self, monkeypatch): + """Repair restores the FIRER, never the on/off state the user chose.""" + monkeypatch.delenv("PG_SCHEDULER_ENABLED", raising=False) + self._mocks(stale=True) + with ( + patch("scheduler.ownership.PeriodicTask") as PT, + patch("scheduler.ownership.PeriodicTasks"), + ): + ownership.reconcile_ownership_for(_PID, _ORG, active=False) + beat = PT.objects.filter.return_value.update.call_args.kwargs + assert beat["enabled"] is False + # Baselined even though it stays paused: if it is resumed later, Beat must not + # then replay the backlog accrued while PG owned it. + assert "last_run_at" in beat + + def test_flipt_is_never_consulted_while_the_gate_is_off(self, monkeypatch): + """The repair resolves to Beat unconditionally. + + Asking Flipt could return True and re-hand the schedule to a PG scheduler + that is not running — turning a double-fire into no firer at all. + """ + monkeypatch.delenv("PG_SCHEDULER_ENABLED", raising=False) + monkeypatch.setenv("FLIPT_SERVICE_AVAILABLE", "true") + self._mocks(stale=True) + with ( + patch("scheduler.ownership.PeriodicTask"), + patch("scheduler.ownership.PeriodicTasks"), + patch( + "scheduler.ownership.check_feature_flag_status", return_value=True + ) as flag, + ): + ownership.reconcile_ownership_for(_PID, _ORG, active=True) + flag.assert_not_called() + + def test_a_db_error_on_the_check_falls_back_to_writing_nothing(self, monkeypatch): + """A repair that cannot confirm it is needed must not run.""" + monkeypatch.delenv("PG_SCHEDULER_ENABLED", raising=False) + sched = patch("scheduler.ownership.PgPeriodicSchedule").start() + sched.objects.filter.side_effect = Exception("db down") + with ( + patch("scheduler.ownership.PeriodicTask") as PT, + patch("scheduler.ownership.PeriodicTasks") as PTs, + ): + assert ownership.reconcile_ownership_for(_PID, _ORG, active=True) is False + PT.objects.filter.assert_not_called() + PTs.update_changed.assert_not_called() + + class TestReconcileAtomicityRealDB: """The load-bearing invariant: the pg_owned write and the PeriodicTask write are ONE transaction — if the PeriodicTask update fails, pg_owned rolls back @@ -227,3 +421,138 @@ def test_periodictask_update_failure_rolls_back_pg_owned(self): finally: RealPeriodicTask.objects.filter(name=pid).delete() PgPeriodicSchedule.objects.filter(pipeline_id=pid).delete() + + +class TestNextRunBaselineOnTransition: + """An ownership or on/off change must never cause an immediate unscheduled run. + + The PG tick selects `WHERE pg_owned AND enabled AND (next_run_at IS NULL OR + next_run_at <= now())`. A next_run_at left over from an earlier PG-ownership period + is already in the past, so the row fires on the very next pass. NULL is the guard — + "record a baseline next tick, don't fire this cycle" (pg_queue/models.py:366). + + Observed on integration 2026-08-14: gallh_load_test carried + next_run_at=2026-08-12 06:08 and fired ~2s after being re-enabled — two days late, + on top of the operator's own manual run. + + CORRECTION (2026-08-24): this docstring used to add "Beat never did this: + DatabaseScheduler keeps no persisted next_run_at and recomputes due-ness from the + crontab each tick." That is wrong. It recomputes from ``PeriodicTask.last_run_at``, + which is precisely why it DOES catch up — releasing the fleet fired four pipelines + and three periodics inside 30 ms. Beat has the same failure mode from the other + direction; see TestBeatClockBaselineOnRelease below. + """ + + def _reconcile(self, monkeypatch, *, was_pg_owned, now_pg_owned): + monkeypatch.setenv("PG_SCHEDULER_ENABLED", "true") + monkeypatch.setenv("FLIPT_SERVICE_AVAILABLE", "true") + with ( + patch("scheduler.ownership.PgPeriodicSchedule") as Sched, + patch("scheduler.ownership.PeriodicTask"), + patch("scheduler.ownership.PeriodicTasks"), + patch( + "scheduler.ownership.transaction.atomic", + return_value=contextlib.nullcontext(), + ), + patch( + "scheduler.ownership.check_feature_flag_status", + return_value=now_pg_owned, + ), + ): + qs = Sched.objects.filter.return_value + qs.values_list.return_value.first.return_value = was_pg_owned + qs.update.return_value = 1 + ownership.reconcile_ownership_for(_PID, _ORG, active=True) + return qs.update.call_args.kwargs + + def test_handing_over_to_pg_baselines(self, monkeypatch): + # The bug: a stale next_run_at surviving adoption fires the pipeline at once. + updates = self._reconcile(monkeypatch, was_pg_owned=False, now_pg_owned=True) + assert updates["pg_owned"] is True + assert updates["next_run_at"] is None + + def test_releasing_to_beat_still_baselines(self, monkeypatch): + updates = self._reconcile(monkeypatch, was_pg_owned=True, now_pg_owned=False) + assert updates["next_run_at"] is None + + def test_an_unchanged_owner_is_NOT_re_baselined(self, monkeypatch): + """The reason this is scoped to the transition. + + reconcile runs on every pipeline save. Clearing unconditionally would let a + save at 12:07:59 re-baseline a 12:08 next_run_at to 13:08 — silently skipping + that fire. Idempotent re-runs (every deploy) must leave the schedule alone. + """ + updates = self._reconcile(monkeypatch, was_pg_owned=True, now_pg_owned=True) + assert "next_run_at" not in updates + + +class TestBeatClockBaselineOnRelease: + """Handing a schedule BACK to Beat must reset Beat's clock, or Beat replays. + + The symmetric half of the class above, and the one that was missing. + DatabaseScheduler stores no next_run_at; it derives due-ness from + ``PeriodicTask.last_run_at`` against the crontab. A pipeline that spent days + PG-owned still carries the last_run_at from before the hand-over, so the instant + ``enabled`` flips back it is overdue by every interval it missed — and Beat fires + them all at once. + + Observed on integration 2026-08-24: `converge_pg_scheduler` released 23 schedules + and Beat dispatched four pipelines plus three dashboard_metrics.* periodics within + 30 ms of logging "Released to Beat". + """ + + def _reconcile(self, monkeypatch, *, was_pg_owned, now_pg_owned, active=True): + monkeypatch.setenv("PG_SCHEDULER_ENABLED", "true") + monkeypatch.setenv("FLIPT_SERVICE_AVAILABLE", "true") + with ( + patch("scheduler.ownership.PgPeriodicSchedule") as Sched, + patch("scheduler.ownership.PeriodicTask") as Beat, + patch("scheduler.ownership.PeriodicTasks"), + patch( + "scheduler.ownership.transaction.atomic", + return_value=contextlib.nullcontext(), + ), + patch( + "scheduler.ownership.check_feature_flag_status", + return_value=now_pg_owned, + ), + ): + qs = Sched.objects.filter.return_value + qs.values_list.return_value.first.return_value = was_pg_owned + qs.update.return_value = 1 + ownership.reconcile_ownership_for(_PID, _ORG, active=active) + return Beat.objects.filter.return_value.update.call_args.kwargs + + def test_releasing_to_beat_stamps_last_run_at(self, monkeypatch): + beat = self._reconcile(monkeypatch, was_pg_owned=True, now_pg_owned=False) + assert beat["enabled"] is True + assert "last_run_at" in beat, ( + "release must baseline Beat's clock, else every interval missed while " + "PG owned the schedule is overdue and Beat replays them at once" + ) + assert beat["last_run_at"] is not None + + def test_handing_over_to_pg_does_NOT_touch_beats_clock(self, monkeypatch): + """Adoption switches Beat OFF, so its clock is irrelevant — and overwriting it + would destroy the value the eventual release needs to restore from.""" + beat = self._reconcile(monkeypatch, was_pg_owned=False, now_pg_owned=True) + assert beat["enabled"] is False + assert "last_run_at" not in beat + + def test_an_unchanged_owner_is_NOT_re_stamped(self, monkeypatch): + """Same transition-scoping reason as next_run_at: reconcile runs on every + pipeline save, and stamping unconditionally would push Beat's clock forward + on an ordinary edit and silently skip a due fire.""" + beat = self._reconcile(monkeypatch, was_pg_owned=False, now_pg_owned=False) + assert "last_run_at" not in beat + + def test_a_paused_pipeline_is_released_disabled_but_still_baselined( + self, monkeypatch + ): + """A paused schedule comes back paused — but if it is later resumed, Beat must + not then replay the backlog it accrued while PG owned it.""" + beat = self._reconcile( + monkeypatch, was_pg_owned=True, now_pg_owned=False, active=False + ) + assert beat["enabled"] is False + assert "last_run_at" in beat diff --git a/backend/workflow_manager/execution/tests/test_pg_finalization_fixes.py b/backend/workflow_manager/execution/tests/test_pg_finalization_fixes.py index c3d7252167..d9cc5ba4da 100644 --- a/backend/workflow_manager/execution/tests/test_pg_finalization_fixes.py +++ b/backend/workflow_manager/execution/tests/test_pg_finalization_fixes.py @@ -22,9 +22,20 @@ def _age(execution, seconds): - """Backdate modified_at (auto_now) via a direct UPDATE that bypasses auto_now.""" - WorkflowExecution.objects.filter(pk=execution.pk).update( - modified_at=timezone.now() - timedelta(seconds=seconds) + """Backdate modified_at (auto_now) via a direct UPDATE that bypasses auto_now. + + Ages the FILE rows too, because that is where ``recover_stuck_pg_executions`` + measures staleness: the execution row's ``modified_at`` is effectively its start + time (file completions write the file row, not the execution row), so ageing only + the execution would describe a run that started long ago and finished a moment + ago — a live execution mid-callback, which the endpoint must NOT touch. A real + strand has quiet files, and that is what this reproduces. Tests that specifically + need the two timestamps to diverge override the file rows after calling this. + """ + stale = timezone.now() - timedelta(seconds=seconds) + WorkflowExecution.objects.filter(pk=execution.pk).update(modified_at=stale) + WorkflowFileExecution.objects.filter(workflow_execution=execution).update( + modified_at=stale ) @@ -46,9 +57,9 @@ def _exec(self, status, pg=True, files=()): ) return ex - def _call(self, stuck_seconds=60): + def _call(self, stuck_seconds=60, limit=100): req = MagicMock() - req.data = {"stuck_seconds": stuck_seconds, "limit": 100} + req.data = {"stuck_seconds": stuck_seconds, "limit": limit} return self.view.recover_stuck_pg_executions(req).data def test_all_files_completed_recovers_to_completed(self): @@ -102,28 +113,81 @@ def test_negative_stuck_seconds_does_not_match_live_execution(self): assert out["scanned"] == 0 assert ex.status == ExecutionStatus.EXECUTING.value - def test_file_less_stuck_is_skipped_not_failed(self): - # A file-less PG exec may still be QUEUED (a backlog/outage can outlast the - # stuck window) — it must NOT be failed, so a delayed worker can still - # finalize it (the one-way guard would otherwise block recovery). + def test_file_less_stuck_is_left_alone_and_no_longer_consumes_the_window(self): + """A file-less PG exec may still be QUEUED (a backlog/outage can outlast the + stuck window) — it must NOT be failed, so a delayed worker can still finalize + it (the one-way guard would otherwise block recovery). That invariant is + unchanged and is what the status assertion below pins. + + The COUNTER expectation is deliberately updated: this used to assert + ``skipped == 1``, i.e. the row was selected and then rejected downstream. + It is now filtered out at selection instead, so it is never scanned. The + mechanism moved because a downstream skip leaves ``modified_at`` untouched, + so these rows were re-selected on every sweep forever and crowded genuinely + recoverable executions out of the window — see the starvation test below. + """ ex = self._exec(ExecutionStatus.PENDING, files=[]) _age(ex, 9999) out = self._call() ex.refresh_from_db() - assert out["skipped"] == 1 + assert out["scanned"] == 0 + assert out["skipped"] == 0 + assert out["failed"] == 0 assert ex.status == ExecutionStatus.PENDING.value - def test_celery_execution_never_scanned(self): + def test_celery_execution_IS_now_recovered(self): + """INVERTED deliberately (UN-3796). This asserted `scanned == 0` — that a + Celery execution is never touched. + + That exclusion is the bug. A PG cutover removes the Celery workers, and an + execution whose files all finished but whose chord callback never fired then + sits EXECUTING forever: the callback is gone, there is no pg_barrier_state to + recover from, and this endpoint refused to look. Nothing else could close it. + + The finalization logic was always transport-agnostic — it reads file statuses + and recomputes the terminal status. Only the selection was narrow. + """ ex = self._exec( ExecutionStatus.EXECUTING, pg=False, files=[ExecutionStatus.COMPLETED] ) _age(ex, 9999) out = self._call() ex.refresh_from_db() - assert out["scanned"] == 0 # queue_message_id IS NULL → PG filter excludes it + assert out["recovered"] == 1 + assert ex.status == ExecutionStatus.COMPLETED.value + + def test_a_NEVER_DISPATCHED_execution_is_left_to_the_undispatched_sweep(self): + """The disjointness guard. Both handles NULL means the request died before + dispatch — undispatched_sweep.py's row, which marks it ERROR with "did not + start, safe to re-run". + + Both sweeps run on the same reaper cadence against the same table, so an + overlap would be live from the first tick: one marking it ERROR while the + other finalized it from whatever files happened to exist. + """ + ex = WorkflowExecution.objects.create( + workflow=self.wf, + status=ExecutionStatus.EXECUTING.value, + queue_message_id=None, + task_id=None, + ) + WorkflowFileExecution.objects.create( + workflow_execution=ex, file_name="f", status=ExecutionStatus.COMPLETED.value + ) + _age(ex, 9999) + out = self._call() + ex.refresh_from_db() + assert out["scanned"] == 0 assert ex.status == ExecutionStatus.EXECUTING.value - def test_still_processing_is_skipped(self): + def test_still_processing_is_never_finalized(self): + """A non-terminal file means work may still be in flight — finalizing would + terminalize a live execution, and the one-way guard then blocks any correction. + That invariant is unchanged. + + Counter expectation updated for the same reason as the file-less case: the row + is now excluded at selection rather than skipped after being scanned. + """ ex = self._exec( ExecutionStatus.EXECUTING, files=[ExecutionStatus.COMPLETED, ExecutionStatus.EXECUTING], @@ -131,7 +195,8 @@ def test_still_processing_is_skipped(self): _age(ex, 9999) out = self._call() ex.refresh_from_db() - assert out["skipped"] == 1 + assert out["scanned"] == 0 + assert out["failed"] == 0 assert ex.status == ExecutionStatus.EXECUTING.value def test_recently_modified_not_recovered(self): @@ -141,7 +206,160 @@ def test_recently_modified_not_recovered(self): assert out["scanned"] == 0 assert ex.status == ExecutionStatus.EXECUTING.value + def test_backlog_of_unrecoverable_rows_does_not_starve_a_recoverable_one(self): + """THE regression. Selection is oldest-first with a hard ``limit``, so any row + that is selected but never finalized occupies a slot on every future sweep — + a skip does not advance ``modified_at``, so the same rows come back forever. + A recoverable execution behind them is never reached: not scanned, not + skipped, invisible. + + Found live during the UN-3796 cutover rehearsal: 1964 candidates, 1122 of them + permanently unrecoverable (476 file-less, 646 with a non-terminal file), and a + freshly stranded execution at rank 1964 that every sweep failed to see while + logging a healthy-looking "scanned 100, recovered 0". + + Mutation check: drop either ``Exists`` clause from the selection query and this + fails — the backlog fills the window and ``recovered`` is 0. + """ + limit = 3 + # Older than the victim, and permanently unrecoverable: the two shapes that + # made up the real backlog. + for i in range(limit * 2): + junk = self._exec( + ExecutionStatus.EXECUTING, + files=[] if i % 2 else [ExecutionStatus.EXECUTING], + ) + _age(junk, 9000) + victim = self._exec(ExecutionStatus.EXECUTING, files=[ExecutionStatus.COMPLETED]) + _age(victim, 100) # NEWEST → last in oldest-first order + + out = self._call(limit=limit) + + victim.refresh_from_db() + assert out["recovered"] == 1 + assert victim.status == ExecutionStatus.COMPLETED.value + + def test_window_is_spent_only_on_rows_that_can_finalize(self): + """The drain property that makes oldest-first safe: everything scanned is + finalized, so it leaves the candidate set and the backlog shrinks. If skips + can consume the window the queue stops moving and FIFO stops being fair. + """ + for _ in range(5): + _age(self._exec(ExecutionStatus.EXECUTING, files=[]), 9000) + _age( + self._exec(ExecutionStatus.EXECUTING, files=[ExecutionStatus.EXECUTING]), + 9000, + ) + _age( + self._exec(ExecutionStatus.EXECUTING, files=[ExecutionStatus.COMPLETED]), 9000 + ) + + out = self._call() + + assert out["scanned"] == 1 + assert out["recovered"] == 1 + assert out["skipped"] == 0 + + def test_execution_time_comes_from_the_last_file_not_from_now(self): + """``update_execution()`` stamps ``execution_time = now - created_at`` on any + terminal transition. Correct for a live finalization; wrong here, where the + sweep may finalize work that ended long ago — it would record how late the + reaper was, not how long the run took. The integration backlog would have + written multi-day runtimes onto executions that ran for minutes. + """ + ex = self._exec(ExecutionStatus.EXECUTING, files=[ExecutionStatus.COMPLETED]) + _age(ex, 9000) + # Set the timestamps AFTER _age — it backdates the file rows too, so doing this + # first would have them overwritten and the assertion would read 0.0. + created = timezone.now() - timedelta(seconds=9000) + # modified_at MUST be passed explicitly here. BaseModelManager.update() does + # `kwargs.setdefault("modified_at", timezone.now())`, so updating created_at + # alone silently re-stamps modified_at to NOW and un-ages the row _age() just + # aged — the execution then fails `modified_at__lt=cutoff`, the endpoint reports + # scanned=0, and the test fails on STATUS with no hint that a timestamp moved. + # That cost a CI round trip; the manager's docstring documents the override. + WorkflowExecution.objects.filter(pk=ex.pk).update( + created_at=created, + modified_at=timezone.now() - timedelta(seconds=9000), + ) + WorkflowFileExecution.objects.filter(workflow_execution=ex).update( + modified_at=created + timedelta(seconds=42) + ) + + self._call() + + ex.refresh_from_db() + assert ex.status == ExecutionStatus.COMPLETED.value + assert ex.execution_time == 42.0 + + def test_execution_time_left_alone_when_no_file_timestamp(self): + """No usable timestamp → keep whatever is there rather than write a worse + guess. Pins the early return, which a refactor could silently drop. + + Exercised directly rather than through the endpoint: ``modified_at`` is + ``auto_now`` and never NULL in practice, and a row with no file timestamp can + no longer be SELECTED anyway (staleness is measured on the files now, and NULL + does not satisfy ``< cutoff``). Driving it through the endpoint would assert + the selection filter, not the early return this test exists to pin. + """ + ex = self._exec(ExecutionStatus.EXECUTING, files=[ExecutionStatus.COMPLETED]) + # modified_at is NOT NULL at the database level, so it cannot be forced to NULL + # to reach the early return — an earlier version of this test tried and died on + # an IntegrityError. Patch the aggregate instead, which is the only way the + # helper can legitimately see a missing timestamp. + ex.execution_time = 7.5 + ex.save(update_fields=["execution_time"]) + + with patch( + "workflow_manager.file_execution.models.WorkflowFileExecution.objects" + ) as objects: + objects.filter.return_value.aggregate.return_value = {"last": None} + self.view._restamp_execution_time_from_files(ex) + + ex.refresh_from_db() + assert ex.execution_time == 7.5 + + def test_execution_whose_files_JUST_finished_is_not_recovered(self): + """The callback's window. ``workflow_execution.modified_at`` is effectively the + START time — file completions write the file row, not the execution row — so any + run longer than ``stuck_seconds`` is formally "stuck" while still running. + + The all-files-terminal filter stops that being catastrophic, but on its own it + leaves a live race: between the last file going terminal and the callback + finalizing, a perfectly healthy execution passes every other check. A sweep + landing in that window finalizes it first, the terminal-one-way guard then + refuses the callback's write, and the notification is silently lost. + + Measuring staleness on the files closes it. Mutation check: drop the + ``last_file_at__lt=cutoff`` filter and this fails. + """ + ex = self._exec(ExecutionStatus.EXECUTING, files=[ExecutionStatus.COMPLETED]) + _age(ex, 9000) # execution row looks ancient (it is just the start time) + WorkflowFileExecution.objects.filter(workflow_execution=ex).update( + modified_at=timezone.now() # ...but the work finished a moment ago + ) + + out = self._call(stuck_seconds=600) + + ex.refresh_from_db() + assert out["scanned"] == 0 + assert ex.status == ExecutionStatus.EXECUTING.value + + def test_execution_whose_files_went_quiet_long_ago_IS_recovered(self): + """The other half of the same predicate — moving staleness onto the files must + not stop genuinely abandoned work from being recovered. + """ + ex = self._exec(ExecutionStatus.EXECUTING, files=[ExecutionStatus.COMPLETED]) + _age(ex, 9000) + WorkflowFileExecution.objects.filter(workflow_execution=ex).update( + modified_at=timezone.now() - timedelta(seconds=3600) + ) + + out = self._call(stuck_seconds=600) + ex.refresh_from_db() + assert out["recovered"] == 1 + assert ex.status == ExecutionStatus.COMPLETED.value class TerminalOneWayGuardTests(TestCase): @@ -420,6 +638,8 @@ def test_result_acknowledge_does_not_touch_status_or_counters(self): assert ex.status == ExecutionStatus.COMPLETED.value assert ex.successful_files == 1 assert ex.result_acknowledged is True + + class RetrieveNotFoundTests(TestCase): """A missing execution must return 404, not 500 (UN-3719). The reaper's orphan-claim sweep relies on the deterministic 404 to GC claims for deleted diff --git a/backend/workflow_manager/internal_views.py b/backend/workflow_manager/internal_views.py index 4f9de9f0aa..0aa0fb1f3a 100644 --- a/backend/workflow_manager/internal_views.py +++ b/backend/workflow_manager/internal_views.py @@ -646,12 +646,30 @@ def recover_stuck_pg_executions(self, request): stuck past ``stuck_seconds`` whose files are ALL terminal, recomputes the correct terminal status from those files, and finalizes them. - Scoped to PG by ``queue_message_id`` (the durable per-row transport marker; - Celery rows carry ``task_id`` and ``queue_message_id=NULL``), so it needs no - Flipt context and never touches Celery executions. + **Covers BOTH transports (UN-3796).** It was scoped to PG by + ``queue_message_id__isnull=False``, which made a Celery execution invisible — + so one stranded by a PG cutover (its workers removed while its chord callback + was still to fire) sat ``EXECUTING`` forever with nothing able to close it. The + finalization logic itself was always transport-agnostic: it reads file statuses + and recomputes the terminal status, never touching a queue. Only the selection + was narrow. + + Now requires the row to have been dispatched on **some** transport — + ``queue_message_id`` (PG) or ``task_id`` (Celery). That deliberately excludes + rows with BOTH handles NULL, which are never-dispatched and belong to + ``workflow_v2/undispatched_sweep.py``; the two predicates are disjoint on + ``task_id`` so they can never contend for the same row. (The + ``total == 0 → skipped`` guard below would also catch those, but relying on a + downstream guard for correctness is how overlaps get reintroduced.) + + The method and route keep the ``_pg_`` name despite no longer being PG-only: + the URL is an internal-API contract between backend and workers, and renaming + it would break during a rolling deploy where an older worker still calls the + old path. Misleading name, deliberate trade. """ from datetime import timedelta + from django.db.models import Exists, Max, OuterRef, Q from django.utils import timezone from workflow_manager.workflow_v2.enums import ExecutionStatus @@ -672,15 +690,72 @@ def recover_stuck_pg_executions(self, request): cutoff = timezone.now() - timedelta(seconds=stuck_seconds) terminal = ExecutionStatus.terminal_values() + # Select ONLY rows the per-row guard below can actually finalize: at least one + # file execution, and none of them still non-terminal. + # + # Without this the window fills with rows _recover_one_stuck_pg_execution() + # always skips — and a skip leaves ``modified_at`` untouched, so the SAME rows + # are re-selected on every sweep, forever. A genuinely recoverable execution + # behind them is never reached: not scanned, not skipped, simply invisible. + # Measured on integration during the UN-3796 cutover rehearsal: 1964 candidates, + # of which 476 had no file rows and 646 had a non-terminal file — 57% permanent + # skip fodder — while the freshly stranded execution sat at rank 1964 and every + # sweep logged a healthy-looking "scanned 100, recovered 0". + # + # Filtering here rather than leaning on the downstream guard is also what makes + # the drain real: every selected row finalizes and drops out of the candidate + # set, so the backlog shrinks monotonically. That is precisely what keeps the + # oldest-first ordering below starvation-free — FIFO is only fair if the queue + # actually moves. + files = WorkflowFileExecution.objects.filter(workflow_execution=OuterRef("pk")) stuck_ids = list( WorkflowExecution.objects.filter( - queue_message_id__isnull=False, # PG-only; Celery uses task_id + # Dispatched by ANY evidence: the positive stamp, or either transport + # handle. dispatched_at alone is not enough — pre-0027 rows never got + # one — and the handles alone are not enough either, which is the gap + # that makes this three-way rather than two. + # + # Once undispatched_sweep also requires `dispatched_at IS NULL`, a row + # that WAS dispatched but whose handle write failed stops matching that + # sweep. Without dispatched_at here it would match neither, and an + # execution whose files all finished could sit non-terminal forever with + # nothing able to close it — trading "wrongly terminalised" for "never + # terminalised". The three-way test keeps the two sweeps disjoint AND + # jointly exhaustive: undispatched takes rows where all three are absent, + # this takes rows where any one is present. + Q(dispatched_at__isnull=False) + | Q(queue_message_id__isnull=False) + | Q(task_id__isnull=False), status__in=[ ExecutionStatus.PENDING.value, ExecutionStatus.EXECUTING.value, ], modified_at__lt=cutoff, ) + .filter(Exists(files), ~Exists(files.exclude(status__in=terminal))) + # Staleness must be measured on the FILES, not on the execution row. + # + # ``workflow_execution.modified_at`` is effectively the START time: file + # completions write the file row, not the execution row, so it does not + # advance while work proceeds. ``modified_at < cutoff`` therefore reads as + # "started more than stuck_seconds ago", not "quiet for stuck_seconds" — + # and every run longer than the window is formally eligible for recovery + # WHILE IT IS STILL RUNNING. (Observed: a 14-minute ETL whose execution row + # still read its 12:01:44 start at 12:15, with a 600 s window.) + # + # The all-files-terminal filter above stops that being catastrophic, but it + # leaves a live race: between the last file going terminal and the callback + # finalizing, a healthy execution passes every check here. If a sweep lands + # in that window the reaper finalizes first, the terminal-one-way guard then + # refuses the callback's own write, and the notification never fires. + # + # Requiring the LAST FILE to also be older than the cutoff closes it: an + # execution whose files just finished is not yet stale, so the callback keeps + # its window and only genuinely abandoned work is recovered. Both conditions + # are kept — the indexed execution-row predicate cheaply narrows the scan, + # this one decides correctness. + .annotate(last_file_at=Max("file_executions__modified_at")) + .filter(last_file_at__lt=cutoff) .order_by("modified_at") .values_list("id", flat=True)[:limit] ) @@ -753,6 +828,7 @@ def _recover_one_stuck_pg_execution(self, exec_id, cutoff, terminal) -> str: else: computed = ExecutionStatus.COMPLETED execution.update_execution(status=computed) + self._restamp_execution_time_from_files(execution) self._update_file_aggregates( execution, { @@ -828,6 +904,35 @@ def _truncate_error_message(error_msg: str | None, execution_id) -> str | None: return error_msg[:253] + "..." return error_msg + @staticmethod + def _restamp_execution_time_from_files(execution) -> None: + """Correct the runtime ``update_execution()`` just stamped from ``now()``. + + Both transports' status writers set ``execution_time = now - created_at`` on + any terminal transition (``_apply_legacy_update`` / ``_apply_pg_guarded_update`` + in the execution model). That is right for a live finalization, where "now" IS + when the work ended, and wrong here: this sweep finalizes executions whose files + finished long before anything noticed, so the stamp would record how late the + reaper was rather than how long the work took. On the integration backlog that + meant multi-day runtimes on executions that ran for minutes — every one of them + feeding whatever dashboards read this column. + + Recompute from the last file to reach a terminal state, which is the closest + durable record of when the work actually ended. If no file carries a timestamp, + leave whatever is there alone rather than substituting a worse guess. + """ + from django.db.models import Max + + last_file_at = WorkflowFileExecution.objects.filter( + workflow_execution=execution + ).aggregate(last=Max("modified_at"))["last"] + if not last_file_at or not execution.created_at: + return + execution.execution_time = round( + (last_file_at - execution.created_at).total_seconds(), 3 + ) + execution.save(update_fields=["execution_time"]) + @staticmethod def _update_file_aggregates(execution, validated_data) -> None: """Persist total / successful / failed file counts when present.""" diff --git a/backend/workflow_manager/workflow_v2/execution.py b/backend/workflow_manager/workflow_v2/execution.py index cbe02d2783..213c81302b 100644 --- a/backend/workflow_manager/workflow_v2/execution.py +++ b/backend/workflow_manager/workflow_v2/execution.py @@ -409,6 +409,37 @@ def update_execution_task(execution_id: str, task_id: str) -> None: except WorkflowExecution.DoesNotExist: logger.error(f"execution doesn't exist {execution_id}") + @staticmethod + def mark_dispatched(execution_id: str) -> None: + """Stamp ``dispatched_at`` — the positive record that the orchestrator task + reached its transport. + + Called by the dispatcher the instant dispatch returns, BEFORE any handle + bookkeeping. That ordering is the whole point: ``_record_dispatch_handle`` has + three paths that leave both handles NULL on an execution that IS running (an + exception the caller swallows, an empty handle, an unparseable PG handle), and + the undispatched sweep used to read that absence as "never dispatched" and mark + the live row ERROR. One write upstream of all three closes it. + + Best-effort, matching the handle write it precedes: a failure here must never + abort the caller, which is past the point of no return — the orchestrator is + already running. A failure leaves exactly the pre-existing ambiguity rather than + a new one. + + ``update()`` rather than ``save()``: no auto_now field on this model may be + re-stamped by a bookkeeping write, and this must not race the orchestrator's own + status writes on the same row. + """ + from django.utils import timezone + + updated = WorkflowExecution.objects.filter(id=execution_id).update( + dispatched_at=timezone.now() + ) + if not updated: + logger.error( + f"Could not stamp dispatched_at: execution {execution_id} not found" + ) + @staticmethod def update_execution_queue_message_id( execution_id: str, queue_message_id: int | None diff --git a/backend/workflow_manager/workflow_v2/execution_log_internal_urls.py b/backend/workflow_manager/workflow_v2/execution_log_internal_urls.py index 7f60dc88d4..d59d85268a 100644 --- a/backend/workflow_manager/workflow_v2/execution_log_internal_urls.py +++ b/backend/workflow_manager/workflow_v2/execution_log_internal_urls.py @@ -33,4 +33,9 @@ execution_log_internal_views.ProcessLogHistoryAPIView.as_view(), name="process_log_history", ), + path( + "sweep-undispatched-executions/", + execution_log_internal_views.SweepUndispatchedExecutionsAPIView.as_view(), + name="sweep_undispatched_executions", + ), ] diff --git a/backend/workflow_manager/workflow_v2/execution_log_internal_views.py b/backend/workflow_manager/workflow_v2/execution_log_internal_views.py index 708376dcc3..0c07b04742 100644 --- a/backend/workflow_manager/workflow_v2/execution_log_internal_views.py +++ b/backend/workflow_manager/workflow_v2/execution_log_internal_views.py @@ -16,6 +16,9 @@ process_log_history_from_cache, ) from workflow_manager.workflow_v2.models import WorkflowExecution +from workflow_manager.workflow_v2.undispatched_sweep import ( + sweep_undispatched_executions, +) logger = logging.getLogger(__name__) @@ -171,3 +174,32 @@ def post(self, request: Request) -> Response: return Response( {"error": str(e)}, status=status.HTTP_500_INTERNAL_SERVER_ERROR ) + + +class SweepUndispatchedExecutionsAPIView(APIView): + """Terminalise executions that were created but never dispatched. + + Called by the PG-queue reaper on its sweep tick. The reaper owns *when* — it is + leader-elected, so exactly one instance sweeps — and the backend owns the *logic*, + mirroring ProcessLogHistoryAPIView above. + + It lives here rather than in the reaper because the reaper deliberately never + touches backend tables: it reads execution state through + ``get_workflow_execution`` and writes through ``update_workflow_execution_status``. + A ``WorkflowExecution`` query inside the worker would break that boundary, and the + sweep also needs the rate limiter — both backend-side. + + Transport-agnostic on purpose: the create-then-dispatch window sits upstream of + ``resolve_transport``, so this recovers Celery-path strands too. + """ + + def post(self, request: Request) -> Response: + """Run one sweep. Returns the number of executions terminalised.""" + try: + swept = sweep_undispatched_executions() + return Response({"swept": swept}) + except Exception as e: + logger.error(f"Error sweeping undispatched executions: {e}", exc_info=True) + return Response( + {"error": str(e)}, status=status.HTTP_500_INTERNAL_SERVER_ERROR + ) diff --git a/backend/workflow_manager/workflow_v2/migrations/0026_workflowexecution_undispatched_idx.py b/backend/workflow_manager/workflow_v2/migrations/0026_workflowexecution_undispatched_idx.py new file mode 100644 index 0000000000..c8b7157d49 --- /dev/null +++ b/backend/workflow_manager/workflow_v2/migrations/0026_workflowexecution_undispatched_idx.py @@ -0,0 +1,132 @@ +"""Add a partial index over UNDISPATCHED workflow executions. + +Serves the undispatched-execution sweep (``workflow_v2/undispatched_sweep.py``), which +terminalises rows left ``PENDING`` because the request died between +``create_workflow_execution`` and ``execute_workflow_async`` — committed but never +queued, so no ``pg_barrier_state`` row exists and the reaper's barrier scan cannot see +them. Measured on integration 2026-08-17: 967 such rows from a single 5-minute load +test whose API tier shed 71% of its traffic. + +The sweep runs every 5 minutes (``WORKER_PG_REAPER_SWEEP_SECONDS``) with:: + + WHERE status = 'PENDING' AND created_at < now() - grace + AND task_id IS NULL AND queue_message_id IS NULL + +Design +------ +* PARTIAL — the index covers only rows that are simultaneously PENDING and + undispatched. In steady state that is ~ZERO rows: an execution leaves PENDING within + seconds of creation, and one of ``task_id`` / ``queue_message_id`` is stamped the + moment dispatch succeeds. So the index is effectively empty, costs almost nothing to + maintain, and only grows when something is actually wrong. +* ``created_at`` as the indexed column — the sweep's only range predicate, and the + ORDER BY for the batch limit. +* WHY A NEW INDEX AT ALL. ``we_active_by_workflow_idx`` (migration 0023) *is* usable + here — ``status = 'PENDING'`` provably implies its ``NOT IN (terminal)`` predicate — + so without this the fallback is a full scan of that (small) partial index rather than + a table seq scan. That is survivable but planner-dependent and grows with the count + of active executions. This makes it an exact index scan over a near-empty index. +* CONCURRENTLY + ``atomic = False`` — ``workflow_execution`` is a multi-million-row + table in production; a plain ``AddIndex`` holds a SHARE lock for the whole build and + blocks writes, i.e. blocks every execution in flight. +* FROZEN LITERAL — migrations must not import app enums, so ``'PENDING'`` is hardcoded. + Drift against ``ExecutionStatus.PENDING`` is caught by + ``tests/test_undispatched_execution_index.py`` (model/enum introspection only, no + test DB), mirroring what ``test_active_execution_index.py`` does for 0023. +* INVALID-INDEX GUARD — ``IF NOT EXISTS`` silently no-ops over a leftover INVALID index + from an interrupted CONCURRENTLY build, and Django would then record this migration as + applied while the index is physically unusable (never read, write overhead only). The + second statement RAISEs in that case, so the failure is loud rather than green-but-broken. + +Deployment +---------- +``CREATE INDEX CONCURRENTLY`` scans the table and can run for minutes on a large table — +long enough to time out a deploy's ``migrate`` step. Prefer building it OUT OF BAND +*before* the deploy; the migration then no-ops via ``IF NOT EXISTS``:: + + CREATE INDEX CONCURRENTLY IF NOT EXISTS we_undispatched_idx + ON workflow_execution (created_at) + WHERE status = 'PENDING' + AND task_id IS NULL + AND queue_message_id IS NULL; + +Then confirm the planner uses it (the partial-index proof depends on the literals +reaching the planner as constants):: + + EXPLAIN SELECT id, workflow_id FROM workflow_execution + WHERE status = 'PENDING' AND created_at < now() - interval '15 minutes' + AND task_id IS NULL AND queue_message_id IS NULL + ORDER BY created_at LIMIT 500; + -- expect: Index Scan using we_undispatched_idx + +Recovery +-------- +An interrupted build leaves an INVALID index that adds write overhead but is never read. +``IF NOT EXISTS`` will NOT rebuild over it (and the guard below RAISEs on it), so drop it +first and re-run:: + + DROP INDEX CONCURRENTLY IF EXISTS we_undispatched_idx; +""" + +from django.db import migrations, models +from django.db.models import Q + +INDEX_NAME = "we_undispatched_idx" + +# Frozen — see the FROZEN LITERAL note above. +PENDING_STATUS = "PENDING" + +_ASSERT_INDEX_VALID = f""" +DO $$ +BEGIN + IF EXISTS ( + SELECT 1 FROM pg_class c + JOIN pg_index i ON i.indexrelid = c.oid + WHERE c.relname = '{INDEX_NAME}' AND NOT i.indisvalid + ) THEN + RAISE EXCEPTION 'Index {INDEX_NAME} exists but is INVALID (a prior CREATE INDEX CONCURRENTLY was interrupted). Drop it and re-run: DROP INDEX CONCURRENTLY IF EXISTS {INDEX_NAME};'; + END IF; +END +$$; +""" + + +class Migration(migrations.Migration): + # CREATE / DROP INDEX CONCURRENTLY cannot run inside a transaction block. + atomic = False + + dependencies = [("workflow_v2", "0025_workflow_workflow_org_modified_idx")] + + operations = [ + migrations.SeparateDatabaseAndState( + database_operations=[ + migrations.RunSQL( + sql=( + f"CREATE INDEX CONCURRENTLY IF NOT EXISTS {INDEX_NAME} " + "ON workflow_execution (created_at) " + f"WHERE status = '{PENDING_STATUS}' " + "AND task_id IS NULL " + "AND queue_message_id IS NULL;" + ), + reverse_sql=f"DROP INDEX CONCURRENTLY IF EXISTS {INDEX_NAME};", + ), + migrations.RunSQL( + sql=_ASSERT_INDEX_VALID, reverse_sql=migrations.RunSQL.noop + ), + ], + state_operations=[ + migrations.AddIndex( + model_name="workflowexecution", + index=models.Index( + fields=["created_at"], + name=INDEX_NAME, + condition=Q( + status=PENDING_STATUS, + task_id__isnull=True, + queue_message_id__isnull=True, + ), + ), + ), + ], + ), + ] diff --git a/backend/workflow_manager/workflow_v2/migrations/0027_workflowexecution_dispatched_at.py b/backend/workflow_manager/workflow_v2/migrations/0027_workflowexecution_dispatched_at.py new file mode 100644 index 0000000000..9d483d152d --- /dev/null +++ b/backend/workflow_manager/workflow_v2/migrations/0027_workflowexecution_dispatched_at.py @@ -0,0 +1,69 @@ +"""Record dispatch as a POSITIVE fact instead of inferring it from absent handles. + +The undispatched sweep (``workflow_v2/undispatched_sweep.py``) decides that an +execution "never dispatched" from ``task_id IS NULL AND queue_message_id IS NULL``. +That is an inference, and it is unsound: three paths in ``workflow_helper`` reach +exactly that state AFTER the message is already on its transport — + +* ``_record_dispatch_handle`` raising and the caller swallowing it, under the comment + "continuing — the orchestrator is already running"; +* the handle coming back empty and the recorder returning without writing; +* a PG handle that will not parse as a bigint, same early return. + +So the sweep could claim a RUNNING execution, mark it ERROR and tell the user "You can +safely run it again". On the Celery transport that self-corrected — the orchestrator ran +regardless and its terminal write superseded the ERROR. On PG it does NOT: both worker +entry points stop on a terminal execution (``general/tasks.py`` returns +``skipped_terminal_execution``; ``file_processing/tasks.py`` raises +``_TerminalExecutionSkip``), so the message is acked and the work is silently DROPPED. + +``dispatched_at`` removes the inference. It is stamped by the dispatcher the moment +dispatch succeeds, BEFORE any handle bookkeeping can fail, so all three paths above are +covered by one write. + +Nullable and additive +--------------------- +No backfill, and none is needed. Every pre-existing row has ``dispatched_at IS NULL`` +whatever its true state, so the sweep predicate KEEPS the two handle checks alongside the +new one:: + + WHERE status = 'PENDING' AND created_at < now() - grace + AND dispatched_at IS NULL + AND task_id IS NULL AND queue_message_id IS NULL + +That is what makes this a ONE-PHASE deploy. During a rolling upgrade an old backend pod +dispatches without stamping; the row then matches ``dispatched_at IS NULL`` but is +excluded by ``task_id IS NOT NULL``. Drop the handle checks and you would need to stamp +everywhere first and switch the predicate in a later release, or sweep live work +mid-deploy. + +The index over that predicate is swapped separately in 0028, which needs +``atomic = False`` for ``CREATE INDEX CONCURRENTLY``; keeping the column addition here +means this migration stays a fast, transactional, single-purpose change. +""" + +from django.db import migrations, models + + +class Migration(migrations.Migration): + dependencies = [ + ("workflow_v2", "0026_workflowexecution_undispatched_idx"), + ] + + operations = [ + migrations.AddField( + model_name="workflowexecution", + name="dispatched_at", + field=models.DateTimeField( + null=True, + blank=True, + db_comment=( + "Set when the orchestrator task was handed to its transport. NULL means " + "never dispatched. The POSITIVE fact the undispatched sweep needs: absent " + "task_id/queue_message_id also occurs for executions that ARE running, " + "because three paths in workflow_helper return without recording a handle " + "after the message is already enqueued." + ), + ), + ), + ] diff --git a/backend/workflow_manager/workflow_v2/migrations/0028_undispatched_idx_dispatched_at.py b/backend/workflow_manager/workflow_v2/migrations/0028_undispatched_idx_dispatched_at.py new file mode 100644 index 0000000000..e6cd88710b --- /dev/null +++ b/backend/workflow_manager/workflow_v2/migrations/0028_undispatched_idx_dispatched_at.py @@ -0,0 +1,136 @@ +"""Re-key the undispatched partial index on ``dispatched_at``. + +0027 added the column and explains why the inference it replaces was unsound. This +swaps the index that serves the sweep so the new predicate stays an index scan rather +than degrading to a sequential scan of a multi-million-row table. + +The sweep now runs:: + + WHERE status = 'PENDING' AND created_at < now() - grace + AND dispatched_at IS NULL + AND task_id IS NULL AND queue_message_id IS NULL + +The handle checks are KEPT deliberately — see 0027 for why they are what make the +rollout single-phase — so the index condition carries all three. + +Ordering matters here. The NEW index is created before the OLD one is dropped, so the +predicate is served throughout: a deploy that is interrupted between the two operations +leaves both present, which costs a little write overhead and nothing else. Dropping +first would leave the sweep unindexed for the length of the build. + +Same operational rules as 0026, which this mirrors: + +* ``atomic = False`` — CREATE / DROP INDEX CONCURRENTLY cannot run in a transaction. +* Prefer building OUT OF BAND before the deploy; ``IF NOT EXISTS`` then no-ops:: + + CREATE INDEX CONCURRENTLY IF NOT EXISTS we_undispatched_dispatch_idx + ON workflow_execution (created_at) + WHERE status = 'PENDING' + AND dispatched_at IS NULL + AND task_id IS NULL + AND queue_message_id IS NULL; + +* An interrupted build leaves an INVALID index that is never read but still costs + writes. ``IF NOT EXISTS`` will not rebuild over it and the guard below RAISEs on it, + so drop and re-run:: + + DROP INDEX CONCURRENTLY IF EXISTS we_undispatched_dispatch_idx; + +* Confirm the planner uses it — the partial-index proof depends on the literals + reaching the planner as constants:: + + EXPLAIN SELECT id, workflow_id FROM workflow_execution + WHERE status = 'PENDING' AND created_at < now() - interval '1 hour' + AND dispatched_at IS NULL + AND task_id IS NULL AND queue_message_id IS NULL + ORDER BY created_at LIMIT 500; + -- expect: Index Scan using we_undispatched_dispatch_idx + +A NEW NAME rather than a rebuild under the old one: an index cannot be redefined in +place, and reusing ``we_undispatched_idx`` would mean dropping before creating, leaving +the window this ordering exists to avoid. +""" + +from django.db import migrations, models +from django.db.models import Q + +OLD_INDEX_NAME = "we_undispatched_idx" +INDEX_NAME = "we_undispatched_dispatch_idx" + +# Frozen literal, tied to ExecutionStatus by +# tests/test_undispatched_execution_index.py — see 0026. +PENDING_STATUS = "PENDING" + +_ASSERT_INDEX_VALID = f""" +DO $$ +BEGIN + IF EXISTS ( + SELECT 1 FROM pg_class c + JOIN pg_index i ON i.indexrelid = c.oid + WHERE c.relname = '{INDEX_NAME}' AND NOT i.indisvalid + ) THEN + RAISE EXCEPTION 'Index {INDEX_NAME} exists but is INVALID (a prior CREATE INDEX CONCURRENTLY was interrupted). Drop it and re-run: DROP INDEX CONCURRENTLY IF EXISTS {INDEX_NAME};'; + END IF; +END +$$; +""" + + +class Migration(migrations.Migration): + # CREATE / DROP INDEX CONCURRENTLY cannot run inside a transaction block. + atomic = False + + dependencies = [("workflow_v2", "0027_workflowexecution_dispatched_at")] + + operations = [ + migrations.SeparateDatabaseAndState( + database_operations=[ + # Create the replacement FIRST — the predicate stays served throughout. + migrations.RunSQL( + sql=( + f"CREATE INDEX CONCURRENTLY IF NOT EXISTS {INDEX_NAME} " + "ON workflow_execution (created_at) " + f"WHERE status = '{PENDING_STATUS}' " + "AND dispatched_at IS NULL " + "AND task_id IS NULL " + "AND queue_message_id IS NULL;" + ), + reverse_sql=f"DROP INDEX CONCURRENTLY IF EXISTS {INDEX_NAME};", + ), + migrations.RunSQL( + sql=_ASSERT_INDEX_VALID, reverse_sql=migrations.RunSQL.noop + ), + # Only then retire the old one. Reverse recreates it, so a rollback of + # this migration leaves the pre-0028 predicate indexed. + migrations.RunSQL( + sql=f"DROP INDEX CONCURRENTLY IF EXISTS {OLD_INDEX_NAME};", + reverse_sql=( + f"CREATE INDEX CONCURRENTLY IF NOT EXISTS {OLD_INDEX_NAME} " + "ON workflow_execution (created_at) " + f"WHERE status = '{PENDING_STATUS}' " + "AND task_id IS NULL " + "AND queue_message_id IS NULL;" + ), + ), + ], + state_operations=[ + migrations.RemoveIndex( + model_name="workflowexecution", + name=OLD_INDEX_NAME, + ), + migrations.AddIndex( + model_name="workflowexecution", + index=models.Index( + fields=["created_at"], + name=INDEX_NAME, + condition=Q( + status=PENDING_STATUS, + dispatched_at__isnull=True, + task_id__isnull=True, + queue_message_id__isnull=True, + ), + ), + ), + ], + ), + ] diff --git a/backend/workflow_manager/workflow_v2/models/execution.py b/backend/workflow_manager/workflow_v2/models/execution.py index a0ec40ae64..d6082aa423 100644 --- a/backend/workflow_manager/workflow_v2/models/execution.py +++ b/backend/workflow_manager/workflow_v2/models/execution.py @@ -213,6 +213,17 @@ class Type(models.TextChoices): db_comment="Details of encountered errors", ) attempts = models.IntegerField(default=0, db_comment="number of attempts taken") + dispatched_at = models.DateTimeField( + null=True, + blank=True, + db_comment=( + "Set when the orchestrator task was handed to its transport. NULL means " + "never dispatched. The POSITIVE fact the undispatched sweep needs: absent " + "task_id/queue_message_id also occurs for executions that ARE running, " + "because three paths in workflow_helper return without recording a handle " + "after the message is already enqueued." + ), + ) execution_time = models.FloatField(default=0, db_comment="execution time in seconds") tags = models.ManyToManyField(Tag, related_name="workflow_executions", blank=True) @@ -239,6 +250,28 @@ class Meta: name="we_active_by_workflow_idx", condition=~Q(status__in=["COMPLETED", "STOPPED", "ERROR"]), ), + # Partial index over UNDISPATCHED executions — see migration 0026. + # Serves the undispatched-execution sweep (undispatched_sweep.py), which + # terminalises rows left PENDING because the request died between + # create_workflow_execution and execute_workflow_async. Effectively EMPTY + # in steady state: an execution leaves PENDING within seconds, and one of + # task_id / queue_message_id is stamped the moment dispatch succeeds — so + # it costs almost nothing and only grows when something is wrong. + # we_active_by_workflow_idx above is *usable* for the same predicate + # (PENDING implies NOT IN terminal) but is keyed on workflow_id, which the + # sweep does not filter on — so without this the sweep falls back to a full + # scan of that index. Literals are frozen in the migration and tied to + # ExecutionStatus by tests/test_undispatched_execution_index.py. + models.Index( + fields=["created_at"], + name="we_undispatched_dispatch_idx", + condition=Q( + status="PENDING", + dispatched_at__isnull=True, + task_id__isnull=True, + queue_message_id__isnull=True, + ), + ), ] @property diff --git a/backend/workflow_manager/workflow_v2/tests/test_step_execution_transport.py b/backend/workflow_manager/workflow_v2/tests/test_step_execution_transport.py new file mode 100644 index 0000000000..2e179ad0e1 --- /dev/null +++ b/backend/workflow_manager/workflow_v2/tests/test_step_execution_transport.py @@ -0,0 +1,120 @@ +"""Tests for the PG guard on single-step execution (UN-3445 acceptance gate). + +Single-step is the one entry path whose fan-out was never transport-gated: it +reaches the Celery chord in ``process_input_files`` directly, while the normal +path's PG fan-out lives in the general worker. With the flag on and the Celery +file_processing workers scaled to zero, those batches would never be consumed +and the execution would hang in EXECUTING with no PG reaper aware of it. + +So the guard has to hold in *both* directions, and that is what these pin: +Celery must still run step execution unchanged (it is what staging and +production run today), and PG must refuse it loudly rather than dispatch into a +queue with no consumer. + +DB-free: WorkflowExecution, resolve_transport and run_workflow are all mocked. +""" + +from unittest.mock import MagicMock, patch + +import pytest +from workflow_manager.workflow_v2.exceptions import WorkflowExecutionError +from workflow_manager.workflow_v2.models import Workflow +from workflow_manager.workflow_v2.workflow_helper import WorkflowHelper + +_TRANSPORT = "workflow_manager.workflow_v2.workflow_helper.resolve_transport" +_EXECUTION = "workflow_manager.workflow_v2.workflow_helper.WorkflowExecution" +_RUN = "workflow_manager.workflow_v2.workflow_helper.WorkflowHelper.run_workflow" + +# The production comparison is `execution_action is ...START.value`, so the test +# must pass that exact object rather than an equal string literal. +_START = Workflow.ExecutionAction.START.value +_EXECUTION_ID = "11111111-1111-1111-1111-111111111111" + + +def _workflow(): + workflow = MagicMock() + workflow.organization.organization_id = "org_acme" + return workflow + + +def _patched_execution(): + """Mock WorkflowExecution, keeping DoesNotExist a real exception class. + + A bare MagicMock attribute is not catchable, so the production + ``except WorkflowExecution.DoesNotExist`` would raise TypeError instead. + """ + model = MagicMock() + model.DoesNotExist = type("DoesNotExist", (Exception,), {}) + model.objects.get.return_value = MagicMock() + return model + + +class TestStepExecutionTransportGuard: + def test_pg_transport_refuses_instead_of_dispatching_to_celery(self): + with patch(_EXECUTION, _patched_execution()), patch( + _TRANSPORT, return_value="pg_queue" + ), patch(_RUN) as run: + with pytest.raises(WorkflowExecutionError, match="not supported"): + WorkflowHelper.step_execution( + workflow=_workflow(), + execution_action=_START, + execution_id=_EXECUTION_ID, + ) + # The point of the guard: nothing reaches the chord. Asserting the raise + # alone would still pass if run_workflow had already fanned out. + run.assert_not_called() + + def test_celery_transport_still_runs_step_execution(self): + """Flag-off is what staging and production run — it must be untouched.""" + with patch(_EXECUTION, _patched_execution()), patch( + _TRANSPORT, return_value="celery" + ), patch(_RUN, return_value="ran") as run: + result = WorkflowHelper.step_execution( + workflow=_workflow(), + execution_action=_START, + execution_id=_EXECUTION_ID, + ) + assert result == "ran" + assert run.call_args.kwargs["single_step"] is True + + def test_transport_is_resolved_on_the_execution_id(self): + """entity_id must be the execution id — it is what Flipt buckets on, and + what keeps one execution from re-bucketing across transports.""" + with patch(_EXECUTION, _patched_execution()), patch( + _TRANSPORT, return_value="celery" + ) as resolve, patch(_RUN): + WorkflowHelper.step_execution( + workflow=_workflow(), + execution_action=_START, + execution_id=_EXECUTION_ID, + ) + assert resolve.call_args.kwargs["execution_id"] == _EXECUTION_ID + assert resolve.call_args.kwargs["organization_id"] == "org_acme" + + def test_no_execution_id_creates_a_step_execution_without_resolving(self): + """The START-without-id branch only mints a row; it dispatches nothing, so + it must not consume a Flipt evaluation (and must not be blocked).""" + with patch(_TRANSPORT) as resolve, patch.object( + WorkflowHelper, "create_and_make_execution_response", return_value="created" + ): + result = WorkflowHelper.step_execution( + workflow=_workflow(), execution_action=_START, execution_id=None + ) + assert result == "created" + resolve.assert_not_called() + + def test_missing_execution_falls_back_to_creating_one(self): + """Pre-existing behaviour: a stale execution_id re-mints rather than 404s. + Pinned so the guard's placement inside the try/except cannot change it.""" + model = _patched_execution() + model.objects.get.side_effect = model.DoesNotExist + with patch(_EXECUTION, model), patch(_TRANSPORT) as resolve, patch.object( + WorkflowHelper, "create_and_make_execution_response", return_value="created" + ): + result = WorkflowHelper.step_execution( + workflow=_workflow(), + execution_action=_START, + execution_id=_EXECUTION_ID, + ) + assert result == "created" + resolve.assert_not_called() diff --git a/backend/workflow_manager/workflow_v2/tests/test_undispatched_execution_index.py b/backend/workflow_manager/workflow_v2/tests/test_undispatched_execution_index.py new file mode 100644 index 0000000000..4268a2f7fa --- /dev/null +++ b/backend/workflow_manager/workflow_v2/tests/test_undispatched_execution_index.py @@ -0,0 +1,149 @@ +"""Guard: the undispatched-execution partial index stays in sync with reality. + +``we_undispatched_dispatch_idx`` (``WorkflowExecution.Meta.indexes``) hardcodes the literal +``'PENDING'`` and the two dispatch-handle columns. The same predicate is written a +second time in migration 0026's ``RunSQL``, and a third time as the sweep's WHERE clause +in ``undispatched_sweep.py`` — migrations cannot import app enums, so the literal cannot +simply reference ``ExecutionStatus``. + +Three copies of one predicate is exactly how an index silently stops matching the query +it exists for: the sweep keeps working, just without the index, and nobody notices until +it is scanning a multi-million-row table every 5 minutes. These assert the copies agree. + +Model introspection only — no test database required, so this runs in the unit tier +alongside ``test_active_execution_index.py``, which does the same job for 0023. +""" + +from __future__ import annotations + +import os +import re +from pathlib import Path + +import django +from django.apps import apps + +os.environ.setdefault("DJANGO_SETTINGS_MODULE", "backend.settings.test") +if not apps.ready: + django.setup() + +from unstract.core.data_models import ExecutionStatus # noqa: E402 + +INDEX_NAME = "we_undispatched_dispatch_idx" +_MIGRATION = ( + Path(__file__).resolve().parent.parent + / "migrations" + / "0026_workflowexecution_undispatched_idx.py" +) + + +def _index(): + model = apps.get_model("workflow_v2", "WorkflowExecution") + return next((i for i in model._meta.indexes if i.name == INDEX_NAME), None) + + +def _condition_children() -> dict: + index = _index() + assert index is not None, f"{INDEX_NAME} is missing from WorkflowExecution.Meta" + return {c[0]: c[1] for c in index.condition.children if isinstance(c, tuple)} + + +class TestTheModelIndexMatchesTheEnum: + def test_the_status_literal_is_the_pending_enum_value(self): + """If ExecutionStatus.PENDING's *value* ever changes, the index silently stops + matching every row the sweep looks for.""" + assert _condition_children()["status"] == ExecutionStatus.PENDING.value + + def test_all_three_dispatch_signals_are_in_the_predicate(self): + """Dropping any makes the index non-matching for the sweep's WHERE clause + (Postgres needs the query predicate to imply the index predicate), so the sweep + would quietly fall back to scanning. + + Three, not two: `dispatched_at` is the authoritative test (0027), and the two + handle checks are retained so the rollout is single-phase — a row dispatched by + an old pod mid-deploy has no stamp but does have a handle, and must not be swept. + """ + children = _condition_children() + assert children.get("dispatched_at__isnull") is True + assert children.get("task_id__isnull") is True + assert children.get("queue_message_id__isnull") is True + + def test_it_is_keyed_on_created_at(self): + """The sweep's only range predicate and its ORDER BY for the batch limit.""" + assert _index().fields == ["created_at"] + + +class TestTheMigrationMatchesTheModel: + """The migration's raw SQL is the copy that actually builds the index; the model's + Index() is only Django state. They must not disagree. + """ + + def test_the_migration_sql_uses_the_same_status_literal(self): + sql = _MIGRATION.read_text() + assert f"status = '{ExecutionStatus.PENDING.value}'" in sql + + def test_the_migration_sql_carries_both_handle_conditions(self): + sql = _MIGRATION.read_text() + assert "task_id IS NULL" in sql + assert "queue_message_id IS NULL" in sql + + def test_the_migration_builds_concurrently_and_is_non_atomic(self): + """workflow_execution is multi-million-row in production — a plain AddIndex + holds a SHARE lock for the whole build and blocks every in-flight execution. + """ + sql = _MIGRATION.read_text() + assert "CREATE INDEX CONCURRENTLY IF NOT EXISTS" in sql + assert re.search(r"^\s*atomic\s*=\s*False", sql, re.MULTILINE) + + def test_it_guards_against_a_leftover_invalid_index(self): + """IF NOT EXISTS no-ops over an INVALID index from an interrupted build, and + Django would record the migration as applied over something unusable. + """ + sql = _MIGRATION.read_text() + assert "RAISE EXCEPTION" in sql + assert "indisvalid" in sql + + +class TestTheSweepQueryMatchesTheIndex: + def test_the_sweep_predicate_uses_the_same_columns(self): + """The third copy. If the sweep's SQL drifts from the index predicate, Postgres + can no longer prove the implication and the partial index becomes unusable — + the sweep still returns correct rows, just by scanning. + """ + from workflow_manager.workflow_v2 import undispatched_sweep + + sql = undispatched_sweep._CLAIM_SQL + assert "task_id IS NULL" in sql + assert "queue_message_id IS NULL" in sql + assert "created_at <" in sql + assert "ORDER BY created_at" in sql + + def test_the_predicate_appears_TWICE__inner_select_and_outer_recheck(self): + """The race guard, and the one thing a substring check silently misses. + + The inner SELECT picks and locks candidates; the OUTER WHERE re-carries the same + predicate so Postgres re-evaluates it at write time under those locks. Drop it + from the outer clause and the statement still looks right — the strings are all + still present, courtesy of the inner SELECT — but a row dispatched between the + two is now marked ERROR *while it runs*. + + Verified: removing the outer re-check passed every other test in this file. + """ + from workflow_manager.workflow_v2 import undispatched_sweep + + sql = undispatched_sweep._CLAIM_SQL + for clause in ("task_id IS NULL", "queue_message_id IS NULL"): + assert sql.count(clause) >= 2, ( + f"{clause!r} appears {sql.count(clause)}x — it must be in BOTH the " + "inner SELECT and the outer re-check, or the claim is racy" + ) + # The status check likewise: %s placeholders, so count the column instead. + assert sql.count("status = %s") >= 2, "status must be re-checked in the UPDATE" + + def test_it_locks_candidates_with_skip_locked(self): + """Without FOR UPDATE the outer re-check reads unlocked rows and the race + reopens; without SKIP LOCKED an overlapping sweeper blocks instead of yielding. + """ + from workflow_manager.workflow_v2 import undispatched_sweep + + assert "FOR UPDATE SKIP LOCKED" in undispatched_sweep._CLAIM_SQL diff --git a/backend/workflow_manager/workflow_v2/tests/test_undispatched_sweep.py b/backend/workflow_manager/workflow_v2/tests/test_undispatched_sweep.py new file mode 100644 index 0000000000..e69c2d42d0 --- /dev/null +++ b/backend/workflow_manager/workflow_v2/tests/test_undispatched_sweep.py @@ -0,0 +1,358 @@ +"""PENDING is not terminal — an undispatched execution must not stay that way. + +Guards the sweep that closes the create-then-dispatch window (see the module +docstring). The two failure modes that matter are opposite in direction, and both are +pinned here: + + * FAILING TO SWEEP leaves an execution that can never reach COMPLETED or ERROR — + invisible failure. 967 of them appeared from one 5-minute load test. + * SWEEPING TOO EAGERLY marks a LIVE execution ERROR. Far worse, so the young-row, + dispatched-row and raced-row cases are asserted explicitly rather than assumed. +""" + +from __future__ import annotations + +import uuid +from datetime import timedelta + +import pytest +from django.utils import timezone + +from unstract.core.data_models import ExecutionStatus + +# NOT module-level: TestUserFacingMessage asserts a pure string contract and needs +# no database, so it must stay in the UNIT tier. conftest auto-marks django_db +# tests as `integration`, which the unit lane excludes — a module-level marker +# would have hidden the message assertions behind live Postgres for no reason. + + +def _execution(**overrides): + """A PENDING, never-dispatched execution old enough to sweep. + + A workflow is REQUIRED, not incidental: ``WorkflowExecution.save()`` calls + ``_handle_execution_cache()``, which dereferences ``self.workflow.id``. Built + without one, every test here died with ``AttributeError: 'NoneType' object has no + attribute 'id'`` before reaching its assertion — which is what CI reported the + moment this PR left draft and the integration tier actually ran. + """ + from workflow_manager.workflow_v2.models import Workflow, WorkflowExecution + + fields = { + "id": uuid.uuid4(), + "workflow": Workflow.objects.create( + workflow_name=f"wf-sweep-{uuid.uuid4().hex[:8]}" + ), + "status": ExecutionStatus.PENDING.value, + "task_id": None, + "queue_message_id": None, + } + fields.update(overrides) + created_at = fields.pop("created_at", timezone.now() - timedelta(hours=2)) + ex = WorkflowExecution.objects.create(**fields) + # created_at is auto_now_add; bypass it with a direct UPDATE so age is controllable. + WorkflowExecution.objects.filter(id=ex.id).update(created_at=created_at) + ex.refresh_from_db() + return ex + + +@pytest.mark.django_db +class TestSweepsWhatNothingElseOwns: + def test_an_aged_undispatched_execution_is_marked_error(self): + from workflow_manager.workflow_v2.models import WorkflowExecution + from workflow_manager.workflow_v2.undispatched_sweep import ( + UNDISPATCHED_ERROR_MESSAGE, + sweep_undispatched_executions, + ) + + ex = _execution() + assert sweep_undispatched_executions() == 1 + + ex.refresh_from_db() + assert ex.status == ExecutionStatus.ERROR.value + assert ex.error_message == UNDISPATCHED_ERROR_MESSAGE + assert ( + WorkflowExecution.objects.filter(status=ExecutionStatus.PENDING.value).count() + == 0 + ) + + def test_it_is_idempotent(self): + """Runs on every tick — a second pass must find nothing, not re-write.""" + from workflow_manager.workflow_v2.undispatched_sweep import ( + sweep_undispatched_executions, + ) + + _execution() + assert sweep_undispatched_executions() == 1 + assert sweep_undispatched_executions() == 0 + + +@pytest.mark.django_db +class TestNeverTouchesALiveExecution: + """The dangerous direction. Each case is a way a live run could be killed.""" + + def test_a_young_execution_is_left_alone(self): + """Dispatch follows creation within seconds; age is the only separator.""" + from workflow_manager.workflow_v2.undispatched_sweep import ( + sweep_undispatched_executions, + ) + + ex = _execution(created_at=timezone.now() - timedelta(seconds=30)) + assert sweep_undispatched_executions() == 0 + ex.refresh_from_db() + assert ex.status == ExecutionStatus.PENDING.value + + def test_a_pg_dispatched_execution_is_left_alone(self): + """queue_message_id set => the PG queue holds it; the reaper owns it now.""" + from workflow_manager.workflow_v2.undispatched_sweep import ( + sweep_undispatched_executions, + ) + + ex = _execution(queue_message_id=4242) + assert sweep_undispatched_executions() == 0 + ex.refresh_from_db() + assert ex.status == ExecutionStatus.PENDING.value + + def test_a_celery_dispatched_execution_is_left_alone(self): + """task_id set => Celery holds it. The window is upstream of + resolve_transport, so the predicate must be transport-aware or a flag-off + deployment would have its live executions terminalised. + """ + from workflow_manager.workflow_v2.undispatched_sweep import ( + sweep_undispatched_executions, + ) + + ex = _execution(task_id=uuid.uuid4()) + assert sweep_undispatched_executions() == 0 + ex.refresh_from_db() + assert ex.status == ExecutionStatus.PENDING.value + + @pytest.mark.parametrize( + "status", + [ + ExecutionStatus.EXECUTING.value, + ExecutionStatus.COMPLETED.value, + ExecutionStatus.ERROR.value, + ExecutionStatus.STOPPED.value, + ], + ) + def test_only_pending_is_swept(self, status): + """EXECUTING is running; the terminal three are already resolved. Rewriting + a COMPLETED run's error_message would be corruption, not recovery. + """ + from workflow_manager.workflow_v2.undispatched_sweep import ( + sweep_undispatched_executions, + ) + + ex = _execution(status=status) + assert sweep_undispatched_executions() == 0 + ex.refresh_from_db() + assert ex.status == status + assert ex.error_message == "" + + def test_a_row_dispatched_mid_sweep_is_not_clobbered(self): + """The read-then-write race, which is why the UPDATE re-carries the predicate. + + Simulates dispatch landing between candidate selection and the write: the + row is passed in as a candidate id but no longer matches, so Postgres skips + it. A plain `filter(id__in=ids).update(...)` would mark a live run ERROR. + """ + from workflow_manager.workflow_v2.models import WorkflowExecution + from workflow_manager.workflow_v2.undispatched_sweep import ( + sweep_undispatched_executions, + ) + + ex = _execution() + WorkflowExecution.objects.filter(id=ex.id).update(queue_message_id=99) + + assert sweep_undispatched_executions() == 0 + ex.refresh_from_db() + assert ex.status == ExecutionStatus.PENDING.value + + +@pytest.mark.django_db +class TestDispatchIsAPositiveFactNotAnInference: + """The fix for the defect the whole file's caution was working around. + + The sweep used to infer "never dispatched" from `task_id IS NULL AND + queue_message_id IS NULL`. Three paths in workflow_helper reach that state on a + RUNNING execution — a swallowed exception around _record_dispatch_handle, an empty + handle, an unparseable PG handle — so the sweep could mark live work ERROR and, on + the PG transport, get it silently DROPPED (both worker entry points stop on a + terminal execution rather than superseding it). + + `dispatched_at` is stamped the instant dispatch succeeds, upstream of all three. + """ + + def test_a_dispatched_row_is_never_swept_even_with_no_handles(self): + """THE regression. This is exactly the shape the three failure paths produce: + dispatch succeeded, neither handle was recorded. Before dispatched_at this row + was swept to ERROR while its worker was still running. + """ + from workflow_manager.workflow_v2.undispatched_sweep import ( + sweep_undispatched_executions, + ) + + ex = _execution(dispatched_at=timezone.now() - timedelta(hours=2)) + assert sweep_undispatched_executions() == 0 + ex.refresh_from_db() + assert ex.status == ExecutionStatus.PENDING.value + + def test_an_unstamped_row_with_a_handle_is_still_not_swept(self): + """The single-phase rollout guarantee. During a rolling deploy an OLD backend + pod dispatches without stamping, so the row has no dispatched_at — but it does + have a handle. Keeping the handle checks alongside the new one is what lets this + ship in one release instead of two. + """ + from workflow_manager.workflow_v2.undispatched_sweep import ( + sweep_undispatched_executions, + ) + + ex = _execution(task_id=uuid.uuid4()) + assert sweep_undispatched_executions() == 0 + ex.refresh_from_db() + assert ex.status == ExecutionStatus.PENDING.value + + def test_a_genuinely_undispatched_row_is_still_swept(self): + """The fix must not disable the sweep. A row with no stamp and no handle is the + 967-orphan case this exists for, and it must still terminalise. + """ + from workflow_manager.workflow_v2.undispatched_sweep import ( + sweep_undispatched_executions, + ) + + ex = _execution() + assert sweep_undispatched_executions() == 1 + ex.refresh_from_db() + assert ex.status == ExecutionStatus.ERROR.value + + +class TestGracePeriodIsSizedForTheQueue: + """The grace period is the ONLY thing separating "abandoned" from "dispatched but + not yet recorded", and on PG that distinction has teeth. + + A row this sweep terminalises while its message is still queued does not get run + late — both PG worker entry points STOP on a terminal execution and ack the message + (general/tasks.py returns skipped_terminal_execution; file_processing/tasks.py + raises _TerminalExecutionSkip), so the work is dropped. On Celery the orchestrator + ran regardless and its own terminal write superseded the ERROR, which is why this + was survivable before the PG transport and is not after it. + + An hour is chosen to sit clear of any realistic dequeue latency. Lowering it back + toward the queue's p99 wait re-opens that window, so this pins the value: change it + deliberately, with a latency figure in hand, not as a tidy-up. + """ + + def test_the_grace_period_is_an_hour(self): + from workflow_manager.workflow_v2.undispatched_sweep import ( + DEFAULT_MIN_AGE_SECONDS, + ) + + assert DEFAULT_MIN_AGE_SECONDS == 3600 + + def test_it_stays_clear_of_the_barrier_stuck_timeout(self): + """Must remain well under ~2.5h or this sweep and the barrier reaper start + contending for the same rows — the disjointness the two sweeps depend on. + """ + from workflow_manager.workflow_v2.undispatched_sweep import ( + DEFAULT_MIN_AGE_SECONDS, + ) + + assert DEFAULT_MIN_AGE_SECONDS < 9000 + + +class TestNeverDeletesStagedInput: + """The sweep must not perform any IRREVERSIBLE cleanup, because its "never + dispatched" test is an INFERENCE and the inference is unsound. + + Both handles being NULL is reached by three paths in ``workflow_helper`` AFTER the + message is on its transport: ``_record_dispatch_handle`` raising and being swallowed + by the caller ("continuing — the orchestrator is already running"), an empty handle + returning early, and a PG handle that will not parse returning early. So a claimed + row may be a LIVE execution. + + Marking it ERROR is survivable — the running worker's own terminal write supersedes + it, and error→completed is explicitly permitted by the status guard. Deleting its + staged input is not: the worker is still going to read it. + + This pins the ABSENCE of that call. Re-adding it is only safe once dispatch is a + positive fact (a stamped ``dispatched_at``) rather than an inferred absence. + """ + + def test_the_sweep_module_does_not_delete_api_storage(self): + import inspect + + from workflow_manager.workflow_v2 import undispatched_sweep + + source = inspect.getsource(undispatched_sweep) + executable = "\n".join( + line for line in source.splitlines() if not line.lstrip().startswith("#") + ) + assert "delete_api_storage_dir" not in executable + + def test_releasing_resources_still_frees_the_rate_limit_slot(self): + """The reversible half must survive the removal — a held slot consumes the + org's API-deployment concurrency budget until the limiter TTL expires it. + """ + import inspect + + from workflow_manager.workflow_v2 import undispatched_sweep + + source = inspect.getsource(undispatched_sweep._release_abandoned_resources) + assert "release_slot" in source + + +class TestUserFacingMessage: + """`error_message` is rendered in the UI — ExecutionSerializer uses `exclude`, + not `fields`, so every unlisted model field is serialized to customers. + """ + + def test_it_fits_the_column_with_the_ref_intact(self): + """CharField(max_length=256) truncates silently, which would eat the trailing + ref code and leave support without the greppable handle. + """ + from workflow_manager.workflow_v2.models.execution import ( + EXECUTION_ERROR_LENGTH, + ) + from workflow_manager.workflow_v2.undispatched_sweep import ( + UNDISPATCHED_ERROR_MESSAGE, + ) + + assert len(UNDISPATCHED_ERROR_MESSAGE) <= EXECUTION_ERROR_LENGTH + assert UNDISPATCHED_ERROR_MESSAGE.rstrip().endswith("(ref: EXEC_NOT_STARTED)") + + def test_it_names_no_internals(self): + """A customer reading this must not meet our vocabulary. The existing reaper + strings ("the final aggregating callback never fired before the barrier + expired") are the anti-pattern this guards against. + """ + from workflow_manager.workflow_v2.undispatched_sweep import ( + UNDISPATCHED_ERROR_MESSAGE, + ) + + leaked = [ + w + for w in ( + "queue", + "barrier", + "dispatch", + "gateway", + "celery", + "worker", + "reaper", + "502", + "null", + ) + if w in UNDISPATCHED_ERROR_MESSAGE.lower() + ] + assert not leaked, f"internal term(s) in a user-facing message: {leaked}" + + def test_it_answers_what_the_user_needs(self): + """Did anything run, is my data affected, what do I do.""" + from workflow_manager.workflow_v2.undispatched_sweep import ( + UNDISPATCHED_ERROR_MESSAGE, + ) + + low = UNDISPATCHED_ERROR_MESSAGE.lower() + assert "did not start" in low + assert "no files were processed" in low + assert "run it again" in low diff --git a/backend/workflow_manager/workflow_v2/undispatched_sweep.py b/backend/workflow_manager/workflow_v2/undispatched_sweep.py new file mode 100644 index 0000000000..cc3d77f356 --- /dev/null +++ b/backend/workflow_manager/workflow_v2/undispatched_sweep.py @@ -0,0 +1,303 @@ +"""Terminalise executions that were created but never dispatched to a queue. + +``PENDING`` is not a terminal state, so every execution must eventually reach +``COMPLETED`` or ``ERROR``. One case had no owner: + + deployment_helper.py:236 create_workflow_execution(...) -> commits a PENDING row + ... file handling, validation ... + deployment_helper.py:310 execute_workflow_async(...) -> dispatches + +An abort between those two — client disconnect, gateway timeout, pod eviction, OOM — +leaves a committed ``PENDING`` row that was never queued. Nothing recovers it: + +* the **reaper** recovers stranded work by scanning ``pg_barrier_state``, and a barrier + only exists once a batch has been dispatched. No dispatch, no barrier, invisible. +* ``execute_workflow_async`` marks *dispatch failures* ERROR, but never runs at all if + the request dies before reaching it. + +Observed on integration 2026-08-17: an 800-user load test saturated the API tier +(2 backend pods), 71% of requests got a 502 from the load balancer, and **967 +executions were left PENDING** — 2914 execution rows from 2444 requests, because the +row outlived the request. ``pg_queue_message`` was completely empty, confirming the +work was never enqueued rather than enqueued-and-lost. + +**Not PG-specific.** The create-then-dispatch ordering sits *upstream* of +``resolve_transport``, so the Celery path has the identical window. The predicate here +is transport-aware for exactly that reason. + +**The predicate.** ``workflow_helper.py:566/570`` stamps ``queue_message_id`` (PG) or +``task_id`` (Celery) immediately after a successful dispatch, and the model documents +that the other stays NULL. So ``PENDING`` + both handles NULL + older than the grace +period means the dispatch USUALLY never happened — but see the note in +_release_abandoned_resources: three paths in workflow_helper reach that exact +state AFTER dispatch, which is why this sweep does no irreversible cleanup and +why the real fix is a positive dispatched_at rather than an inferred absence. +Do not shorten the grace period or reinstate cleanup on the strength of this +paragraph alone. Nominally it means the dispatch never happened — no JSONB probing, no cross-table joins, and +correct under either transport. +""" + +from __future__ import annotations + +import logging +import os + +from django.utils import timezone + +from unstract.core.data_models import ExecutionStatus + +logger = logging.getLogger(__name__) + +# Grace period before an undispatched execution is considered abandoned. +# +# Dispatch normally follows row creation within seconds, so this is generous by two +# orders of magnitude. It has to be: the ONLY thing separating "abandoned" from "about +# to be dispatched" is elapsed time, and terminalising a live execution is far worse +# than leaving a dead one a while longer. Well under the barrier stuck-timeout (~2.5h) +# so the two sweeps never contend for the same row. +# RAISED 900 -> 3600 deliberately, and the reason is PG-specific. +# +# The predicate cannot distinguish "never dispatched" from "dispatched, handle not +# recorded, message still sitting in the queue" — see the note in +# _release_abandoned_resources for the three paths that produce the second state. On +# CELERY that mismatch was survivable: the orchestrator runs on the row regardless and +# its terminal write supersedes the sweep's ERROR. On PG it is not. Both worker entry +# points STOP on a terminal execution (general/tasks.py returns +# skipped_terminal_execution; file_processing/tasks.py raises _TerminalExecutionSkip), +# so a row this sweep terminalises while its message is still queued gets acked and +# dropped rather than run. +# +# That failure needs a slow queue AND an unrecorded handle at once. An hour is chosen to +# sit well clear of any realistic dequeue latency, which is what shrinks the overlap to +# near-nothing without new machinery. It is a mitigation, not the fix: the fix is to make +# dispatch a POSITIVE fact (a stamped dispatched_at) so the predicate stops guessing. +# +# Still well under the barrier stuck-timeout (~2.5h), so the two sweeps never contend for +# the same row. Cost of the raise: a genuinely undispatched execution shows PENDING for +# up to an hour before it errors. Override per-environment with the env var if a +# deployment's queue latency justifies something tighter. +_MIN_AGE_ENV = "UNDISPATCHED_EXECUTION_GRACE_SECONDS" +DEFAULT_MIN_AGE_SECONDS = 3600 # 1 hour + +# Bounds one sweep so a large backlog (a 502 storm leaves hundreds) can't hold a long +# transaction open. Whatever is left is picked up by the next tick. +_BATCH_LIMIT_ENV = "UNDISPATCHED_EXECUTION_SWEEP_LIMIT" +DEFAULT_BATCH_LIMIT = 500 + +# USER-FACING. `error_message` is rendered in the UI: ExecutionSerializer uses +# `exclude`, not `fields`, so every unlisted model field is serialized — this string +# reaches customers. It answers their three questions (did anything run? is my data +# affected? what do I do?) and deliberately names no internals: no queues, barriers, +# dispatch or gateways. The machine-readable ref is for support/greppability, and the +# precise technical cause stays in the log line below. +# +# Must fit EXECUTION_ERROR_LENGTH (256); the model truncates silently, which would +# otherwise eat the ref code at the end. +UNDISPATCHED_ERROR_MESSAGE = ( + "This execution did not start. The request was interrupted before any processing " + "began, so no files were processed. You can safely run it again. " + "(ref: EXEC_NOT_STARTED)" +) + + +def _positive_int_from_env(name: str, default: int) -> int: + """Env override, falling back loudly rather than silently on a bad value. + + A shortened grace period is the dangerous direction (it terminalises live + executions), so a typo must not quietly take effect. + """ + raw = os.environ.get(name) + if raw is None: + return default + try: + value = int(raw) + except (TypeError, ValueError): + logger.warning("%s=%r is not an integer; using default %s", name, raw, default) + return default + if value <= 0: + logger.warning("%s=%r must be > 0; using default %s", name, raw, default) + return default + return value + + +# Claim the batch in ONE statement that also reports exactly which rows it won. +# +# `UPDATE ... RETURNING` is the whole design in a single round trip: +# * the inner SELECT bounds the batch and takes row locks (SKIP LOCKED, so an +# overlapping sweeper can never block or double-claim); +# * the OUTER WHERE re-carries the full predicate, so Postgres re-evaluates it at +# write time under those locks — a row dispatched between selection and write is +# silently skipped rather than marked ERROR *while it runs*, which is the one +# failure this sweep must never cause; +# * RETURNING yields precisely the claimed rows, which is what lets the irreversible +# cleanup run for those and only those. (An earlier draft looped one UPDATE per +# row to get that guarantee; this keeps it and drops 500 round trips.) +# +# Django's .update() cannot RETURN rows, hence raw SQL — the same approach the PG-queue +# reaper's own sweeps use. The table name is interpolated from _meta (never user input). +_CLAIM_SQL = """ +UPDATE {table} + SET status = %s, error_message = %s, modified_at = %s + WHERE id IN ( + SELECT id + FROM {table} + WHERE status = %s + AND created_at < %s + AND dispatched_at IS NULL + AND task_id IS NULL + AND queue_message_id IS NULL + ORDER BY created_at + LIMIT %s + FOR UPDATE SKIP LOCKED + ) + AND status = %s + AND dispatched_at IS NULL + AND task_id IS NULL + AND queue_message_id IS NULL +RETURNING id, workflow_id +""" + + +def sweep_undispatched_executions( + min_age_seconds: int | None = None, limit: int | None = None +) -> int: + """Mark aged, never-dispatched PENDING executions ERROR. Returns the count. + + Idempotent and safe to run on every sweep: a row that has since been dispatched or + terminalised no longer matches the predicate, so a second pass finds nothing. + + Race-free by construction — see :data:`_CLAIM_SQL`. + """ + from django.db import connection + + from workflow_manager.workflow_v2.models import WorkflowExecution + + min_age = min_age_seconds or _positive_int_from_env( + _MIN_AGE_ENV, DEFAULT_MIN_AGE_SECONDS + ) + batch_limit = limit or _positive_int_from_env(_BATCH_LIMIT_ENV, DEFAULT_BATCH_LIMIT) + now = timezone.now() + cutoff = now - timezone.timedelta(seconds=min_age) + + sql = _CLAIM_SQL.format(table=WorkflowExecution._meta.db_table) + with connection.cursor() as cursor: + cursor.execute( + sql, + [ + ExecutionStatus.ERROR.value, + UNDISPATCHED_ERROR_MESSAGE, + now, + ExecutionStatus.PENDING.value, + cutoff, + batch_limit, + ExecutionStatus.PENDING.value, + ], + ) + claimed = cursor.fetchall() + + if not claimed: + return 0 + + for execution_id, workflow_id in claimed: + _release_abandoned_resources(str(execution_id), str(workflow_id)) + + # The operator-facing half of the story. The user-facing column says only what a + # customer needs; the cause belongs here, where it can name internals freely. + logger.error( + "Undispatched-execution sweep: marked %s execution(s) ERROR. They were created " + "but never queued (task_id and queue_message_id both NULL) and exceeded the %ss " + "grace period — the request died between create_workflow_execution and " + "execute_workflow_async. Ids: %s", + len(claimed), + min_age, + ", ".join(str(row[0]) for row in claimed[:20]), + ) + return len(claimed) + + +def _release_abandoned_resources(execution_id: str, workflow_id: str) -> None: + """Do what the abort prevented the request's own error path from doing. + + ``deployment_helper`` releases the rate-limit slot and deletes the API storage dir + when staging *raises*. An abort raises nothing — the thread is simply gone — so + neither runs and the execution leaks both. + + **Best-effort, and deliberately so.** The status write already succeeded and is the + part that matters; a failure to tidy up must never propagate and stall the rest of + the batch. + + Only the slot is released here. The staged input is deliberately retained — see the + note in the body for why, and why re-adding that delete needs the predicate fixed + first. + """ + # Slot first: it is the one with a live cost. Held slots consume the org's API + # deployment concurrency budget until the Redis ZSET TTL (6h) expires them, so a + # 502 storm can throttle a tenant for hours. Self-healing, but slowly. + try: + from api_v2.rate_limiter import APIDeploymentRateLimiter + + from workflow_manager.workflow_v2.models import WorkflowExecution + + organization = ( + WorkflowExecution.objects.select_related("workflow__organization") + .get(id=execution_id) + .workflow.organization + ) + # str(...organization_id), NOT the model instance. release_slot builds its Redis + # key by formatting the argument into "api_deployment:rate_limit:org:{org_id}", + # and acquire_slot built it from str(organization.organization_id). Passing the + # instance yields "...:org:Organization object (12)", which matches nothing — the + # ZREM removes a non-member, returns 0, raises nothing, and the slot stays held + # for the full 6h TTL. Silent, and the exact harm this call exists to prevent. + APIDeploymentRateLimiter.release_slot( + str(organization.organization_id), execution_id + ) + except Exception: + logger.warning( + "Undispatched sweep: could not release the rate-limit slot for %s " + "(it expires with the limiter TTL regardless)", + execution_id, + exc_info=True, + ) + + # The staged input is deliberately NOT deleted here. + # + # This sweep infers "never dispatched" from the ABSENCE of both handles + # (task_id IS NULL AND queue_message_id IS NULL). That inference is not sound: + # three paths in workflow_helper leave a RUNNING execution in exactly that state, + # all of them after the message is already on its transport — + # + # * _record_dispatch_handle raises and the caller swallows it, explicitly + # "continuing — the orchestrator is already running" (workflow_helper.py:686) + # * the handle comes back empty and it returns without writing (:544) + # * a PG handle will not parse as a bigint and it returns without writing (:562) + # + # so after the 15-minute grace this sweep can claim a live execution. + # + # Do NOT read the surviving ERROR write as harmless. An earlier version of this note + # claimed it "self-corrects — the running worker's own terminal write supersedes it". + # That is false on the PG transport: both worker entry points treat a terminal + # execution as a reason to STOP, not to overwrite. general/tasks.py returns + # skipped_terminal_execution before it would write EXECUTING, and + # file_processing/tasks.py raises _TerminalExecutionSkip. So a wrongly-claimed row is + # silently DROPPED — the message is acked, no retry is scheduled, and the only trace + # is a worker WARNING. The realistic trigger is not the three no-handle paths above + # but a saturated queue: a message enqueued successfully and still unconsumed at 15 + # minutes gets marked ERROR and then discarded, under exactly the backlog conditions + # that produced the orphans this sweep was built for. + # + # Even where a terminal write does land, error_message is not cleared by it, so a + # wrongly-claimed run can complete while still showing the customer "This execution + # did not start... You can safely run it again." + # + # Deleting the input is worse again and is what this removal addresses: the worker is + # still going to read it, and the message invites a re-run against data that is gone. + # + # Dropping the delete keeps the reversible half of the sweep and removes the + # irreversible one, at the cost of leaking an input directory for executions that + # genuinely never started. That is the right trade while the predicate is unsound; + # the real fix is to make dispatch a POSITIVE fact (stamp dispatched_at in the same + # call that records the handle, including on the three branches above, and key both + # this predicate and the 0026 partial index on it) — tracked separately, since it + # needs a migration and an index rebuild. + _ = workflow_id # retained: the signature is restored with the positive-fact fix diff --git a/backend/workflow_manager/workflow_v2/workflow_helper.py b/backend/workflow_manager/workflow_v2/workflow_helper.py index d22dcc85bb..971fb7a8c1 100644 --- a/backend/workflow_manager/workflow_v2/workflow_helper.py +++ b/backend/workflow_manager/workflow_v2/workflow_helper.py @@ -666,6 +666,28 @@ def execute_workflow_async( # the bookkeeping below must NOT flip the (now-running) row to ERROR. dispatched = True + # Record dispatch as a POSITIVE fact, before any bookkeeping that can + # fail. The undispatched sweep used to infer "never dispatched" from + # `task_id IS NULL AND queue_message_id IS NULL`, which the three paths + # below all reach on a RUNNING execution — the swallowed exception around + # _record_dispatch_handle, and its two early returns for an empty or + # unparseable handle. This single write sits upstream of all three, so the + # sweep can ask a question with a true answer instead of guessing. + # + # Its own try/except for the same reason the handle write has one: we are + # past the point of no return and nothing here may abort the caller. If it + # fails we are back to the old ambiguity for this one execution, which is + # strictly no worse than before. + try: + WorkflowExecutionServiceHelper.mark_dispatched(execution_id) + except Exception: + logger.exception( + f"[{org_schema}] Failed to stamp dispatched_at for execution " + f"'{execution_id}'; continuing — the orchestrator is already " + "running. This row is now indistinguishable from an undispatched " + "one and may be terminalised by the undispatched sweep." + ) + workflow_execution: WorkflowExecution = WorkflowExecution.objects.get( id=execution_id ) @@ -1033,6 +1055,37 @@ def step_execution( ) try: workflow_execution = WorkflowExecution.objects.get(pk=execution_id) + # Single-step is the one entry path that never moved onto PG. Its + # only fan-out is the Celery chord in `process_input_files` — the + # normal path's PG fan-out lives in the general worker instead, so + # nothing here is transport-gated. With the flag on and the Celery + # file_processing workers scaled to zero (the epic's acceptance + # gate), those batches would sit unconsumed and the execution would + # hang in EXECUTING, invisible to the PG reaper. Fail fast instead: + # a 500 with a stated cause beats a silent forever-EXECUTING row. + # + # Unreachable from the UI — the step buttons are gone and both live + # call sites pass isStepExecution=false, so `execution_action` is + # never sent (frontend Agency.jsx). This guards a direct API call. + # + # Resolved here rather than in `process_input_files` so the Celery + # path stays byte-identical and an already-resolved execution is + # never re-resolved. Step executions are created by + # `create_and_make_execution_response`, which resolves no transport, + # so this is their first and only resolution. + transport = resolve_transport( + execution_id=execution_id, + organization_id=workflow.organization.organization_id, + workflow_id=workflow.id, + ) + if is_pg_transport(transport): + raise WorkflowExecutionError( + "Single-step execution is not supported on the Postgres " + "queue transport: it has no PG fan-out, and its Celery " + "batches would never be consumed. Run the full workflow " + "instead, or disable pg_queue_enabled for this " + "organization." + ) return WorkflowHelper.run_workflow( workflow=workflow, single_step=True, diff --git a/docker/docker-compose.yaml b/docker/docker-compose.yaml index 10a7b4e5db..dda8f27a21 100644 --- a/docker/docker-compose.yaml +++ b/docker/docker-compose.yaml @@ -742,6 +742,93 @@ services: profiles: - pg-queue + # Dashboard-metrics periodics over PG (UN-3796) — the execution half of the Beat + # replacement. The PG scheduler fires dashboard_metrics.* onto dashboard_metric_events + # and this drains them, replacing the Celery worker-metrics + # (-A backend worker -Q dashboard_metric_events). + # + # A dedicated service rather than another queue on worker-pg-scheduler: the + # consumer's health heartbeat is stamped per poll and FROZEN while a task runs, so + # HEALTH_STALE doubles as an upper bound on one task's wall clock. The aggregation + # runs for minutes; under the scheduler's 240s bound it would trip the liveness + # probe, restart the pod, and take in-flight pipeline triggers down with it. + # + # No rabbitmq dependency — the tasks are thin HTTP proxies to the backend internal + # API and dispatch nothing onward. MAX_ATTEMPTS=1 because for a periodic the next + # cron tick supersedes a failed run; retrying would let a persistently failing + # aggregation stack up in-flight messages. + worker-pg-metrics: + image: unstract/worker-unified:${VERSION} + container_name: unstract-worker-pg-metrics + restart: unless-stopped + # The GENERIC consumer command, with identity carried by the env below — the + # container entrypoint (run-worker-docker.sh) is env-driven and its `pg-*` + # dispatch accepts only pg-queue-consumer / pg-consumer / pg-queue-reaper / + # pg-reaper / reaper, rejecting anything else with exit 1. The named + # `pg-metrics` role is for the HOST runner (`./run-worker.sh pg-metrics`), + # which is a different script and does understand roles. + command: ["pg-queue-consumer"] + ports: + - "8101:8090" + env_file: + - ../workers/.env + - ./essentials.env + depends_on: + - db + - redis + environment: + - ENVIRONMENT=development + - APPLICATION_NAME=unstract-worker-pg-metrics + - WORKER_BARRIER_BACKEND=pg + - WORKER_PG_QUEUE_CONSUMER_WORKER_TYPE=scheduler + - WORKER_PG_QUEUE_CONSUMER_QUEUE=dashboard_metric_events + - WORKER_PG_QUEUE_CONSUMER_HEALTH_PORT=8090 + # Sized for a minutes-long aggregation: stale > vt > the task's own ceiling. + - WORKER_PG_QUEUE_CONSUMER_VT_SECONDS=${WORKER_PG_METRICS_VT_SECONDS:-900} + - WORKER_PG_QUEUE_CONSUMER_HEALTH_STALE_SECONDS=${WORKER_PG_METRICS_HEALTH_STALE_SECONDS:-960} + # One at a time — these are global singletons, not parallel work. + - WORKER_PG_QUEUE_CONSUMER_CONCURRENCY=1 + - WORKER_PG_QUEUE_CONSUMER_MAX_ATTEMPTS=1 + labels: + - traefik.enable=false + profiles: + - pg-queue + + # Log streaming without Celery (UN-3755) — the flag-on replacement for + # worker-log-consumer-v2. Note the transport here is a REDIS LIST, not + # pg_queue_message: the log consumer's own sinks are already Redis (log_history_queue + # + the Socket.IO emit), so routing through Postgres would detour a lossy, high-rate + # stream through the platform's most contended resource to buy durability that logs + # explicitly discard at their 10k cap. It sits in the pg-queue profile because it is + # part of the same flag-on fleet, not because it speaks to Postgres. + # + # LOG_TRANSPORT must be set to "redis" on every PUBLISHER too (backend, workers, + # runner) or this drains an empty list while logs keep going to RabbitMQ. + worker-log-stream-consumer: + image: unstract/worker-unified:${VERSION} + container_name: unstract-worker-log-stream-consumer + restart: unless-stopped + command: ["log-stream-consumer"] + env_file: + - ../workers/.env + - ./essentials.env + depends_on: + - redis + environment: + - ENVIRONMENT=development + - APPLICATION_NAME=unstract-worker-log-stream-consumer + - WORKER_TYPE=log_consumer + - WORKER_NAME=log-stream-consumer + - LOG_TRANSPORT=redis + - LOG_STREAM_QUEUE_NAME=${LOG_STREAM_QUEUE_NAME:-log_stream_queue} + # Downstream sinks are unchanged — same durable buffer the Celery worker wrote to. + - LOG_HISTORY_QUEUE_NAME=${LOG_HISTORY_QUEUE_NAME:-log_history_queue} + - ENABLE_LOG_HISTORY=${ENABLE_LOG_HISTORY:-true} + labels: + - traefik.enable=false + profiles: + - pg-queue + # Executor RPC over PG — runs execute_extraction as a request-reply: claims # work from Postgres and writes the ExecutionResult to pg_task_result for the # blocking caller. Same heavy executor runtime as worker-executor-v2 (tool diff --git a/docker/sample.env b/docker/sample.env index e5e1ca7c68..6c96821a42 100644 --- a/docker/sample.env +++ b/docker/sample.env @@ -133,6 +133,17 @@ ENABLE_METRICS=true # WORKER_PG_REAPER_HEALTH_PORT - opt-in reaper liveness port (unset = off). WORKER_PG_REAPER_INTERVAL_SECONDS=5 # Reaper sweep interval (seconds) # +# worker-pg-metrics (UN-3796) — drains the dashboard_metrics.* periodics the PG +# scheduler fires, replacing Celery Beat + worker-metrics. Its knobs are sized for a +# task that runs for MINUTES, unlike every other consumer here: +# +# WORKER_PG_METRICS_VT_SECONDS - claim/drain bound (default 900). +# WORKER_PG_METRICS_HEALTH_STALE_SECONDS- liveness bound (default 960). The heartbeat +# freezes while a task runs, so this is effectively an +# upper bound on ONE task's wall clock — it must stay +# above the vt, which must stay above the task itself. +# above the vt, which must stay above the task itself. +# # NOTE: routing executions to PG is a SEPARATE, later step. Running these # services does NOT move any traffic — the single `pg_queue_enabled` Flipt flag # (default off, fail-closed) decides per-execution transport, and stays off until diff --git a/runner/src/unstract/runner/constants.py b/runner/src/unstract/runner/constants.py index b46971d618..e4f7449fb7 100644 --- a/runner/src/unstract/runner/constants.py +++ b/runner/src/unstract/runner/constants.py @@ -44,3 +44,10 @@ class Env: CELERY_BROKER_BASE_URL = "CELERY_BROKER_BASE_URL" CELERY_BROKER_USER = "CELERY_BROKER_USER" CELERY_BROKER_PASS = "CELERY_BROKER_PASS" + # Log-streaming transport (UN-3755). Must reach the SIDECAR: it is the process + # that calls LogPublisher.publish for tool logs + # (tool_sidecar/log_processor.py:165), and the sidecar's environment is a + # hand-picked allowlist rather than an inherited one — so anything absent here is + # silently absent there, and it falls back to publishing on Celery/RabbitMQ. + LOG_TRANSPORT = "LOG_TRANSPORT" + LOG_STREAM_QUEUE_NAME = "LOG_STREAM_QUEUE_NAME" diff --git a/runner/src/unstract/runner/runner.py b/runner/src/unstract/runner/runner.py index 1125bf69e8..2414ac7548 100644 --- a/runner/src/unstract/runner/runner.py +++ b/runner/src/unstract/runner/runner.py @@ -245,6 +245,16 @@ def _get_sidecar_container_config( "CELERY_BROKER_BASE_URL": os.getenv(Env.CELERY_BROKER_BASE_URL), "CELERY_BROKER_USER": os.getenv(Env.CELERY_BROKER_USER), "CELERY_BROKER_PASS": os.getenv(Env.CELERY_BROKER_PASS), + # Log transport (UN-3755). The sidecar is a LogPublisher producer + # (log_processor.py:165), so it must agree with every other publisher or + # its tool logs go to RabbitMQ while the rest go to Redis — and with the + # Celery log consumer scaled to zero, those logs are simply lost. This + # dict is an allowlist, not inherited env, so omitting it fails silently. + # Needs no new credentials: REDIS_* is already passed above. + Env.LOG_TRANSPORT: os.getenv(Env.LOG_TRANSPORT, "celery"), + Env.LOG_STREAM_QUEUE_NAME: os.getenv( + Env.LOG_STREAM_QUEUE_NAME, "log_stream_queue" + ), "CONTAINER_NAME": container_name, } sidecar_config = self.client.get_container_run_config( diff --git a/runner/tests/test_sidecar_log_transport.py b/runner/tests/test_sidecar_log_transport.py new file mode 100644 index 0000000000..01c46280b1 --- /dev/null +++ b/runner/tests/test_sidecar_log_transport.py @@ -0,0 +1,84 @@ +"""The sidecar must inherit the log transport (UN-3755). + +``_get_sidecar_container_config`` builds the sidecar's environment as a hand-picked +**allowlist**, not inherited env. So a variable the sidecar needs but that nobody added +here is silently absent at runtime — no error, no warning. + +That matters for ``LOG_TRANSPORT`` specifically: the sidecar is a ``LogPublisher`` +producer (``tool_sidecar/log_processor.py:165``), and ``LogPublisher`` defaults to the +Celery/RabbitMQ transport when the variable is unset. Miss it and tool logs keep going +to ``celery_log_task_queue`` while every other publisher moves to Redis — and once the +Celery log consumer is scaled to zero those logs are simply dropped, with live +streaming and ``execution_log`` rows both silently missing for container-based tools. +""" + +from __future__ import annotations + +from unittest.mock import MagicMock + +import pytest + +from unstract.runner.constants import Env +from unstract.runner.runner import UnstractRunner + + +@pytest.fixture +def sidecar_env(monkeypatch): + """Build a sidecar env dict with the client mocked out.""" + + def _build(**env): + for key, value in env.items(): + monkeypatch.setenv(key, value) + runner = UnstractRunner.__new__(UnstractRunner) + runner.client = MagicMock() + runner.client.get_container_run_config.side_effect = ( + lambda **kwargs: kwargs # return the call kwargs so we can read `envs` + ) + config = runner._get_sidecar_container_config( + container_name="c", + shared_log_dir="/d", + shared_log_file="/d/log", + organization_id="org", + execution_id="exec", + file_execution_id="fe", + messaging_channel="chan", + tool_instance_id="ti", + ) + return config["envs"] + + return _build + + +class TestSidecarLogTransport: + def test_defaults_to_celery_when_unset(self, sidecar_env, monkeypatch): + monkeypatch.delenv(Env.LOG_TRANSPORT, raising=False) + envs = sidecar_env() + # Present and explicit, not merely absent — the sidecar's LogPublisher reads + # this key, and an explicit default documents the flag-off state. + assert envs[Env.LOG_TRANSPORT] == "celery" + + def test_forwards_redis_transport_to_the_sidecar(self, sidecar_env): + envs = sidecar_env(LOG_TRANSPORT="redis") + assert envs[Env.LOG_TRANSPORT] == "redis" + + def test_forwards_the_stream_queue_name(self, sidecar_env): + # If a deployment renames the queue, the sidecar must push to the same list + # the consumer drains, or its logs land somewhere nobody reads. + envs = sidecar_env(LOG_TRANSPORT="redis", LOG_STREAM_QUEUE_NAME="custom_stream") + assert envs[Env.LOG_STREAM_QUEUE_NAME] == "custom_stream" + + def test_redis_credentials_are_already_present(self, sidecar_env): + """The Redis transport needs no NEW credential in this allowlist. + + This is the concrete payoff of choosing a Redis list over the PG queue for the + log hop: REDIS_* is already forwarded, whereas PG would have required adding + database credentials to every spawned sidecar. + """ + envs = sidecar_env(LOG_TRANSPORT="redis", REDIS_HOST="r", REDIS_PORT="6379") + assert envs["REDIS_HOST"] == "r" + assert envs["REDIS_PORT"] == "6379" + + def test_celery_broker_still_forwarded_for_the_flag_off_path(self, sidecar_env): + # Flag-off must stay intact: the sidecar still publishes over AMQP. + envs = sidecar_env(CELERY_BROKER_BASE_URL="amqp://x") + assert envs["CELERY_BROKER_BASE_URL"] == "amqp://x" diff --git a/tests/groups.yaml b/tests/groups.yaml index 0fef1e82a5..d2c35e34d6 100644 --- a/tests/groups.yaml +++ b/tests/groups.yaml @@ -111,6 +111,21 @@ groups: install_editable: true coverage_source: src + # runner/tests was registered in NO group, so it was collected by nothing. Its + # only suite guards that LOG_TRANSPORT and LOG_STREAM_QUEUE_NAME reach the tool + # sidecar's hand-picked env allowlist — a list that drops anything not named in + # it, so a missing entry means the sidecar silently logs to the wrong transport. + # Nothing else covers that, and it went in alongside the Redis log transport. + # install_editable because the suite imports `unstract.runner`, which resolves + # only from this package's own src layout (packages = ["src/unstract"]). + unit-runner: + tier: unit + workdir: runner + paths: [tests] + uv_sync_group: test + install_editable: true + coverage_source: src + # ── Integration tier: needs infra but not full platform ──────────────────── integration-backend: tier: integration diff --git a/unstract/core/src/unstract/core/data_models.py b/unstract/core/src/unstract/core/data_models.py index eba226743e..bc2fd6dbc9 100644 --- a/unstract/core/src/unstract/core/data_models.py +++ b/unstract/core/src/unstract/core/data_models.py @@ -286,13 +286,23 @@ class QueueMessageState(StrEnum): set, mirroring the ``priority`` (fairness) precedent. ``str`` Enum → serialises to its value and compares equal to the bare string. - - ``READY`` — claimable: the dequeue's partial claim index holds only these. - - ``CLAIMED`` — in-flight: a consumer holds it, ``vt`` is its renewable lease; + - ``READY`` — claimable: the dequeue's partial claim index holds only these. + - ``CLAIMED`` — in-flight: a consumer holds it, ``vt`` is its renewable lease; re-armed back to ``READY`` by the reaper when the lease expires (crash). + - ``SCHEDULED`` — deferred (UN-3843): enqueued with a future ``available_at`` + (Celery ``countdown``/``eta`` parity) and **deliberately absent from the claim + index**, so a not-yet-due row costs the hot claim path nothing. The reaper + promotes it to ``READY`` once ``available_at <= now()`` + (``reaper.promote_due_scheduled``), which is why delivery is "not before + ``available_at``" rather than exactly at it — granularity is the reaper tick. + Filtering ``available_at`` in the claim instead would put every not-yet-due row + back inside ``pg_queue_message_claim_idx`` to be walked and discarded on every + claim — precisely the scan-past cost this enum exists to remove. """ READY = "ready" CLAIMED = "claimed" + SCHEDULED = "scheduled" # Fairness L3 priority bounds (1..10, higher = claimed sooner). Single source of diff --git a/unstract/core/src/unstract/core/pubsub_helper.py b/unstract/core/src/unstract/core/pubsub_helper.py index 3bec8475e0..d42c487183 100644 --- a/unstract/core/src/unstract/core/pubsub_helper.py +++ b/unstract/core/src/unstract/core/pubsub_helper.py @@ -12,6 +12,35 @@ from unstract.core.cache.redis_client import create_redis_client from unstract.core.constants import LogEventArgument, LogProcessingTask +#: Transport for the log-streaming hop between this publisher and the log consumer. +#: ``celery`` (default) publishes a Celery-protocol message to RabbitMQ; ``redis`` +#: RPUSHes onto a Redis list drained by the PG-era log consumer. +#: +#: Deliberately an env var rather than the ``pg_queue_enabled`` Flipt flag. The flag +#: routes *per execution*, which is meaningless here: the consumer is a single +#: deployment reading from exactly one place, so a per-org split would strand half the +#: logs in a queue nobody drains. Producer and consumer must agree cluster-wide, which +#: makes this deployment config. It must be flipped in step with the flag rollout. +_LOG_TRANSPORT_ENV = "LOG_TRANSPORT" +_LOG_TRANSPORT_REDIS = "redis" + +#: Redis list used when ``LOG_TRANSPORT=redis``. Distinct from ``log_history_queue``, +#: which is the *downstream* durable buffer the consumer writes to — this one is the +#: transport itself, and is normally near-empty. +_LOG_STREAM_QUEUE_ENV = "LOG_STREAM_QUEUE_NAME" +_LOG_STREAM_QUEUE_DEFAULT = "log_stream_queue" + +#: Cap mirroring ``store_execution_log``'s (unstract/core/log_utils.py): without it a +#: stopped consumer grows this list until Redis OOMs, which would take down far more +#: than logging. Dropping is the correct behaviour for this payload — logs are already +#: best-effort and are discarded at the same cap one hop downstream. +_LOG_STREAM_MAX_SIZE_ENV = "LOG_STREAM_QUEUE_MAX_SIZE" +_LOG_STREAM_MAX_SIZE_DEFAULT = 10000 + + +def _use_redis_log_transport() -> bool: + return os.getenv(_LOG_TRANSPORT_ENV, "celery").strip().lower() == _LOG_TRANSPORT_REDIS + class LogPublisher: broker_url = str( @@ -160,27 +189,30 @@ def publish(cls, channel_id: str, payload: dict[str, Any]) -> bool: """Publish a message to the queue.""" try: event = f"logs:{channel_id}" - with cls.kombu_conn.Producer(serializer="json") as producer: - task_message = cls._get_task_message( - user_session_id=channel_id, - event=event, - message=payload, - ) - headers = cls._get_task_header(LogProcessingTask.TASK_NAME) - # Publish the message to the queue - producer.publish( - body=task_message, - exchange="", - headers=headers, - routing_key=LogProcessingTask.QUEUE_NAME, - compression=None, - retry=True, - ) - logging.debug(f"Published '{channel_id}' <= {payload}") - - # Persisting messages for unified notification - if payload.get("type") == "LOG": - cls.store_for_unified_notification(event, payload) + task_message = cls._get_task_message( + user_session_id=channel_id, + event=event, + message=payload, + ) + if _use_redis_log_transport(): + cls._publish_via_redis(task_message) + else: + with cls.kombu_conn.Producer(serializer="json") as producer: + headers = cls._get_task_header(LogProcessingTask.TASK_NAME) + # Publish the message to the queue + producer.publish( + body=task_message, + exchange="", + headers=headers, + routing_key=LogProcessingTask.QUEUE_NAME, + compression=None, + retry=True, + ) + logging.debug(f"Published '{channel_id}' <= {payload}") + + # Persisting messages for unified notification + if payload.get("type") == "LOG": + cls.store_for_unified_notification(event, payload) except Exception as e: logging.error( f"Failed to publish '{channel_id}' <= {payload}" @@ -189,6 +221,42 @@ def publish(cls, channel_id: str, payload: dict[str, Any]) -> bool: return False return True + @classmethod + def _publish_via_redis(cls, task_message: dict[str, Any]) -> None: + """RPUSH the log envelope onto the Redis transport list. + + The envelope carries the task name alongside the kwargs so the consumer + dispatches by name exactly as the Celery header did — a bare kwargs blob would + leave the consumer guessing, and a name mismatch is the silent failure mode + (message read, no handler, dropped with nothing at the publish site to trace). + + Raises on Redis failure; ``publish()`` owns the swallow, so a logging fault can + never break an execution. + """ + queue_name = os.getenv(_LOG_STREAM_QUEUE_ENV, _LOG_STREAM_QUEUE_DEFAULT) + max_size = int( + os.getenv(_LOG_STREAM_MAX_SIZE_ENV, str(_LOG_STREAM_MAX_SIZE_DEFAULT)) + ) + redis_client = cls._get_redis_client() + + # O(1), and the same llen-then-push shape store_execution_log already uses one + # hop downstream. Two Redis round trips per log line is the price of not letting + # a stopped consumer OOM Redis. + if redis_client.llen(queue_name) >= max_size: + logging.warning( + f"Log stream queue '{queue_name}' at capacity ({max_size}), " + "dropping current log - log consumer may be down or falling behind" + ) + return + + envelope = json.dumps( + { + "task": LogProcessingTask.TASK_NAME, + "kwargs": task_message["kwargs"], + } + ) + redis_client.rpush(queue_name, envelope) + @classmethod def store_for_unified_notification(cls, event: str, payload: dict[str, Any]) -> None: """Helps persist messages for unified notification. diff --git a/workers/log_consumer/redis_stream_consumer.py b/workers/log_consumer/redis_stream_consumer.py new file mode 100644 index 0000000000..a2242f11f2 --- /dev/null +++ b/workers/log_consumer/redis_stream_consumer.py @@ -0,0 +1,163 @@ +"""Redis-list consumer for the log stream (UN-3755). + +Drains the list ``LogPublisher.publish()`` writes to when ``LOG_TRANSPORT=redis`` and +runs the **existing** ``logs_consumer`` task body against each envelope. Nothing about +what a log *does* changes here: the task still RPUSHes to ``log_history_queue`` and +emits over Socket.IO, and the log-history scheduler still drains that to +``execution_log``. This module replaces only the transport hop that RabbitMQ used to +provide, so the log consumer no longer needs a Celery deployment. + +**Why a reliable-queue pattern rather than a bare BLPOP.** The Celery consumer runs with +``task_acks_late=True`` (``shared/infrastructure/config/worker_config.py:545``, and +``backend/celery_config.py:63``), so a worker that dies mid-task gets the message +redelivered. ``BLPOP`` deletes on read, which would silently make crashes lossy — a real +regression, not a no-op. ``BLMOVE`` instead parks the envelope on a per-pod *processing* +list and removes it only after the handler returns, and startup re-queues whatever the +previous incarnation left behind. + +Bound worth knowing: recovery is keyed on ``HOSTNAME``, so a **container restart** (crash +loop, OOM kill — the dominant failure mode) recovers fully, while a **pod replacement** +strands that pod's in-flight envelopes, at most one per concurrent loop. That is strictly +better than BLPOP and, for a stream already discarded wholesale at its 10k cap, close +enough to the Celery behaviour it replaces. Sweeping other pods' lists is deliberately not +attempted: with multiple replicas it cannot distinguish a dead owner from a live one, and +would duplicate log lines on every start. +""" + +from __future__ import annotations + +import json +import os +import signal +import socket +import sys +from types import FrameType +from typing import Any + +from shared.enums.worker_enums import WorkerType +from shared.infrastructure.config.builder import WorkerBuilder +from shared.infrastructure.logging import WorkerLogger + +from unstract.core.cache.redis_client import create_redis_client +from unstract.core.constants import LogProcessingTask + +logger = WorkerLogger.setup(WorkerType.LOG_CONSUMER) + +# Build the Celery app before importing tasks: ``@worker_task`` is ``shared_task``, which +# binds to the current app at call time. Same ordering as ``log_consumer/worker.py``. +app, config = WorkerBuilder.build_celery_app(WorkerType.LOG_CONSUMER) + +from log_consumer.tasks import logs_consumer # noqa: E402 + +_QUEUE_NAME = os.getenv("LOG_STREAM_QUEUE_NAME", "log_stream_queue") +# BLMOVE blocks up to this long before returning None, which is the loop's only chance to +# notice a shutdown signal. Keep it well under the pod's terminationGracePeriodSeconds. +_BLOCK_TIMEOUT_SECONDS = int(os.getenv("LOG_STREAM_BLOCK_TIMEOUT", "5")) +# redis-py enforces ``socket_timeout`` on the BLMOVE read itself, so it MUST exceed the +# server-side block or every call aborts mid-block with ``redis.TimeoutError``. +# ``create_redis_client`` defaults it to 5s — exactly ``_BLOCK_TIMEOUT_SECONDS`` — and the +# socket won that race in integration: a traceback every ~5.01s while idle, plus a +# disconnect/reconnect each time. Worse, an envelope BLMOVE had already moved to the +# processing list was stranded there until this pod restarted, because the reply never +# reached the client. Identical trap to the one documented at +# ``workers/queue_backend/pg_queue/result_backend.py:152``. +_SOCKET_TIMEOUT_SECONDS = _BLOCK_TIMEOUT_SECONDS + 5 + +_shutdown = False + + +def _processing_list_name() -> str: + """Per-pod parking list, so one pod never reclaims another's in-flight envelope.""" + return f"{_QUEUE_NAME}:processing:{os.getenv('HOSTNAME') or socket.gethostname()}" + + +def _handle_signal(signum: int, _frame: FrameType | None) -> None: + """Finish the envelope in hand, then stop. Never abandon work mid-flight.""" + global _shutdown + _shutdown = True + logger.info("Signal %s received; finishing current log then shutting down", signum) + + +def _recover_in_flight(redis_client: Any, processing: str) -> None: + """Re-queue anything this pod was mid-way through when it last died. + + LMOVE back to the *head* of the source list so recovered envelopes are re-processed + before newer ones, preserving rough log ordering. + """ + recovered = 0 + while redis_client.lmove(processing, _QUEUE_NAME, "RIGHT", "LEFT"): + recovered += 1 + if recovered: + logger.warning( + "Recovered %d in-flight log envelope(s) from a previous run of this pod", + recovered, + ) + + +def _dispatch(raw: bytes | str) -> None: + """Run one envelope through the existing task body. + + Dispatches **by name** so a producer/consumer mismatch fails loudly here rather than + silently discarding the message. + """ + envelope = json.loads(raw) + task_name = envelope.get("task") + if task_name != LogProcessingTask.TASK_NAME: + raise ValueError( + f"Unexpected task {task_name!r} on '{_QUEUE_NAME}'; " + f"expected {LogProcessingTask.TASK_NAME!r}" + ) + logs_consumer(**envelope.get("kwargs", {})) + + +def run() -> int: + signal.signal(signal.SIGTERM, _handle_signal) + signal.signal(signal.SIGINT, _handle_signal) + + # Not ``RedisQueueClient.from_env()``: that hard-codes the 5s socket timeout, which + # cannot outlive this loop's block. Built directly so the two stay related by + # construction — see _SOCKET_TIMEOUT_SECONDS. + redis_client = create_redis_client( + decode_responses=True, + socket_timeout=_SOCKET_TIMEOUT_SECONDS, + ) + processing = _processing_list_name() + logger.info( + "Log stream consumer starting: queue='%s' processing='%s'", + _QUEUE_NAME, + processing, + ) + _recover_in_flight(redis_client, processing) + + while not _shutdown: + try: + raw = redis_client.blmove( + _QUEUE_NAME, processing, _BLOCK_TIMEOUT_SECONDS, "LEFT", "RIGHT" + ) + except Exception: + # Connection blips must not kill the pod — the next iteration reconnects via + # the client's own retry. Sleeping is unnecessary: BLMOVE already blocks. + logger.error("Log stream read failed; retrying", exc_info=True) + continue + + if raw is None: # timeout, no work — loop so shutdown can be observed + continue + + try: + _dispatch(raw) + except Exception: + # Match the Celery consumer's posture: a poison envelope is logged and + # dropped, never retried forever. logs_consumer already swallows its own + # sink failures, so reaching here means a malformed envelope. + logger.error("Discarding unprocessable log envelope", exc_info=True) + finally: + # Remove exactly one copy, whether it succeeded or was discarded — leaving it + # parked would have it re-queued on the next restart and replayed forever. + redis_client.lrem(processing, 1, raw) + + logger.info("Log stream consumer stopped") + return 0 + + +if __name__ == "__main__": + sys.exit(run()) diff --git a/workers/queue_backend/pg_queue/metrics.py b/workers/queue_backend/pg_queue/metrics.py index 5678b38198..d34c7d885d 100644 --- a/workers/queue_backend/pg_queue/metrics.py +++ b/workers/queue_backend/pg_queue/metrics.py @@ -252,6 +252,20 @@ def __init__( "faults that share pg_reaper_tick_failures_total)", registry=self.registry, ) + self.queue_promoted = Counter( + "pg_reaper_queue_promoted_total", + "Due scheduled queue messages promoted to 'ready' (delayed-visibility " + "delivery for countdown/eta dispatches — the reaper is on the DELIVERY " + "path for these, not just recovery)", + registry=self.registry, + ) + self.queue_promote_failures = Counter( + "pg_reaper_queue_promote_failures_total", + "Promotion sweep attempts that raised (delayed messages did not become " + "claimable this tick; alert on a sustained non-zero rate — the symptom is " + "silent non-delivery, not an error at the enqueue site)", + registry=self.registry, + ) self.claim_recovered = Counter( "pg_reaper_claim_recovered_total", "Orphan orchestration claims recovered (crash-window execution " @@ -268,6 +282,21 @@ def __init__( "Orphan-claim recovery attempts that raised (row left for retry)", registry=self.registry, ) + self.undispatched_swept = Counter( + "pg_reaper_undispatched_swept_total", + "Executions terminalised because they were created but never dispatched " + "(no task_id and no queue_message_id past the grace period). These have no " + "barrier, so barrier recovery cannot see them; a sustained non-zero rate " + "means requests are dying between create_workflow_execution and dispatch", + registry=self.registry, + ) + self.undispatched_sweep_failures = Counter( + "pg_reaper_undispatched_sweep_failures_total", + "Undispatched-execution sweep calls that raised (swallowed so the tick " + "still dispatches schedules; a sustained non-zero rate means PENDING rows " + "are accumulating unrecovered)", + registry=self.registry, + ) self.sweep_failures = Counter( "pg_reaper_sweep_failures_total", "Whole-sweep failures, by swept table (see the reaper fail-streak log)", diff --git a/workers/queue_backend/pg_queue/pg_scheduler.py b/workers/queue_backend/pg_queue/pg_scheduler.py index 7ccfe40be1..6b87acdaf1 100644 --- a/workers/queue_backend/pg_queue/pg_scheduler.py +++ b/workers/queue_backend/pg_queue/pg_scheduler.py @@ -56,8 +56,14 @@ class _DueSchedule(NamedTuple): - """One row from the due-schedules scan — names bound to columns at one site - so a future reorder of the SELECT can't silently misassign fields. + """One row from the due-schedules scan. + + THE FIELD NAMES ARE THE QUERY. They are emitted verbatim as the SELECT's column + list (see the scan below), so every one MUST be a column of + ``pg_periodic_schedule`` — renaming a field here rewrites the SQL and needs a + matching migration in ``backend/pg_queue/models.py``. That is what makes a reorder + harmless (there is no second list to drift from) and a RENAME dangerous, which is + the opposite of what an earlier version of this docstring implied. """ pipeline_id: uuid.UUID @@ -143,9 +149,8 @@ def dispatch_due_schedules(conn: PgConnection) -> int: base = cur.fetchone()[0] cur.execute( f""" - SELECT pipeline_id, organization_id, workflow_id, pipeline_name, - cron_string, next_run_at - FROM {qualified('pg_periodic_schedule')} + SELECT {", ".join(f'"{f}"' for f in _DueSchedule._fields)} + FROM {qualified("pg_periodic_schedule")} WHERE pg_owned AND enabled AND (next_run_at IS NULL OR next_run_at <= %s) """, @@ -230,3 +235,153 @@ def dispatch_due_schedules(conn: PgConnection) -> int: ) return fired + + +class _DuePeriodicTask(NamedTuple): + """One row from the generic-periodic due scan (UN-3796). + + Sibling of :class:`_DueSchedule`, and carries the same rule: the field names are + emitted verbatim as the SELECT's column list, so each MUST be a column of + ``pg_periodic_task``. + + Note this type says ``org_id`` where its sibling says ``organization_id``. Tempting + to unify — do not, without a migration. The two back different tables, and renaming + this one would silently select a column ``pg_periodic_task`` does not have; the + resulting UndefinedColumn propagates out of the dispatch and takes down the whole + leader tick, retention sweep and gauge refresh included. + """ + + name: str + task_name: str + queue: str + task_args: list + task_kwargs: dict + org_id: str + cron_string: str + next_run_at: datetime | None + + +def _quiesce_invalid_periodic_cron(conn: PgConnection, row: _DuePeriodicTask) -> None: + """Disable a generic periodic whose cron won't parse, so it stops being + re-selected (and re-logging a traceback) every tick. Mirrors + :func:`_quiesce_invalid_cron` for the pipeline table. + """ + logger.exception( + "PG scheduler: invalid cron %r for periodic %r — disabling the row", + row.cron_string, + row.name, + ) + try: + with conn.cursor() as cur: + cur.execute( + f"UPDATE {qualified('pg_periodic_task')} " + "SET enabled = FALSE WHERE name = %s", + (row.name,), + ) + conn.commit() + except Exception: + with contextlib.suppress(Exception): + conn.rollback() + + +def dispatch_due_periodic_tasks(conn: PgConnection) -> int: + """Fire PG-owned, enabled, due **non-pipeline** periodics; return the count fired. + + The Beat-replacement half for everything that isn't a pipeline trigger + (UN-3796): ``dashboard_metrics.*``, log-history, audit, and anything an operator + adds. Structurally identical to :func:`dispatch_due_schedules` — leader-gated by + the caller, DB clock throughout, per-row transaction, enqueue and ``next_run_at`` + advance in ONE transaction so a crash between them cannot double-fire, first + observation of a freshly-owned row records a baseline without firing. + + The one real difference: each row carries its **own** target queue and its own + task/args/kwargs, where the pipeline dispatcher rebuilds one fixed argument list + onto one fixed queue. That is the whole reason for the second table. + """ + try: + with conn.cursor() as cur: + cur.execute("SELECT now()") + base = cur.fetchone()[0] + cur.execute( + f""" + SELECT {", ".join(f'"{f}"' for f in _DuePeriodicTask._fields)} + FROM {qualified("pg_periodic_task")} + WHERE pg_owned AND enabled + AND (next_run_at IS NULL OR next_run_at <= %s) + """, + (base,), + ) + due = [_DuePeriodicTask(*row) for row in cur.fetchall()] + conn.commit() + except Exception: + with contextlib.suppress(Exception): + conn.rollback() + raise + + fired = 0 + for row in due: + try: + nxt = compute_next_run(row.cron_string, base) + except Exception: + _quiesce_invalid_periodic_cron(conn, row) + continue + + try: + if row.next_run_at is None: + with conn.cursor() as cur: + cur.execute( + f"UPDATE {qualified('pg_periodic_task')} " + "SET next_run_at = %s WHERE name = %s", + (nxt, row.name), + ) + conn.commit() + logger.info( + "PG scheduler: baselined periodic %r (next_run_at=%s, not fired)", + row.name, + nxt, + ) + continue + + payload = to_payload( + row.task_name, + args=list(row.task_args or []), + kwargs=dict(row.task_kwargs or {}), + queue=row.queue, + fairness=None, + ) + # Enqueue + advance in ONE transaction (see the module docstring). + with conn.cursor() as cur: + cur.execute( + insert_message_sql(), + ( + row.queue, + json.dumps(payload), + row.org_id or "", + DEFAULT_PRIORITY, + ), + ) + cur.execute( + f"UPDATE {qualified('pg_periodic_task')} " + "SET last_run_at = %s, next_run_at = %s WHERE name = %s", + (base, nxt, row.name), + ) + conn.commit() + except Exception: + with contextlib.suppress(Exception): + conn.rollback() + logger.exception( + "PG scheduler: failed to fire periodic %r — leaving for next tick", + row.name, + ) + continue + + fired += 1 + logger.info( + "PG scheduler: fired periodic %r (%s) → %s (next_run_at=%s)", + row.name, + row.task_name, + row.queue, + nxt, + ) + + return fired diff --git a/workers/queue_backend/pg_queue/reaper.py b/workers/queue_backend/pg_queue/reaper.py index f31a86a93c..b43438c44f 100644 --- a/workers/queue_backend/pg_queue/reaper.py +++ b/workers/queue_backend/pg_queue/reaper.py @@ -70,7 +70,7 @@ from .leader_election import LeaderLease, default_worker_id from .liveness import LivenessServer as _BaseLivenessServer from .metrics import ReaperMetrics -from .pg_scheduler import dispatch_due_schedules +from .pg_scheduler import dispatch_due_periodic_tasks, dispatch_due_schedules from .recovery import mark_execution_error from .schema import qualified @@ -175,6 +175,46 @@ def reaper_interval_from_env() -> float: return value +#: How long an execution must sit non-terminal with EVERY file already terminal before +#: the safety net finalizes it (``WORKER_PG_STUCK_EXECUTION_RECOVERY_SECONDS``). +#: +#: Ten minutes, and deliberately NOT the barrier stuck-timeout it used to default to. +#: Those answer different questions: the barrier timeout bounds how long a batch may +#: make no progress (hours, because a single file can legitimately take that long), +#: whereas this must outlast the worst-case callback latency after the last file goes +#: terminal (see below — that is what sizes it). +#: +#: **The threshold IS load-bearing — do not lower it casually.** It was not, once: the +#: only guard was all-files-terminal (``internal_views.py``: ``total == 0 or +#: terminal < total → skipped``), and against that guard the value genuinely did not +#: matter. It does now. The selection also requires ``last_file_at < cutoff``, because +#: all-files-terminal alone leaves a live race: between the last file going terminal and +#: the callback finalizing, a HEALTHY execution passes every other check, and a sweep +#: landing there finalizes it first — the status guard then refuses the callback's own +#: write and the notification is lost. This threshold is what sizes that window, so it +#: must exceed the worst-case callback latency after the last file completes. An earlier +#: version of this comment said "seconds" and "no matter how short this is"; acting on +#: that today reopens the race. +#: +#: What remains true is that inheriting the multi-hour barrier timeout is wrong — it +#: answers a different question (how long a batch may make no progress, where one file +#: can legitimately take hours) and it cost a lot: a +#: stranded execution — the common shape being a Celery run whose callback worker was +#: removed by a deploy (its grace is 300s while file-processing gets 7200s) — sat dead +#: for ~2.5 h before anything noticed. +#: +#: A DEFAULT rather than something operators set, because production and on-prem have +#: no operator and no Flipt. A value that only exists as an env override is a value +#: those environments never get, and they are exactly the ones that cannot diagnose a +#: hung execution themselves. +#: +#: Not lowered further because ``total`` counts the file rows that EXIST, not +#: ``execution.total_files`` — while discovery is still creating rows, a brief +#: "all existing rows terminal" moment is possible. Ten minutes clears that comfortably. +#: Comparing against ``total_files`` would close the window properly and allow less. +_DEFAULT_STUCK_RECOVERY_SECONDS = 600 + + def _positive_duration_from_env(name: str, default: _N, cast: Callable[[str], _N]) -> _N: """Read a positive duration env var (default on unset; raise on invalid/<=0). @@ -339,6 +379,47 @@ def rearm_expired_claims(conn: PgConnection) -> int: raise +def promote_due_scheduled(conn: PgConnection) -> int: + """Make deferred queue messages claimable: ``scheduled`` + due -> ``ready``. + + Delayed visibility (UN-3843, Celery ``countdown``/``eta`` parity). A deferred + row is enqueued ``state='scheduled'`` with a future ``available_at`` and is + **absent from the claim's partial index**, so consumers cannot see it and it + costs the hot claim path nothing while it waits. This sweep is the only thing + that promotes it, which makes the reaper part of the delivery path for delayed + messages — not just the recovery path. It is already the mandatory singleton for + crash redelivery, so no new deployment dependency, but a queue that uses delays + inherits its liveness alert. + + Consequence to hold onto: delivery is **"not before ``available_at``"**, never + early, with granularity of the reaper tick (default 5s) — not exact ETA. Celery's + countdown is likewise approximate, and every consumer of this (staggered sends, + retry backoff) needs a floor rather than an instant. + + Cheap and bounded: ``pg_queue_message_scheduled_idx`` (partial, ``scheduled`` + only, keyed on ``available_at``) means a backlog of far-future rows costs one + index seek rather than a scan. Queue-agnostic — one statement promotes every due + row across all queues. Idempotent (a promoted row leaves the predicate); rolls + back on error. Runs every **leader** tick alongside + :func:`rearm_expired_claims`; see :meth:`PgReaper.tick`. + """ + ready = QueueMessageState.READY.value + scheduled = QueueMessageState.SCHEDULED.value + try: + with conn.cursor() as cur: + cur.execute( + f"UPDATE {qualified('pg_queue_message')} " + f"SET state = '{ready}' " + f"WHERE state = '{scheduled}' AND available_at <= now()" + ) + promoted = cur.rowcount + conn.commit() + return promoted + except Exception: + _rollback_after_sweep_failure(conn, "pg_queue_message") + raise + + def _execution_status( api_client: InternalAPIClient, execution_id: str, organization_id: str ) -> str | object | None: @@ -1038,15 +1119,13 @@ def __init__( ) if self._dedup_retention <= 0: raise ValueError("dedup_retention_seconds must be positive") - # Safety-net recovery of PG executions stranded non-terminal after all files + # Safety-net recovery of executions stranded non-terminal after all files # completed (barrier gone → invisible to the PG-table sweeps). Enablement + - # kill-switch semantics live in stuck_recovery_enabled_from_env(); the stuck - # window (below) defaults to the barrier stuck-timeout so it never races a - # legitimately long execution. + # kill-switch semantics live in stuck_recovery_enabled_from_env(). self._stuck_recovery_enabled = stuck_recovery_enabled_from_env() self._stuck_recovery_seconds = _positive_duration_from_env( "WORKER_PG_STUCK_EXECUTION_RECOVERY_SECONDS", - self._stuck_timeout_seconds, + _DEFAULT_STUCK_RECOVERY_SECONDS, int, ) # None → "never swept", so the first leader tick sweeps immediately; set to @@ -1186,6 +1265,31 @@ def tick(self) -> TickOutcome: ) self._discard_owned_sweep_conn() raise + # Delayed-visibility delivery (UN-3843): promote due 'scheduled' rows so + # consumers can claim them. Placed with the re-arm sweep because it shares + # its cadence requirement — this is the DELIVERY path for delayed messages, + # so a slower interval would directly add latency to every countdown/eta + # dispatch. Same failure posture as the re-arm above (dedicated counter, + # re-raise, discard the conn): a stalled promotion sweep means delayed + # messages silently never fire, which must not be swallowed. + try: + promoted = promote_due_scheduled(self._get_sweep_conn()) + if promoted: + self._metrics.queue_promoted.inc(promoted) + logger.info( + "Reaper: promoted %s due scheduled queue message(s) to 'ready' " + "(delayed-visibility delivery)", + promoted, + ) + except Exception: + self._metrics.queue_promote_failures.inc() + logger.exception( + "Reaper: promotion sweep failed — delayed (countdown/eta) queue " + "messages will not become claimable this tick " + "(see pg_reaper_queue_promote_failures_total)" + ) + self._discard_owned_sweep_conn() + raise # Orchestrator's second job: fire due PG-owned schedules (Beat # replacement). Ordered AFTER recovery so this cycle's recovery has # already completed before any scheduler error can propagate (the except @@ -1196,6 +1300,17 @@ def tick(self) -> TickOutcome: except Exception: self._discard_owned_sweep_conn() raise + # ...and the non-pipeline periodics (UN-3796): dashboard_metrics.*, + # log-history, audit, anything an operator adds. Separate call because each + # row carries its own task/args/queue rather than the pipeline trigger's one + # fixed shape; same leader gating, same dark-by-default posture (nothing + # fires until a row is pg_owned). Ordered after the pipeline dispatch so a + # fault here cannot stop pipelines, which are the customer-visible ones. + try: + dispatch_due_periodic_tasks(self._get_sweep_conn()) + except Exception: + self._discard_owned_sweep_conn() + raise # Orchestrator's third job: retention cleanup (cadence-gated, so it does # NOT run every tick). Last so a sweep error can't skip recovery/schedules. self._maybe_sweep() @@ -1249,8 +1364,59 @@ def _maybe_sweep(self) -> None: dedup, claims, ) + self._sweep_undispatched_executions() self._maybe_recover_stuck_executions() + def _sweep_undispatched_executions(self) -> None: + """Terminalise executions created but never dispatched (backend-side). + + The gap this closes: an abort between ``create_workflow_execution`` and + ``execute_workflow_async`` commits a PENDING row that was never queued. No + queue message and no barrier are ever created, so + :func:`recover_expired_barriers` — which scans ``pg_barrier_state`` — cannot + see it, and PENDING is not terminal, so nothing else resolves it either. + Observed as 967 orphans from one load test whose API tier shed 71% of traffic. + + Delegated to the backend rather than done here: this reaper deliberately never + touches backend tables (it reads/writes execution state through the internal + API), and the sweep also needs the rate limiter. + The reaper contributes what it uniquely has — leader election, so exactly one + instance sweeps — while the backend owns the logic. + + Cadence-gated with the retention sweeps: these rows are already older than the + grace period, so there is nothing to gain from running it every tick. + + **Swallowed on failure**, unlike barrier recovery. This cleans up work that is + already dead; a fault here must not discard the connection or abort the tick + and thereby defer schedule dispatch, which serves live traffic. + """ + try: + response = self._get_api_client().sweep_undispatched_executions() + except Exception: + self._metrics.undispatched_sweep_failures.inc() + logger.warning( + "Reaper: undispatched-execution sweep failed; retrying next sweep", + exc_info=True, + ) + return + # sweep_undispatched_executions() is typed `-> dict[str, Any]` and returns the + # parsed body verbatim — {"swept": N}, no {"data": ...} envelope. Reading + # `.data` off a plain dict yields None, so this counter was always 0: the + # metric never incremented and the "terminalised N" line never fired, leaving + # the 967-orphan condition this sweep exists to detect completely invisible. + # The sibling recover_stuck_pg_executions() hits the same flat-body shape and + # documents it (execution_client.py:352-358); it wraps there instead, which is + # why that one works. Read the dict directly here rather than add a second + # wrapping convention. + swept = (response if isinstance(response, dict) else {}).get("swept", 0) + if swept: + self._metrics.undispatched_swept.inc(swept) + logger.info( + "Reaper: terminalised %s undispatched execution(s) " + "(created but never queued)", + swept, + ) + def _maybe_recover_stuck_executions(self) -> None: """Opt-in safety-net: finalize PG executions stranded non-terminal after all files completed. diff --git a/workers/queue_backend/pg_queue/schema.py b/workers/queue_backend/pg_queue/schema.py index 92273ae026..98976efd76 100644 --- a/workers/queue_backend/pg_queue/schema.py +++ b/workers/queue_backend/pg_queue/schema.py @@ -53,6 +53,7 @@ "pg_orchestration_claim", "pg_orchestrator_lock", "pg_periodic_schedule", + "pg_periodic_task", } ) diff --git a/workers/run-worker-docker.sh b/workers/run-worker-docker.sh index be5aa86789..05d2b9398f 100755 --- a/workers/run-worker-docker.sh +++ b/workers/run-worker-docker.sh @@ -536,6 +536,19 @@ run_pg_consumer() { exec "$PG_QUEUE_PYTHON_BIN" -m pg_queue_consumer } +run_log_stream_consumer() { + ensure_pg_interpreter + export WORKER_NAME="${WORKER_NAME:-log-stream-consumer}" + + # Named for its transport, not the `pg-` family: this drains a Redis list, not + # pg_queue_message. It is nonetheless the flag-on replacement for the Celery + # worker-log-consumer, which is why it lives beside the PG components here. + print_status $GREEN "Starting log stream consumer (Redis transport)..." + print_status $BLUE "Queue: ${LOG_STREAM_QUEUE_NAME:-log_stream_queue}" + + exec "$PG_QUEUE_PYTHON_BIN" -m log_consumer.redis_stream_consumer +} + run_pg_reaper() { ensure_pg_interpreter export WORKER_NAME="${WORKER_NAME:-pg-reaper}" @@ -567,6 +580,9 @@ case "${1:-}" in pg-queue-reaper|pg-reaper|reaper) run_pg_reaper ;; + log-stream-consumer) + run_log_stream_consumer + ;; pg-*) # Obviously-PG-intended but unrecognized (e.g. a typo'd command) — fail # loudly instead of silently coercing it into a default Celery worker. diff --git a/workers/run-worker.sh b/workers/run-worker.sh index 0ec399deec..5c37466bff 100755 --- a/workers/run-worker.sh +++ b/workers/run-worker.sh @@ -57,6 +57,7 @@ readonly PG_ROLE_FILEPROC="pg-fileproc" readonly PG_ROLE_CALLBACK="pg-callback" readonly PG_ROLE_SCHEDULER="pg-scheduler" readonly PG_ROLE_EXECUTOR="pg-executor" +readonly PG_ROLE_METRICS="pg-metrics" declare -rA PG_CONSUMER_ROLES=( ["$PG_ROLE_ORCH_API"]="api_deployment;celery_api_deployments" ["$PG_ROLE_ORCH_GENERAL"]="general;celery" @@ -70,6 +71,14 @@ declare -rA PG_CONSUMER_ROLES=( # the result to pg_task_result for the blocking caller. Queues mirror the # Celery executor's CELERY_QUEUES_EXECUTOR. ["$PG_ROLE_EXECUTOR"]="executor;celery_executor_legacy,celery_executor_agentic,celery_executor_agentic_table" + # Runs the dashboard_metrics.* periodics fired by the PG scheduler tick, + # replacing the Celery 'workerMetrics' (-Q dashboard_metric_events). Its own + # role rather than a queue bolted onto pg-scheduler: the consumer's health + # heartbeat freezes for the duration of a task, so HEALTH_STALE is an upper + # bound on one task's wall clock. The aggregation runs for minutes; sharing + # pg-scheduler's tight 240s bound would restart that pod and take in-flight + # pipeline triggers down with it. Sized in docker-compose.yaml. + ["$PG_ROLE_METRICS"]="scheduler;dashboard_metric_events" ) declare -rA PG_QUEUE_MEMBERS=( ["$PG_QUEUE_CONSUMER_TYPE"]=1 @@ -80,6 +89,7 @@ declare -rA PG_QUEUE_MEMBERS=( ["$PG_ROLE_CALLBACK"]=1 ["$PG_ROLE_SCHEDULER"]=1 ["$PG_ROLE_EXECUTOR"]=1 + ["$PG_ROLE_METRICS"]=1 ) # The Celery transport set: every worker EXCEPT the PG-queue members — the # *complement* of the 'pg-queue' set, so the two transports' logs can be tailed @@ -119,6 +129,7 @@ declare -A WORKERS=( ["$PG_ROLE_CALLBACK"]="$PG_ROLE_CALLBACK" ["$PG_ROLE_SCHEDULER"]="$PG_ROLE_SCHEDULER" ["$PG_ROLE_EXECUTOR"]="$PG_ROLE_EXECUTOR" + ["$PG_ROLE_METRICS"]="$PG_ROLE_METRICS" # PG Queue reaper — leader-elected recovery loop (barrier-orphan sweep) ["reaper"]="$PG_QUEUE_REAPER_TYPE" ["pg-queue-reaper"]="$PG_QUEUE_REAPER_TYPE" diff --git a/workers/scheduler/dashboard_metrics_tasks.py b/workers/scheduler/dashboard_metrics_tasks.py new file mode 100644 index 0000000000..44bbe50440 --- /dev/null +++ b/workers/scheduler/dashboard_metrics_tasks.py @@ -0,0 +1,126 @@ +"""Thin PG-side tasks for the dashboard-metrics periodics (UN-3796). + +These replace ``workerMetrics`` (``celery -A backend worker -Q dashboard_metric_events``) +once the matching schedules are adopted by the PG scheduler. They do **no** work +themselves: each POSTs to a backend internal endpoint that calls the real Django +function, mirroring ``log_consumer/process_log_history.py`` and +``process_notification_buffer.py``. The workers image has no Django, so the ORM-heavy +aggregation cannot run here. + +**Registered under the exact Beat task names.** That is deliberate: the mirror copies a +Beat row verbatim (task, queue, args, kwargs), so matching names mean no remap table and +``--release`` stays a true inverse. The same name therefore exists in two registries — +the backend image (the Django implementation, consumed by Celery) and this one (the HTTP +proxy, consumed by ``worker-pg-metrics``). They are separate processes with separate +registries and neither imports the other; the name is a logical contract, not a symbol +clash. A reader who assumes otherwise will be very confused, hence this paragraph. + +Redelivery is safe without extra guards: the backend's upserts are +``INSERT … ON CONFLICT DO UPDATE SET`` (overwrite with recomputed values, not +increment) and the cleanups are ``DELETE … WHERE ts < cutoff``, so a double-run costs +duplicate DB work, never wrong numbers. +""" + +from __future__ import annotations + +import os +from typing import Any + +import httpx +from queue_backend import worker_task +from shared.infrastructure.logging import WorkerLogger + +logger = WorkerLogger.get_logger(__name__) + +# Sits just above gunicorn's --timeout 600, which is the ceiling that actually applies +# here (the task's Celery soft_time_limit=600 / time_limit=660 govern the Celery path, +# where a Celery worker runs the body in-process — not this one). Above it so a long +# aggregation surfaces the SERVER's error; a shorter client timeout would abort first +# and read as a network fault while the server kept working. +DEFAULT_HTTP_TIMEOUT_SECONDS = 630.0 + +_AGGREGATE_PATH = "v1/dashboard-metrics/aggregate/" +_CLEANUP_HOURLY_PATH = "v1/dashboard-metrics/cleanup/hourly/" +_CLEANUP_DAILY_PATH = "v1/dashboard-metrics/cleanup/daily/" + + +def _call_internal( + path: str, + *, + method: str = "POST", + body: dict[str, Any] | None = None, + timeout: float = DEFAULT_HTTP_TIMEOUT_SECONDS, +) -> dict[str, Any]: + """Call a backend internal endpoint and return its decoded body. + + Single seam so tests have one thing to patch. + + **Raises** on any failure, unlike ``process_log_history.py`` which returns False. + That script is driven by a bash loop with no other channel; this runs inside a PG + consumer, where raising is what marks the message failed and gets it logged loudly + (with ``MAX_ATTEMPTS=1`` it is then dropped rather than retried — the next cron tick + supersedes it). + + Never sends ``X-Organization-ID``: these are global aggregations, and the middleware + would otherwise scope every ORM read in them to a single tenant. + """ + base_url = os.getenv("INTERNAL_API_BASE_URL") + api_key = os.getenv("INTERNAL_SERVICE_API_KEY") + if not base_url: + raise RuntimeError("INTERNAL_API_BASE_URL environment variable not set") + if not api_key: + raise RuntimeError("INTERNAL_SERVICE_API_KEY environment variable not set") + + url = f"{base_url.rstrip('/')}/{path}" + # Transport-level retries only — these re-establish a connection that never + # delivered the request. They do NOT re-send after the server received it, which + # matters: a retry on a slow-but-live aggregation would run it twice concurrently. + transport = httpx.HTTPTransport(retries=3) + with httpx.Client(transport=transport) as client: + response = client.request( + method, + url, + headers={"Authorization": f"Bearer {api_key}"}, + json=body, + timeout=timeout, + ) + if response.status_code != 200: + raise RuntimeError( + f"{method} {path} failed: HTTP {response.status_code} {response.text[:500]}" + ) + return response.json() + + +def _log_if_skipped(name: str, result: dict[str, Any]) -> None: + """Surface a lock-held no-op. + + The backend returns success with ``skipped=True`` when the Redis lock is held. That + is correct behaviour, but left at INFO a permanently leaked lock looks like 96 + successful runs a day that did nothing. + """ + if result.get("skipped"): + logger.warning( + "%s did no work: %s", name, result.get("reason", "reported skipped=True") + ) + + +@worker_task(name="dashboard_metrics.aggregate_from_sources") +def dashboard_metrics_aggregate() -> dict[str, Any]: + """Aggregate source tables into the hourly/daily/monthly metrics tables.""" + result = _call_internal(_AGGREGATE_PATH) + _log_if_skipped("dashboard_metrics.aggregate_from_sources", result) + return result + + +@worker_task(name="dashboard_metrics.cleanup_hourly_data") +def dashboard_metrics_cleanup_hourly(retention_days: int | None = None) -> dict[str, Any]: + """Delete hourly metrics older than the retention window.""" + body = {"retention_days": retention_days} if retention_days is not None else None + return _call_internal(_CLEANUP_HOURLY_PATH, body=body) + + +@worker_task(name="dashboard_metrics.cleanup_daily_data") +def dashboard_metrics_cleanup_daily(retention_days: int | None = None) -> dict[str, Any]: + """Delete daily metrics older than the retention window.""" + body = {"retention_days": retention_days} if retention_days is not None else None + return _call_internal(_CLEANUP_DAILY_PATH, body=body) diff --git a/workers/scheduler/tasks.py b/workers/scheduler/tasks.py index 3c98fda9b9..b30742c450 100644 --- a/workers/scheduler/tasks.py +++ b/workers/scheduler/tasks.py @@ -9,6 +9,21 @@ from queue_backend import FairnessKey, QueueBackend, dispatch, worker_task from queue_backend.fairness import WorkloadType + +# Register the dashboard-metrics proxy tasks on this worker type (UN-3796). Imported +# purely for the side effect. +# +# ABSOLUTE package import, and it must stay that way — this module is reached by TWO +# different mechanisms and only this form survives both: +# 1. worker.py loads /app/scheduler/tasks.py BY PATH with the worker dir on sys.path. +# No parent package is set, so `from . import ...` cannot resolve. +# 2. general/tasks.py does `from scheduler.tasks import execute_pipeline_task_v2`, +# importing it as a package module. Then /app/scheduler is NOT on sys.path, so a +# bare `import dashboard_metrics_tasks` raises ModuleNotFoundError and crash-loops +# worker-general (this happened — the bare form shipped and broke the general +# worker at flag-off, where PG is not even involved). +# `/app` is on PYTHONPATH (run-worker-docker.sh), so `scheduler.` resolves under both. +from scheduler import dashboard_metrics_tasks # noqa: F401, E402 (side-effect import) from shared.enums.status_enums import PipelineStatus from shared.enums.worker_enums import QueueName from shared.infrastructure.config import WorkerConfig diff --git a/workers/shared/api/internal_client.py b/workers/shared/api/internal_client.py index a81287ccff..bba4af5296 100644 --- a/workers/shared/api/internal_client.py +++ b/workers/shared/api/internal_client.py @@ -1445,6 +1445,21 @@ def get( """Make GET request.""" return self.base_client.get(endpoint, params, organization_id) + def sweep_undispatched_executions(self) -> dict[str, Any]: + """Ask the backend to terminalise executions that were never dispatched. + + Closes a gap no worker can see: an abort between ``create_workflow_execution`` + and ``execute_workflow_async`` commits a PENDING execution row without ever + creating a queue message or a barrier, so the reaper's barrier scan cannot + find it and PENDING — a non-terminal state — is never resolved. + + Backend-side because the sweep touches ``WorkflowExecution``, the API-deployment + rate limiter; the reaper only supplies leader + election so exactly one instance runs it. Org-agnostic: the sweep spans all + organizations, so no ``organization_id`` is sent. + """ + return self.post("v1/execution-logs/sweep-undispatched-executions/", data={}) + def post( self, endpoint: str, data: dict[str, Any], organization_id: str | None = None ) -> dict[str, Any]: diff --git a/workers/tests/test_dashboard_metrics_tasks.py b/workers/tests/test_dashboard_metrics_tasks.py new file mode 100644 index 0000000000..ce8ff853ac --- /dev/null +++ b/workers/tests/test_dashboard_metrics_tasks.py @@ -0,0 +1,147 @@ +"""Tests for the thin dashboard-metrics proxy tasks (UN-3796). + +The tasks themselves do nothing but call a backend internal endpoint, so what is worth +pinning is the contract around that call: the registered names (a mismatch means the PG +consumer drops the message as an unknown task — the failure mode with no error at the +enqueue site), the request shape, and the failure posture. +""" + +from __future__ import annotations + +import sys +from pathlib import Path +from unittest.mock import MagicMock, patch + +import pytest + +# Import via the PACKAGE, matching `scheduler/tasks.py`'s +# `from scheduler import dashboard_metrics_tasks` (OSS 6397dd578) — which is the form +# BOTH runtime mechanisms converge on: worker.py's by-path load of tasks.py still +# resolves that import against /app, so `scheduler.dashboard_metrics_tasks` is the +# single module object either way. +# +# The bare `import dashboard_metrics_tasks` used here previously mirrored only the +# by-path load, and once tasks.py started importing the package form the two coexisted +# as SEPARATE module objects with separate Celery registrations. Patching one left the +# other live, so a test that believed it had mocked the HTTP client made a REAL request +# and failed on DNS — but only when some earlier test in the run had already imported +# `scheduler.tasks`, which is why it looked like flakiness rather than a wiring bug. +_WORKERS_ROOT = Path(__file__).resolve().parent.parent +if str(_WORKERS_ROOT) not in sys.path: + sys.path.insert(0, str(_WORKERS_ROOT)) + +from scheduler import dashboard_metrics_tasks as dmt # noqa: E402 + +_ENV = { + "INTERNAL_API_BASE_URL": "http://backend:8000/internal", + "INTERNAL_SERVICE_API_KEY": "test-key", +} + + +@pytest.fixture(autouse=True) +def _env(monkeypatch): + for k, v in _ENV.items(): + monkeypatch.setenv(k, v) + + +class TestRegistration: + """The names must match the Beat rows exactly, or the mirror's verbatim copy + produces a message no consumer can resolve — silently dropped as poison. + """ + + @pytest.mark.parametrize( + "name,func", + [ + ("dashboard_metrics.aggregate_from_sources", "dashboard_metrics_aggregate"), + ("dashboard_metrics.cleanup_hourly_data", "dashboard_metrics_cleanup_hourly"), + ("dashboard_metrics.cleanup_daily_data", "dashboard_metrics_cleanup_daily"), + ], + ) + def test_task_is_registered_under_the_beat_name(self, name, func): + assert getattr(dmt, func).name == name + + +class TestCallContract: + def test_aggregate_posts_to_the_aggregate_endpoint(self): + with patch.object(dmt, "_call_internal", return_value={"success": True}) as call: + dmt.dashboard_metrics_aggregate() + assert call.call_args[0][0] == "v1/dashboard-metrics/aggregate/" + + @pytest.mark.parametrize( + "func,path", + [ + ("dashboard_metrics_cleanup_hourly", "v1/dashboard-metrics/cleanup/hourly/"), + ("dashboard_metrics_cleanup_daily", "v1/dashboard-metrics/cleanup/daily/"), + ], + ) + def test_cleanup_passes_retention_through(self, func, path): + with patch.object(dmt, "_call_internal", return_value={"deleted": 1}) as call: + getattr(dmt, func)(retention_days=45) + assert call.call_args[0][0] == path + assert call.call_args.kwargs["body"] == {"retention_days": 45} + + def test_cleanup_omits_body_when_no_retention_given(self): + # The backend then applies the same default the Beat kwargs carry, so an + # unspecified call matches the Celery path rather than inventing a value here. + with patch.object(dmt, "_call_internal", return_value={}) as call: + dmt.dashboard_metrics_cleanup_hourly() + assert call.call_args.kwargs["body"] is None + + def test_lock_held_result_is_surfaced_not_swallowed(self, caplog): + # A permanently leaked lock otherwise looks like 96 successful no-op runs a day. + with patch.object( + dmt, + "_call_internal", + return_value={"success": True, "skipped": True, "reason": "lock_held"}, + ): + result = dmt.dashboard_metrics_aggregate() + assert result["skipped"] is True + + +class TestInternalCall: + def _response(self, status_code=200, payload=None): + r = MagicMock() + r.status_code = status_code + r.json.return_value = payload if payload is not None else {"ok": True} + r.text = "boom" + return r + + def test_sends_bearer_auth_and_never_an_org_header(self): + # X-Organization-ID would make the middleware scope every ORM read in these + # global aggregations to one tenant. + with patch.object(dmt.httpx, "Client") as client_cls: + client = client_cls.return_value.__enter__.return_value + client.request.return_value = self._response() + dmt._call_internal("v1/x/") + headers = client.request.call_args.kwargs["headers"] + assert headers == {"Authorization": "Bearer test-key"} + assert not any(h.lower() == "x-organization-id" for h in headers) + + def test_timeout_outlasts_the_server_side_ceiling(self): + # The ceiling that matters is gunicorn's --timeout 600, not the task's Celery + # time_limit=660 (no Celery worker runs it on this path). The client must sit + # ABOVE 600 so a long run surfaces the server's error rather than our own + # timeout — which would look like a network fault and hide the real cause. + with patch.object(dmt.httpx, "Client") as client_cls: + client = client_cls.return_value.__enter__.return_value + client.request.return_value = self._response() + dmt._call_internal("v1/x/") + assert client.request.call_args.kwargs["timeout"] > 600 + + def test_non_200_raises(self): + with patch.object(dmt.httpx, "Client") as client_cls: + client = client_cls.return_value.__enter__.return_value + client.request.return_value = self._response(status_code=500) + with pytest.raises(RuntimeError, match="HTTP 500"): + dmt._call_internal("v1/x/") + + @pytest.mark.parametrize( + "missing", ["INTERNAL_API_BASE_URL", "INTERNAL_SERVICE_API_KEY"] + ) + def test_missing_config_raises_rather_than_returning_falsy(self, monkeypatch, missing): + # Deliberately different from process_log_history.py, which returns False: that + # runs under a bash loop with no other channel. Here raising is what marks the + # message failed and gets it logged. + monkeypatch.delenv(missing) + with pytest.raises(RuntimeError, match=missing): + dmt._call_internal("v1/x/") diff --git a/workers/tests/test_log_stream_consumer.py b/workers/tests/test_log_stream_consumer.py new file mode 100644 index 0000000000..8c3f3091e6 --- /dev/null +++ b/workers/tests/test_log_stream_consumer.py @@ -0,0 +1,212 @@ +"""Tests for the Redis log-stream consumer's delivery semantics (UN-3755). + +The consumer replaces a Celery worker running with ``task_acks_late=True``, where a +crash mid-task means redelivery. A bare ``BLPOP`` would have quietly made crashes lossy, +so the loop parks each envelope on a per-pod processing list and removes it only after +the handler returns. These pin that contract — it is the part that is easy to regress +into "logs vanish when a pod restarts" without any test noticing. + +The module is loaded over faked worker-framework imports so the test needs neither a +Celery app nor a live Redis. +""" + +from __future__ import annotations + +import importlib.util +import json +import sys +import types +from pathlib import Path +from unittest.mock import MagicMock, patch + +import pytest + +_MODULE = ( + Path(__file__).resolve().parent.parent / "log_consumer" / "redis_stream_consumer.py" +) + + +def _load(monkeypatch): + """Import the consumer with its framework + task imports stubbed out.""" + monkeypatch.setenv("LOG_STREAM_QUEUE_NAME", "log_stream_queue") + monkeypatch.setenv("HOSTNAME", "pod-abc") + + def _mod(name, **attrs): + m = types.ModuleType(name) + for k, v in attrs.items(): + setattr(m, k, v) + return m + + logs_consumer = MagicMock(name="logs_consumer") + stubs = { + "shared": _mod("shared"), + "shared.enums": _mod("shared.enums"), + "shared.enums.worker_enums": _mod( + "shared.enums.worker_enums", + WorkerType=types.SimpleNamespace(LOG_CONSUMER="log_consumer"), + ), + "shared.infrastructure": _mod("shared.infrastructure"), + "shared.infrastructure.config": _mod("shared.infrastructure.config"), + "shared.infrastructure.config.builder": _mod( + "shared.infrastructure.config.builder", + WorkerBuilder=types.SimpleNamespace( + build_celery_app=lambda _t: (MagicMock(), MagicMock()) + ), + ), + "shared.infrastructure.logging": _mod( + "shared.infrastructure.logging", + WorkerLogger=types.SimpleNamespace(setup=lambda _t: MagicMock()), + ), + "log_consumer": _mod("log_consumer"), + "log_consumer.tasks": _mod("log_consumer.tasks", logs_consumer=logs_consumer), + } + for name, mod in stubs.items(): + monkeypatch.setitem(sys.modules, name, mod) + + spec = importlib.util.spec_from_file_location( + "log_consumer.redis_stream_consumer", _MODULE + ) + mod = importlib.util.module_from_spec(spec) + spec.loader.exec_module(mod) + mod._test_logs_consumer = logs_consumer + return mod + + +@pytest.fixture +def consumer(monkeypatch): + return _load(monkeypatch) + + +def _envelope(task="logs_consumer", **kwargs): + return json.dumps({"task": task, "kwargs": kwargs or {"event": "logs:c", "room": "c"}}) + + +class TestDispatch: + def test_runs_the_existing_task_body_with_the_envelope_kwargs(self, consumer): + consumer._dispatch(_envelope(event="logs:c1", user_session_id="c1")) + consumer._test_logs_consumer.assert_called_once_with( + event="logs:c1", user_session_id="c1" + ) + + def test_rejects_an_unexpected_task_name_loudly(self, consumer): + # Dispatching blind would run the log handler on a foreign payload; raising here + # surfaces a producer/consumer mismatch instead of corrupting the stream. + with pytest.raises(ValueError, match="Unexpected task"): + consumer._dispatch(_envelope(task="something_else")) + consumer._test_logs_consumer.assert_not_called() + + +class TestAtLeastOnceDelivery: + def test_processing_list_is_scoped_to_this_pod(self, consumer): + # A shared list would let one pod reclaim another's in-flight envelope and + # replay it while the owner is still working on it. + assert consumer._processing_list_name() == "log_stream_queue:processing:pod-abc" + + def test_startup_requeues_what_the_previous_run_left_in_flight(self, consumer): + redis = MagicMock() + redis.lmove.side_effect = [b"a", b"b", None] + consumer._recover_in_flight(redis, "proc") + assert redis.lmove.call_count == 3 + # Back to the HEAD of the source list, so recovered logs precede newer ones. + assert redis.lmove.call_args_list[0][0] == ("proc", "log_stream_queue", "RIGHT", "LEFT") + + def _one_shot_redis(self, consumer, raw): + """A redis mock that yields exactly one envelope, then ends the loop. + + ``lmove`` must return None or startup recovery spins forever — a real Redis + returns nil on an empty list, but a bare MagicMock is truthy. + """ + redis = MagicMock() + redis.lmove.return_value = None + + def _blmove(*_a, **_k): + if redis.blmove.call_count == 1: + return raw + consumer._shutdown = True + return None + + redis.blmove.side_effect = _blmove + return redis + + def test_envelope_is_removed_only_after_the_handler_returns(self, consumer): + raw = _envelope() + redis = self._one_shot_redis(consumer, raw) + order = [] + redis.lrem.side_effect = lambda *a: order.append("lrem") + consumer._test_logs_consumer.side_effect = lambda **_: order.append("handled") + + with patch.object(consumer, "create_redis_client", return_value=redis): + consumer.run() + + # Order is the whole point: lrem before the handler would lose the envelope on + # a crash, which is exactly the acks_late behaviour this replaces. + assert order == ["handled", "lrem"] + redis.lrem.assert_called_once_with("log_stream_queue:processing:pod-abc", 1, raw) + + def test_a_poison_envelope_is_dropped_not_replayed_forever(self, consumer): + raw = b"not-json" + redis = self._one_shot_redis(consumer, raw) + with patch.object(consumer, "create_redis_client", return_value=redis): + consumer.run() + + # Still removed from the processing list — otherwise startup recovery would + # re-queue it on every restart and the loop would never drain. + redis.lrem.assert_called_once_with("log_stream_queue:processing:pod-abc", 1, raw) + + +class TestSocketTimeoutOutlivesTheBlock: + """redis-py applies ``socket_timeout`` to the BLMOVE read itself. + + Shipped equal to the block (both 5s) and the socket won the race in integration: + ``redis.exceptions.TimeoutError: Timeout reading from socket`` every ~5.01s while + idle, each one tearing down the connection. The quiet part is worse — BLMOVE is + atomic server-side, so an envelope could be moved onto the processing list and its + reply then lost with the socket, stranding that log until the pod restarted. + """ + + def test_socket_timeout_strictly_exceeds_the_block_timeout(self, consumer): + assert consumer._SOCKET_TIMEOUT_SECONDS > consumer._BLOCK_TIMEOUT_SECONDS + + def test_the_margin_holds_when_the_block_is_tuned_up(self, monkeypatch): + # The two must stay related by construction, not by both happening to be + # defaults — a deployment raising the block alone would resurrect the bug. + monkeypatch.setenv("LOG_STREAM_BLOCK_TIMEOUT", "45") + mod = _load(monkeypatch) + assert mod._BLOCK_TIMEOUT_SECONDS == 45 + assert mod._SOCKET_TIMEOUT_SECONDS > 45 + + def test_the_client_is_actually_built_with_that_timeout(self, consumer): + """The constant is inert unless it reaches ``create_redis_client``.""" + redis = MagicMock() + redis.lmove.return_value = None + + def _blmove(*_a, **_k): + consumer._shutdown = True + return None + + redis.blmove.side_effect = _blmove + + with patch.object( + consumer, "create_redis_client", return_value=redis + ) as factory: + consumer.run() + + assert factory.call_args.kwargs["socket_timeout"] == ( + consumer._SOCKET_TIMEOUT_SECONDS + ) + + def test_blmove_is_called_with_the_block_timeout(self, consumer): + """Pins the other half of the pair the invariant is about.""" + redis = MagicMock() + redis.lmove.return_value = None + + def _blmove(*_a, **_k): + consumer._shutdown = True + return None + + redis.blmove.side_effect = _blmove + + with patch.object(consumer, "create_redis_client", return_value=redis): + consumer.run() + + assert redis.blmove.call_args[0][2] == consumer._BLOCK_TIMEOUT_SECONDS diff --git a/workers/tests/test_log_stream_transport.py b/workers/tests/test_log_stream_transport.py new file mode 100644 index 0000000000..34d8cf6f1c --- /dev/null +++ b/workers/tests/test_log_stream_transport.py @@ -0,0 +1,111 @@ +"""Tests for the log-streaming transport branch (UN-3755). + +``LogPublisher.publish`` is the hop that lets the log consumer stop being a Celery +worker. What matters is that the flag-off path is untouched (it is what staging and +production run), that the flag-on envelope carries the task **name** so the consumer can +dispatch it, and that a logging fault can never break an execution. +""" + +from __future__ import annotations + +import json +from unittest.mock import MagicMock, patch + +import pytest + +from unstract.core.constants import LogProcessingTask +from unstract.core.pubsub_helper import LogPublisher + +_PAYLOAD = {"type": "LOG", "level": "INFO", "log": "hello", "timestamp": 1.0} + + +@pytest.fixture +def redis_client(): + client = MagicMock() + client.llen.return_value = 0 + with patch.object(LogPublisher, "_get_redis_client", return_value=client): + yield client + + +@pytest.fixture +def kombu(): + with patch.object(LogPublisher, "kombu_conn") as conn: + yield conn.Producer.return_value.__enter__.return_value + + +class TestFlagOff: + """Default transport must remain exactly what it is on main.""" + + def test_publishes_to_amqp_and_never_touches_the_redis_stream( + self, monkeypatch, kombu, redis_client + ): + monkeypatch.delenv("LOG_TRANSPORT", raising=False) + assert LogPublisher.publish("chan-1", _PAYLOAD) is True + + kombu.publish.assert_called_once() + kwargs = kombu.publish.call_args.kwargs + assert kwargs["routing_key"] == LogProcessingTask.QUEUE_NAME + assert kwargs["headers"] == {"task": LogProcessingTask.TASK_NAME} + redis_client.rpush.assert_not_called() + + def test_an_unrelated_transport_value_still_uses_celery( + self, monkeypatch, kombu, redis_client + ): + # Fail closed: only the exact opt-in switches transport. + monkeypatch.setenv("LOG_TRANSPORT", "rabbit") + LogPublisher.publish("chan-1", _PAYLOAD) + kombu.publish.assert_called_once() + redis_client.rpush.assert_not_called() + + +class TestFlagOn: + @pytest.fixture(autouse=True) + def _enable(self, monkeypatch): + monkeypatch.setenv("LOG_TRANSPORT", "redis") + monkeypatch.setenv("LOG_STREAM_QUEUE_NAME", "log_stream_queue") + + def test_pushes_an_envelope_and_never_touches_amqp(self, kombu, redis_client): + assert LogPublisher.publish("chan-1", _PAYLOAD) is True + kombu.publish.assert_not_called() + + queue, raw = redis_client.rpush.call_args[0] + assert queue == "log_stream_queue" + envelope = json.loads(raw) + # The task name is what the consumer dispatches on. Without it a rename would + # silently drop every log with nothing at the publish site to trace it from. + assert envelope["task"] == LogProcessingTask.TASK_NAME + assert envelope["kwargs"]["message"] == _PAYLOAD + assert envelope["kwargs"]["user_session_id"] == "chan-1" + assert envelope["kwargs"]["event"] == "logs:chan-1" + + def test_drops_at_capacity_rather_than_growing_until_redis_ooms(self, redis_client): + redis_client.llen.return_value = 10_000 + assert LogPublisher.publish("chan-1", _PAYLOAD) is True + redis_client.rpush.assert_not_called() + + def test_a_redis_failure_is_swallowed_not_raised(self, redis_client): + # A logging fault must never surface into the execution that emitted it. + redis_client.rpush.side_effect = RuntimeError("redis down") + assert LogPublisher.publish("chan-1", _PAYLOAD) is False + + def test_unified_notification_still_stored(self, redis_client): + with patch.object(LogPublisher, "store_for_unified_notification") as store: + LogPublisher.publish("chan-1", _PAYLOAD) + store.assert_called_once() + assert store.call_args[0][0] == "logs:chan-1" + + +class TestEnvelopeParity: + """Both transports must carry the same kwargs, or the consumer behaves differently + depending on how the log arrived.""" + + def test_kwargs_match_across_transports(self, monkeypatch, kombu, redis_client): + monkeypatch.delenv("LOG_TRANSPORT", raising=False) + LogPublisher.publish("chan-1", _PAYLOAD) + celery_kwargs = kombu.publish.call_args.kwargs["body"]["kwargs"] + + monkeypatch.setenv("LOG_TRANSPORT", "redis") + LogPublisher.publish("chan-1", _PAYLOAD) + redis_kwargs = json.loads(redis_client.rpush.call_args[0][1])["kwargs"] + + assert celery_kwargs == redis_kwargs diff --git a/workers/tests/test_pg_finalization_fixes.py b/workers/tests/test_pg_finalization_fixes.py index 14981a1f5c..0645244595 100644 --- a/workers/tests/test_pg_finalization_fixes.py +++ b/workers/tests/test_pg_finalization_fixes.py @@ -117,6 +117,57 @@ def test_non_dict_body_yields_empty_data(self): _RECOVERY_ENV = "WORKER_PG_STUCK_EXECUTION_RECOVERY_ENABLED" +_RECOVERY_SECONDS_ENV = "WORKER_PG_STUCK_EXECUTION_RECOVERY_SECONDS" + + +class TestStuckRecoveryWindowDefault: + """The window must be a DEFAULT, not something an operator sets. + + It used to inherit the barrier stuck-timeout (~2.5 h), which answers a different + question: that one bounds how long a batch may make no progress, where a single + file legitimately can take hours. This one only needs to outlast a callback that + is about to fire — seconds. + + The threshold is not what makes recovery safe; the endpoint skips unless EVERY + file is terminal, so a running execution is never a candidate however short this + is. Inheriting hours therefore bought nothing and left a stranded execution dead + for ~2.5 h — the common shape being a Celery run whose callback worker was removed + by a deploy (300 s grace, against file-processing's 7200 s). + + Pinned as a default because production and on-prem have **no operator and no + Flipt**: a value that only exists as an env override is one those environments + never receive, and they are precisely the ones that cannot diagnose a hung + execution themselves. + """ + + def _seconds(self, monkeypatch, value=None): + from queue_backend.pg_queue import reaper as R + + monkeypatch.delenv(_RECOVERY_SECONDS_ENV, raising=False) + if value is not None: + monkeypatch.setenv(_RECOVERY_SECONDS_ENV, value) + return R._positive_duration_from_env( + _RECOVERY_SECONDS_ENV, R._DEFAULT_STUCK_RECOVERY_SECONDS, int + ) + + def test_default_is_ten_minutes(self, monkeypatch): + assert self._seconds(monkeypatch) == 600 + + def test_default_is_NOT_the_barrier_stuck_timeout(self, monkeypatch): + """The regression this replaces. Any multi-hour value here means a stranded + execution stays dead for hours in an environment with nobody to notice.""" + assert self._seconds(monkeypatch) < 3600 + + def test_env_still_overrides_for_tuning(self, monkeypatch): + assert self._seconds(monkeypatch, "120") == 120 + + def test_a_nonsense_value_raises_rather_than_silently_defaulting(self, monkeypatch): + """Loud-on-misconfig, matching the other duration knobs: silently falling back + would hide a typo'd ConfigMap behind behaviour that looks deliberate.""" + import pytest + + with pytest.raises(ValueError): + self._seconds(monkeypatch, "not-a-number") class TestStuckRecoveryDefaultOn: diff --git a/workers/tests/test_pg_metrics.py b/workers/tests/test_pg_metrics.py index 55f4d47750..e21bc2b67f 100644 --- a/workers/tests/test_pg_metrics.py +++ b/workers/tests/test_pg_metrics.py @@ -401,6 +401,8 @@ def test_leader_tick_refreshes_gauges(self): with ( patch.object(reaper_mod, "recover_expired_barriers", return_value=[]), patch.object(reaper_mod, "rearm_expired_claims", return_value=0), + patch.object(reaper_mod, "promote_due_scheduled", return_value=0), + patch.object(reaper_mod, "dispatch_due_periodic_tasks", return_value=0), patch.object(reaper_mod, "refresh_queue_gauges") as refresh, ): reaper.tick() @@ -414,6 +416,8 @@ def test_refresh_is_cadence_gated_and_resumes(self): with ( patch.object(reaper_mod, "recover_expired_barriers", return_value=[]), patch.object(reaper_mod, "rearm_expired_claims", return_value=0), + patch.object(reaper_mod, "promote_due_scheduled", return_value=0), + patch.object(reaper_mod, "dispatch_due_periodic_tasks", return_value=0), patch.object(reaper_mod, "refresh_queue_gauges") as refresh, ): reaper.tick() @@ -432,6 +436,8 @@ def test_refresh_failure_never_fails_the_tick_and_consumes_the_interval(self): with ( patch.object(reaper_mod, "recover_expired_barriers", return_value=[]), patch.object(reaper_mod, "rearm_expired_claims", return_value=0), + patch.object(reaper_mod, "promote_due_scheduled", return_value=0), + patch.object(reaper_mod, "dispatch_due_periodic_tasks", return_value=0), patch.object( reaper_mod, "refresh_queue_gauges", side_effect=RuntimeError("db") ) as refresh, @@ -458,6 +464,8 @@ def test_lost_leadership_clears_snapshot_and_resets_cadence(self): with ( patch.object(reaper_mod, "recover_expired_barriers", return_value=[]), patch.object(reaper_mod, "rearm_expired_claims", return_value=0), + patch.object(reaper_mod, "promote_due_scheduled", return_value=0), + patch.object(reaper_mod, "dispatch_due_periodic_tasks", return_value=0), patch.object(reaper_mod, "refresh_queue_gauges") as refresh, ): reaper.tick() # becomes leader, refresh #1 @@ -479,6 +487,8 @@ def test_renew_raise_clears_snapshot_before_propagating(self): with ( patch.object(reaper_mod, "recover_expired_barriers", return_value=[]), patch.object(reaper_mod, "rearm_expired_claims", return_value=0), + patch.object(reaper_mod, "promote_due_scheduled", return_value=0), + patch.object(reaper_mod, "dispatch_due_periodic_tasks", return_value=0), patch.object(reaper_mod, "refresh_queue_gauges"), ): reaper.tick() # becomes leader @@ -499,6 +509,8 @@ def test_sweep_failure_increments_labeled_counter(self, monkeypatch): with ( patch.object(reaper_mod, "recover_expired_barriers", return_value=[]), patch.object(reaper_mod, "rearm_expired_claims", return_value=0), + patch.object(reaper_mod, "promote_due_scheduled", return_value=0), + patch.object(reaper_mod, "dispatch_due_periodic_tasks", return_value=0), patch.object(reaper_mod, "refresh_queue_gauges"), ): reaper.tick() # first leader tick sweeps immediately @@ -513,6 +525,8 @@ def test_sweep_failure_increments_labeled_counter(self, monkeypatch): with ( patch.object(reaper_mod, "recover_expired_barriers", return_value=[]), patch.object(reaper_mod, "rearm_expired_claims", return_value=0), + patch.object(reaper_mod, "promote_due_scheduled", return_value=0), + patch.object(reaper_mod, "dispatch_due_periodic_tasks", return_value=0), patch.object(reaper_mod, "refresh_queue_gauges"), ): reaper.tick() diff --git a/workers/tests/test_pg_queue_client.py b/workers/tests/test_pg_queue_client.py index a79ae82088..3fa1a4baac 100644 --- a/workers/tests/test_pg_queue_client.py +++ b/workers/tests/test_pg_queue_client.py @@ -26,8 +26,9 @@ from queue_backend.pg_queue import PgQueueClient, QueueMessage from queue_backend.pg_queue.client import _SEND_RETRY_BACKOFF_SECONDS from queue_backend.pg_queue.connection import create_pg_connection -from queue_backend.pg_queue.reaper import rearm_expired_claims +from queue_backend.pg_queue.reaper import promote_due_scheduled, rearm_expired_claims from queue_backend.pg_queue.schema import qualified +from unstract.core.data_models import QueueMessageState # --- Unit: SQL shape against a mocked connection --- @@ -670,6 +671,82 @@ def test_reaper_does_not_rearm_a_live_lease(self, pg_conn, queue_name): assert rearm_expired_claims(pg_conn) == 0 # vt not expired → untouched assert client.read(queue_name, vt_seconds=30, qty=10) == [] # still claimed + def _schedule(self, pg_conn, queue_name, seconds_from_now): + """Insert a deferred row the way the backend producer writes one.""" + state = QueueMessageState.SCHEDULED.value + with pg_conn.cursor() as cur: + cur.execute( + f"INSERT INTO {qualified('pg_queue_message')} " + "(queue_name, message, org_id, priority, enqueued_at, vt, read_ct, " + " state, available_at) " + "VALUES (%s, %s::jsonb, '', 5, now(), now(), 0, %s, " + " now() + make_interval(secs => %s)) " + "RETURNING msg_id", + (queue_name, '{"n": 1}', state, seconds_from_now), + ) + msg_id = cur.fetchone()[0] + pg_conn.commit() + return msg_id + + def test_scheduled_row_is_invisible_to_the_claim(self, pg_conn, queue_name): + # The core delayed-visibility guarantee (UN-3843): a not-yet-due row must + # never be handed to a consumer. It is enforced structurally — 'scheduled' + # is absent from the claim's partial index — not by a time predicate. + client = PgQueueClient(conn=pg_conn) + self._schedule(pg_conn, queue_name, 3600) + assert client.read(queue_name, vt_seconds=30, qty=10) == [] + + def test_promotion_makes_a_due_row_claimable(self, pg_conn, queue_name): + # ...and the reaper sweep is what releases it. Negative offset = already due. + client = PgQueueClient(conn=pg_conn) + msg_id = self._schedule(pg_conn, queue_name, -1) + assert client.read(queue_name, vt_seconds=30, qty=10) == [] # still deferred + assert promote_due_scheduled(pg_conn) == 1 + claimed = client.read(queue_name, vt_seconds=30, qty=10) + assert [m.msg_id for m in claimed] == [msg_id] + + def test_promotion_leaves_not_yet_due_rows_alone(self, pg_conn, queue_name): + # The sweep must not release early — "not before available_at" is the whole + # contract. A bug here fires a staggered send all at once. + self._schedule(pg_conn, queue_name, 3600) + assert promote_due_scheduled(pg_conn) == 0 + assert PgQueueClient(conn=pg_conn).read(queue_name, vt_seconds=30, qty=10) == [] + + def test_promotion_is_idempotent(self, pg_conn, queue_name): + # A promoted row leaves the predicate, so a second tick is a no-op rather + # than re-arming an already-claimed message back to 'ready' (double-run). + client = PgQueueClient(conn=pg_conn) + self._schedule(pg_conn, queue_name, -1) + assert promote_due_scheduled(pg_conn) == 1 + assert promote_due_scheduled(pg_conn) == 0 + client.read(queue_name, vt_seconds=30, qty=10) # now claimed + assert promote_due_scheduled(pg_conn) == 0 # claimed rows are not re-promoted + + def test_rearm_does_not_touch_scheduled_rows(self, pg_conn, queue_name): + # The two sweeps share a table and both write state='ready'; re-arm keys on + # 'claimed' only, so a deferred row can't be released by the crash path. + self._schedule(pg_conn, queue_name, 3600) + assert rearm_expired_claims(pg_conn) == 0 + assert PgQueueClient(conn=pg_conn).read(queue_name, vt_seconds=30, qty=10) == [] + + def test_default_available_at_keeps_the_raw_insert_working(self, pg_conn, queue_name): + # Deploy-safety guard: the workers' enqueue SQL does not mention + # available_at, so the column MUST keep a DB-level default (migration 0002 + # re-adds the one Django's AddField drops). Without it every worker enqueue + # fails with a not-null violation, in any deploy order. + client = PgQueueClient(conn=pg_conn) + msg_id = client.send(queue_name, {"n": 1}) + with pg_conn.cursor() as cur: + cur.execute( + f"SELECT available_at, state FROM {qualified('pg_queue_message')} " + "WHERE msg_id = %s", + (msg_id,), + ) + available_at, state = cur.fetchone() + assert state == QueueMessageState.READY.value + assert available_at is not None + assert client.read(queue_name, vt_seconds=30, qty=10) # immediately claimable + def test_state_check_constraint_matches_enum(self, pg_conn, queue_name): # Drift guard (mirrors test_db_check_constraint_matches_fairness_bounds): # the DB CheckConstraint must accept exactly the QueueMessageState values diff --git a/workers/tests/test_pg_reaper.py b/workers/tests/test_pg_reaper.py index 4a6f56f55c..add6d8c5ce 100644 --- a/workers/tests/test_pg_reaper.py +++ b/workers/tests/test_pg_reaper.py @@ -31,6 +31,7 @@ dedup_retention_from_env, reaper_interval_from_env, reaper_sweep_interval_from_env, + promote_due_scheduled, rearm_expired_claims, recover_expired_barriers, sweep_expired_results, @@ -49,6 +50,7 @@ def stub_scheduler_tick(monkeypatch): mock = MagicMock(return_value=0) monkeypatch.setattr(reaper_mod, "dispatch_due_schedules", mock) + monkeypatch.setattr(reaper_mod, "dispatch_due_periodic_tasks", MagicMock(return_value=0)) return mock @@ -63,6 +65,16 @@ def stub_queue_rearm(monkeypatch): return mock +# The leader tick also promotes due scheduled queue messages (UN-3843 delayed +# visibility). Same reason as the re-arm stub above: the leadership / connection +# tests drive dummy connections that would blow up on a real UPDATE. +@pytest.fixture(autouse=True) +def stub_queue_promote(monkeypatch): + mock = MagicMock(return_value=0) + monkeypatch.setattr(reaper_mod, "promote_due_scheduled", mock) + return mock + + # The leader tick also runs the retention sweep (UN-3610). Stub the two sweep # helpers by default so the leadership / connection tests don't hit a real DELETE # on their dummy connections; the SQL-contract tests import the real helpers @@ -417,6 +429,82 @@ def test_rearm_error_discards_conn_and_counts_failure(self, stub_queue_rearm): assert reaper.metrics.queue_rearm_failures._value.get() == 1 +class TestQueuePromoteTick: + """UN-3843: the leader promotes due scheduled queue messages each cycle, next to + the re-arm sweep. Promotion SQL is in TestRetentionSweepSql; here we assert the + wiring, gating, metric guard and error posture (the autouse + ``stub_queue_promote`` patches the helper). + + Cadence matters more here than for the other sweeps: this is the DELIVERY path + for every countdown/eta dispatch, so running it anywhere but the per-tick + recovery block would add latency to each delayed message. + """ + + def _reaper(self, lease): + return PgReaper( + lease, interval_seconds=0.01, sweep_conn=object(), api_client=object() + ) + + def test_leader_runs_promotion(self, stub_queue_promote): + reaper = self._reaper(_FakeLease(acquires=True, renews=True)) + with patch.object(reaper_mod, "recover_expired_barriers", return_value=[]): + reaper.tick() + stub_queue_promote.assert_called_once() + + def test_standby_does_not_promote(self, stub_queue_promote): + # Two promoters would race on the same rows; leadership is what serialises it. + reaper = self._reaper(_FakeLease(acquires=False)) + with patch.object(reaper_mod, "recover_expired_barriers"): + reaper.tick() + stub_queue_promote.assert_not_called() + + def test_promotion_runs_after_rearm_before_schedule( + self, stub_queue_rearm, stub_queue_promote, stub_scheduler_tick + ): + order = [] + reaper = self._reaper(_FakeLease(acquires=True, renews=True)) + stub_queue_rearm.side_effect = lambda *_: order.append("rearm") or 0 + stub_queue_promote.side_effect = lambda *_: order.append("promote") or 0 + stub_scheduler_tick.side_effect = lambda *_: order.append("schedule") + with patch.object( + reaper_mod, + "recover_expired_barriers", + side_effect=lambda *_, **__: order.append("recover") or [], + ): + reaper.tick() + assert order == ["recover", "rearm", "promote", "schedule"] + + def test_metric_incremented_only_when_nonzero(self, stub_queue_promote): + reaper = self._reaper(_FakeLease(acquires=True, renews=True)) + counter = reaper.metrics.queue_promoted + with patch.object(reaper_mod, "recover_expired_barriers", return_value=[]): + stub_queue_promote.return_value = 0 # nothing due + reaper.tick() + assert counter._value.get() == 0 # guarded by `if promoted:` + stub_queue_promote.return_value = 2 + reaper.tick() + assert counter._value.get() == 2 + + def test_promote_error_discards_conn_and_counts_failure(self, stub_queue_promote): + # Failure must not be swallowed: the symptom of a silently-stalled promotion + # sweep is delayed messages that simply never fire, with no error at the + # enqueue site to trace it back from. Dedicated counter, discard, re-raise. + reaper = PgReaper( + _FakeLease(acquires=True, renews=True), + interval_seconds=0.01, + api_client=object(), + ) + owned = MagicMock() + owned.closed = False + reaper._sweep_conn = owned + stub_queue_promote.side_effect = psycopg2.OperationalError("db gone") + with patch.object(reaper_mod, "recover_expired_barriers", return_value=[]): + with pytest.raises(psycopg2.OperationalError): + reaper.tick() + assert reaper._sweep_conn is None # discarded + assert reaper.metrics.queue_promote_failures._value.get() == 1 + + class TestRetentionSweepSql: """The sweep helpers' SQL contract (mock cursor, no DB). These call the real helpers (imported at module load), unaffected by the autouse stub which patches @@ -462,17 +550,31 @@ def test_rearm_expired_claims_sql(self): assert "WHERE state = 'claimed' AND vt <= now()" in sql conn.commit.assert_called_once() + def test_promote_due_scheduled_sql(self): + # UN-3843 delayed-visibility helper. The predicate is the whole contract: + # 'scheduled' scopes it to deferred rows (so it can never re-arm a claimed + # message), and `available_at <= now()` is what makes delivery "not before" + # the requested time. Its integration test is Postgres-gated. + conn, cur = self._conn_cur(2) + assert promote_due_scheduled(conn) == 2 + sql = cur.execute.call_args[0][0] + assert f"UPDATE {qualified('pg_queue_message')}" in sql + assert "SET state = 'ready'" in sql + assert "WHERE state = 'scheduled' AND available_at <= now()" in sql + conn.commit.assert_called_once() + @pytest.mark.parametrize( "sweep", [ lambda conn: sweep_expired_results(conn), lambda conn: sweep_orphan_dedup(conn, 60), lambda conn: rearm_expired_claims(conn), + lambda conn: promote_due_scheduled(conn), ], - ids=["expired_results", "orphan_dedup", "rearm_claims"], + ids=["expired_results", "orphan_dedup", "rearm_claims", "promote_scheduled"], ) def test_sweep_rolls_back_on_error(self, sweep): - # Both helpers have their own try/except/rollback — exercise each. + # Each helper owns its try/except/rollback — exercise every one. cur = MagicMock() cur.execute.side_effect = psycopg2.OperationalError("dead") conn = MagicMock() diff --git a/workers/tests/test_pg_scheduler.py b/workers/tests/test_pg_scheduler.py index 5086d3c121..e5ea53bc08 100644 --- a/workers/tests/test_pg_scheduler.py +++ b/workers/tests/test_pg_scheduler.py @@ -7,6 +7,7 @@ """ import datetime +import json import uuid import psycopg2 @@ -16,11 +17,47 @@ SCHEDULER_QUEUE_NAME, _build_trigger_payload, compute_next_run, + dispatch_due_periodic_tasks, dispatch_due_schedules, ) class TestPureHelpers: + def test_the_field_names_ARE_the_select_column_lists(self): + """The one guarantee the _fields-derived SELECT cannot give itself. + + Deriving the column list from the type makes a REORDER harmless — there is no + second list to drift from. It makes a RENAME silent in a new way: the query + simply starts asking for a different column. Every test that would execute + these queries needs live Postgres and SKIPS without it, so a rename would + otherwise reach production unchallenged. + + Pinning the tuples turns that into a red test in the default lane, which forces + whoever renames a field to confirm a matching migration in + backend/pg_queue/models.py. Do NOT update these literals to match a rename + without doing so — that is the mistake this exists to stop. + """ + from queue_backend.pg_queue.pg_scheduler import _DuePeriodicTask, _DueSchedule + + assert _DueSchedule._fields == ( + "pipeline_id", + "organization_id", + "workflow_id", + "pipeline_name", + "cron_string", + "next_run_at", + ) + assert _DuePeriodicTask._fields == ( + "name", + "task_name", + "queue", + "task_args", + "task_kwargs", + "org_id", + "cron_string", + "next_run_at", + ) + def test_next_run_is_strictly_after_base(self): base = datetime.datetime(2026, 6, 18, 10, 0, 0) # 09:00 daily — already past at 10:00, so the next is tomorrow 09:00. @@ -80,8 +117,16 @@ def _seed(conn, *, pg_owned, enabled, next_run_at, cron="0 9 * * *"): created_at, updated_at) VALUES (%s, %s, %s, %s, %s, %s, %s, NULL, %s, now(), now()) """, - (pid, _MARKER, str(uuid.uuid4()), "Test ETL", cron, enabled, pg_owned, - next_run_at), + ( + pid, + _MARKER, + str(uuid.uuid4()), + "Test ETL", + cron, + enabled, + pg_owned, + next_run_at, + ), ) conn.commit() return pid @@ -111,7 +156,9 @@ def clean(pg_conn): """Remove any rows this test created (the tick commits, so teardown must).""" yield pg_conn with pg_conn.cursor() as cur: - cur.execute("DELETE FROM pg_periodic_schedule WHERE organization_id = %s", (_MARKER,)) + cur.execute( + "DELETE FROM pg_periodic_schedule WHERE organization_id = %s", (_MARKER,) + ) cur.execute("DELETE FROM pg_queue_message WHERE org_id = %s", (_MARKER,)) pg_conn.commit() @@ -275,3 +322,188 @@ def execute(self, sql, params=None): if self._fail_when(sql): raise psycopg2.OperationalError("forced failure") return self._cur.execute(sql, params) + + +# --- generic (non-pipeline) periodics: UN-3796 --- + +_PT_MARKER = f"test_pgperiodic_{uuid.uuid4().hex[:8]}" +_PT_QUEUE = f"{_PT_MARKER}_queue" + + +def _seed_periodic( + conn, + *, + pg_owned, + enabled, + next_run_at, + cron="0 9 * * *", + queue=_PT_QUEUE, + task_args=None, + task_kwargs=None, + org_id="", +): + """Insert one pg_periodic_task row; returns its name.""" + name = f"{_PT_MARKER}_{uuid.uuid4().hex[:6]}" + with conn.cursor() as cur: + cur.execute( + """ + INSERT INTO pg_periodic_task + (name, task_name, queue, task_args, task_kwargs, org_id, + cron_string, enabled, pg_owned, last_run_at, next_run_at, + created_at, updated_at) + VALUES (%s, %s, %s, %s::jsonb, %s::jsonb, %s, %s, %s, %s, NULL, %s, + now(), now()) + """, + ( + name, + "dashboard_metrics.aggregate_from_sources", + queue, + json.dumps(task_args if task_args is not None else []), + json.dumps(task_kwargs if task_kwargs is not None else {}), + org_id, + cron, + enabled, + pg_owned, + next_run_at, + ), + ) + conn.commit() + return name + + +def _periodic_row(conn, name): + with conn.cursor() as cur: + cur.execute( + "SELECT last_run_at, next_run_at, enabled FROM pg_periodic_task " + "WHERE name = %s", + (name,), + ) + return cur.fetchone() + + +def _messages_on(conn, queue): + with conn.cursor() as cur: + cur.execute( + "SELECT message FROM pg_queue_message WHERE queue_name = %s", (queue,) + ) + return [r[0] for r in cur.fetchall()] + + +@pytest.fixture +def clean_periodic(pg_conn): + yield pg_conn + with pg_conn.cursor() as cur: + cur.execute( + "DELETE FROM pg_periodic_task WHERE name LIKE %s", (f"{_PT_MARKER}%",) + ) + cur.execute( + "DELETE FROM pg_queue_message WHERE queue_name LIKE %s", (f"{_PT_MARKER}%",) + ) + pg_conn.commit() + + +class TestDispatchDuePeriodicTasks: + """UN-3796 — the Beat replacement for everything that isn't a pipeline.""" + + def test_due_owned_row_fires_onto_its_own_queue_and_advances(self, clean_periodic): + conn = clean_periodic + past = datetime.datetime(2020, 1, 1, tzinfo=datetime.timezone.utc) + name = _seed_periodic( + conn, + pg_owned=True, + enabled=True, + next_run_at=past, + task_kwargs={"retention_days": 30}, + ) + + assert dispatch_due_periodic_tasks(conn) == 1 + + # Enqueued onto the row's OWN queue — the reason this can't reuse the + # pipeline dispatcher, which targets one fixed queue. + msgs = _messages_on(conn, _PT_QUEUE) + assert len(msgs) == 1 + assert msgs[0]["task_name"] == "dashboard_metrics.aggregate_from_sources" + assert msgs[0]["kwargs"] == {"retention_days": 30} + assert msgs[0]["queue"] == _PT_QUEUE + + last_run, next_run, _ = _periodic_row(conn, name) + assert last_run is not None and next_run > last_run + + def test_org_id_is_carried_onto_the_queue_message(self, clean_periodic): + # org_id is empty for every periodic today (these are global jobs), but the + # column exists so a later unification with pg_periodic_schedule — which is + # org-scoped — stays a data migration. Pin that the dispatcher actually + # threads it through rather than hardcoding "". + conn = clean_periodic + past = datetime.datetime(2020, 1, 1, tzinfo=datetime.timezone.utc) + _seed_periodic( + conn, pg_owned=True, enabled=True, next_run_at=past, org_id="org_acme" + ) + assert dispatch_due_periodic_tasks(conn) == 1 + with conn.cursor() as cur: + cur.execute( + "SELECT org_id FROM pg_queue_message WHERE queue_name = %s", + (_PT_QUEUE,), + ) + assert cur.fetchone()[0] == "org_acme" + + def test_first_observation_baselines_without_firing(self, clean_periodic): + # No burst when a row is handed over: Beat fires a new schedule at its next + # cron match, not immediately, and the hand-over must match that. + conn = clean_periodic + name = _seed_periodic(conn, pg_owned=True, enabled=True, next_run_at=None) + + assert dispatch_due_periodic_tasks(conn) == 0 + assert _messages_on(conn, _PT_QUEUE) == [] + last_run, next_run, _ = _periodic_row(conn, name) + assert last_run is None and next_run is not None + + @pytest.mark.parametrize( + "pg_owned,enabled", [(False, True), (True, False), (False, False)] + ) + def test_unowned_or_disabled_rows_never_fire(self, clean_periodic, pg_owned, enabled): + # Dark by default: until a row is explicitly adopted, Beat is the only firer. + conn = clean_periodic + past = datetime.datetime(2020, 1, 1, tzinfo=datetime.timezone.utc) + _seed_periodic(conn, pg_owned=pg_owned, enabled=enabled, next_run_at=past) + assert dispatch_due_periodic_tasks(conn) == 0 + assert _messages_on(conn, _PT_QUEUE) == [] + + def test_not_yet_due_row_does_not_fire(self, clean_periodic): + conn = clean_periodic + future = datetime.datetime(2999, 1, 1, tzinfo=datetime.timezone.utc) + _seed_periodic(conn, pg_owned=True, enabled=True, next_run_at=future) + assert dispatch_due_periodic_tasks(conn) == 0 + assert _messages_on(conn, _PT_QUEUE) == [] + + def test_invalid_cron_disables_the_row_instead_of_looping(self, clean_periodic): + # Otherwise it is re-selected and re-logs a traceback every tick, forever. + conn = clean_periodic + past = datetime.datetime(2020, 1, 1, tzinfo=datetime.timezone.utc) + name = _seed_periodic( + conn, pg_owned=True, enabled=True, next_run_at=past, cron="not a cron" + ) + assert dispatch_due_periodic_tasks(conn) == 0 + assert _periodic_row(conn, name)[2] is False # enabled -> False + assert _messages_on(conn, _PT_QUEUE) == [] + + def test_second_tick_does_not_refire_an_advanced_row(self, clean_periodic): + # The enqueue and the next_run_at advance share one transaction, so a row + # that fired is not due again until its next cron match. + conn = clean_periodic + past = datetime.datetime(2020, 1, 1, tzinfo=datetime.timezone.utc) + _seed_periodic(conn, pg_owned=True, enabled=True, next_run_at=past) + assert dispatch_due_periodic_tasks(conn) == 1 + assert dispatch_due_periodic_tasks(conn) == 0 + assert len(_messages_on(conn, _PT_QUEUE)) == 1 + + def test_rows_are_isolated_from_each_other(self, clean_periodic): + # A bad cron on one row must not stop the others firing that tick. + conn = clean_periodic + past = datetime.datetime(2020, 1, 1, tzinfo=datetime.timezone.utc) + _seed_periodic( + conn, pg_owned=True, enabled=True, next_run_at=past, cron="garbage" + ) + _seed_periodic(conn, pg_owned=True, enabled=True, next_run_at=past) + assert dispatch_due_periodic_tasks(conn) == 1 + assert len(_messages_on(conn, _PT_QUEUE)) == 1 diff --git a/workers/tests/test_pg_schema_drift.py b/workers/tests/test_pg_schema_drift.py index 2dbcbe2a38..384ae357a0 100644 --- a/workers/tests/test_pg_schema_drift.py +++ b/workers/tests/test_pg_schema_drift.py @@ -45,6 +45,11 @@ "read_ct", "priority", "state", + # Delayed visibility (countdown/eta), migration 0002. Absent from this set + # until now, while workers depend on it in raw SQL — so renaming or dropping + # it would have kept this guard green and broken delayed delivery at runtime, + # which is exactly the drift this file exists to catch. + "available_at", }, "pg_task_result": { "task_id", @@ -79,6 +84,24 @@ "created_at", "updated_at", }, + # UN-3796 — generic (non-pipeline) Beat periodics. Sibling of the table above; + # each row carries its own task/queue/args rather than the pipeline trigger's + # one fixed shape. + "pg_periodic_task": { + "name", + "task_name", + "queue", + "task_args", + "task_kwargs", + "org_id", + "cron_string", + "enabled", + "pg_owned", + "last_run_at", + "next_run_at", + "created_at", + "updated_at", + }, } diff --git a/workers/tests/test_reaper_undispatched_sweep.py b/workers/tests/test_reaper_undispatched_sweep.py new file mode 100644 index 0000000000..099c1eb82a --- /dev/null +++ b/workers/tests/test_reaper_undispatched_sweep.py @@ -0,0 +1,94 @@ +"""The reaper's trigger for the backend undispatched-execution sweep. + +The sweep's *logic* lives in the backend (it touches WorkflowExecution, the rate +limiter and API storage — none of which a worker may reach). What the reaper +contributes is leader election, so exactly one instance runs it. These pin the +wiring, which is the part the backend tests cannot see: + + * it is CADENCE-GATED, not per-tick — the rows are already older than the grace + period, so running it every 5s would be pure DB load; + * a failure is SWALLOWED — unlike barrier recovery, this cleans up work that is + already dead, and discarding the connection would defer schedule dispatch, which + serves live traffic. +""" + +from __future__ import annotations + +from types import SimpleNamespace +from unittest.mock import MagicMock, patch + +from queue_backend.pg_queue import reaper as reaper_mod + + +def _reaper_with(api_client): + """A PgReaper stand-in exposing only what _sweep_undispatched_executions uses. + + Built with __new__ so no DB/lease/metrics registry is required; the two collaborators + the method actually touches are injected. + """ + r = reaper_mod.PgReaper.__new__(reaper_mod.PgReaper) + r._get_api_client = lambda: api_client # type: ignore[method-assign] + r._metrics = MagicMock() + return r + + +class TestTheTriggerIsWiredToTheBackend: + def test_it_calls_the_backend_sweep(self): + api = MagicMock() + # The REAL client returns the parsed body verbatim (typed `-> dict[str, Any]`), + # not an object with `.data`. This fixture used to be SimpleNamespace(data=...), + # a shape nothing in the call chain produces — so it asserted inc(3) while + # production read `.data` off a dict and always incremented by 0. The test + # passed against a fiction and hid a dead metric. + api.sweep_undispatched_executions.return_value = {"swept": 3} + r = _reaper_with(api) + r._sweep_undispatched_executions() + api.sweep_undispatched_executions.assert_called_once_with() + # Counted by how many were recovered, not "one sweep ran" — the rate is the + # signal that requests are dying before dispatch. + r._metrics.undispatched_swept.inc.assert_called_once_with(3) + + def test_a_zero_result_is_not_an_error(self): + """The steady state. Must stay quiet, not log every 5 minutes forever.""" + api = MagicMock() + api.sweep_undispatched_executions.return_value = {"swept": 0} + r = _reaper_with(api) + r._sweep_undispatched_executions() # no raise + # A zero must NOT touch the counter: inc(0) is harmless but a nonzero rate is + # the alert signal, so keep the series clean. + r._metrics.undispatched_swept.inc.assert_not_called() + + def test_a_missing_or_odd_payload_does_not_raise(self): + """A backend on an older image returns no `swept` key — that must not take + down a leader tick that also dispatches schedules. + """ + for payload in (None, {}, {"unexpected": 1}, SimpleNamespace(data={"swept": 9})): + api = MagicMock() + api.sweep_undispatched_executions.return_value = payload + _reaper_with(api)._sweep_undispatched_executions() # no raise + + +class TestAFailureCannotBreakTheTick: + def test_an_api_error_is_swallowed(self): + """Barrier recovery re-raises and discards the connection because a stranded + RUNNING execution is urgent. This is the opposite case — the executions are + already dead — and raising here would abort the tick and defer schedule + dispatch for live traffic. Cleanup must never outrank scheduling. + """ + api = MagicMock() + api.sweep_undispatched_executions.side_effect = RuntimeError("backend down") + r = _reaper_with(api) + r._sweep_undispatched_executions() # must not raise + # Swallowed, but counted — otherwise a persistently failing sweep is invisible + # and PENDING rows accumulate unrecovered. + r._metrics.undispatched_sweep_failures.inc.assert_called_once() + + def test_the_failure_is_logged_so_it_is_not_silent(self): + """Swallowed is not the same as hidden — a persistently failing sweep means + PENDING rows accumulate, and that has to be visible. + """ + api = MagicMock() + api.sweep_undispatched_executions.side_effect = RuntimeError("backend down") + with patch.object(reaper_mod.logger, "warning") as warn: + _reaper_with(api)._sweep_undispatched_executions() + assert warn.called diff --git a/workers/tests/test_worker_task_module_imports.py b/workers/tests/test_worker_task_module_imports.py new file mode 100644 index 0000000000..df54943500 --- /dev/null +++ b/workers/tests/test_worker_task_module_imports.py @@ -0,0 +1,156 @@ +"""Every worker's tasks.py must import under BOTH mechanisms worker.py uses. + +This exists because a bare ``import dashboard_metrics_tasks`` in scheduler/tasks.py +shipped and crash-looped **worker-general** in integration: + + File "/app/general/tasks.py", line 16, in + from scheduler.tasks import execute_pipeline_task_v2 + File "/app/scheduler/tasks.py", line 14, in + import dashboard_metrics_tasks + ModuleNotFoundError: No module named 'dashboard_metrics_tasks' + +The bare form resolves only when ``worker.py`` loads ``scheduler/tasks.py`` BY PATH with +``/app/scheduler`` appended to ``sys.path``. It does not resolve when another worker +imports it as a package module — ``general/tasks.py`` does exactly that, and then +``/app/scheduler`` is not on the path. Flag-off; PG is not involved. + +The entire workers suite passed while this was broken, because nothing exercised the two +import mechanisms. That is the gap these close. + +**Why subprocesses.** Importing the same file under two names (``dashboard_metrics_tasks`` +and ``scheduler.dashboard_metrics_tasks``) creates two module objects and duplicate +Celery task registrations, which breaks other suites that patch one copy — doing this +in-process made test_dashboard_metrics_tasks fail. A fresh interpreter per mechanism is +both properly isolated and a truer reproduction of what a booting worker does. +""" + +from __future__ import annotations + +import json +import subprocess +import sys +from pathlib import Path + +import pytest + +_WORKERS_ROOT = Path(__file__).resolve().parent.parent + +# Discovered, not hardcoded, so a new worker is covered automatically. +_TASK_MODULES = sorted( + p.parent.name + for p in _WORKERS_ROOT.glob("*/tasks.py") + if not p.parent.name.startswith((".", "_")) and p.parent.name != "tests" +) + +# Mechanism 2 — `from .tasks import ...`, the way general/tasks.py reaches +# scheduler/tasks.py. Only the workers root is on sys.path, never the worker's own +# directory: that is the situation that broke. +_PACKAGE_IMPORT_PROBE = """ +import importlib, json, sys, traceback +sys.path.insert(0, {root!r}) +failures = {{}} +for name in {workers!r}: + try: + importlib.import_module(name + ".tasks") + except Exception: + failures[name] = traceback.format_exc().strip().splitlines()[-1] +print("RESULT" + json.dumps(failures)) +""" + +# Mechanism 1 — worker.py's spec_from_file_location load with the worker dir appended, +# i.e. how a worker loads its own tasks.py at boot. +_PATH_IMPORT_PROBE = """ +import importlib.util, json, sys, traceback +sys.path.insert(0, {root!r}) +failures = {{}} +for name in {workers!r}: + wd = {root!r} + "/" + name + sys.path.append(wd) + try: + spec = importlib.util.spec_from_file_location("tasks_" + name.replace("-", "_"), + wd + "/tasks.py") + mod = importlib.util.module_from_spec(spec) + spec.loader.exec_module(mod) + except Exception: + failures[name] = traceback.format_exc().strip().splitlines()[-1] + finally: + sys.path.remove(wd) +print("RESULT" + json.dumps(failures)) +""" + + +def _run_probe(source: str, workers: list[str] | None = None) -> dict[str, str]: + proc = subprocess.run( + [sys.executable, "-c", + source.format(root=str(_WORKERS_ROOT), workers=workers or _TASK_MODULES)], + capture_output=True, + text=True, + cwd=str(_WORKERS_ROOT), + timeout=300, + ) + marker = [ln for ln in proc.stdout.splitlines() if ln.startswith("RESULT")] + if not marker: + pytest.fail(f"probe did not complete.\nstdout:\n{proc.stdout}\nstderr:\n{proc.stderr}") + return json.loads(marker[-1][len("RESULT") :]) + + +def test_discovery_found_the_workers() -> None: + """Guard against the glob matching nothing and every case passing vacuously.""" + assert len(_TASK_MODULES) >= 5, _TASK_MODULES + assert "scheduler" in _TASK_MODULES + assert "general" in _TASK_MODULES + + +# A hyphenated directory (api-deployment) can never be imported package-style — it is +# not a valid Python identifier — so only the file-path mechanism applies to it. +_PACKAGE_IMPORTABLE = [w for w in _TASK_MODULES if w.isidentifier()] + + +def test_every_tasks_module_imports_as_a_package() -> None: + """The mechanism that broke: worker's own dir NOT on sys.path.""" + failures = _run_probe(_PACKAGE_IMPORT_PROBE, workers=_PACKAGE_IMPORTABLE) + assert not failures, "package-style import failed:\n" + "\n".join( + f" {k}: {v}" for k, v in sorted(failures.items()) + ) + + +def test_every_tasks_module_imports_by_file_path() -> None: + """The mechanism worker.py uses for a worker's own tasks.py.""" + failures = _run_probe(_PATH_IMPORT_PROBE) + assert not failures, "file-path import failed:\n" + "\n".join( + f" {k}: {v}" for k, v in sorted(failures.items()) + ) + + +def test_scheduler_registers_the_metrics_proxies_under_their_wire_names() -> None: + """The broken import is a SIDE-EFFECT import — it must actually register. + + Asserting only that scheduler.tasks imports would still pass if someone 'fixed' it + by deleting the import, silently unregistering the three dashboard-metrics proxies + that the PG metrics consumer resolves BY NAME. + """ + source = """ +import importlib, json, sys +sys.path.insert(0, {root!r}) +m = importlib.import_module("scheduler.dashboard_metrics_tasks") +print("RESULT" + json.dumps({{ + "aggregate": m.dashboard_metrics_aggregate.name, + "hourly": m.dashboard_metrics_cleanup_hourly.name, + "daily": m.dashboard_metrics_cleanup_daily.name, +}})) +""" + proc = subprocess.run( + [sys.executable, "-c", source.format(root=str(_WORKERS_ROOT))], + capture_output=True, + text=True, + cwd=str(_WORKERS_ROOT), + timeout=300, + ) + marker = [ln for ln in proc.stdout.splitlines() if ln.startswith("RESULT")] + assert marker, f"probe failed:\n{proc.stdout}\n{proc.stderr}" + names = json.loads(marker[-1][len("RESULT") :]) + assert names == { + "aggregate": "dashboard_metrics.aggregate_from_sources", + "hourly": "dashboard_metrics.cleanup_hourly_data", + "daily": "dashboard_metrics.cleanup_daily_data", + }