From b7590230532d097df89b82d397fe31e68b13b2b3 Mon Sep 17 00:00:00 2001 From: ali Date: Wed, 5 Aug 2026 10:59:14 +0530 Subject: [PATCH 01/33] =?UTF-8?q?UN-3843=20[GATED-FEAT]=20PG=20queue=20?= =?UTF-8?q?=E2=80=94=20delayed-visibility=20primitive=20(countdown/eta)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The PG queue delivered every message immediately, so any dispatch relying on Celery's countdown/eta had to stay Celery-only. Adds the deferral primitive. A deferred row is written state='scheduled' with a future available_at and is absent from the claim's partial index (which covers only 'ready'), so a pending delay costs the hot claim path nothing. The reaper promotes it to 'ready' once available_at passes, on the same per-tick cadence as the crash re-arm. Deliberately NOT `available_at <= now()` in the claim, which is what the ticket originally proposed: that parks every not-yet-due row inside pg_queue_message_claim_idx to be walked and discarded on every claim, which is exactly the scan-past cost the state-machine claim was introduced to remove. The trade is granularity — delivery is "not before available_at", never early, at reaper-tick resolution (default 5s) rather than exact ETA. Celery's countdown is likewise approximate, and the consumers of this (staggered sends, retry backoff) need a floor, not an instant. - unstract.core: QueueMessageState.SCHEDULED (shared enum, drift-tested). - backend: available_at column, widened state constraint, partial pg_queue_message_scheduled_idx, and countdown/eta on enqueue_task. Non-positive countdown / past eta resolve to the immediate path so a computed-zero stagger step doesn't pay a tick. - workers: promote_due_scheduled() + tick wiring, with a dedicated failure counter and re-raise — a silently stalled sweep means delayed messages never fire, with nothing at the enqueue site to trace it back from. Migration 0002 re-adds a persistent `DEFAULT now()` that Django's AddField drops. Without it the workers' raw enqueue (explicit column list, no available_at) would fail with a not-null violation in ANY deploy order. Covered by test_default_available_at_keeps_the_raw_insert_working. Additive and inert: existing rows and every existing enqueue resolve to available_at=now() / state='ready', unchanged. Flag-off untouched. Tests: 22 new (8 producer DB-free, 6 worker integration on live Postgres, 6 reaper wiring, SQL contract + rollback parametrisation). 1464 workers pass, 47 backend pg_queue pass. Co-Authored-By: Claude Opus 5 (1M context) --- .../0002_pgqueuemessage_available_at.py | 67 +++++++++++ backend/pg_queue/models.py | 34 +++++- backend/pg_queue/producer.py | 57 ++++++++++ backend/pg_queue/tests/test_producer.py | 89 +++++++++++++++ .../core/src/unstract/core/data_models.py | 14 ++- workers/queue_backend/pg_queue/metrics.py | 14 +++ workers/queue_backend/pg_queue/reaper.py | 66 +++++++++++ workers/tests/test_pg_metrics.py | 7 ++ workers/tests/test_pg_queue_client.py | 79 ++++++++++++- workers/tests/test_pg_reaper.py | 105 +++++++++++++++++- 10 files changed, 526 insertions(+), 6 deletions(-) create mode 100644 backend/pg_queue/migrations/0002_pgqueuemessage_available_at.py 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/models.py b/backend/pg_queue/models.py index abd5bab447..c6992dee4b 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", + ), ] 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_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/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/workers/queue_backend/pg_queue/metrics.py b/workers/queue_backend/pg_queue/metrics.py index 5678b38198..a3d77c43fb 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 " diff --git a/workers/queue_backend/pg_queue/reaper.py b/workers/queue_backend/pg_queue/reaper.py index f31a86a93c..c60cd90789 100644 --- a/workers/queue_backend/pg_queue/reaper.py +++ b/workers/queue_backend/pg_queue/reaper.py @@ -339,6 +339,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: @@ -1186,6 +1227,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 diff --git a/workers/tests/test_pg_metrics.py b/workers/tests/test_pg_metrics.py index 55f4d47750..4dd3550a4f 100644 --- a/workers/tests/test_pg_metrics.py +++ b/workers/tests/test_pg_metrics.py @@ -401,6 +401,7 @@ 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, "refresh_queue_gauges") as refresh, ): reaper.tick() @@ -414,6 +415,7 @@ 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, "refresh_queue_gauges") as refresh, ): reaper.tick() @@ -432,6 +434,7 @@ 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, "refresh_queue_gauges", side_effect=RuntimeError("db") ) as refresh, @@ -458,6 +461,7 @@ 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, "refresh_queue_gauges") as refresh, ): reaper.tick() # becomes leader, refresh #1 @@ -479,6 +483,7 @@ 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, "refresh_queue_gauges"), ): reaper.tick() # becomes leader @@ -499,6 +504,7 @@ 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, "refresh_queue_gauges"), ): reaper.tick() # first leader tick sweeps immediately @@ -513,6 +519,7 @@ 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, "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..4127b4f1aa 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, @@ -63,6 +64,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 +428,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 +549,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() From 9e406be4fcd2a2bc460bbfa3de06466c4ffe8fdf Mon Sep 17 00:00:00 2001 From: ali Date: Wed, 5 Aug 2026 14:09:57 +0530 Subject: [PATCH 02/33] UN-3796 [GATED-FEAT] PG scheduler: fire non-pipeline Beat periodics (scheduling half) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Beat is the only thing still firing periodic tasks, so it blocks the "scale every Celery deployment to zero" gate. The PG scheduler already existed but mirrored and fired ONLY pipeline schedules. This adds the generic half. - PgPeriodicTask: a sibling of PgPeriodicSchedule keyed on PeriodicTask.name, carrying task_name/queue/args/kwargs/org_id + cron and run-state. A sibling rather than widening the existing table: that one is live, mid-ramp and dual-written from five sites, and the two carry different ownership POLICIES (per-pipeline Flipt percentage vs all-or-nothing operator adopt). org_id is present from the start so a later unification is a data migration, not a redesign. - dispatch_due_periodic_tasks(): same shape as the pipeline dispatcher — leader gated, per-row txn, enqueue+advance in ONE transaction, baseline-without-firing on first observation, invalid-cron quiesce. Differs only in that each row carries its own task/args/queue, which is the whole reason for the second table. - mirror_pg_periodic_tasks: mirrors generically from Beat (the live PeriodicTask table is the only authority — DatabaseScheduler keeps schedules as rows, not code). --adopt flips pg_owned AND disables the Beat row in one transaction; a row PG-owned while Beat still has it enabled fires twice. Running it against a real Beat table surfaced three things source could not: celery.backend_cleanup (Celery's own result-backend housekeeping — excluded, it retires with Celery); a legacy execute_pipeline_task_v2 row (a pipeline trigger under a second task path, which would have given one pipeline two owners — now excluded by task path, not by luck); and workflow_log_history_v2 at a 30-second interval, which has no cron expression and is skipped loudly rather than rounded up to a minute. Scale, per review: excluded paths are filtered in SQL, not Python — there is one PeriodicTask row per scheduled pipeline, so the earlier .all() dragged the whole pipeline population through memory to discard it. Both commands now use .iterator(chunk_size=) with --batch-size; reconcile_pg_schedules also stops materialising every mirrored id as a set. Its five existing tests pass unchanged (behaviour is identical); new tests pin the chunking so it can't be silently undone. Inert: rows land pg_owned=False, so the PG scheduler fires nothing and Beat keeps firing everything until an operator adopts a row. Flag-off untouched. Tests: 1475 workers, 72 backend pg_queue. Co-Authored-By: Claude Opus 5 (1M context) --- .../commands/mirror_pg_periodic_tasks.py | 349 ++++++++++++++++++ .../commands/reconcile_pg_schedules.py | 57 ++- .../migrations/0003_pgperiodictask.py | 40 ++ backend/pg_queue/models.py | 78 ++++ .../tests/test_mirror_pg_periodic_tasks.py | 157 ++++++++ .../test_reconcile_pg_schedules_command.py | 121 +++++- .../queue_backend/pg_queue/pg_scheduler.py | 144 ++++++++ workers/queue_backend/pg_queue/reaper.py | 13 +- workers/queue_backend/pg_queue/schema.py | 1 + workers/tests/test_pg_metrics.py | 7 + workers/tests/test_pg_reaper.py | 1 + workers/tests/test_pg_scheduler.py | 180 +++++++++ workers/tests/test_pg_schema_drift.py | 18 + 13 files changed, 1139 insertions(+), 27 deletions(-) create mode 100644 backend/pg_queue/management/commands/mirror_pg_periodic_tasks.py create mode 100644 backend/pg_queue/migrations/0003_pgperiodictask.py create mode 100644 backend/pg_queue/tests/test_mirror_pg_periodic_tasks.py 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..370c7efd43 --- /dev/null +++ b/backend/pg_queue/management/commands/mirror_pg_periodic_tasks.py @@ -0,0 +1,349 @@ +"""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_celery_beat.models import IntervalSchedule, PeriodicTask + +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 where one exists exactly + (``*/N`` minutes / hours / days). + + 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. + """ + 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: + return f"*/{every} * * * *" + if interval.period == IntervalSchedule.HOURS and every < 24: + return f"0 */{every} * * *" + if interval.period == IntervalSchedule.DAYS and every < 32: + return f"0 0 */{every} * *" + # SECONDS, or a step too large to express as a single cron field. + 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: + PgPeriodicTask.objects.update_or_create( + name=plan.name, defaults=plan.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 + verb = "adopt" if to_pg else "release" + self.stdout.write(f"{verb} {row.name!r} (beat enabled -> {not to_pg})") + 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"]) + PeriodicTask.objects.filter(name=row.name).update(enabled=not to_pg) + 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..53c77df1f6 100644 --- a/backend/pg_queue/management/commands/reconcile_pg_schedules.py +++ b/backend/pg_queue/management/commands/reconcile_pg_schedules.py @@ -23,6 +23,10 @@ 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,11 +55,25 @@ 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." + ), + ) 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"] + if batch_size < 1: + raise CommandError("--batch-size must be >= 1") + backfilled = self._backfill_mirrors(dry_run, batch_size) + reconciled, pg_owned, failed = self._reconcile_all(dry_run, batch_size) prefix = "[dry-run] " if dry_run else "" summary = ( @@ -93,15 +111,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 +154,18 @@ def _backfill_mirrors(self, dry_run: bool) -> int: backfilled += 1 return backfilled - def _reconcile_all(self, dry_run: bool) -> tuple[int, int, int]: + 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/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 c6992dee4b..500c6d9084 100644 --- a/backend/pg_queue/models.py +++ b/backend/pg_queue/models.py @@ -382,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/tests/test_mirror_pg_periodic_tasks.py b/backend/pg_queue/tests/test_mirror_pg_periodic_tasks.py new file mode 100644 index 0000000000..9945a30277 --- /dev/null +++ b/backend/pg_queue/tests/test_mirror_pg_periodic_tasks.py @@ -0,0 +1,157 @@ +"""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 + +import pytest + +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 + (2, HOURS, "0 */2 * * *"), + (3, DAYS, "0 0 */3 * *"), + ], + ) + 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))) == "" + + 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 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..a02fc3464b 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,10 +75,10 @@ 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 @@ -64,9 +94,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 +112,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 +127,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 +145,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/workers/queue_backend/pg_queue/pg_scheduler.py b/workers/queue_backend/pg_queue/pg_scheduler.py index 7ccfe40be1..47c185d0d7 100644 --- a/workers/queue_backend/pg_queue/pg_scheduler.py +++ b/workers/queue_backend/pg_queue/pg_scheduler.py @@ -230,3 +230,147 @@ 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`. Same reason for existing: the field names are + bound to the SELECT's column order at exactly one site. + """ + + 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 name, task_name, queue, task_args, task_kwargs, org_id, + cron_string, next_run_at + 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 c60cd90789..c0ab3d1317 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 @@ -1262,6 +1262,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() 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/tests/test_pg_metrics.py b/workers/tests/test_pg_metrics.py index 4dd3550a4f..e21bc2b67f 100644 --- a/workers/tests/test_pg_metrics.py +++ b/workers/tests/test_pg_metrics.py @@ -402,6 +402,7 @@ def test_leader_tick_refreshes_gauges(self): 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() @@ -416,6 +417,7 @@ def test_refresh_is_cadence_gated_and_resumes(self): 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() @@ -435,6 +437,7 @@ def test_refresh_failure_never_fails_the_tick_and_consumes_the_interval(self): 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, @@ -462,6 +465,7 @@ def test_lost_leadership_clears_snapshot_and_resets_cadence(self): 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 @@ -484,6 +488,7 @@ def test_renew_raise_clears_snapshot_before_propagating(self): 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 @@ -505,6 +510,7 @@ def test_sweep_failure_increments_labeled_counter(self, monkeypatch): 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 @@ -520,6 +526,7 @@ def test_sweep_failure_increments_labeled_counter(self, monkeypatch): 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_reaper.py b/workers/tests/test_pg_reaper.py index 4127b4f1aa..add6d8c5ce 100644 --- a/workers/tests/test_pg_reaper.py +++ b/workers/tests/test_pg_reaper.py @@ -50,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 diff --git a/workers/tests/test_pg_scheduler.py b/workers/tests/test_pg_scheduler.py index 5086d3c121..91b4af04bd 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,6 +17,7 @@ SCHEDULER_QUEUE_NAME, _build_trigger_payload, compute_next_run, + dispatch_due_periodic_tasks, dispatch_due_schedules, ) @@ -275,3 +277,181 @@ 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..2fbd84cdb9 100644 --- a/workers/tests/test_pg_schema_drift.py +++ b/workers/tests/test_pg_schema_drift.py @@ -79,6 +79,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", + }, } From c6f77e4255a1b0340ed7d613b3ba7f622d1eb69e Mon Sep 17 00:00:00 2001 From: ali Date: Wed, 5 Aug 2026 14:23:41 +0530 Subject: [PATCH 03/33] UN-3796 [GATED-FEAT] PG metrics execution path: worker-pg-metrics + internal API MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The scheduling half (previous commit) fires dashboard_metrics.* onto dashboard_metric_events, but nothing on PG could execute them: the three tasks are Django (ORM, cache, a Redis lock) and the PG consumer bootstraps its tasks from workers/ with no Django. Firing them would have stranded the messages — nothing drains that queue on PG. Execution follows the pattern this repo already uses for Django-side periodic work (process_log_history.py / process_notification_buffer.py -> backend internal API): - backend: dashboard_metrics/internal_views.py + internal_urls.py, registered under /internal/v1/dashboard-metrics/. The three task bodies are plain functions that happen to carry @shared_task, so they are called VERBATIM — aggregation windows and the Redis lock reused, not reimplemented. - workers: three thin @worker_task proxies registered under the EXACT Beat task names, so the mirror's verbatim copy needs no remap and --release stays a true inverse. The same name now exists in two registries (backend = the Django implementation for Celery; workers = the HTTP proxy for PG) — separate processes, separate registries, neither imports the other. Documented loudly in both. - deployment: a pg-metrics role + worker-pg-metrics compose service, NOT another queue on worker-pg-scheduler. The consumer's health heartbeat freezes while a task runs, so HEALTH_STALE is an upper bound on one task's wall clock; the scheduler's 240s bound against a minutes-long aggregation would trip the liveness probe, restart the pod and take in-flight pipeline triggers with it. Sized 900/960, concurrency 1, MAX_ATTEMPTS=1 (for a periodic the next cron tick supersedes a failed run; retrying stacks in-flight messages). Two correctness notes: - Org scoping: the Celery path runs with no organization set (hence _base_manager throughout tasks.py). The proxies never send X-Organization-ID and each view clears StateStore defensively — it is a thread-local and gunicorn reuses threads, so a leftover value would silently scope a global aggregation to one tenant. - Redelivery is safe: the 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 work, never wrong numbers. That is also why chunking needs no cross-call lock. Chunking seam ships OFF: _run_aggregation(org_ids=None) + _active_org_ids() extracted and GET aggregate/orgs/ exposed, with DASHBOARD_METRICS_ORG_CHUNK_SIZE default 0. Gunicorn caps a request at 600s and a SIGKILLed request leaks the Redis lock until its 900s self-heal, so this is the escape hatch — flipping it is env-only. _run_aggregation() with no args is unchanged, so the Celery path is byte-identical. Inert until an operator adopts a schedule; flag-off untouched. Tests: 1494 workers (19 new, incl. the registry check that catches "unknown task -> message silently dropped"), 72 backend pg_queue. The 9 errors in dashboard_metrics/tests/test_tasks.py are pre-existing — verified identical with these changes stashed (backend DB tests can't run without the tenant schema). Co-Authored-By: Claude Opus 5 (1M context) --- backend/backend/internal_base_urls.py | 7 + backend/dashboard_metrics/internal_urls.py | 34 +++ backend/dashboard_metrics/internal_views.py | 155 +++++++++++++ backend/dashboard_metrics/tasks.py | 61 +++-- docker/docker-compose.yaml | 49 ++++ docker/sample.env | 15 ++ workers/run-worker.sh | 11 + workers/scheduler/dashboard_metrics_tasks.py | 176 +++++++++++++++ workers/scheduler/tasks.py | 5 + workers/tests/test_dashboard_metrics_tasks.py | 209 ++++++++++++++++++ 10 files changed, 703 insertions(+), 19 deletions(-) create mode 100644 backend/dashboard_metrics/internal_urls.py create mode 100644 backend/dashboard_metrics/internal_views.py create mode 100644 workers/scheduler/dashboard_metrics_tasks.py create mode 100644 workers/tests/test_dashboard_metrics_tasks.py 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..23d1e0291a --- /dev/null +++ b/backend/dashboard_metrics/internal_urls.py @@ -0,0 +1,34 @@ +"""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( + "aggregate/orgs/", + internal_views.ActiveOrgsAPIView.as_view(), + name="aggregate_active_orgs", + ), + 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..bd01220bb5 --- /dev/null +++ b/backend/dashboard_metrics/internal_views.py @@ -0,0 +1,155 @@ +"""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 ( + _active_org_ids, + _run_aggregation, + 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 + +# Window used to decide which orgs are "active" — must match the daily window +# _run_aggregation uses, or the chunked path would process a different set than the +# in-process one. +ACTIVE_ORG_WINDOW_DAYS = 7 + + +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. + + With no body this is the whole job, lock included — byte-identical to what the + Celery task does. With ``{"org_ids": [...]}`` it runs only that slice and skips the + lock: the caller is splitting one logical run across several requests, and the lock + is a load guard rather than a correctness one (the upserts overwrite with recomputed + values, so slices cannot double-count). + """ + + def post(self, request: Request) -> Response: + org_ids = request.data.get("org_ids") if isinstance(request.data, dict) else None + if org_ids is None: + return self._run(aggregate_metrics_from_sources) + if not isinstance(org_ids, list): + return Response( + {"error": f"org_ids must be a list, got {type(org_ids).__name__}"}, + status=status.HTTP_400_BAD_REQUEST, + ) + return self._run(_run_aggregation, org_ids=org_ids) + + +class ActiveOrgsAPIView(_MetricsTaskAPIView): + """List the orgs an aggregation would process — the chunking seam. + + A caller fetches this, splits it, and posts each slice to ``aggregate/``. + """ + + def get(self, request: Request) -> Response: + _clear_org_context() + try: + from datetime import timedelta + + from django.utils import timezone + + since = timezone.now() - timedelta(days=ACTIVE_ORG_WINDOW_DAYS) + return Response({"org_ids": sorted(str(x) for x in _active_org_ids(since))}) + except Exception as exc: + logger.error("dashboard-metrics active-orgs failed: %s", exc, exc_info=True) + return Response( + {"error": str(exc)}, status=status.HTTP_500_INTERNAL_SERVER_ERROR + ) + + +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/tasks.py b/backend/dashboard_metrics/tasks.py index 181c985137..7365215734 100644 --- a/backend/dashboard_metrics/tasks.py +++ b/backend/dashboard_metrics/tasks.py @@ -395,10 +395,39 @@ def _aggregate_llm_combined( _upsert_agg(monthly_agg, key, metric_type, value) -def _run_aggregation() -> dict[str, Any]: +def _active_org_ids(since: datetime) -> set[str]: + """Org ids with workflow activity since ``since``. + + Pre-filter to orgs with recent activity to reduce DB load. Uses the 7-day daily + window rather than the 2-month monthly one because hourly/daily queries only need + recent data, and monthly totals for dormant orgs were already written when the org + was active — re-running just overwrites the same values. Avoids 28 queries per + dormant org that had activity 2-8 weeks ago. + + Extracted so the PG path can fetch the work list up front and process it in + chunks (UN-3796); the in-process caller below is unchanged. + """ + return set( + WorkflowExecution.objects.filter(created_at__gte=since) + .values_list("workflow__organization_id", flat=True) + .distinct() + ) + + +def _run_aggregation(org_ids: list[str] | None = None) -> dict[str, Any]: """Execute the actual aggregation logic. Separated from the task function to keep the lock management clean. + + ``org_ids`` restricts the run to a caller-supplied slice instead of discovering + active orgs itself (UN-3796). That lets the PG path split one aggregation across + several bounded calls, which matters because gunicorn caps a request at 600s while + this task declares ``time_limit=660``. Passing ``None`` — every existing caller — + behaves exactly as before. + + Chunking needs no cross-call lock: the upserts below are + ``INSERT … ON CONFLICT DO UPDATE SET``, i.e. they overwrite with recomputed values + rather than incrementing, so slices cannot double-count. """ end_date = timezone.now() @@ -459,25 +488,19 @@ def _run_aggregation() -> dict[str, Any]: "orgs_processed": 0, } - # Pre-filter to orgs with recent activity to reduce DB load. - # Uses daily_start (7 days) instead of monthly_start (2 months) because: - # - Hourly/daily queries only need recent data (24h / 7d windows) - # - Monthly totals for dormant orgs were already written by previous - # runs when the org was active — re-running just overwrites same values - # - This avoids 28 queries per dormant org that had activity 2-8 weeks ago - active_org_ids = set( - WorkflowExecution.objects.filter( - created_at__gte=daily_start, + # A caller-supplied slice skips discovery entirely (see the docstring); otherwise + # pre-filter to orgs with recent activity to reduce DB load. + if org_ids is not None: + active_org_ids = set(org_ids) + logger.info("Aggregation: %d org(s) supplied by caller", len(active_org_ids)) + else: + active_org_ids = _active_org_ids(daily_start) + total_orgs = Organization.objects.count() + logger.info( + "Aggregation: %d active orgs out of %d total", + len(active_org_ids), + total_orgs, ) - .values_list("workflow__organization_id", flat=True) - .distinct() - ) - total_orgs = Organization.objects.count() - logger.info( - "Aggregation: %d active orgs out of %d total", - len(active_org_ids), - total_orgs, - ) if not active_org_ids: return { diff --git a/docker/docker-compose.yaml b/docker/docker-compose.yaml index 9a8db90afe..fa57c8e13f 100644 --- a/docker/docker-compose.yaml +++ b/docker/docker-compose.yaml @@ -732,6 +732,55 @@ 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 + command: ["pg-metrics"] + 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 + # Split the aggregation into slices of N orgs if one run would exceed + # gunicorn's 600s request ceiling. 0 = one unsliced call (default). + - DASHBOARD_METRICS_ORG_CHUNK_SIZE=${DASHBOARD_METRICS_ORG_CHUNK_SIZE:-0} + 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..6cc186075a 100644 --- a/docker/sample.env +++ b/docker/sample.env @@ -133,6 +133,21 @@ 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. +# DASHBOARD_METRICS_ORG_CHUNK_SIZE - split the aggregation into slices of N orgs +# (0 = one unsliced call, the default). Raise this if +# one aggregation would exceed gunicorn's 600s request +# ceiling; the upserts overwrite rather than increment, +# so slices cannot double-count. +# # 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/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..06cbb5aca0 --- /dev/null +++ b/workers/scheduler/dashboard_metrics_tasks.py @@ -0,0 +1,176 @@ +"""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 + +# Optional chunking (default off). When > 0 the aggregation is split into slices of +# this many orgs, each its own request — the escape hatch for installations where one +# aggregation would exceed gunicorn's 600s request ceiling. Off by default because +# splitting costs extra round trips and most installations do not need it. +_ORG_CHUNK_SIZE_ENV = "DASHBOARD_METRICS_ORG_CHUNK_SIZE" + +_AGGREGATE_PATH = "v1/dashboard-metrics/aggregate/" +_ACTIVE_ORGS_PATH = "v1/dashboard-metrics/aggregate/orgs/" +_CLEANUP_HOURLY_PATH = "v1/dashboard-metrics/cleanup/hourly/" +_CLEANUP_DAILY_PATH = "v1/dashboard-metrics/cleanup/daily/" + + +def _org_chunk_size() -> int: + raw = os.getenv(_ORG_CHUNK_SIZE_ENV, "0") + try: + return max(0, int(raw)) + except ValueError: + logger.warning( + "%s=%r is not an integer; chunking disabled", _ORG_CHUNK_SIZE_ENV, raw + ) + return 0 + + +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.""" + chunk_size = _org_chunk_size() + if not chunk_size: + result = _call_internal(_AGGREGATE_PATH) + _log_if_skipped("dashboard_metrics.aggregate_from_sources", result) + return result + + org_ids = _call_internal(_ACTIVE_ORGS_PATH, method="GET", timeout=60.0)["org_ids"] + if not org_ids: + logger.info("dashboard_metrics: no active orgs; nothing to aggregate") + return {"success": True, "organizations_processed": 0, "chunks": 0} + + chunks = [org_ids[i : i + chunk_size] for i in range(0, len(org_ids), chunk_size)] + processed = failures = 0 + for index, chunk in enumerate(chunks, start=1): + try: + _call_internal(_AGGREGATE_PATH, body={"org_ids": chunk}) + processed += len(chunk) + except Exception: + # One bad slice must not lose the rest — the whole point of chunking is + # that a run is made of independent, idempotent pieces. + failures += 1 + logger.exception( + "dashboard_metrics: chunk %s/%s failed (%s orgs)", + index, + len(chunks), + len(chunk), + ) + if failures == len(chunks): + raise RuntimeError(f"dashboard_metrics: all {len(chunks)} chunks failed") + return { + "success": failures == 0, + "organizations_processed": processed, + "chunks": len(chunks), + "failed_chunks": failures, + } + + +@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..885ecdcf31 100644 --- a/workers/scheduler/tasks.py +++ b/workers/scheduler/tasks.py @@ -7,6 +7,11 @@ import traceback from typing import Any +# Register the dashboard-metrics proxy tasks on this worker type (UN-3796). Imported +# purely for the side effect. A BARE module import, not a relative one: worker.py loads +# this file by path with the worker directory on sys.path, so there is no parent package +# for `from . import ...` to resolve against. +import dashboard_metrics_tasks # noqa: F401, E402 (side-effect import) from queue_backend import FairnessKey, QueueBackend, dispatch, worker_task from queue_backend.fairness import WorkloadType from shared.enums.status_enums import PipelineStatus diff --git a/workers/tests/test_dashboard_metrics_tasks.py b/workers/tests/test_dashboard_metrics_tasks.py new file mode 100644 index 0000000000..b42f396d21 --- /dev/null +++ b/workers/tests/test_dashboard_metrics_tasks.py @@ -0,0 +1,209 @@ +"""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 + +# worker.py puts the worker directory on sys.path and loads tasks.py by file path, so +# the module is importable by bare name at runtime. Mirror that here. +_SCHEDULER_DIR = Path(__file__).resolve().parent.parent / "scheduler" +if str(_SCHEDULER_DIR) not in sys.path: + sys.path.insert(0, str(_SCHEDULER_DIR)) + +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) + monkeypatch.delenv("DASHBOARD_METRICS_ORG_CHUNK_SIZE", raising=False) + + +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/") + + +class TestChunking: + """Off by default; the escape hatch for installations where one aggregation would + exceed gunicorn's 600s request ceiling. + """ + + def test_disabled_by_default_makes_one_unsliced_call(self): + with patch.object(dmt, "_call_internal", return_value={"success": True}) as call: + dmt.dashboard_metrics_aggregate() + assert call.call_count == 1 + assert call.call_args.kwargs.get("body") is None + + def test_slices_cover_every_org_exactly_once(self, monkeypatch): + monkeypatch.setenv("DASHBOARD_METRICS_ORG_CHUNK_SIZE", "2") + orgs = ["a", "b", "c", "d", "e"] + calls = [] + + def fake(path, **kw): + if path.endswith("orgs/"): + return {"org_ids": orgs} + calls.append(kw["body"]["org_ids"]) + return {"success": True} + + with patch.object(dmt, "_call_internal", side_effect=fake): + result = dmt.dashboard_metrics_aggregate() + + assert len(calls) == 3 # 2 + 2 + 1 + assert [o for c in calls for o in c] == orgs # disjoint, complete, in order + assert result["organizations_processed"] == 5 + + def test_one_failing_slice_does_not_abort_the_rest(self, monkeypatch): + monkeypatch.setenv("DASHBOARD_METRICS_ORG_CHUNK_SIZE", "1") + + def fake(path, **kw): + if path.endswith("orgs/"): + return {"org_ids": ["a", "b", "c"]} + if kw["body"]["org_ids"] == ["b"]: + raise RuntimeError("chunk blew up") + return {"success": True} + + with patch.object(dmt, "_call_internal", side_effect=fake): + result = dmt.dashboard_metrics_aggregate() + + assert result["failed_chunks"] == 1 + assert result["organizations_processed"] == 2 + assert result["success"] is False + + def test_all_slices_failing_raises(self, monkeypatch): + monkeypatch.setenv("DASHBOARD_METRICS_ORG_CHUNK_SIZE", "1") + + def fake(path, **kw): + if path.endswith("orgs/"): + return {"org_ids": ["a", "b"]} + raise RuntimeError("down") + + with patch.object(dmt, "_call_internal", side_effect=fake): + with pytest.raises(RuntimeError, match="all 2 chunks failed"): + dmt.dashboard_metrics_aggregate() + + def test_no_active_orgs_is_a_clean_no_op(self, monkeypatch): + monkeypatch.setenv("DASHBOARD_METRICS_ORG_CHUNK_SIZE", "10") + with patch.object(dmt, "_call_internal", return_value={"org_ids": []}): + result = dmt.dashboard_metrics_aggregate() + assert result == {"success": True, "organizations_processed": 0, "chunks": 0} + + def test_garbage_chunk_size_disables_chunking_rather_than_crashing(self, monkeypatch): + monkeypatch.setenv("DASHBOARD_METRICS_ORG_CHUNK_SIZE", "not-a-number") + with patch.object(dmt, "_call_internal", return_value={"success": True}) as call: + dmt.dashboard_metrics_aggregate() + assert call.call_count == 1 From 6419dd1188fd147cb0346058977f7a2bab01bcb0 Mon Sep 17 00:00:00 2001 From: ali Date: Wed, 5 Aug 2026 16:01:19 +0530 Subject: [PATCH 04/33] UN-3796 [FIX] Bump PeriodicTasks.update_changed() on adopt/release or Beat never reloads MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Found via UN-2898, which describes exactly this failure mode for the pipeline path. --adopt sets pg_owned=True and disables the Beat row in one transaction so the two can never both fire. But the disable is a bulk `PeriodicTask.objects.filter(...).update(...)`, which 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 read "disabled" while Beat carried on firing alongside the PG scheduler: a DOUBLE FIRE, which is precisely what the atomic transaction exists to prevent. For cleanup_* that means concurrent deletes; for aggregate_* double work. --release fails the mirror way — Beat never resumes. Fix mirrors scheduler/ownership.py:132, which already does this on the pipeline path with the same rationale in a comment. Test pins both directions and was verified to go RED with the call removed — the symptom is invisible in the DB (the row really is disabled), so only a mutation check proves the guard works. Co-Authored-By: Claude Opus 5 (1M context) --- .../commands/mirror_pg_periodic_tasks.py | 11 ++++- .../tests/test_mirror_pg_periodic_tasks.py | 44 +++++++++++++++++++ 2 files changed, 54 insertions(+), 1 deletion(-) diff --git a/backend/pg_queue/management/commands/mirror_pg_periodic_tasks.py b/backend/pg_queue/management/commands/mirror_pg_periodic_tasks.py index 370c7efd43..09585e4c82 100644 --- a/backend/pg_queue/management/commands/mirror_pg_periodic_tasks.py +++ b/backend/pg_queue/management/commands/mirror_pg_periodic_tasks.py @@ -31,7 +31,7 @@ from django.core.management.base import BaseCommand, CommandError from django.db import transaction -from django_celery_beat.models import IntervalSchedule, PeriodicTask +from django_celery_beat.models import IntervalSchedule, PeriodicTask, PeriodicTasks from pg_queue.models import PgPeriodicTask @@ -345,5 +345,14 @@ def _set_ownership( row.next_run_at = None row.save(update_fields=["pg_owned", "next_run_at", "updated_at"]) PeriodicTask.objects.filter(name=row.name).update(enabled=not to_pg) + # 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/tests/test_mirror_pg_periodic_tasks.py b/backend/pg_queue/tests/test_mirror_pg_periodic_tasks.py index 9945a30277..8e1afdf6d7 100644 --- a/backend/pg_queue/tests/test_mirror_pg_periodic_tasks.py +++ b/backend/pg_queue/tests/test_mirror_pg_periodic_tasks.py @@ -13,8 +13,10 @@ 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, @@ -155,3 +157,45 @@ 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): + row = SimpleNamespace(name="h", pg_owned=(flag == "--release"), next_run_at=None) + 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() From 81be693c8f3610b82f313760fbc64996ddcd32ce Mon Sep 17 00:00:00 2001 From: ali Date: Wed, 5 Aug 2026 16:36:46 +0530 Subject: [PATCH 05/33] UN-3796 [FIX] worker-pg-metrics compose command would crash-loop the container MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit `command: ["pg-metrics"]` reached run-worker-docker.sh, whose pg-* dispatch (:563-578) accepts only pg-queue-consumer / pg-consumer / pg-queue-reaper / pg-reaper / reaper and rejects everything else with a loud exit 1: pg-*) print_status $RED "Unrecognized PG-queue command: '$1' ..."; exit 1 ;; So the service would have exited immediately and crash-looped. I had assumed the `pg-metrics` role added to run-worker.sh would be honoured in the container. It isn't — that script is the HOST runner. Containers go through run-worker-docker.sh, which is purely env-driven (`run_pg_consumer` reads WORKER_PG_QUEUE_CONSUMER_WORKER_TYPE / _QUEUE at :515-516). Every sibling PG compose service already uses the generic command for exactly this reason. Switch to ["pg-queue-consumer"]; the env already carries type=scheduler and queue=dashboard_metric_events, so nothing else changes. The run-worker.sh role stays — `./run-worker.sh pg-metrics` is a real local-dev path, same as pg-scheduler and pg-executor — and the comment now records that the two paths differ. Co-Authored-By: Claude Opus 5 (1M context) --- docker/docker-compose.yaml | 8 +++++++- 1 file changed, 7 insertions(+), 1 deletion(-) diff --git a/docker/docker-compose.yaml b/docker/docker-compose.yaml index fa57c8e13f..d80f0b980f 100644 --- a/docker/docker-compose.yaml +++ b/docker/docker-compose.yaml @@ -751,7 +751,13 @@ services: image: unstract/worker-unified:${VERSION} container_name: unstract-worker-pg-metrics restart: unless-stopped - command: ["pg-metrics"] + # 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: From 4c2d452a0e237c9f26053783661fa9575cf93f88 Mon Sep 17 00:00:00 2001 From: ali Date: Thu, 6 Aug 2026 12:21:27 +0530 Subject: [PATCH 06/33] =?UTF-8?q?UN-3796=20[FIX]=20Remove=20the=20org-chun?= =?UTF-8?q?king=20seam=20=E2=80=94=20it=20was=20a=20new=20feature,=20not?= =?UTF-8?q?=20a=20migration?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit DASHBOARD_METRICS_ORG_CHUNK_SIZE and everything behind it had no counterpart in the Celery worker, which runs the whole aggregation in one task with no slicing. I built it on a concern I then disproved myself: - I framed gunicorn's --timeout 600 as a NEW ceiling the PG path introduces. It is not. The Celery task already declares soft_time_limit=600 and gunicorn declares --timeout 600 — the same 600s. A run exceeding it would already be failing on Celery with SoftTimeLimitExceeded. - The "~120s p95" threshold attached to it was invented outright. And a >600s run is not a PG-only loss either: SoftTimeLimitExceeded is not in autoretry_for, so Celery loses the run too. Both recover on the next tick, because _acquire_aggregation_lock already self-heals a SIGKILLed run (its docstring says so) and AGGREGATION_LOCK_TIMEOUT equals the schedule interval. So the seam guarded an unchanged ceiling, nothing used it (default 0 = off), and part of it edited a file on the flag-off Celery path — which this branch's own rule says must stay byte-identical to main. Removed: the org_ids parameter and _active_org_ids extraction (dashboard_metrics/ tasks.py reverted to main outright), ActiveOrgsAPIView and the aggregate/orgs route, the chunking branch in the worker proxy, the env var from compose and sample.env, and the six TestChunking cases. What remains is a strict 1:1 with Celery: 3 Beat periodics -> PG scheduler; 3 tasks -> 3 thin proxies calling 3 endpoints that invoke the same functions verbatim, lock included; workerMetrics -> worker-pg-metrics. Verified: `git diff origin/main -- backend/dashboard_metrics/tasks.py` is empty, which is the check that the Celery path is untouched. 1488 workers pass, 74 backend pg_queue, lint clean. Co-Authored-By: Claude Opus 5 (1M context) --- backend/dashboard_metrics/internal_urls.py | 5 -- backend/dashboard_metrics/internal_views.py | 46 +----------- backend/dashboard_metrics/tasks.py | 61 +++++----------- docker/docker-compose.yaml | 3 - docker/sample.env | 6 +- workers/scheduler/dashboard_metrics_tasks.py | 56 +-------------- workers/tests/test_dashboard_metrics_tasks.py | 72 ------------------- 7 files changed, 26 insertions(+), 223 deletions(-) diff --git a/backend/dashboard_metrics/internal_urls.py b/backend/dashboard_metrics/internal_urls.py index 23d1e0291a..bc9d698f75 100644 --- a/backend/dashboard_metrics/internal_urls.py +++ b/backend/dashboard_metrics/internal_urls.py @@ -16,11 +16,6 @@ internal_views.AggregateMetricsAPIView.as_view(), name="aggregate_metrics", ), - path( - "aggregate/orgs/", - internal_views.ActiveOrgsAPIView.as_view(), - name="aggregate_active_orgs", - ), path( "cleanup/hourly/", internal_views.CleanupHourlyMetricsAPIView.as_view(), diff --git a/backend/dashboard_metrics/internal_views.py b/backend/dashboard_metrics/internal_views.py index bd01220bb5..f776633944 100644 --- a/backend/dashboard_metrics/internal_views.py +++ b/backend/dashboard_metrics/internal_views.py @@ -34,8 +34,6 @@ from utils.local_context import StateStore from dashboard_metrics.tasks import ( - _active_org_ids, - _run_aggregation, aggregate_metrics_from_sources, cleanup_daily_metrics, cleanup_hourly_metrics, @@ -49,11 +47,6 @@ DEFAULT_HOURLY_RETENTION_DAYS = 30 DEFAULT_DAILY_RETENTION_DAYS = 365 -# Window used to decide which orgs are "active" — must match the daily window -# _run_aggregation uses, or the chunked path would process a different set than the -# in-process one. -ACTIVE_ORG_WINDOW_DAYS = 7 - def _clear_org_context() -> None: """Drop any organization left in this thread's StateStore. See the module docstring. @@ -96,45 +89,12 @@ def _run(self, fn, *args: Any, **kwargs: Any) -> Response: class AggregateMetricsAPIView(_MetricsTaskAPIView): """Run the metrics aggregation. - With no body this is the whole job, lock included — byte-identical to what the - Celery task does. With ``{"org_ids": [...]}`` it runs only that slice and skips the - lock: the caller is splitting one logical run across several requests, and the lock - is a load guard rather than a correctness one (the upserts overwrite with recomputed - values, so slices cannot double-count). + 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: - org_ids = request.data.get("org_ids") if isinstance(request.data, dict) else None - if org_ids is None: - return self._run(aggregate_metrics_from_sources) - if not isinstance(org_ids, list): - return Response( - {"error": f"org_ids must be a list, got {type(org_ids).__name__}"}, - status=status.HTTP_400_BAD_REQUEST, - ) - return self._run(_run_aggregation, org_ids=org_ids) - - -class ActiveOrgsAPIView(_MetricsTaskAPIView): - """List the orgs an aggregation would process — the chunking seam. - - A caller fetches this, splits it, and posts each slice to ``aggregate/``. - """ - - def get(self, request: Request) -> Response: - _clear_org_context() - try: - from datetime import timedelta - - from django.utils import timezone - - since = timezone.now() - timedelta(days=ACTIVE_ORG_WINDOW_DAYS) - return Response({"org_ids": sorted(str(x) for x in _active_org_ids(since))}) - except Exception as exc: - logger.error("dashboard-metrics active-orgs failed: %s", exc, exc_info=True) - return Response( - {"error": str(exc)}, status=status.HTTP_500_INTERNAL_SERVER_ERROR - ) + return self._run(aggregate_metrics_from_sources) class CleanupHourlyMetricsAPIView(_MetricsTaskAPIView): diff --git a/backend/dashboard_metrics/tasks.py b/backend/dashboard_metrics/tasks.py index 7365215734..181c985137 100644 --- a/backend/dashboard_metrics/tasks.py +++ b/backend/dashboard_metrics/tasks.py @@ -395,39 +395,10 @@ def _aggregate_llm_combined( _upsert_agg(monthly_agg, key, metric_type, value) -def _active_org_ids(since: datetime) -> set[str]: - """Org ids with workflow activity since ``since``. - - Pre-filter to orgs with recent activity to reduce DB load. Uses the 7-day daily - window rather than the 2-month monthly one because hourly/daily queries only need - recent data, and monthly totals for dormant orgs were already written when the org - was active — re-running just overwrites the same values. Avoids 28 queries per - dormant org that had activity 2-8 weeks ago. - - Extracted so the PG path can fetch the work list up front and process it in - chunks (UN-3796); the in-process caller below is unchanged. - """ - return set( - WorkflowExecution.objects.filter(created_at__gte=since) - .values_list("workflow__organization_id", flat=True) - .distinct() - ) - - -def _run_aggregation(org_ids: list[str] | None = None) -> dict[str, Any]: +def _run_aggregation() -> dict[str, Any]: """Execute the actual aggregation logic. Separated from the task function to keep the lock management clean. - - ``org_ids`` restricts the run to a caller-supplied slice instead of discovering - active orgs itself (UN-3796). That lets the PG path split one aggregation across - several bounded calls, which matters because gunicorn caps a request at 600s while - this task declares ``time_limit=660``. Passing ``None`` — every existing caller — - behaves exactly as before. - - Chunking needs no cross-call lock: the upserts below are - ``INSERT … ON CONFLICT DO UPDATE SET``, i.e. they overwrite with recomputed values - rather than incrementing, so slices cannot double-count. """ end_date = timezone.now() @@ -488,19 +459,25 @@ def _run_aggregation(org_ids: list[str] | None = None) -> dict[str, Any]: "orgs_processed": 0, } - # A caller-supplied slice skips discovery entirely (see the docstring); otherwise - # pre-filter to orgs with recent activity to reduce DB load. - if org_ids is not None: - active_org_ids = set(org_ids) - logger.info("Aggregation: %d org(s) supplied by caller", len(active_org_ids)) - else: - active_org_ids = _active_org_ids(daily_start) - total_orgs = Organization.objects.count() - logger.info( - "Aggregation: %d active orgs out of %d total", - len(active_org_ids), - total_orgs, + # Pre-filter to orgs with recent activity to reduce DB load. + # Uses daily_start (7 days) instead of monthly_start (2 months) because: + # - Hourly/daily queries only need recent data (24h / 7d windows) + # - Monthly totals for dormant orgs were already written by previous + # runs when the org was active — re-running just overwrites same values + # - This avoids 28 queries per dormant org that had activity 2-8 weeks ago + active_org_ids = set( + WorkflowExecution.objects.filter( + created_at__gte=daily_start, ) + .values_list("workflow__organization_id", flat=True) + .distinct() + ) + total_orgs = Organization.objects.count() + logger.info( + "Aggregation: %d active orgs out of %d total", + len(active_org_ids), + total_orgs, + ) if not active_org_ids: return { diff --git a/docker/docker-compose.yaml b/docker/docker-compose.yaml index d80f0b980f..7c98d52a13 100644 --- a/docker/docker-compose.yaml +++ b/docker/docker-compose.yaml @@ -779,9 +779,6 @@ services: # One at a time — these are global singletons, not parallel work. - WORKER_PG_QUEUE_CONSUMER_CONCURRENCY=1 - WORKER_PG_QUEUE_CONSUMER_MAX_ATTEMPTS=1 - # Split the aggregation into slices of N orgs if one run would exceed - # gunicorn's 600s request ceiling. 0 = one unsliced call (default). - - DASHBOARD_METRICS_ORG_CHUNK_SIZE=${DASHBOARD_METRICS_ORG_CHUNK_SIZE:-0} labels: - traefik.enable=false profiles: diff --git a/docker/sample.env b/docker/sample.env index 6cc186075a..6c96821a42 100644 --- a/docker/sample.env +++ b/docker/sample.env @@ -142,11 +142,7 @@ WORKER_PG_REAPER_INTERVAL_SECONDS=5 # Reaper sweep interval (seconds) # 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. -# DASHBOARD_METRICS_ORG_CHUNK_SIZE - split the aggregation into slices of N orgs -# (0 = one unsliced call, the default). Raise this if -# one aggregation would exceed gunicorn's 600s request -# ceiling; the upserts overwrite rather than increment, -# so slices cannot double-count. +# 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 diff --git a/workers/scheduler/dashboard_metrics_tasks.py b/workers/scheduler/dashboard_metrics_tasks.py index 06cbb5aca0..44bbe50440 100644 --- a/workers/scheduler/dashboard_metrics_tasks.py +++ b/workers/scheduler/dashboard_metrics_tasks.py @@ -39,29 +39,11 @@ # and read as a network fault while the server kept working. DEFAULT_HTTP_TIMEOUT_SECONDS = 630.0 -# Optional chunking (default off). When > 0 the aggregation is split into slices of -# this many orgs, each its own request — the escape hatch for installations where one -# aggregation would exceed gunicorn's 600s request ceiling. Off by default because -# splitting costs extra round trips and most installations do not need it. -_ORG_CHUNK_SIZE_ENV = "DASHBOARD_METRICS_ORG_CHUNK_SIZE" - _AGGREGATE_PATH = "v1/dashboard-metrics/aggregate/" -_ACTIVE_ORGS_PATH = "v1/dashboard-metrics/aggregate/orgs/" _CLEANUP_HOURLY_PATH = "v1/dashboard-metrics/cleanup/hourly/" _CLEANUP_DAILY_PATH = "v1/dashboard-metrics/cleanup/daily/" -def _org_chunk_size() -> int: - raw = os.getenv(_ORG_CHUNK_SIZE_ENV, "0") - try: - return max(0, int(raw)) - except ValueError: - logger.warning( - "%s=%r is not an integer; chunking disabled", _ORG_CHUNK_SIZE_ENV, raw - ) - return 0 - - def _call_internal( path: str, *, @@ -125,41 +107,9 @@ def _log_if_skipped(name: str, result: dict[str, Any]) -> None: @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.""" - chunk_size = _org_chunk_size() - if not chunk_size: - result = _call_internal(_AGGREGATE_PATH) - _log_if_skipped("dashboard_metrics.aggregate_from_sources", result) - return result - - org_ids = _call_internal(_ACTIVE_ORGS_PATH, method="GET", timeout=60.0)["org_ids"] - if not org_ids: - logger.info("dashboard_metrics: no active orgs; nothing to aggregate") - return {"success": True, "organizations_processed": 0, "chunks": 0} - - chunks = [org_ids[i : i + chunk_size] for i in range(0, len(org_ids), chunk_size)] - processed = failures = 0 - for index, chunk in enumerate(chunks, start=1): - try: - _call_internal(_AGGREGATE_PATH, body={"org_ids": chunk}) - processed += len(chunk) - except Exception: - # One bad slice must not lose the rest — the whole point of chunking is - # that a run is made of independent, idempotent pieces. - failures += 1 - logger.exception( - "dashboard_metrics: chunk %s/%s failed (%s orgs)", - index, - len(chunks), - len(chunk), - ) - if failures == len(chunks): - raise RuntimeError(f"dashboard_metrics: all {len(chunks)} chunks failed") - return { - "success": failures == 0, - "organizations_processed": processed, - "chunks": len(chunks), - "failed_chunks": failures, - } + result = _call_internal(_AGGREGATE_PATH) + _log_if_skipped("dashboard_metrics.aggregate_from_sources", result) + return result @worker_task(name="dashboard_metrics.cleanup_hourly_data") diff --git a/workers/tests/test_dashboard_metrics_tasks.py b/workers/tests/test_dashboard_metrics_tasks.py index b42f396d21..eb43c9641b 100644 --- a/workers/tests/test_dashboard_metrics_tasks.py +++ b/workers/tests/test_dashboard_metrics_tasks.py @@ -32,7 +32,6 @@ def _env(monkeypatch): for k, v in _ENV.items(): monkeypatch.setenv(k, v) - monkeypatch.delenv("DASHBOARD_METRICS_ORG_CHUNK_SIZE", raising=False) class TestRegistration: @@ -136,74 +135,3 @@ def test_missing_config_raises_rather_than_returning_falsy(self, monkeypatch, mi monkeypatch.delenv(missing) with pytest.raises(RuntimeError, match=missing): dmt._call_internal("v1/x/") - - -class TestChunking: - """Off by default; the escape hatch for installations where one aggregation would - exceed gunicorn's 600s request ceiling. - """ - - def test_disabled_by_default_makes_one_unsliced_call(self): - with patch.object(dmt, "_call_internal", return_value={"success": True}) as call: - dmt.dashboard_metrics_aggregate() - assert call.call_count == 1 - assert call.call_args.kwargs.get("body") is None - - def test_slices_cover_every_org_exactly_once(self, monkeypatch): - monkeypatch.setenv("DASHBOARD_METRICS_ORG_CHUNK_SIZE", "2") - orgs = ["a", "b", "c", "d", "e"] - calls = [] - - def fake(path, **kw): - if path.endswith("orgs/"): - return {"org_ids": orgs} - calls.append(kw["body"]["org_ids"]) - return {"success": True} - - with patch.object(dmt, "_call_internal", side_effect=fake): - result = dmt.dashboard_metrics_aggregate() - - assert len(calls) == 3 # 2 + 2 + 1 - assert [o for c in calls for o in c] == orgs # disjoint, complete, in order - assert result["organizations_processed"] == 5 - - def test_one_failing_slice_does_not_abort_the_rest(self, monkeypatch): - monkeypatch.setenv("DASHBOARD_METRICS_ORG_CHUNK_SIZE", "1") - - def fake(path, **kw): - if path.endswith("orgs/"): - return {"org_ids": ["a", "b", "c"]} - if kw["body"]["org_ids"] == ["b"]: - raise RuntimeError("chunk blew up") - return {"success": True} - - with patch.object(dmt, "_call_internal", side_effect=fake): - result = dmt.dashboard_metrics_aggregate() - - assert result["failed_chunks"] == 1 - assert result["organizations_processed"] == 2 - assert result["success"] is False - - def test_all_slices_failing_raises(self, monkeypatch): - monkeypatch.setenv("DASHBOARD_METRICS_ORG_CHUNK_SIZE", "1") - - def fake(path, **kw): - if path.endswith("orgs/"): - return {"org_ids": ["a", "b"]} - raise RuntimeError("down") - - with patch.object(dmt, "_call_internal", side_effect=fake): - with pytest.raises(RuntimeError, match="all 2 chunks failed"): - dmt.dashboard_metrics_aggregate() - - def test_no_active_orgs_is_a_clean_no_op(self, monkeypatch): - monkeypatch.setenv("DASHBOARD_METRICS_ORG_CHUNK_SIZE", "10") - with patch.object(dmt, "_call_internal", return_value={"org_ids": []}): - result = dmt.dashboard_metrics_aggregate() - assert result == {"success": True, "organizations_processed": 0, "chunks": 0} - - def test_garbage_chunk_size_disables_chunking_rather_than_crashing(self, monkeypatch): - monkeypatch.setenv("DASHBOARD_METRICS_ORG_CHUNK_SIZE", "not-a-number") - with patch.object(dmt, "_call_internal", return_value={"success": True}) as call: - dmt.dashboard_metrics_aggregate() - assert call.call_count == 1 From 873bfa7a6f973f52e3a69b4a3e3b99dce4ceecbe Mon Sep 17 00:00:00 2001 From: ali Date: Thu, 6 Aug 2026 15:51:02 +0530 Subject: [PATCH 07/33] UN-3445 [GATED-FEAT] Refuse single-step execution on the PG transport MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Single-step is the last entry path whose fan-out was never transport-gated. `step_execution` -> `run_workflow` -> `process_input_files` builds a Celery chord directly in the Django web process; the normal path never reaches that code with the flag on, because the PG dispatcher enqueues async_execute_bin and the *workers'* general task does its own PG-gated fan-out (workers/general/tasks.py:875). So the backend chord is Celery-only residue reachable solely via step. With the flag on and the Celery file_processing workers scaled to zero (the epic's acceptance gate), those batches sit unconsumed and the execution hangs in EXECUTING forever, invisible to the PG reaper. Fail fast instead: a 500 naming the cause beats a silent forever-EXECUTING row, which is the exact failure class this epic exists to eliminate. The guard lives in step_execution, NOT in process_input_files. The latter is also on the Celery hot path, where the transport is already resolved upstream — re-resolving there would break the "resolved once per execution" invariant stated in workflow_v2/transport.py and add a Flipt call to every execution. Step executions are created by create_and_make_execution_response, which resolves no transport, so this is their first and only resolution. Unreachable from the UI: every write of `execution_action` in Agency.jsx sits behind an isStepExecution guard, and both live call sites pass false (the START/NEXT/STOP/ CONTINUE buttons are gone; only "Run Workflow" remains). Cloud has no frontend override. This guards a direct API call. Known trade-off: during the rollout window (flag on, Celery still up) step execution would work today and now fails. Accepted because the gate says nothing may fall back to Celery and the UI cannot generate the call. Not confirmed against production data — no Superset session available — so this rests on code reading. Flag-off is untouched and byte-identical; two of the five tests pin that specifically, since it is what staging and production run. Tests: 5 new, mutation-checked (stubbing the guard to `if False and ...` turns exactly one red). 139 pass across pg_queue/, workflow_manager/, notification_v2/, pipeline_v2/ at -m "not integration"; the DB-bound tests are deselected here and run in CI. Co-Authored-By: Claude Opus 5 (1M context) --- .../tests/test_step_execution_transport.py | 120 ++++++++++++++++++ .../workflow_v2/workflow_helper.py | 31 +++++ 2 files changed, 151 insertions(+) create mode 100644 backend/workflow_manager/workflow_v2/tests/test_step_execution_transport.py 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/workflow_helper.py b/backend/workflow_manager/workflow_v2/workflow_helper.py index d22dcc85bb..66b59ee715 100644 --- a/backend/workflow_manager/workflow_v2/workflow_helper.py +++ b/backend/workflow_manager/workflow_v2/workflow_helper.py @@ -1033,6 +1033,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, From 2fa627282d58c8f9a5faeb7136e157fd8f5bc8e5 Mon Sep 17 00:00:00 2001 From: ali Date: Fri, 7 Aug 2026 11:37:40 +0530 Subject: [PATCH 08/33] UN-3755 [GATED-FEAT] Log streaming without Celery: Redis-list transport MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit worker-log-consumer is the last worker forcing a Celery deployment to stay up, so it blocks the acceptance gate. This adds a second transport for the hop between LogPublisher and the consumer, selected by LOG_TRANSPORT. Redis, not the PG queue — the decision, recorded so it is not relitigated: The PG option never removed Redis. Both of the consumer's sinks are already Redis (RPUSH log_history_queue + the Socket.IO emit), so routing through pg_queue_message inserts Postgres into the MIDDLE of a Redis-to-Redis path. It buys mechanism consistency, not durability: logs stay lossy (dropped at the 10k cap in log_utils.py:107, every exception swallowed) and durability still begins at execution_log via the existing batched drain. Nor is Redis a second mechanism. Logs already use a Redis list drained by a loop — log_history_queue + worker-log-history-scheduler-v2, the very next stage of this same pipeline. This extends that pattern one hop upstream. Cost sealed it. LogPublisher lives in unstract/core with 14 call sites across 6 deployables, including tool-sidecar, spawned per file execution (live, not legacy: shared/workflow/execution/service.py:52 -> WorkflowExecutionService -> ToolSandbox -> runner -> sidecar). PG would need psycopg2 in unstract/core inherited by every importer, a third producer there (it can import neither Django nor workers.queue_backend), DB credentials in runner.py:229's sidecar env allowlist, and one unpooled psycopg2.connect() per live sidecar. At 30+ logs per execution that is ~3 DB ops per line against the platform's most contended resource. BLMOVE, not BLPOP. I expected the Celery consumer to run acks_late=False, making loss-on-crash a non-regression. It does not — task_acks_late defaults to true (shared/infrastructure/config/worker_config.py:545, backend/celery_config.py:63), so a crash today REDELIVERS. BLPOP deletes on read and would have quietly made crashes lossy. The loop parks each envelope on a per-pod processing list and removes it only after the handler returns; startup re-queues what the previous incarnation left. Residual gap, stated plainly: a container restart (crash loop, OOM — the dominant mode) recovers fully; a pod REPLACEMENT strands that pod's in-flight envelope. Sweeping other pods' lists cannot distinguish a dead owner from a live one and would duplicate logs on every start, so it is not attempted. Business logic is untouched. The pipeline is 8 stages; this changes 1-3 only. The consumer runs the SAME logs_consumer body, writes the same two Redis sinks, and the same 5s scheduler drains to the same execution_log table. Flag-off is unchanged: LOG_TRANSPORT defaults to celery and only the exact string "redis" opts in. The Kombu publish block is byte-identical; task_message construction moved above the branch (pure, no behaviour change). The publisher side must flip too, or this drains an empty list — wired in the cloud chart as one global.logTransport knob (backend, workers, runner; runner forwards it to each sidecar). Tests: 13 new; 1501 pass across the full workers suite (no regressions). Three mutations verified red — capacity guard, task-name in envelope, ack-before-handle ordering. NOT verified: no end-to-end run. The gate (flag on, worker-log-consumer at zero, logs still streaming and landing in execution_log) needs a live stack. Co-Authored-By: Claude Opus 5 (1M context) --- docker/docker-compose.yaml | 35 ++++ .../core/src/unstract/core/pubsub_helper.py | 110 +++++++++--- workers/log_consumer/redis_stream_consumer.py | 148 +++++++++++++++++ workers/run-worker-docker.sh | 16 ++ workers/tests/test_log_stream_consumer.py | 156 ++++++++++++++++++ workers/tests/test_log_stream_transport.py | 111 +++++++++++++ 6 files changed, 555 insertions(+), 21 deletions(-) create mode 100644 workers/log_consumer/redis_stream_consumer.py create mode 100644 workers/tests/test_log_stream_consumer.py create mode 100644 workers/tests/test_log_stream_transport.py diff --git a/docker/docker-compose.yaml b/docker/docker-compose.yaml index 7c98d52a13..d4b6c79d83 100644 --- a/docker/docker-compose.yaml +++ b/docker/docker-compose.yaml @@ -784,6 +784,41 @@ services: 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/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..acf02aa4de --- /dev/null +++ b/workers/log_consumer/redis_stream_consumer.py @@ -0,0 +1,148 @@ +"""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_queue_client import RedisQueueClient +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")) + +_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) + + redis_client = RedisQueueClient.from_env().redis_client + 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/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/tests/test_log_stream_consumer.py b/workers/tests/test_log_stream_consumer.py new file mode 100644 index 0000000000..7fa76bbe83 --- /dev/null +++ b/workers/tests/test_log_stream_consumer.py @@ -0,0 +1,156 @@ +"""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, "RedisQueueClient") as rq: + rq.from_env.return_value.redis_client = 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, "RedisQueueClient") as rq: + rq.from_env.return_value.redis_client = 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) 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 From c7882e27473aa6f527b224c8cf1200bb44e4cd6a Mon Sep 17 00:00:00 2001 From: ali Date: Fri, 7 Aug 2026 16:13:14 +0530 Subject: [PATCH 09/33] UN-3796 [GATED-FEAT] Declare the metrics periodics for the PG scheduler MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The three dashboard_metrics schedules are static, code-declared rows — Beat gets them from a data migration (0002_setup_periodic_tasks, update_or_create), so the PG scheduler gets them the same way rather than by mirroring Beat's table at runtime. That removes the mirror command from the rollout path for these entirely. Placed in dashboard_metrics/ rather than pg_queue/ deliberately: the failure mode that matters is the two declarations drifting apart, and a reviewer editing a schedule only sees both if they sit side by side. Why a migration here when the pipeline mirror stays a management command: three fixed rows known at build time is exactly what a data migration is for — tiny, idempotent, 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. The PG scheduler's due scan is WHERE pg_owned AND enabled, so nothing is selectable; and NULL next_run_at means "baseline next tick", not "overdue, fire now", so enabling the flag later cannot produce 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. The drift test compares them after json.loads; a silent mismatch there would change the cleanup retention. Safe to ship flag-off, verified rather than assumed: - pg_queue/0003 (the table) is branch-only, so table and seed rows ship together. - Single schema — django_tenants is commented out and TENANT_APPS is empty — so 3 rows total, not 3 per org. - Only reader is pg_scheduler.py, gated on pg_owned; the invalid-cron UPDATE at :265 is reachable only for rows that scan already selected. - workerPgScheduler.enabled is false in base values with no env override, so the PG scheduler is not even deployed. Inert twice over. - Both pg_queue and dashboard_metrics are in settings/base.py, which every variant imports, so apps.get_model resolves everywhere the migration runs. - Beat is untouched: the migration writes only pg_periodic_task. Caveat: this is static analysis. The migration has not been applied against a live database here (no DB; DB-bound tests are deselected as integration). Integration is the real confirmation. Tests: 12, DB-free — the Beat migration's forward function is run against fakes and its declarations compared to the PG specs, so the test cannot drift alongside a hardcoded copy. Mutation-checked: changing a cron (2am->4am) and a retention kwarg (365->90) each turn it red. 151 pass across dashboard_metrics/, pg_queue/, scheduler/, workflow_v2/. Co-Authored-By: Claude Opus 5 (1M context) --- .../migrations/0004_pg_periodic_tasks.py | 97 +++++++++++++++ .../test_pg_periodic_task_declarations.py | 113 ++++++++++++++++++ 2 files changed, 210 insertions(+) create mode 100644 backend/dashboard_metrics/migrations/0004_pg_periodic_tasks.py create mode 100644 backend/dashboard_metrics/tests/test_pg_periodic_task_declarations.py 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 From 9dc960401220d426349796314d5caa1d3b171ec1 Mon Sep 17 00:00:00 2001 From: ali Date: Fri, 7 Aug 2026 17:10:12 +0530 Subject: [PATCH 10/33] UN-3796 [GATED-FEAT] Mirror-only mode + automatic backfill on deploy (OSS) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Pipeline schedules created before the PG mirror existed have no pg_periodic_schedule row, so the PG scheduler has nothing to fire for them. Turning the flag on without backfilling first means those pipelines silently stop running. Until now the backfill was a manual command nobody was reminded to run. --mirror-only, and why it is load-bearing: reconcile_pg_schedules has no --adopt flag — handle() runs _reconcile_all unconditionally. With the rollout off that is inert (resolve_schedule_owner fails closed, so enabled = active AND NOT False leaves Beat untouched), but with the rollout ON it flips pg_owned AND disables the matching Beat PeriodicTask. That is a behaviour change no unattended job may make on its own. Backfilling is additive at every flag state, so that is all automation does; the ownership hand-over stays an explicit operator action. entrypoint.sh runs it inside the existing --migrate branch, after migrate returns. It is an ordinary management command, NOT a migration — sequenced after migrate only because pg_periodic_schedule must exist first. Only the single backend service passes --migrate, so there is no concurrent-replica race, and every other compose service is restart: unless-stopped, so a one-shot service would be a new convention for no gain. Best-effort by design (`|| echo WARNING`): entrypoint.sh has no `set -e`, but don't depend on that. A mirror failure must never stop the backend booting — Beat keeps firing everything in that case, which is the safe state. Runs on EVERY start rather than once. Schedules created while an older backend was deployed, and rows previously skipped for malformed PeriodicTask.args, are only picked up by a re-run; it is idempotent (already-mirrored pipelines are skipped), so there is nothing to retire until Celery is decommissioned. Not a gap this closes and not one it needs to: schedules created or edited SINCE the mirror shipped are already dual-written by SchedulerHelper on every save (helper.py:70), unconditionally and flag-independently. This is pre-existing rows only. Tests: 12 in the command suite (1 new), 117 across pg_queue/, dashboard_metrics/, scheduler/. Mutation-checked: making --mirror-only reconcile anyway turns it red. Co-Authored-By: Claude Opus 5 (1M context) --- backend/entrypoint.sh | 19 +++++++++++++++ .../commands/reconcile_pg_schedules.py | 24 ++++++++++++++++++- .../test_reconcile_pg_schedules_command.py | 19 +++++++++++++++ 3 files changed, 61 insertions(+), 1 deletion(-) diff --git a/backend/entrypoint.sh b/backend/entrypoint.sh index 5d44345fda..b7f82c2ca4 100755 --- a/backend/entrypoint.sh +++ b/backend/entrypoint.sh @@ -27,6 +27,25 @@ done if [ "$migrate" = true ]; then echo "Migration initiated" .venv/bin/python manage.py migrate + + # Backfill PG-scheduler mirror rows for pipeline schedules created before the + # mirror existed (UN-3796). NOT a migration — an ordinary management command, + # sequenced after migrate only because pg_periodic_schedule must exist first. + # + # --mirror-only: purely additive, writes only pg_periodic_schedule and never a + # Beat PeriodicTask, so it is safe at any flag state. Ownership hand-over stays + # an explicit operator action. + # + # Runs on EVERY start, not once: schedules created while an older backend was + # deployed, and rows previously skipped for malformed args, are only picked up by + # a re-run. It is idempotent (already-mirrored pipelines are skipped), so there + # is nothing to retire until Celery is decommissioned. + # + # Best-effort by design: a mirror failure must never stop the backend from + # starting. Beat keeps firing everything in that case, which is the safe state. + echo "PG schedule mirror backfill initiated" + .venv/bin/python manage.py reconcile_pg_schedules --mirror-only \ + || echo "WARNING: PG schedule mirror backfill failed; continuing startup (schedules stay on Beat)" fi # Configure Gunicorn based on --dev flag diff --git a/backend/pg_queue/management/commands/reconcile_pg_schedules.py b/backend/pg_queue/management/commands/reconcile_pg_schedules.py index 53c77df1f6..d20ce8663a 100644 --- a/backend/pg_queue/management/commands/reconcile_pg_schedules.py +++ b/backend/pg_queue/management/commands/reconcile_pg_schedules.py @@ -66,20 +66,42 @@ def add_arguments(self, parser: Any) -> None: "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 — see handle()." + ), + ) def handle(self, *args: Any, **options: Any) -> None: dry_run = options["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) - reconciled, pg_owned, failed = self._reconcile_all(dry_run, batch_size) + + # --mirror-only exists because the deploy-time automation must be safe to run + # at ANY flag state. The reconcile below is fail-closed (rollout off → every + # schedule stays on Beat, nothing written), but with the rollout ON it flips + # ownership *and disables the matching Beat PeriodicTask* — a behaviour change + # no unattended job should make on its own. Backfilling is inert at every flag + # state, so that is what automation runs; ownership stays an operator action. + if mirror_only: + reconciled, pg_owned, failed = 0, 0, 0 + else: + reconciled, pg_owned, failed = self._reconcile_all(dry_run, batch_size) prefix = "[dry-run] " if dry_run else "" summary = ( f"{prefix}backfilled={backfilled} reconciled={reconciled} " f"pg_owned={pg_owned} 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)) 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 a02fc3464b..04d8d07753 100644 --- a/backend/pg_queue/tests/test_reconcile_pg_schedules_command.py +++ b/backend/pg_queue/tests/test_reconcile_pg_schedules_command.py @@ -85,6 +85,25 @@ def test_backfills_only_unmirrored_and_reconciles(self): 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"]') From 6352186cea43169f224dbb71d9e1b6c69e040410 Mon Sep 17 00:00:00 2001 From: ali Date: Fri, 7 Aug 2026 18:02:17 +0530 Subject: [PATCH 11/33] UN-3796 [GATED-FEAT] Implement the PG_SCHEDULER_ENABLED gate so Beat stays sole scheduler MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Beat remains the ONLY scheduler for the whole PG rollout; the PG scheduler is activated later, as part of retiring Beat. That works because pipelines already reach PG without it: execute_pipeline_task_v2 -> WorkflowHelper.complete_execution -> execute_workflow_async -> resolve_transport (workflow_helper.py:631). The PG scheduler exists to retire the Beat *deployment*, not to get work onto PG. One scheduler at all times means the duplicate-trigger class never arises. pg_scheduler_enabled was already referenced by scheduler/helper.py:79 and twice in reconcile_pg_schedules (docstring + help text) — but no code implemented it. resolve_schedule_owner gated on pg_queue_enabled alone, welding the two flips together. This builds the gate that was designed and documented but never written. Without it, turning pg_queue_enabled on hands 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 — a user saving a schedule is enough. Two call sites, both required: - resolve_schedule_owner returns False immediately, short-circuiting BEFORE the Flipt call so there is no evaluation per schedule save. - reconcile_ownership_for returns early, writing NEITHER Beat table. This one is load-bearing beyond the obvious: merely resolving to Beat would still issue PeriodicTask.update(enabled=active) and bump PeriodicTasks.update_changed() on every save — writing back a value Beat already had and forcing a Beat reload, on tables we promised not to touch. (Both exist: ..._periodictask holds the schedules; ..._periodictasks is a single-row reload signal.) It also stops pg_owned being set. Otherwise it would drift to True across the rollout and the day the PG scheduler is switched on it would immediately fire everything Beat is also firing. Flag-off is unchanged. Flag-on now leaves Beat's tables untouched, so turning the flag back off is a pure Flipt flip with nothing to restore. Tests: 5 new for the gate (incl. reconcile touching neither Beat table); the 7 existing hand-over tests still pin the FINAL-phase behaviour and are scoped behind an autouse fixture that turns the gate on, rather than weakened. 167 pass across scheduler/, pg_queue/, dashboard_metrics/, workflow_v2/, pipeline_v2/. Both gates mutation-checked. Co-Authored-By: Claude Opus 5 (1M context) --- backend/scheduler/ownership.py | 58 +++++++++++++- .../tests/test_pg_schedule_ownership.py | 75 +++++++++++++++++++ 2 files changed, 129 insertions(+), 4 deletions(-) diff --git a/backend/scheduler/ownership.py b/backend/scheduler/ownership.py index 9c1e1f2878..7bd76f776d 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 " @@ -94,7 +135,16 @@ 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. 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. """ + if not pg_scheduler_enabled(): + return False pg_owned = resolve_schedule_owner(pipeline_id, organization_id) try: with transaction.atomic(): diff --git a/backend/scheduler/tests/test_pg_schedule_ownership.py b/backend/scheduler/tests/test_pg_schedule_ownership.py index bb4c1af48a..aff3b6225d 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(); @@ -167,6 +187,61 @@ 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), + ): + 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 set either — otherwise switching the PG scheduler on later + # would fire every already-owned schedule while Beat is still firing it too. + Sched.objects.filter.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 From 5e759c2a917d13da82e64f61780fb2a59c904491 Mon Sep 17 00:00:00 2001 From: ali Date: Mon, 10 Aug 2026 10:12:52 +0530 Subject: [PATCH 12/33] UN-3755 [FIX] Forward LOG_TRANSPORT into the tool sidecar MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Bug in my own UN-3755 change, found while checking whether tool containers use Celery. They do not — but the SIDECAR does publish logs, and it never received the transport switch. tool-sidecar/log_processor.py:165 calls LogPublisher.publish. LogPublisher reads LOG_TRANSPORT and defaults to Celery when unset. The sidecar's environment comes from the hand-picked allowlist in runner.py's _get_sidecar_container_config — it does NOT inherit the runner's env — and LOG_TRANSPORT was not in it. Effect with the flag on: every other publisher moves to the Redis list while the sidecar keeps publishing to celery_log_task_queue. Once worker-log-consumer is scaled to zero (the acceptance gate) those tool logs are dropped outright — no live streaming, no execution_log rows — silently, for container-based tool workflows. The cloud values comment asserted "runner forwards it to each tool sidecar". That was not true; I wrote it without checking. Corrected there to say the forwarding is explicit, not inherited, and to point at the test that now pins it. LOG_STREAM_QUEUE_NAME is forwarded too: if a deployment renames the queue, the sidecar must push to the same list the consumer drains. Needs NO new credentials — REDIS_* is already in that allowlist. That is the concrete payoff of choosing a Redis list over the PG queue for this hop; PG would have required database credentials in every spawned sidecar. Scope note, since it also corrects something I overstated earlier: the sidecar path is NOT every file execution. Structure-tool workflows now run in-process via the executor (file_processing/structure_tool_task.py: "Replaces the Docker-container-based StructureTool.run()"), so runner + tool container + sidecar remain only for non-structure tools. I had cited per-execution sidecars as a load-bearing argument for Redis over PG; it is real but narrower than I presented. The decision stands on the other grounds (PG would sit in the middle of a Redis-to-Redis path, and unstract/core has no psycopg2). Tests: 5 new, DB-free. Mutation-checked — removing the LOG_TRANSPORT line reproduces the bug and turns 2 red. 16 pass across the runner suite. Also adds runner/tests/ as a package root; the runner venv needs `uv sync --group test`. Co-Authored-By: Claude Opus 5 (1M context) --- runner/src/unstract/runner/constants.py | 7 ++ runner/src/unstract/runner/runner.py | 10 +++ runner/tests/test_sidecar_log_transport.py | 84 ++++++++++++++++++++++ 3 files changed, 101 insertions(+) create mode 100644 runner/tests/test_sidecar_log_transport.py 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" From 6397dd57869c212f519223e213aeebbc747ae8ff Mon Sep 17 00:00:00 2001 From: ali Date: Tue, 11 Aug 2026 11:29:30 +0530 Subject: [PATCH 13/33] UN-3796 [FIX] Bare import in scheduler/tasks.py crash-looped worker-general MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Regression I introduced in c6f77e425. The dashboard-metrics side-effect import was written as a BARE module import: import dashboard_metrics_tasks and my comment on it explicitly justified that form. The justification covered one entry path and missed the other. worker.py reaches a tasks.py two different ways: 1. BY PATH — spec_from_file_location with the worker's own dir appended to sys.path. A bare import resolves, because /app/scheduler is on the path. 2. AS A PACKAGE — general/tasks.py:16 does `from scheduler.tasks import execute_pipeline_task_v2`. Now /app/scheduler is NOT on sys.path, only /app. The bare import raises ModuleNotFoundError. So worker-general crash-looped 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' This is a FLAG-OFF regression — PG is not involved, the flag was never enabled, and no PG worker was deployed. It breaks the Celery path we promised to leave untouched. Fix: `from scheduler import dashboard_metrics_tasks`. /app is on PYTHONPATH (run-worker-docker.sh:571), so it resolves under both mechanisms. Comment rewritten to record why it must stay absolute. Why the existing 1501 tests missed it: nothing exercised either import mechanism. The new suite closes that — it discovers every */tasks.py and imports each one BOTH ways. Mutation-checked: restoring the bare import reproduces the exact production ModuleNotFoundError. The tests run each mechanism in a SUBPROCESS. In-process, importing the same file as both `dashboard_metrics_tasks` and `scheduler.dashboard_metrics_tasks` creates two module objects and duplicate Celery registrations, which made test_dashboard_metrics_tasks fail on patches that hit the wrong copy. A fresh interpreter is properly isolated and is also a truer reproduction of a booting worker. Also asserts the three proxies register under their exact wire names — deleting the import would otherwise "fix" the crash while silently unregistering the tasks the PG metrics consumer resolves by name. api-deployment is excluded from the package-import check only: a hyphen is not a valid Python identifier, so it can never be imported that way. Tests: 4 new; 1505 pass across the workers suite (was 1501 + 4, no regressions). Co-Authored-By: Claude Opus 5 (1M context) --- workers/scheduler/tasks.py | 20 ++- .../tests/test_worker_task_module_imports.py | 156 ++++++++++++++++++ 2 files changed, 171 insertions(+), 5 deletions(-) create mode 100644 workers/tests/test_worker_task_module_imports.py diff --git a/workers/scheduler/tasks.py b/workers/scheduler/tasks.py index 885ecdcf31..b30742c450 100644 --- a/workers/scheduler/tasks.py +++ b/workers/scheduler/tasks.py @@ -7,13 +7,23 @@ import traceback from typing import Any -# Register the dashboard-metrics proxy tasks on this worker type (UN-3796). Imported -# purely for the side effect. A BARE module import, not a relative one: worker.py loads -# this file by path with the worker directory on sys.path, so there is no parent package -# for `from . import ...` to resolve against. -import dashboard_metrics_tasks # noqa: F401, E402 (side-effect import) 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/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", + } From 02fce4424213f110f5b511fdb0c9168fad44d4ca Mon Sep 17 00:00:00 2001 From: ali Date: Tue, 11 Aug 2026 17:19:53 +0530 Subject: [PATCH 14/33] UN-3755 [FIX] Log stream consumer socket timeout must outlive its BLMOVE block MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Regression I introduced in 2fa627282, found running in integration: ERROR {module:redis_stream_consumer} :- Log stream read failed; retrying redis.exceptions.TimeoutError: Timeout reading from socket every ~5.01 seconds while idle. _BLOCK_TIMEOUT_SECONDS was 5, and create_redis_client defaults socket_timeout to 5 (unstract/core/.../redis_client.py:44). redis-py enforces socket_timeout on the blocking read itself, so the socket expired at the same instant the server-side block did — and won the race. Hence the exact 5.01s cadence. This is the SAME trap already documented at workers/queue_backend/pg_queue/result_backend.py:152, where the identical pairing is spelled out for the PG result signal. I wrote that warning and then reproduced the bug three files over. Two consequences, one loud and one quiet: 1. A traceback every 5s while idle, each tearing down the connection (_disconnect_raise) and reconnecting on the next pass. Log noise plus needless connection churn. 2. The one that matters: BLMOVE is atomic server-side. When the socket timed out after Redis had already moved an envelope onto the processing list but before the reply reached the client, that envelope was stranded there — recovered only by _recover_in_flight on a restart of the same pod. A real, if narrow, log-loss window. Fix: build the client with socket_timeout DERIVED from the block _SOCKET_TIMEOUT_SECONDS = _BLOCK_TIMEOUT_SECONDS + 5 so the two cannot drift apart. RedisQueueClient.from_env() is dropped because it hard-codes the 5s timeout and exposes no override; create_redis_client takes one. Every method used here (blmove/lmove/lrem) is native redis-py, so nothing is lost. _BLOCK_TIMEOUT_SECONDS deliberately stays 5s: it is the loop's only chance to observe a shutdown signal and must stay well under the pod's 60s terminationGracePeriodSeconds. This removes the exception, not the wakeup. Tests: 4 new, covering the invariant, its survival when LOG_STREAM_BLOCK_TIMEOUT is raised to 45 (proving it is structural rather than two defaults coinciding), that the constant actually reaches create_redis_client, and that BLMOVE still gets the block value. Two existing tests re-pointed off the removed RedisQueueClient. Mutation-checked: reverting to `= _BLOCK_TIMEOUT_SECONDS` fails 2 of them. Flag-gated: this worker only runs when workerLogStreamConsumer is enabled. Co-Authored-By: Claude Opus 5 (1M context) --- workers/log_consumer/redis_stream_consumer.py | 19 +++++- workers/tests/test_log_stream_consumer.py | 64 +++++++++++++++++-- 2 files changed, 77 insertions(+), 6 deletions(-) diff --git a/workers/log_consumer/redis_stream_consumer.py b/workers/log_consumer/redis_stream_consumer.py index acf02aa4de..a2242f11f2 100644 --- a/workers/log_consumer/redis_stream_consumer.py +++ b/workers/log_consumer/redis_stream_consumer.py @@ -38,7 +38,7 @@ from shared.infrastructure.config.builder import WorkerBuilder from shared.infrastructure.logging import WorkerLogger -from unstract.core.cache.redis_queue_client import RedisQueueClient +from unstract.core.cache.redis_client import create_redis_client from unstract.core.constants import LogProcessingTask logger = WorkerLogger.setup(WorkerType.LOG_CONSUMER) @@ -53,6 +53,15 @@ # 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 @@ -105,7 +114,13 @@ def run() -> int: signal.signal(signal.SIGTERM, _handle_signal) signal.signal(signal.SIGINT, _handle_signal) - redis_client = RedisQueueClient.from_env().redis_client + # 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'", diff --git a/workers/tests/test_log_stream_consumer.py b/workers/tests/test_log_stream_consumer.py index 7fa76bbe83..8c3f3091e6 100644 --- a/workers/tests/test_log_stream_consumer.py +++ b/workers/tests/test_log_stream_consumer.py @@ -135,8 +135,7 @@ def test_envelope_is_removed_only_after_the_handler_returns(self, consumer): redis.lrem.side_effect = lambda *a: order.append("lrem") consumer._test_logs_consumer.side_effect = lambda **_: order.append("handled") - with patch.object(consumer, "RedisQueueClient") as rq: - rq.from_env.return_value.redis_client = redis + 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 @@ -147,10 +146,67 @@ def test_envelope_is_removed_only_after_the_handler_returns(self, consumer): 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, "RedisQueueClient") as rq: - rq.from_env.return_value.redis_client = redis + 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 From a6ce3ef1e80cb5191b20a0d0a613f485f1bb4f59 Mon Sep 17 00:00:00 2001 From: ali Date: Fri, 14 Aug 2026 15:12:24 +0530 Subject: [PATCH 15/33] UN-3796 [GATED-FEAT] Converge schedule ownership both ways so rollback is a values change MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Beat→PG hand-over had no working reverse. Found on integration 2026-08-12, where Beat and the PG scheduler fired one pipeline 2.4s apart; a single execution resulted only because the Celery consumers happened to be scaled to zero. Three defects, one theme: the reverse direction was never reachable. 1. PIPELINES HAD NO REVERSE AT ALL. reconcile_ownership_for returned early whenever PG_SCHEDULER_ENABLED was off, writing NEITHER table. That is correct for a clean environment, but it assumed pg_owned could never be True while the gate is off. A build predating the gate breaks that: there, resolve_schedule_owner keys on the pg_queue_enabled Flipt flag ALONE, so flipping it hands schedules to PG. Roll forward onto the gate and the row is stranded — nothing can clear it, not even `reconcile_pg_schedules`, which routes back through the same early return. Only hand-written SQL could, which is not a deploy procedure. 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. The gate-off path now ASSERTS its invariant instead of assuming it: one indexed existence check, and still zero writes when the row is already correct — so "Beat's tables stay untouched for the whole rollout" holds for every environment that never had a contradictory row. 2. --release RESURRECTED DELIBERATELY-DISABLED PERIODICS. It wrote enabled=True unconditionally, so a job an operator had switched off in Beat before the migration came back ON after a rollback. A rollback restores the previous state; it does not invent a new one. The pre-migration value was already recorded — plan_mirror copies task.enabled — so release now restores row.enabled. 3. RE-MIRRORING AFTER ADOPTION SILENTLY KILLED THE ROW. `enabled` is the one mirrored field that stops tracking Beat once PG owns it: after --adopt, Beat's copy is False by definition. Copying that back left 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 stops with no error, and the value needed to release it correctly is gone. Cron/args/queue still track Beat, so a schedule edit made while PG owns the row is still picked up. Rather than leave the reverse as a second command an operator must remember, a single converge_pg_scheduler command routes on PG_SCHEDULER_ENABLED — adopt when on, release when off — and entrypoint.sh runs it on every start. Reverting becomes a values change like any other, including on-prem where a deploy cannot reach `manage.py`. Idempotent in both directions; an unset gate converges to Beat, so the default is the safe direction. Periodics stay opt-in behind --periodics (PG_SCHEDULER_ADOPT_PERIODICS), because adopting them needs workerPgMetrics deployed first. What this deliberately does NOT do: converging to Beat restores ownership, not capacity. Beat publishes over RabbitMQ, so released schedules only fire again if workerSchedulerV2 (and workerMetrics) are running — flip those in the SAME change that turns the gate off, or the outage merely moves. The release path logs that warning. Flag-gated: with PG_SCHEDULER_ENABLED unset — every environment today — the only behaviour change is that a contradictory row gets repaired instead of stranded. Tests: 18 new across three files, covering the released-back-to-Beat transaction, a paused pipeline surviving repair unchanged, Flipt never consulted while the gate is off (asking could re-hand a schedule to a scheduler that is not running, turning a double-fire into no firer), a DB error falling back to writing nothing, enabled preserved through a full adopt→release cycle, and the converge routing in both directions. Mutation-checked: restoring the blanket early return fails 2, the blanket re-enable fails 1, the mirror clobber fails 1, inverting the converge direction fails 5. One existing fixture gained an `enabled` attribute a real PgPeriodicTask always has. 370→378 unit-backend pass. NOT verified against a real database — every test here is mocked, and the atomicity test that needs Postgres skips; the rig's integration-backend tier is what exercises the actual transactions. Co-Authored-By: Claude Opus 5 (1M context) --- backend/entrypoint.sh | 22 +++- .../commands/converge_pg_scheduler.py | 111 ++++++++++++++++++ .../commands/mirror_pg_periodic_tasks.py | 29 ++++- .../commands/reconcile_pg_schedules.py | 66 ++++++++++- .../tests/test_converge_pg_scheduler.py | 91 ++++++++++++++ .../tests/test_mirror_pg_periodic_tasks.py | 84 ++++++++++++- backend/scheduler/ownership.py | 67 +++++++++-- .../tests/test_pg_schedule_ownership.py | 105 ++++++++++++++++- 8 files changed, 553 insertions(+), 22 deletions(-) create mode 100644 backend/pg_queue/management/commands/converge_pg_scheduler.py create mode 100644 backend/pg_queue/tests/test_converge_pg_scheduler.py diff --git a/backend/entrypoint.sh b/backend/entrypoint.sh index b7f82c2ca4..e7dd9d85dc 100755 --- a/backend/entrypoint.sh +++ b/backend/entrypoint.sh @@ -43,9 +43,25 @@ if [ "$migrate" = true ]; then # # Best-effort by design: a mirror failure must never stop the backend from # starting. Beat keeps firing everything in that case, which is the safe state. - echo "PG schedule mirror backfill initiated" - .venv/bin/python manage.py reconcile_pg_schedules --mirror-only \ - || echo "WARNING: PG schedule mirror backfill failed; continuing startup (schedules stay on Beat)" + # Converge schedule ownership to whatever PG_SCHEDULER_ENABLED declares — adopt + # when on, release back to Beat when off. 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. + # + # 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. Whatever fired the schedules before keeps firing them in that case. + 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..6c4d4378ed --- /dev/null +++ b/backend/pg_queue/management/commands/converge_pg_scheduler.py @@ -0,0 +1,111 @@ +"""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 — ``_set_ownership`` skips rows already in the + target state, and the pipeline path no-ops when ownership already matches. +* 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 index 09585e4c82..f39cf0dde3 100644 --- a/backend/pg_queue/management/commands/mirror_pg_periodic_tasks.py +++ b/backend/pg_queue/management/commands/mirror_pg_periodic_tasks.py @@ -302,9 +302,18 @@ def _mirror(self, dry_run: bool, *, batch_size: int, limit: int) -> tuple[int, i f"enabled={plan.fields['enabled']}" ) if not dry_run: - PgPeriodicTask.objects.update_or_create( - name=plan.name, defaults=plan.fields - ) + 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 @@ -333,8 +342,16 @@ def _set_ownership( 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 -> {not to_pg})") + self.stdout.write(f"{verb} {row.name!r} (beat enabled -> {beat_enabled})") if not dry_run: with transaction.atomic(): row.pg_owned = to_pg @@ -344,7 +361,9 @@ def _set_ownership( if not to_pg: row.next_run_at = None row.save(update_fields=["pg_owned", "next_run_at", "updated_at"]) - PeriodicTask.objects.filter(name=row.name).update(enabled=not to_pg) + PeriodicTask.objects.filter(name=row.name).update( + enabled=beat_enabled + ) # 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 diff --git a/backend/pg_queue/management/commands/reconcile_pg_schedules.py b/backend/pg_queue/management/commands/reconcile_pg_schedules.py index d20ce8663a..63045b0c92 100644 --- a/backend/pg_queue/management/commands/reconcile_pg_schedules.py +++ b/backend/pg_queue/management/commands/reconcile_pg_schedules.py @@ -18,7 +18,11 @@ 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 @@ -75,6 +79,17 @@ def add_arguments(self, parser: Any) -> None: "PeriodicTask. This is the mode automation runs — 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"] @@ -95,10 +110,15 @@ def handle(self, *args: Any, **options: Any) -> None: 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)" @@ -176,6 +196,48 @@ def _backfill_mirrors(self, dry_run: bool, batch_size: int) -> int: backfilled += 1 return backfilled + 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). 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 index 8e1afdf6d7..52a4c186cf 100644 --- a/backend/pg_queue/tests/test_mirror_pg_periodic_tasks.py +++ b/backend/pg_queue/tests/test_mirror_pg_periodic_tasks.py @@ -176,7 +176,11 @@ class TestBeatReloadSignal: _CMD = "pg_queue.management.commands.mirror_pg_periodic_tasks" def _run(self, flag): - row = SimpleNamespace(name="h", pg_owned=(flag == "--release"), next_run_at=None) + # `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, @@ -199,3 +203,81 @@ 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 _release(self, mirrored_enabled: bool): + row = SimpleNamespace( + name="h", pg_owned=True, 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", "--release") + return Beat.objects.filter.return_value.update.call_args.kwargs + + def test_a_row_that_was_enabled_is_restored_enabled(self): + assert self._release(mirrored_enabled=True) == {"enabled": 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": False} + + +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/scheduler/ownership.py b/backend/scheduler/ownership.py index 7bd76f776d..bd219d7042 100644 --- a/backend/scheduler/ownership.py +++ b/backend/scheduler/ownership.py @@ -118,6 +118,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: @@ -136,16 +157,46 @@ def reconcile_ownership_for( 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. 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. + **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. """ if not pg_scheduler_enabled(): - return False - pg_owned = resolve_schedule_owner(pipeline_id, organization_id) + # 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(): # queryset .update() doesn't fire auto_now, so bump updated_at diff --git a/backend/scheduler/tests/test_pg_schedule_ownership.py b/backend/scheduler/tests/test_pg_schedule_ownership.py index aff3b6225d..d60dcaf49e 100644 --- a/backend/scheduler/tests/test_pg_schedule_ownership.py +++ b/backend/scheduler/tests/test_pg_schedule_ownership.py @@ -233,13 +233,112 @@ def test_reconcile_writes_NEITHER_beat_table(self, monkeypatch): 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 set either — otherwise switching the PG scheduler on later - # would fire every already-owned schedule while Beat is still firing it too. - Sched.objects.filter.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. + PT.objects.filter.return_value.update.assert_called_once_with(enabled=True) + # ...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) + PT.objects.filter.return_value.update.assert_called_once_with(enabled=False) + + 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: From 0be13a707b10a153e6a4cb84d398cf2dae626728 Mon Sep 17 00:00:00 2001 From: ali Date: Fri, 14 Aug 2026 15:40:48 +0530 Subject: [PATCH 16/33] UN-3796 [FIX] Import dashboard_metrics_tasks via the package so tests patch the live module MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Fallout from 6397dd578, exposed by backmerging main. That commit changed scheduler/tasks.py to `from scheduler import dashboard_metrics_tasks` (the bare form crash-looped worker-general). This test still imported the module by BARE name, mirroring only worker.py's by-path load. Once both forms were in play the two coexisted as SEPARATE module objects with separate Celery registrations: the test patched one while the other stayed live, so a test that believed it had mocked the HTTP client issued a REAL request and died on DNS. Order-dependent, which is why it read as flakiness — it needed some earlier test in the run to have imported scheduler.tasks first. Passing the file alone was green; `tests/` as a directory was not. Latent since 6397dd578; main's new tests shifted collection order enough to surface it. Verified against pristine origin/main (1487 passed) to confirm the merge itself was not the cause. The sys.path insertion was the second half of the damage: putting `scheduler/` on the path let scheduler/worker.py SHADOW the top-level worker module, so shared/tests/test_session_lifecycle.py failed with `module 'worker' has no attribute 'on_task_postrun'`. Same root cause, different symptom; both go away by importing the package form and putting the workers ROOT on the path instead. The package form is right rather than merely convenient: worker.py's by-path load of tasks.py still resolves `from scheduler import ...` against /app, so scheduler.dashboard_metrics_tasks is the single module object under BOTH runtime mechanisms. The test now patches what the worker actually runs. Test-only; no production change. unit-workers 1387, unit-core 33, unit-backend 737, all green. Co-Authored-By: Claude Opus 5 (1M context) --- workers/tests/test_dashboard_metrics_tasks.py | 24 +++++++++++++------ 1 file changed, 17 insertions(+), 7 deletions(-) diff --git a/workers/tests/test_dashboard_metrics_tasks.py b/workers/tests/test_dashboard_metrics_tasks.py index eb43c9641b..ce8ff853ac 100644 --- a/workers/tests/test_dashboard_metrics_tasks.py +++ b/workers/tests/test_dashboard_metrics_tasks.py @@ -14,13 +14,23 @@ import pytest -# worker.py puts the worker directory on sys.path and loads tasks.py by file path, so -# the module is importable by bare name at runtime. Mirror that here. -_SCHEDULER_DIR = Path(__file__).resolve().parent.parent / "scheduler" -if str(_SCHEDULER_DIR) not in sys.path: - sys.path.insert(0, str(_SCHEDULER_DIR)) - -import dashboard_metrics_tasks as dmt # noqa: E402 +# 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", From ec0362fe526d05960aef7203789af50ad63b14c7 Mon Sep 17 00:00:00 2001 From: ali Date: Fri, 14 Aug 2026 17:44:16 +0530 Subject: [PATCH 17/33] UN-3796 [FIX] Baseline next_run_at on hand-over and resume, so neither fires a catch-up MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Observed on integration 2026-08-14. gallh_load_test was re-enabled in the UI at 11:14:37; the PG scheduler fired it at 11:14:41 — two seconds later, against next_run_at=2026-08-12 06:08, a timestamp two days old. The operator's own manual run was already in flight, so the pipeline ran twice within 2s. The tick selects: WHERE pg_owned AND enabled AND (next_run_at IS NULL OR next_run_at <= now()) so a next_run_at left over from an earlier PG-ownership period is already in the past and matches on the very next pass. NULL is the guard — "record a baseline next tick, don't fire this cycle" (pg_queue/models.py:366) — but it was only applied on release, not on the two transitions that re-arm a row: * ADOPT — reconcile_ownership_for cleared next_run_at only when handing BACK to Beat, so a stale value survived the hand-over. * RESUME — _mirror_periodic_schedule_set_enabled wrote `enabled` and left next_run_at alone. While a schedule is paused its next_run_at keeps drifting into the past, so re-enabling fires immediately. The longer the pause, the more certain. Stated generally: pausing a PG-owned schedule and later re-enabling it caused an immediate unscheduled run. Beat never behaved this way — DatabaseScheduler keeps no persisted next_run_at and recomputes due-ness from the crontab each tick — so this was a PG-path regression against it, not a cosmetic difference. The adopt fix is scoped to the TRANSITION (was_pg_owned False → True), not to every call. reconcile_ownership_for 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 push it to 13:08 and SKIP that fire. The resume fix is unconditional because that helper is reached only from enable_task/disable_task — an explicit pause/resume — never from the per-save path. Pause deliberately leaves next_run_at alone: nothing fires while disabled, and clearing there would discard the value resume baselines against. An existing test was NARROWED rather than kept: test_pg_owned_does_not_clear_next_run asserted "adopt never clears next_run_at", which is now the opposite of correct. It was also passing for the wrong reason — a bare MagicMock is truthy, so was_pg_owned read as "already owned" by accident. It is now test_an_ALREADY_pg_owned_schedule_does_not_clear_next_run, sets was_pg_owned explicitly, and documents that it covers the True→True case. The guarantee it defends — no re-baselining on an ordinary save — is real and still enforced. NOT the whole story: next_run_at still overloads one nullable timestamp with three meanings (NULL = baseline, past = fire now, future = wait), so a stale value is indistinguishable from a legitimately due one. THREE writers have now forgotten to clear it (the original mirror, adopt, resume). The durable fix is a staleness guard in the tick — if next_run_at is more than one cron period old, baseline instead of firing — enforced once where it fires rather than in every path that touches ownership. Filed separately. unit-backend 740 → 742, unit-workers 1387, pre-commit clean. Mutation-checked one test each: reverting the adopt baseline fails test_handing_over_to_pg_baselines, reverting the resume baseline fails test_resume_clears_next_run_at. Co-Authored-By: Claude Opus 5 (1M context) --- backend/scheduler/ownership.py | 24 +++++- backend/scheduler/tasks.py | 21 ++++- .../tests/test_pg_periodic_schedule_mirror.py | 45 ++++++++++ .../tests/test_pg_schedule_ownership.py | 82 +++++++++++++++++-- 4 files changed, 160 insertions(+), 12 deletions(-) diff --git a/backend/scheduler/ownership.py b/backend/scheduler/ownership.py index bd219d7042..e799ec0624 100644 --- a/backend/scheduler/ownership.py +++ b/backend/scheduler/ownership.py @@ -199,12 +199,30 @@ def reconcile_ownership_for( 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. diff --git a/backend/scheduler/tasks.py b/backend/scheduler/tasks.py index 7db1cad534..601584cf39 100644 --- a/backend/scheduler/tasks.py +++ b/backend/scheduler/tasks.py @@ -123,8 +123,27 @@ 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. + # + # This is Beat parity, not a new rule: DatabaseScheduler holds no persisted + # next_run_at and recomputes due-ness from the crontab each tick, so + # re-enabling never produced a catch-up run there. Observed on integration + # 2026-08-14: gallh_load_test fired ~2s after being re-enabled, against a + # next_run_at two days old. + # + # 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..11a2ac69f9 100644 --- a/backend/scheduler/tests/test_pg_periodic_schedule_mirror.py +++ b/backend/scheduler/tests/test_pg_periodic_schedule_mirror.py @@ -336,3 +336,48 @@ 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). + + Beat parity, not a new rule — DatabaseScheduler keeps no persisted next_run_at and + recomputes due-ness from the crontab each tick, so resume never produced a catch-up + run there. 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. + """ + + 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 d60dcaf49e..b9bb03d583 100644 --- a/backend/scheduler/tests/test_pg_schedule_ownership.py +++ b/backend/scheduler/tests/test_pg_schedule_ownership.py @@ -124,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) @@ -401,3 +409,61 @@ 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. Beat never did this: DatabaseScheduler + keeps no persisted next_run_at and recomputes due-ness from the crontab each tick, + so this is a PG-path regression against it, not a cosmetic difference. + """ + + 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 From d97e26df01cc117f17152d30628271c1824dc716 Mon Sep 17 00:00:00 2001 From: ali Date: Wed, 19 Aug 2026 09:43:48 +0530 Subject: [PATCH 18/33] UN-3796 [FIX] Terminalise executions created but never dispatched MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit PENDING is not a terminal state, so every execution must eventually reach COMPLETED or ERROR. One case had no owner at all. deployment_helper.py commits the execution row at :236 and dispatches at :310, with file staging in between. If the request dies in that window the row is orphaned: * the reaper recovers strands by scanning pg_barrier_state, and a barrier only exists once a batch is dispatched — no dispatch, no barrier, invisible; * execute_workflow_async marks dispatch FAILURES ERROR, but never runs at all; * nothing else looks at PENDING, so the row is a permanent lie in the UI and in every completion metric. Observed on integration 2026-08-17: an 800-user load test left 1947 COMPLETED and 967 PENDING — 2914 rows from 2444 requests, because rows outlived their requests. The backend access log shows 1153 completions and ZERO 4xx/5xx, so 1761 handlers created a row and never returned; 794 had already dispatched (workers finished them) and 967 had not. pg_queue_message was completely empty, proving the queue consumed everything it was given. NOT PG-specific. The window sits upstream of resolve_transport, so the Celery path has it too and the predicate is transport-aware for that reason. But note the TRIGGER is the PG reaper, so this only runs where workerPgReaper is deployed — prod-on-Celery is NOT covered until it migrates. Deliberate: the alternatives were a second caller in the log-history scheduler or a Beat periodic, and neither belongs in this change. 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 the other stays NULL. So PENDING + both handles NULL + older than the grace period means dispatch never happened — no JSONB probing, no joins, correct under either transport. CLAIMING is one UPDATE ... RETURNING. The inner SELECT bounds the batch and takes row locks (SKIP LOCKED); the OUTER WHERE re-carries the full predicate so Postgres re-evaluates at write time under those locks — a row dispatched between selection and write is skipped rather than marked ERROR *while it runs*, the one failure this must never cause. RETURNING names exactly the claimed rows, which is what lets the irreversible cleanup run for those and only those. (A first draft looped one UPDATE per row to get that guarantee; this keeps it and drops 500 round trips.) CLEANUP does what the abort prevented: releases the API-deployment rate-limit slot (held slots consume the org's concurrency budget until the 6h Redis TTL) and deletes the API storage dir (scoped by workflow_id + execution_id, exists()-guarded, so a no-op for a row that died before staging). Best-effort and isolated — the status write already succeeded. TIMING: 15-minute grace, 5-minute sweep cadence, so ~20 min worst case to terminal. Generous on purpose — elapsed time is the ONLY thing separating "abandoned" from "about to dispatch", and terminalising a live execution is far worse than leaving a dead one longer. The error_message is USER-FACING: ExecutionSerializer uses `exclude`, not `fields`, so every unlisted model field is serialized to customers. It names no internals and answers the three things a user needs — did anything run, is my data affected, what do I do: 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) 170/256 chars so the ref survives truncation. The precise cause goes to logger.error with the ids. The no-internals test earned its keep immediately: it caught this commit's own first ref code, EXEC_NOT_DISPATCHED, which leaked "dispatch". INDEX: migration 0026 adds a partial index built CONCURRENTLY, copying 0023's structure on this same table (atomic=False, SeparateDatabaseAndState, IF NOT EXISTS, INVALID-index guard, out-of-band build instructions). workflow_execution is multi-million-row in production, where a plain AddIndex would hold a SHARE lock for the whole build and block every in-flight execution. The index is near-EMPTY in steady state — executions leave PENDING within seconds — so it costs almost nothing and only grows when something is wrong. we_active_by_workflow_idx is usable for this 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. METRICS: pg_reaper_undispatched_swept_total + _failures_total on the reaper's EXISTING liveness server (:8086/metrics) — no new process or scrape target. Follows the recovery sweeps (barrier_recovered, claim_recovered) rather than the retention ones, which have no success counter. Nothing scrapes these in cloud yet (OPERATIONS.md); they are for incident curl. Tests: 28 new. unit-backend 745 -> 755, unit-workers 1387 -> 1392. Mutation-checked: drifting the index status literal fails 1; dropping CONCURRENTLY fails 1; re-raising in the reaper instead of swallowing fails 28 (confirming the swallow is load-bearing — a fault here must not abort a tick that also dispatches schedules). One mutation ESCAPED first time: removing the outer re-check from the claim broke nothing, because the substring assertions were satisfied by the inner SELECT. Fixed by counting occurrences instead — test_the_predicate_appears_TWICE__inner_select_and_outer_recheck now fails on it. NOT verified locally: the 10 django_db tests and the raw SQL itself only run in the rig's integration-backend tier (no Postgres here). The 3 message-contract and 10 drift tests are deliberately kept unit-tier by scoping the marker to the DB-bound classes. Recovery only. The mechanism that killed those handlers is still unconfirmed — two hypotheses were tested and both failed (gunicorn crashes: the messages are graceful and pair with HPA churn, no OOMKilled; short termination grace: it is actually 900s with a 60s preStop). Closing the window itself, and any capacity change, are deliberately out of scope; a sweep is needed regardless, since a node eviction can always kill a request mid-flight. Co-Authored-By: Claude Opus 5 (1M context) --- .../execution_log_internal_urls.py | 5 + .../execution_log_internal_views.py | 32 +++ ...0026_workflowexecution_undispatched_idx.py | 132 ++++++++++ .../workflow_v2/models/execution.py | 21 ++ .../test_undispatched_execution_index.py | 144 +++++++++++ .../tests/test_undispatched_sweep.py | 215 ++++++++++++++++ .../workflow_v2/undispatched_sweep.py | 238 ++++++++++++++++++ workers/queue_backend/pg_queue/metrics.py | 15 ++ workers/queue_backend/pg_queue/reaper.py | 42 ++++ workers/shared/api/internal_client.py | 15 ++ .../tests/test_reaper_undispatched_sweep.py | 95 +++++++ 11 files changed, 954 insertions(+) create mode 100644 backend/workflow_manager/workflow_v2/migrations/0026_workflowexecution_undispatched_idx.py create mode 100644 backend/workflow_manager/workflow_v2/tests/test_undispatched_execution_index.py create mode 100644 backend/workflow_manager/workflow_v2/tests/test_undispatched_sweep.py create mode 100644 backend/workflow_manager/workflow_v2/undispatched_sweep.py create mode 100644 workers/tests/test_reaper_undispatched_sweep.py 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..c7094e2d8e 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 and the API storage connector — 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/models/execution.py b/backend/workflow_manager/workflow_v2/models/execution.py index dc63a19ac7..5b0412707c 100644 --- a/backend/workflow_manager/workflow_v2/models/execution.py +++ b/backend/workflow_manager/workflow_v2/models/execution.py @@ -234,6 +234,27 @@ 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_idx", + condition=Q( + status="PENDING", + task_id__isnull=True, + queue_message_id__isnull=True, + ), + ), ] @property 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..9708ffbf6a --- /dev/null +++ b/backend/workflow_manager/workflow_v2/tests/test_undispatched_execution_index.py @@ -0,0 +1,144 @@ +"""Guard: the undispatched-execution partial index stays in sync with reality. + +``we_undispatched_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_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_both_dispatch_handles_are_in_the_predicate(self): + """Dropping either 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. + """ + children = _condition_children() + 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..9fb1cf92e6 --- /dev/null +++ b/backend/workflow_manager/workflow_v2/tests/test_undispatched_sweep.py @@ -0,0 +1,215 @@ +"""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.""" + from workflow_manager.workflow_v2.models import WorkflowExecution + + fields = { + "id": uuid.uuid4(), + "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 + + +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..b45e9fc074 --- /dev/null +++ b/backend/workflow_manager/workflow_v2/undispatched_sweep.py @@ -0,0 +1,238 @@ +"""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 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. +_MIN_AGE_ENV = "UNDISPATCHED_EXECUTION_GRACE_SECONDS" +DEFAULT_MIN_AGE_SECONDS = 900 # 15 minutes + +# 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 task_id IS NULL + AND queue_message_id IS NULL + ORDER BY created_at + LIMIT %s + FOR UPDATE SKIP LOCKED + ) + AND status = %s + 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. Each side is isolated so one failing does not skip the other. + """ + # 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 + ) + APIDeploymentRateLimiter.release_slot(organization, 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, + ) + + # Then the staged input. Scoped by workflow_id + execution_id, and guarded by an + # exists() check inside, so it is a clean no-op for an execution that died BEFORE + # staging — nothing else's files can be reached from here. + try: + from workflow_manager.endpoint_v2.destination import DestinationConnector + + DestinationConnector.delete_api_storage_dir( + workflow_id=workflow_id, execution_id=execution_id + ) + except Exception: + logger.warning( + "Undispatched sweep: could not delete the API storage dir for %s " + "(orphaned input files remain)", + execution_id, + exc_info=True, + ) diff --git a/workers/queue_backend/pg_queue/metrics.py b/workers/queue_backend/pg_queue/metrics.py index a3d77c43fb..d34c7d885d 100644 --- a/workers/queue_backend/pg_queue/metrics.py +++ b/workers/queue_backend/pg_queue/metrics.py @@ -282,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/reaper.py b/workers/queue_backend/pg_queue/reaper.py index c0ab3d1317..4322d0c8b5 100644 --- a/workers/queue_backend/pg_queue/reaper.py +++ b/workers/queue_backend/pg_queue/reaper.py @@ -1326,8 +1326,50 @@ 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 and the API storage connector. + 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 + swept = (getattr(response, "data", None) or {}).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/shared/api/internal_client.py b/workers/shared/api/internal_client.py index a81287ccff..2342688c35 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 and the API storage connector; 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_reaper_undispatched_sweep.py b/workers/tests/test_reaper_undispatched_sweep.py new file mode 100644 index 0000000000..591044264c --- /dev/null +++ b/workers/tests/test_reaper_undispatched_sweep.py @@ -0,0 +1,95 @@ +"""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() + api.sweep_undispatched_executions.return_value = SimpleNamespace( + data={"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 = SimpleNamespace( + data={"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}): + api = MagicMock() + api.sweep_undispatched_executions.return_value = SimpleNamespace( + data=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 From 2088d6962843aa47471f02baeae7de121df7202c Mon Sep 17 00:00:00 2001 From: ali Date: Thu, 20 Aug 2026 12:33:23 +0530 Subject: [PATCH 19/33] UN-3796 [FIX] Baseline next_run_at on the UI resume path, not just enable_task Re-enabling a paused pipeline fired one spurious catch-up run within ~5s. Observed three times on integration (2026-08-12, -14, -20) against gallh_load_test; the 08-12 occurrence collided with the operator's manual run over the same 100 files and blew the shared Azure gpt-4o quota, failing 83 and 88 files across the two executions. ec0362fe5 fixed this for enable_task/disable_task, but the UI resume does not take that path: it re-saves the pipeline, reaching _schedule_task_job -> mirror_periodic_schedule_upsert and reconcile_ownership_for. The latter baselines only on a Beat->PG hand-over (pg_owned and not was_pg_owned), so an already-pg_owned row kept the next_run_at it held when paused -- by then in the past -- and the tick's `next_run_at <= now()` fired it on the next pass. The upsert could not express the fix: _retargeted_next_run_at returned None both for "leave the column alone" and would have for "write NULL", and the call site read None as the former. Split by a _LEAVE_NEXT_RUN_AT sentinel so NULL can be written deliberately, and clear on the enabled False->True transition only -- the same transition-scoping reconcile_ownership_for uses, so a save at 12:07:59 against a 12:08 next_run_at still fires. Resume is checked before the cron comparison: a plain pause/resume does not change the cron, so the equal-cron early-out would otherwise swallow the common case. Mutation-checked -- dropping the branch, collapsing the sentinel to None, and reordering the two checks each turn tests red. Co-Authored-By: Claude Opus 5 --- backend/scheduler/tasks.py | 96 ++++++++++----- .../tests/test_pg_periodic_schedule_mirror.py | 111 ++++++++++++++++-- 2 files changed, 166 insertions(+), 41 deletions(-) diff --git a/backend/scheduler/tasks.py b/backend/scheduler/tasks.py index 601584cf39..0f6373a9b6 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, diff --git a/backend/scheduler/tests/test_pg_periodic_schedule_mirror.py b/backend/scheduler/tests/test_pg_periodic_schedule_mirror.py index 11a2ac69f9..c256ba5030 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 From 229beab3cdfb4558fb252f923ecc4c06e71a840d Mon Sep 17 00:00:00 2001 From: ali Date: Thu, 20 Aug 2026 15:35:01 +0530 Subject: [PATCH 20/33] UN-3796 [FIX] Only mirror an interval periodic when */N is exactly equivalent cron_from_periodic_task's docstring promised "a step cron where one exists exactly", but the code range-checked (every < 60 / < 24 / < 32). `*/N` restarts at each field boundary, so it means "every N" only when N divides that field's range. Where it does not, the last step of one period runs into the first of the next and the task fires early -- silently, and only at the boundary: */7 minutes -> :00 :07 ... :56 then :00 -- a 4-minute gap, not 7 */45 minutes -> :00 :45 then :00 -- roughly twice as often 0 */5 hours -> 0 5 10 15 20 then 0 -- a 4-hour gap, not 5 Days are worse: `*/N` on day-of-month restarts every month and months are 28-31 days, so `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" beyond N=1, so only every==1 maps, to a plain daily. Now requires 60 % every == 0 / 24 % every == 0 / every == 1. Anything else falls through to plan_mirror's existing skip-and-explain path, the same one that already refuses second-resolution intervals: the periodic stays on Beat visibly rather than being adopted at a frequency nobody chose. Found while verifying the integration cutover, where dashboard_metrics_aggregate_from_sources turned out to be an IntervalSchedule (interval_id=8) rather than a crontab. It was unaffected -- 15 divides 60 -- but Beat schedules are per-environment DB rows that exist in no source file, so staging or production can carry an interval that does not. The (3, DAYS, "0 0 */3 * *") case in test_interval_maps_to_an_exact_step_cron encoded the bug as an expectation; replaced deliberately, since that cron does not fire every 3 days. Mutation-checked: restoring the range check fails 13 tests, dropping only the days restriction fails 6. Co-Authored-By: Claude Opus 5 --- .../commands/mirror_pg_periodic_tasks.py | 38 +++++++++++++--- .../tests/test_mirror_pg_periodic_tasks.py | 44 ++++++++++++++++++- 2 files changed, 74 insertions(+), 8 deletions(-) diff --git a/backend/pg_queue/management/commands/mirror_pg_periodic_tasks.py b/backend/pg_queue/management/commands/mirror_pg_periodic_tasks.py index f39cf0dde3..2a9b58bdc1 100644 --- a/backend/pg_queue/management/commands/mirror_pg_periodic_tasks.py +++ b/backend/pg_queue/management/commands/mirror_pg_periodic_tasks.py @@ -71,13 +71,37 @@ def cron_from_periodic_task(task: PeriodicTask) -> str: 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 where one exists exactly - (``*/N`` minutes / hours / days). + * ``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 @@ -88,13 +112,13 @@ def cron_from_periodic_task(task: PeriodicTask) -> str: every = interval.every if every < 1: return "" - if interval.period == IntervalSchedule.MINUTES and every < 60: + if interval.period == IntervalSchedule.MINUTES and every < 60 and 60 % every == 0: return f"*/{every} * * * *" - if interval.period == IntervalSchedule.HOURS and every < 24: + if interval.period == IntervalSchedule.HOURS and every < 24 and 24 % every == 0: return f"0 */{every} * * *" - if interval.period == IntervalSchedule.DAYS and every < 32: - return f"0 0 */{every} * *" - # SECONDS, or a step too large to express as a single cron field. + if interval.period == IntervalSchedule.DAYS and every == 1: + return "0 0 * * *" + # SECONDS, a step that does not divide its field, or anything else. return "" diff --git a/backend/pg_queue/tests/test_mirror_pg_periodic_tasks.py b/backend/pg_queue/tests/test_mirror_pg_periodic_tasks.py index 52a4c186cf..5605310150 100644 --- a/backend/pg_queue/tests/test_mirror_pg_periodic_tasks.py +++ b/backend/pg_queue/tests/test_mirror_pg_periodic_tasks.py @@ -74,8 +74,13 @@ def test_crontab_is_reconstructed_field_for_field(self): "every,period,expected", [ (15, MINUTES, "*/15 * * * *"), # dashboard_metrics.aggregate_from_sources + (1, MINUTES, "*/1 * * * *"), + (30, MINUTES, "*/30 * * * *"), (2, HOURS, "0 */2 * * *"), - (3, DAYS, "0 0 */3 * *"), + (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): @@ -94,6 +99,43 @@ 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)) == "" From 1c05a3599771e65078a717296724356e93eceada Mon Sep 17 00:00:00 2001 From: ali Date: Mon, 24 Aug 2026 10:48:27 +0530 Subject: [PATCH 21/33] UN-3796 [FIX] Baseline Beat's clock on release, or Beat replays every missed interval Releasing schedules back to Celery Beat fired a burst of catch-up runs. Observed on integration 2026-08-24: converge_pg_scheduler released 23 schedules and Beat dispatched 4 pipelines plus all 3 dashboard_metrics.* periodics within 30 ms of logging "Released to Beat". Release restored PeriodicTask.enabled but left last_run_at untouched. DatabaseScheduler stores no next_run_at -- it derives due-ness from last_run_at against the crontab -- so rows that had been PG-owned since 2026-08-21 were overdue by every interval they had missed, and Beat replayed them the instant `enabled` flipped back. This is the exact mirror of the adopt-side bug fixed in 2088d6962, which baselines next_run_at so a hand-over cannot fire a catch-up. That half was done; this half was not. Fixed at both release sites, scoped to the TRANSITION like its counterpart: scheduler/ownership.py pipelines, on `was_pg_owned and not pg_owned` mirror_pg_periodic_tasks.py periodics, on `not to_pg` Stamping unconditionally would push Beat's clock forward on an ordinary pipeline save (reconcile runs on every save) and silently skip a due fire, so both are gated on the release edge. Adopt deliberately does NOT touch last_run_at: Beat is being switched off, its clock is irrelevant, and overwriting it would destroy the value the eventual release restores from. `timezone` was not imported in mirror_pg_periodic_tasks.py -- the fix would have raised NameError on the very path it repairs, so that import is load-bearing. CORRECTS A CLAIM I HAD PROPAGATED. ec0362fe5 asserted "Beat parity, not a new rule: DatabaseScheduler holds no persisted next_run_at and recomputes due-ness from the crontab each tick, so re-enabling never produced a catch-up run there." That is wrong: it recomputes from last_run_at, which is precisely what makes it catch up. I used that claim to argue the catch-up was a PG-only regression against Beat; it is in fact a hazard both schedulers share. Corrected in all three places it had spread to (ownership.py, scheduler/tasks.py, and two test docstrings), so the next reader does not inherit it. Why it matters beyond tidiness: release is the ROLLBACK path. Staging and production would use it under pressure, and a rollback that immediately replays every overdue schedule -- burning an LLM pass per pipeline -- is a poor thing to discover mid-incident. 782 backend + 1392 worker tests pass; pre-commit clean. Mutation-checked, each red for a distinct reason: dropping the pipeline stamp fails 4, stamping unconditionally instead of on the transition fails 2, dropping the periodics stamp fails 1. Four expectations changed deliberately, not loosened: they asserted the whole kwargs dict (`== {"enabled": True}`), which now carries a second key. Each is per-key now plus new cases, pinning strictly more than the exact-dict form did -- including that a paused pipeline is released disabled but still baselined, so resuming it later does not replay the backlog it accrued while PG owned it. Co-Authored-By: Claude Opus 5 --- .../commands/mirror_pg_periodic_tasks.py | 19 +++- .../tests/test_mirror_pg_periodic_tasks.py | 56 ++++++++++- backend/scheduler/ownership.py | 25 ++++- backend/scheduler/tasks.py | 16 ++- .../tests/test_pg_periodic_schedule_mirror.py | 11 ++- .../tests/test_pg_schedule_ownership.py | 99 ++++++++++++++++++- 6 files changed, 202 insertions(+), 24 deletions(-) diff --git a/backend/pg_queue/management/commands/mirror_pg_periodic_tasks.py b/backend/pg_queue/management/commands/mirror_pg_periodic_tasks.py index 2a9b58bdc1..4e4766b7ba 100644 --- a/backend/pg_queue/management/commands/mirror_pg_periodic_tasks.py +++ b/backend/pg_queue/management/commands/mirror_pg_periodic_tasks.py @@ -31,6 +31,7 @@ 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 @@ -385,9 +386,21 @@ def _set_ownership( if not to_pg: row.next_run_at = None row.save(update_fields=["pg_owned", "next_run_at", "updated_at"]) - PeriodicTask.objects.filter(name=row.name).update( - enabled=beat_enabled - ) + 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() + PeriodicTask.objects.filter(name=row.name).update(**beat_updates) # 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 diff --git a/backend/pg_queue/tests/test_mirror_pg_periodic_tasks.py b/backend/pg_queue/tests/test_mirror_pg_periodic_tasks.py index 5605310150..d5fe342907 100644 --- a/backend/pg_queue/tests/test_mirror_pg_periodic_tasks.py +++ b/backend/pg_queue/tests/test_mirror_pg_periodic_tasks.py @@ -259,9 +259,12 @@ class TestReleaseRestoresRatherThanResurrects: _CMD = "pg_queue.management.commands.mirror_pg_periodic_tasks" - def _release(self, mirrored_enabled: bool): + def _flip(self, flag: str, mirrored_enabled: bool = True): row = SimpleNamespace( - name="h", pg_owned=True, next_run_at=None, enabled=mirrored_enabled + name="h", + pg_owned=(flag == "--release"), + next_run_at=None, + enabled=mirrored_enabled, ) with ( patch(f"{self._CMD}.PgPeriodicTask") as Model, @@ -276,15 +279,58 @@ def _release(self, mirrored_enabled: bool): 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", "--release") + 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": True} + 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": False} + 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: diff --git a/backend/scheduler/ownership.py b/backend/scheduler/ownership.py index e799ec0624..4e64d37b19 100644 --- a/backend/scheduler/ownership.py +++ b/backend/scheduler/ownership.py @@ -240,9 +240,28 @@ 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() + PeriodicTask.objects.filter(name=pipeline_id).update(**beat_updates) # 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 0f6373a9b6..755d878b20 100644 --- a/backend/scheduler/tasks.py +++ b/backend/scheduler/tasks.py @@ -168,11 +168,17 @@ def _mirror_periodic_schedule_set_enabled(pipeline_id: str, enabled: bool) -> No # 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. # - # This is Beat parity, not a new rule: DatabaseScheduler holds no persisted - # next_run_at and recomputes due-ness from the crontab each tick, so - # re-enabling never produced a catch-up run there. Observed on integration - # 2026-08-14: gallh_load_test fired ~2s after being re-enabled, against a - # next_run_at two days old. + # 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 diff --git a/backend/scheduler/tests/test_pg_periodic_schedule_mirror.py b/backend/scheduler/tests/test_pg_periodic_schedule_mirror.py index c256ba5030..8b4439f7a8 100644 --- a/backend/scheduler/tests/test_pg_periodic_schedule_mirror.py +++ b/backend/scheduler/tests/test_pg_periodic_schedule_mirror.py @@ -435,10 +435,15 @@ class TestResumeBaselinesInsteadOfCatchingUp: 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). - Beat parity, not a new rule — DatabaseScheduler keeps no persisted next_run_at and - recomputes due-ness from the crontab each tick, so resume never produced a catch-up - run there. Observed on integration 2026-08-14: gallh_load_test fired ~2s after being + 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): diff --git a/backend/scheduler/tests/test_pg_schedule_ownership.py b/backend/scheduler/tests/test_pg_schedule_ownership.py index b9bb03d583..4f05e8d150 100644 --- a/backend/scheduler/tests/test_pg_schedule_ownership.py +++ b/backend/scheduler/tests/test_pg_schedule_ownership.py @@ -301,7 +301,15 @@ def test_a_stale_row_is_handed_back_to_beat(self, monkeypatch): 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. - PT.objects.filter.return_value.update.assert_called_once_with(enabled=True) + # + # 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() @@ -314,7 +322,11 @@ def test_a_paused_pipeline_is_not_resurrected_by_the_repair(self, monkeypatch): patch("scheduler.ownership.PeriodicTasks"), ): ownership.reconcile_ownership_for(_PID, _ORG, active=False) - PT.objects.filter.return_value.update.assert_called_once_with(enabled=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. @@ -421,9 +433,14 @@ class TestNextRunBaselineOnTransition: 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. Beat never did this: DatabaseScheduler - keeps no persisted next_run_at and recomputes due-ness from the crontab each tick, - so this is a PG-path regression against it, not a cosmetic difference. + 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): @@ -467,3 +484,75 @@ def test_an_unchanged_owner_is_NOT_re_baselined(self, monkeypatch): """ 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 From eabb7889a11b7b1edc6fa47dc6cf674c421f48e9 Mon Sep 17 00:00:00 2001 From: ali Date: Mon, 24 Aug 2026 14:38:51 +0530 Subject: [PATCH 22/33] UN-3796 [FIX] Let the stuck-execution finalizer see Celery executions too An execution stranded by a PG cutover sat EXECUTING forever with nothing able to close it. recover_stuck_pg_executions filtered `queue_message_id__isnull=False`, and a Celery row carries task_id with queue_message_id NULL -- so it was invisible. The other recovery paths could not help either: the chord callback is gone with its worker, and the reaper's barrier sweeps scan pg_barrier_state, which a Celery execution never had. The mechanism is a grace-period asymmetry, not a queue defect: worker-file-processing-v2 grace 7200s worker-file-processing-callback-v2 grace 300s On a cutover both get SIGTERM. File processing has two hours and finishes its batches; the callback worker dies after five minutes. The batches then dispatch process_batch_callback to a worker that is already gone. Every file reaches COMPLETED and the execution never leaves EXECUTING. Only the SELECT was PG-scoped -- _recover_one_stuck_pg_execution was always transport-agnostic, reading file statuses and recomputing the terminal status without touching a queue. So this widens the filter and changes nothing else: Q(queue_message_id__isnull=False) | Q(task_id__isnull=False) Rows with BOTH handles NULL are deliberately still excluded. Those were never dispatched and belong to undispatched_sweep.py, whose claim requires `task_id IS NULL AND queue_message_id IS NULL`. The two predicates are now disjoint on task_id, which matters because both run on the same reaper cadence against the same table -- an overlap would have one marking a row ERROR while the other finalized it. The `total == 0 -> skipped` guard would also have caught it, 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, documented in the docstring. test_celery_execution_never_scanned asserted exactly the behaviour being changed -- INVERTED deliberately, not loosened, with the reason recorded in the test. Added test_a_NEVER_DISPATCHED_execution_is_left_to_the_undispatched_sweep as the disjointness guard. VERIFICATION IS INCOMPLETE AND CI IS THE GATE. 782 unit tests pass, syntax and pre-commit clean -- but the tests that exercise this change are DB-backed and did not run locally: the rig needs testcontainers (unavailable here), and running against the local unstract-db container fails on multi-tenant schema setup. Both new/changed tests are CI-verified only; I have not seen them execute. Co-Authored-By: Claude Opus 5 --- .../tests/test_pg_finalization_fixes.py | 39 ++++++++++++++++++- backend/workflow_manager/internal_views.py | 28 +++++++++++-- 2 files changed, 61 insertions(+), 6 deletions(-) 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..87a69403a2 100644 --- a/backend/workflow_manager/execution/tests/test_pg_finalization_fixes.py +++ b/backend/workflow_manager/execution/tests/test_pg_finalization_fixes.py @@ -113,14 +113,49 @@ def test_file_less_stuck_is_skipped_not_failed(self): assert out["skipped"] == 1 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): diff --git a/backend/workflow_manager/internal_views.py b/backend/workflow_manager/internal_views.py index 4f9de9f0aa..d8bb37cee0 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 Q from django.utils import timezone from workflow_manager.workflow_v2.enums import ExecutionStatus @@ -674,7 +692,9 @@ def recover_stuck_pg_executions(self, request): stuck_ids = list( WorkflowExecution.objects.filter( - queue_message_id__isnull=False, # PG-only; Celery uses task_id + # Dispatched on EITHER transport. Both-NULL means never dispatched — + # undispatched_sweep.py's row, not ours. + Q(queue_message_id__isnull=False) | Q(task_id__isnull=False), status__in=[ ExecutionStatus.PENDING.value, ExecutionStatus.EXECUTING.value, From 4571e0e4823d7f5b4144e9b1708804196e82ff80 Mon Sep 17 00:00:00 2001 From: ali Date: Mon, 24 Aug 2026 16:17:48 +0530 Subject: [PATCH 23/33] UN-3796 [FIX] Default the stuck-recovery window to 10 min, not the barrier timeout The safety net that finalizes a stranded execution waited ~2.5 hours before acting, because WORKER_PG_STUCK_EXECUTION_RECOVERY_SECONDS defaulted to the barrier stuck-timeout. The two answer different questions: barrier stuck-timeout how long a batch may make NO progress -> hours, since a single file can legitimately take that long recovery window debounce against a callback that is -> seconds about to fire Reusing one for the other was convenience, not design, and it bought nothing: the threshold is NOT what makes recovery safe. internal_views.py skips unless EVERY file is terminal (`total == 0 or terminal < total -> skipped`), so a legitimately running execution is never a candidate however short the window is. The long default only meant a genuinely dead execution stayed dead for 2.5 h. The shape that motivates this is a grace-period asymmetry, not a queue defect: worker-file-processing-v2 grace 7200s worker-file-processing-callback-v2 grace 300s On a cutover both get SIGTERM. File processing has two hours and finishes its batches; the callback worker dies after five minutes; the batches then dispatch process_batch_callback to a worker that is gone. Every file is COMPLETED and the execution never leaves EXECUTING. eabb7889a lets the finalizer SEE those; this makes it act in minutes instead of hours. A DEFAULT rather than an operator setting, deliberately. Production and on-prem have no Flipt and no operator: a value that exists only as an env override is one those environments never receive, and they are exactly the ones that cannot diagnose a hung execution themselves. The env still overrides for tuning. Not lowered further because `total` counts the file rows that EXIST, not execution.total_files -- while discovery is still creating rows there is a brief "all existing rows terminal" moment. Ten minutes clears it comfortably. Comparing against total_files would close that window properly and allow less; noted in the constant's docstring as a follow-up. 4 tests, mutation-checked: restoring the 9000s default fails test_default_is_ten_ minutes and test_default_is_NOT_the_barrier_stuck_timeout. 782 backend + 1396 worker tests pass, pre-commit clean. Co-Authored-By: Claude Opus 5 --- workers/queue_backend/pg_queue/reaper.py | 36 +++++++++++++-- workers/tests/test_pg_finalization_fixes.py | 51 +++++++++++++++++++++ 2 files changed, 82 insertions(+), 5 deletions(-) diff --git a/workers/queue_backend/pg_queue/reaper.py b/workers/queue_backend/pg_queue/reaper.py index 4322d0c8b5..ff5f23b185 100644 --- a/workers/queue_backend/pg_queue/reaper.py +++ b/workers/queue_backend/pg_queue/reaper.py @@ -175,6 +175,34 @@ 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 only needs to outlast a callback that is about to fire — seconds. +#: +#: **The threshold is not what makes this safe.** The endpoint skips unless every file +#: is terminal (``internal_views.py``: ``total == 0 or terminal < total → skipped``), so +#: a legitimately running execution is never a candidate no matter how short this is. +#: Inheriting the multi-hour barrier timeout therefore bought nothing and 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). @@ -1079,15 +1107,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 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: From 07bf352102e59187a522223010cf53b08c2bbe76 Mon Sep 17 00:00:00 2001 From: ali Date: Tue, 25 Aug 2026 10:04:10 +0530 Subject: [PATCH 24/33] UN-3796 [FIX] Make the stranded-execution sweep reach recoverable rows MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The safety net that finalizes executions stranded by a transport cutover was selecting rows it could never act on, and never reaching the ones it could. Found by rehearsing the cutover on integration: a 30-file ETL was stranded exactly as designed (every file COMPLETED, execution stuck EXECUTING, no duplicate work, no data loss) and then sat there while every sweep logged a healthy-looking "scanned 100, recovered 0". Three changes, all in the selection path; the finalization logic is untouched. 1. Select only rows that can actually finalize. Candidates were any dispatched PENDING/EXECUTING row past the cutoff, and the all-files-terminal decision happened per-row afterwards. A rejected row is skipped WITHOUT touching modified_at, so under oldest-first ordering the same rows were re-selected on every sweep, forever, and anything behind them was invisible — not scanned, not skipped. Integration held 1964 candidates, 1122 of them permanently unrecoverable (476 with no file rows, 646 with a non-terminal file); the freshly stranded execution ranked 1964th and would never have been reached. Pushing the guard into the query makes the drain real: every selected row finalizes and leaves the candidate set. That is what keeps the existing oldest-first ordering safe — FIFO is only fair if the queue moves — so the ordering is deliberately left alone. 2. Measure staleness on the files, not the execution row. workflow_execution.modified_at is effectively the START time: file completions write the file row, not the execution row. "modified_at < cutoff" therefore meant "started more than stuck_seconds ago", so any run longer than the window was formally stuck while still running — a 14-minute ETL was eligible from minute ten against a 600s window. The all-files-terminal filter kept that from being catastrophic, but left a live race: between the last file going terminal and the callback finalizing, a healthy execution passes every check. A sweep landing there finalizes it first, the terminal-one-way guard then refuses the callback's own write, and the notification is silently lost. Change (1) makes that race reachable for the first time by removing the starvation, so this ships with it rather than after it. 3. Stop fabricating execution_time on retroactive finalization. update_execution() stamps execution_time = now - created_at on any terminal transition. Right for a live finalization; wrong here, where it records how late the reaper was rather than how long the run took — multi-day runtimes on executions that ran for minutes, feeding whatever reads that column. Recomputed from the last file to finish, and left alone when no file carries a usable timestamp. Tests pin the starvation regression (a backlog larger than `limit` composed of permanently-skippable rows must not hide a recoverable one), the drain property, both halves of the file-staleness predicate, and the execution_time behaviour; each states its mutation check. Two existing tests moved from asserting "skipped == 1" to "scanned == 0" — those rows are now excluded at selection rather than rejected downstream. The invariant they protect is unchanged and still asserted: neither is ever finalized or failed. NOT VERIFIED LOCALLY. The backend suite cannot start in this checkout — every test in the file errors at setup with "no schema has been selected to create in", including ones this commit does not touch, and settings.cloud will not import without the enterprise pluggable_apps. Confirmed instead: both files compile, ruff reports exactly the findings HEAD already had, ruff-format clean, scoped pre-commit passed. The tests need CI or a provisioned stack. Co-Authored-By: Claude Opus 5 --- .../tests/test_pg_finalization_fixes.py | 195 ++++++++++++++++-- backend/workflow_manager/internal_views.py | 74 ++++++- 2 files changed, 256 insertions(+), 13 deletions(-) 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 87a69403a2..294b659515 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,15 +113,26 @@ 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_IS_now_recovered(self): @@ -158,7 +180,14 @@ def test_a_NEVER_DISPATCHED_execution_is_left_to_the_undispatched_sweep(self): 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], @@ -166,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): @@ -176,7 +206,146 @@ 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) + WorkflowExecution.objects.filter(pk=ex.pk).update(created_at=created) + 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]) + WorkflowFileExecution.objects.filter(workflow_execution=ex).update( + modified_at=None + ) + ex.execution_time = 7.5 + ex.save(update_fields=["execution_time"]) + + 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): @@ -455,6 +624,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 d8bb37cee0..ef817977e5 100644 --- a/backend/workflow_manager/internal_views.py +++ b/backend/workflow_manager/internal_views.py @@ -669,7 +669,7 @@ def recover_stuck_pg_executions(self, request): """ from datetime import timedelta - from django.db.models import Q + from django.db.models import Exists, Max, OuterRef, Q from django.utils import timezone from workflow_manager.workflow_v2.enums import ExecutionStatus @@ -690,6 +690,24 @@ 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( # Dispatched on EITHER transport. Both-NULL means never dispatched — @@ -701,6 +719,30 @@ def recover_stuck_pg_executions(self, request): ], 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] ) @@ -773,6 +815,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, { @@ -848,6 +891,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.""" From bb772503b65273209b7effbf0f5508f39141284d Mon Sep 17 00:00:00 2001 From: ali Date: Tue, 25 Aug 2026 10:59:57 +0530 Subject: [PATCH 25/33] UN-3796 [FIX] Stop the undispatched sweep deleting a running execution's input The sweep infers "never dispatched" from the ABSENCE of both handles (task_id IS NULL AND queue_message_id IS NULL). That inference is unsound. Three paths in workflow_helper reach exactly that state AFTER the message is already on its transport: * _record_dispatch_handle raises and the caller swallows it, under the comment "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) workflow_helper.py:665-666 states the invariant in so many words: past dispatch, a bookkeeping failure must not flip the now-running row. The sweep then does exactly that 15 minutes later - and, worse, deleted the execution's staged input while the worker was still going to read it, telling the user "You can safely run it again". Marking the row ERROR is survivable: the running worker's own terminal write supersedes it, and error->completed is explicitly permitted by the status guard. Deleting the input is not survivable. This drops the delete and keeps the reversible half (the rate-limit slot release, which has a live cost - a held slot consumes the org's concurrency budget until the limiter TTL expires it). The cost is a leaked 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 because it needs a migration and an index rebuild. Tests pin the ABSENCE of the delete, so it cannot be reinstated without revisiting the predicate, and that the slot release survives the removal. Found by the standardized 16-lens review on PR #2254 (C1, the sole blocking finding). Verified against the source rather than accepted on report: all three no-handle paths confirmed present, and no existing test pinned the deletion. Co-Authored-By: Claude Opus 5 --- .../tests/test_undispatched_sweep.py | 48 +++++++++++++++++-- .../workflow_v2/undispatched_sweep.py | 42 +++++++++------- 2 files changed, 71 insertions(+), 19 deletions(-) diff --git a/backend/workflow_manager/workflow_v2/tests/test_undispatched_sweep.py b/backend/workflow_manager/workflow_v2/tests/test_undispatched_sweep.py index 9fb1cf92e6..f0d7ec3c54 100644 --- a/backend/workflow_manager/workflow_v2/tests/test_undispatched_sweep.py +++ b/backend/workflow_manager/workflow_v2/tests/test_undispatched_sweep.py @@ -60,9 +60,10 @@ def test_an_aged_undispatched_execution_is_marked_error(self): 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 + 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.""" @@ -158,6 +159,47 @@ def test_a_row_dispatched_mid_sweep_is_not_clobbered(self): assert ex.status == ExecutionStatus.PENDING.value +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. diff --git a/backend/workflow_manager/workflow_v2/undispatched_sweep.py b/backend/workflow_manager/workflow_v2/undispatched_sweep.py index b45e9fc074..9d98423aab 100644 --- a/backend/workflow_manager/workflow_v2/undispatched_sweep.py +++ b/backend/workflow_manager/workflow_v2/undispatched_sweep.py @@ -220,19 +220,29 @@ def _release_abandoned_resources(execution_id: str, workflow_id: str) -> None: exc_info=True, ) - # Then the staged input. Scoped by workflow_id + execution_id, and guarded by an - # exists() check inside, so it is a clean no-op for an execution that died BEFORE - # staging — nothing else's files can be reached from here. - try: - from workflow_manager.endpoint_v2.destination import DestinationConnector - - DestinationConnector.delete_api_storage_dir( - workflow_id=workflow_id, execution_id=execution_id - ) - except Exception: - logger.warning( - "Undispatched sweep: could not delete the API storage dir for %s " - "(orphaned input files remain)", - 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. Marking it + # ERROR is wrong but self-corrects — the running worker's own terminal write + # supersedes it, and error→completed is explicitly permitted by the status guard. + # Deleting its input does not self-correct: the worker is still going to read it, + # and the row's user-facing 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 From 0c6b1d1e6cc4bac767b8fbd5118f8dd615cd3f52 Mon Sep 17 00:00:00 2001 From: ali Date: Tue, 25 Aug 2026 11:01:12 +0530 Subject: [PATCH 26/33] UN-3796 [FIX] The undispatched-sweep metric could never increment `swept = (getattr(response, "data", None) or {}).get("swept", 0)` reads `.data` off the return of `sweep_undispatched_executions()`, which is typed `-> dict[str, Any]` and returns the parsed body verbatim: {"swept": N}, with no {"data": ...} envelope. `getattr` on a plain dict yields None, so the counter was always 0. Consequence: `pg_reaper_undispatched_swept_total` never incremented and the "terminalised N undispatched execution(s)" line never fired. The sweep itself worked; every operator-facing signal for it reported nothing. The metric's help text says a sustained non-zero rate means requests are dying between create_workflow_execution and dispatch - the 967-orphan condition this feature was built to detect - and that was invisible. The sibling recover_stuck_pg_executions() meets the same flat-body shape and documents the trap explicitly (execution_client.py:352-358), wrapping the body so `.data` works. That is why that path's counters are correct and this one's were not. Fixed by reading the dict directly rather than introducing a second wrapping convention. The test fixture was the reason this shipped: it returned SimpleNamespace(data={"swept": 3}), a shape nothing in the call chain produces, so it asserted inc(3) against a fiction while production incremented by 0. CI could not have caught it - the test asserted the bug away. Fixtures now return the real contract, and the tolerance test additionally feeds an object with `.data` to prove that shape is ignored rather than silently trusted. Found by the standardized 16-lens review on PR #2254; verified against the client's declared return type and the sibling's comment before accepting. Co-Authored-By: Claude Opus 5 --- workers/queue_backend/pg_queue/reaper.py | 11 ++++++++++- .../tests/test_reaper_undispatched_sweep.py | 19 +++++++++---------- 2 files changed, 19 insertions(+), 11 deletions(-) diff --git a/workers/queue_backend/pg_queue/reaper.py b/workers/queue_backend/pg_queue/reaper.py index ff5f23b185..d94cb485a9 100644 --- a/workers/queue_backend/pg_queue/reaper.py +++ b/workers/queue_backend/pg_queue/reaper.py @@ -1387,7 +1387,16 @@ def _sweep_undispatched_executions(self) -> None: exc_info=True, ) return - swept = (getattr(response, "data", None) or {}).get("swept", 0) + # 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( diff --git a/workers/tests/test_reaper_undispatched_sweep.py b/workers/tests/test_reaper_undispatched_sweep.py index 591044264c..099c1eb82a 100644 --- a/workers/tests/test_reaper_undispatched_sweep.py +++ b/workers/tests/test_reaper_undispatched_sweep.py @@ -35,9 +35,12 @@ def _reaper_with(api_client): class TestTheTriggerIsWiredToTheBackend: def test_it_calls_the_backend_sweep(self): api = MagicMock() - api.sweep_undispatched_executions.return_value = SimpleNamespace( - data={"swept": 3} - ) + # 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() @@ -48,9 +51,7 @@ def test_it_calls_the_backend_sweep(self): 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 = SimpleNamespace( - data={"swept": 0} - ) + 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 @@ -61,11 +62,9 @@ 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}): + for payload in (None, {}, {"unexpected": 1}, SimpleNamespace(data={"swept": 9})): api = MagicMock() - api.sweep_undispatched_executions.return_value = SimpleNamespace( - data=payload - ) + api.sweep_undispatched_executions.return_value = payload _reaper_with(api)._sweep_undispatched_executions() # no raise From 9e502f7da4b387f30563de1a6a3ce66da38009b5 Mon Sep 17 00:00:00 2001 From: ali Date: Tue, 25 Aug 2026 11:25:47 +0530 Subject: [PATCH 27/33] UN-3796 [FIX] Make the reorder guard real, and correct three comments that misdirect Four findings from the standardized review on PR #2254, each verified against source before accepting. 1. _DueSchedule / _DuePeriodicTask did not provide the safety they promise. Both docstrings say the type binds field names to columns "so a future reorder of the SELECT can't silently misassign fields". Both were built by positional unpacking - `_DuePeriodicTask(*row)` - against a hardcoded column list, which is exactly what cannot catch a reorder. Six of the eight periodic fields are same-typed TEXT, so swapping task_name and queue constructs cleanly and then fires the queue name as a task onto a queue named after the task: a periodic that silently stops running, with a MAX_ATTEMPTS=1 message dropped on an unwatched queue. Now derived: `SELECT {", ".join(_DuePeriodicTask._fields)}`. Verified the field names match the previous column lists exactly and in order, so the emitted SQL is unchanged - the claim is simply true now. 2. reaper.py's recovery-window comment told the next maintainer the threshold is not load-bearing. This branch is what made it load-bearing. It argued the value "only needs to outlast a callback that is about to fire - seconds" and that a running execution "is never a candidate no matter how short this is". That held when all-files-terminal was the only guard. The selection now also requires last_file_at < cutoff, precisely because all-files-terminal is NOT sufficient. Anyone trusting the old wording and dropping the window to 30s reopens the callback race. 3. entrypoint.sh documented a command it does not run, and asserted the opposite of what the container does. A block described `--mirror-only` and stated "Ownership hand-over stays an explicit operator action". What runs is converge_pg_scheduler, which with PG_SCHEDULER_ENABLED=true flips pg_owned and DISABLES Beat PeriodicTask rows. An SRE auditing "can a rolling restart move schedules off Beat?" got the wrong answer for the single riskiest behaviour here. Two contradictory "Best-effort by design" paragraphs sat side by side; the stale one is gone. 4. reconcile_pg_schedules.py carried the same stale claim ("that is what automation runs; ownership stays an operator action") and is corrected to say the entrypoint now runs the non-mirror-only mode unattended. Tests: workers/tests/test_pg_scheduler.py 24 passed, test_reaper_undispatched_sweep.py 5 passed. Shell syntax checked with bash -n. Co-Authored-By: Claude Opus 5 --- backend/entrypoint.sh | 35 +++++++++---------- .../commands/reconcile_pg_schedules.py | 18 ++++++---- .../queue_backend/pg_queue/pg_scheduler.py | 10 +++--- workers/queue_backend/pg_queue/reaper.py | 19 +++++++--- 4 files changed, 48 insertions(+), 34 deletions(-) diff --git a/backend/entrypoint.sh b/backend/entrypoint.sh index 9fec4d8f09..b3236d7dda 100755 --- a/backend/entrypoint.sh +++ b/backend/entrypoint.sh @@ -28,26 +28,25 @@ if [ "$migrate" = true ]; then echo "Migration initiated" .venv/bin/python manage.py migrate - # Backfill PG-scheduler mirror rows for pipeline schedules created before the - # mirror existed (UN-3796). NOT a migration — an ordinary management command, - # sequenced after migrate only because pg_periodic_schedule must exist first. + # 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. # - # --mirror-only: purely additive, writes only pg_periodic_schedule and never a - # Beat PeriodicTask, so it is safe at any flag state. Ownership hand-over stays - # an explicit operator action. + # 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. 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.) # - # Runs on EVERY start, not once: schedules created while an older backend was - # deployed, and rows previously skipped for malformed args, are only picked up by - # a re-run. It is idempotent (already-mirrored pipelines are skipped), so there - # is nothing to retire until Celery is decommissioned. - # - # Best-effort by design: a mirror failure must never stop the backend from - # starting. Beat keeps firing everything in that case, which is the safe state. - # Converge schedule ownership to whatever PG_SCHEDULER_ENABLED declares — adopt - # when on, release back to Beat when off. 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. + # 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, and a no-op for rows whose ownership already + # matches, so a restart that changes nothing writes nothing. # # PG_SCHEDULER_ADOPT_PERIODICS additionally moves the dashboard_metrics.* rows; # it is separate because adopting them needs workerPgMetrics deployed. diff --git a/backend/pg_queue/management/commands/reconcile_pg_schedules.py b/backend/pg_queue/management/commands/reconcile_pg_schedules.py index 63045b0c92..7f5a1e78e1 100644 --- a/backend/pg_queue/management/commands/reconcile_pg_schedules.py +++ b/backend/pg_queue/management/commands/reconcile_pg_schedules.py @@ -99,12 +99,18 @@ def handle(self, *args: Any, **options: Any) -> None: raise CommandError("--batch-size must be >= 1") backfilled = self._backfill_mirrors(dry_run, batch_size) - # --mirror-only exists because the deploy-time automation must be safe to run - # at ANY flag state. The reconcile below is fail-closed (rollout off → every - # schedule stays on Beat, nothing written), but with the rollout ON it flips - # ownership *and disables the matching Beat PeriodicTask* — a behaviour change - # no unattended job should make on its own. Backfilling is inert at every flag - # state, so that is what automation runs; ownership stays an operator action. + # --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: diff --git a/workers/queue_backend/pg_queue/pg_scheduler.py b/workers/queue_backend/pg_queue/pg_scheduler.py index 47c185d0d7..caa08c56bd 100644 --- a/workers/queue_backend/pg_queue/pg_scheduler.py +++ b/workers/queue_backend/pg_queue/pg_scheduler.py @@ -143,9 +143,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(_DueSchedule._fields)} + FROM {qualified("pg_periodic_schedule")} WHERE pg_owned AND enabled AND (next_run_at IS NULL OR next_run_at <= %s) """, @@ -292,9 +291,8 @@ def dispatch_due_periodic_tasks(conn: PgConnection) -> int: base = cur.fetchone()[0] cur.execute( f""" - SELECT name, task_name, queue, task_args, task_kwargs, org_id, - cron_string, next_run_at - FROM {qualified('pg_periodic_task')} + SELECT {", ".join(_DuePeriodicTask._fields)} + FROM {qualified("pg_periodic_task")} WHERE pg_owned AND enabled AND (next_run_at IS NULL OR next_run_at <= %s) """, diff --git a/workers/queue_backend/pg_queue/reaper.py b/workers/queue_backend/pg_queue/reaper.py index d94cb485a9..ffa10b8629 100644 --- a/workers/queue_backend/pg_queue/reaper.py +++ b/workers/queue_backend/pg_queue/reaper.py @@ -183,10 +183,21 @@ def reaper_interval_from_env() -> float: #: make no progress (hours, because a single file can legitimately take that long), #: whereas this only needs to outlast a callback that is about to fire — seconds. #: -#: **The threshold is not what makes this safe.** The endpoint skips unless every file -#: is terminal (``internal_views.py``: ``total == 0 or terminal < total → skipped``), so -#: a legitimately running execution is never a candidate no matter how short this is. -#: Inheriting the multi-hour barrier timeout therefore bought nothing and cost a lot: a +#: **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. From 7f3cc38abdf4d503d3a822cd83e6d4bbc9f7ae5a Mon Sep 17 00:00:00 2001 From: ali Date: Tue, 25 Aug 2026 12:00:34 +0530 Subject: [PATCH 28/33] UN-3796 [FIX] Collect runner/tests, pin available_at, and stop silent orphaning on release Four findings from the standardized review on PR #2254. All are test wiring, contract or logging - no control flow or dispatch behaviour changes. 1. runner/tests was in no rig group, so it was collected by nothing. tests/groups.yaml defines no runner group and tox.ini has no alias, so runner/tests/test_sidecar_log_transport.py never ran in any lane. It is the sole guard 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. It shipped alongside the Redis log transport and was verified by nothing. Added `unit-runner`, shaped like unit-core: install_editable, because the suite imports `unstract.runner` which resolves only from this package's own src layout. 2. available_at was missing from the worker schema-drift contract. Migration 0002 adds the column and worker raw SQL depends on it, but WORKER_SCHEMA_CONTRACT's pg_queue_message set never listed it - so renaming or dropping it would keep the guard green while delayed-visibility delivery broke at runtime. That is precisely the drift the file exists to catch. 9 passed. 3. Both Beat-side `.update()` sites discarded their match count. A bulk update returning 0 is indistinguishable from success. On the RELEASE direction 0 rows means PG has let go of a schedule and Beat has no row to take it over - a schedule with no firer at all, on the rollback path, reported as success. Both sites now log when nothing matched. Log-only: no branch, no raise, no change to what is written. 4. converge_pg_scheduler's docstring claimed a no-op that does not happen. "the pipeline path no-ops when ownership already matches" is false: the write and PeriodicTasks.update_changed() are unconditional, so every backend start costs 2N row writes and N Beat reloads regardless. Corrected to say idempotent in outcome but not free, which is what actually justifies running it unattended. Co-Authored-By: Claude Opus 5 --- .../commands/converge_pg_scheduler.py | 8 ++++++-- .../commands/mirror_pg_periodic_tasks.py | 12 +++++++++++- backend/scheduler/ownership.py | 17 ++++++++++++++--- tests/groups.yaml | 15 +++++++++++++++ workers/tests/test_pg_schema_drift.py | 5 +++++ 5 files changed, 51 insertions(+), 6 deletions(-) diff --git a/backend/pg_queue/management/commands/converge_pg_scheduler.py b/backend/pg_queue/management/commands/converge_pg_scheduler.py index 6c4d4378ed..a4341f5827 100644 --- a/backend/pg_queue/management/commands/converge_pg_scheduler.py +++ b/backend/pg_queue/management/commands/converge_pg_scheduler.py @@ -15,8 +15,12 @@ **Safe to run unattended**, which is what lets ``entrypoint.sh`` call it on every start: -* Both directions are idempotent — ``_set_ownership`` skips rows already in the - target state, and the pipeline path no-ops when ownership already matches. +* 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 diff --git a/backend/pg_queue/management/commands/mirror_pg_periodic_tasks.py b/backend/pg_queue/management/commands/mirror_pg_periodic_tasks.py index 4e4766b7ba..2db4958496 100644 --- a/backend/pg_queue/management/commands/mirror_pg_periodic_tasks.py +++ b/backend/pg_queue/management/commands/mirror_pg_periodic_tasks.py @@ -400,7 +400,17 @@ def _set_ownership( # release needs to restore. if not to_pg: beat_updates["last_run_at"] = timezone.now() - PeriodicTask.objects.filter(name=row.name).update(**beat_updates) + matched = PeriodicTask.objects.filter(name=row.name).update( + **beat_updates + ) + if not matched: + # See ownership.py: a bulk .update() returning 0 is + # indistinguishable from success. On release that means a + # periodic PG has let go of and Beat cannot pick up. + self.stderr.write( + f"WARNING: no Beat PeriodicTask named {row.name!r} to " + f"{verb} — 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 diff --git a/backend/scheduler/ownership.py b/backend/scheduler/ownership.py index 4e64d37b19..452a5e2581 100644 --- a/backend/scheduler/ownership.py +++ b/backend/scheduler/ownership.py @@ -109,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, ) @@ -261,7 +260,19 @@ def reconcile_ownership_for( # pipeline save and silently skip a due fire. if was_pg_owned and not pg_owned: beat_updates["last_run_at"] = timezone.now() - PeriodicTask.objects.filter(name=pipeline_id).update(**beat_updates) + matched = PeriodicTask.objects.filter(name=pipeline_id).update(**beat_updates) + if not matched: + # 0 rows means Beat has no row for this pipeline — on the RELEASE + # direction that is a schedule with no firer at all: PG has let go + # and Beat has nothing to take over. Silent before this: a bulk + # .update() reports success by returning 0, so a rollback could + # orphan a schedule and log nothing. + logger.warning( + "reconcile_ownership_for: no Beat PeriodicTask named %s to " + "%s — the schedule may now have no firer", + pipeline_id, + "release to" if not pg_owned else "disable for adoption", + ) # 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/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/workers/tests/test_pg_schema_drift.py b/workers/tests/test_pg_schema_drift.py index 2fbd84cdb9..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", From 6c658c060f8fc7b2591c108821d4763a969d07d2 Mon Sep 17 00:00:00 2001 From: ali Date: Tue, 25 Aug 2026 12:57:43 +0530 Subject: [PATCH 29/33] UN-3796 [FIX] Remediation round 1: fix what the fixes broke Findings from a standardized 16-lens FOLLOWUP review of the four preceding commits, run through five specialist agents. Several are defects the previous round introduced or left behind; each was verified against source before being accepted. CI-SURFACED, and the reason they were invisible until now --------------------------------------------------------- Marking this PR ready for review started the test tier for the first time, and it immediately failed 12 tests. Eight of them were broken from the day they were written: test_undispatched_sweep.py's _execution() built a WorkflowExecution with no workflow, and save() calls _handle_execution_cache() which dereferences self.workflow.id. They could never have passed. This is the prior review's own High - "both PRs are drafts, so no test tier is running" - cashing out exactly as predicted. Fixed by giving the helper a real workflow. Two more were mine from the previous round: modified_at is NOT NULL, so a test that forced it to NULL to reach an early return died on an IntegrityError. It now patches the aggregate instead, which is the only way the helper can legitimately see a missing timestamp. A SILENT NO-OP IN THE HALF THE LAST COMMIT CLAIMED TO PRESERVE -------------------------------------------------------------- release_slot(cls, organization_id: str, ...) was being passed an Organization INSTANCE, so its Redis key formatted to "...:org:Organization object (12)" while acquire_slot had built the key from str(organization.organization_id). The ZREM removed a non-member, returned 0, raised nothing. The previous commit justified dropping the storage delete by saying it "keeps the reversible half (the rate-limit slot release, which has a live cost)" - that half did nothing. Now passes the identifier string, matching the one call site that had it right. A COMMENT CORRECTION THAT INTRODUCED A FRESH FALSE CLAIM -------------------------------------------------------- The previous commit rewrote entrypoint.sh to end "a restart that changes nothing writes nothing", and the NEXT commit corrected converge_pg_scheduler's docstring to say the exact opposite: the pipeline path rewrites both tables and bumps the Beat reload marker for every schedule regardless. The commit whose stated purpose was "correct three comments that misdirect" introduced a new misdirection at the riskiest site. Corrected, along with two other claims in the same file: ownership is additionally Flipt-gated and fails closed (so the env var alone moves nothing when Flipt is blind), and convergence commits per schedule, so a failure leaves a partial hand-over rather than the untouched state the warning promised. "SELF-CORRECTS" WAS FALSE ON THE PG TRANSPORT ---------------------------------------------- The C1 fix reasoned that the surviving ERROR write is tolerable because "the running worker's own terminal write supersedes it". It does not. Both PG worker entry points treat a terminal execution as a reason to STOP: general/tasks.py returns skipped_terminal_execution before it would write EXECUTING, and file_processing/tasks.py raises _TerminalExecutionSkip. A wrongly-claimed row is therefore silently DROPPED - message acked, no retry, only a worker WARNING. The realistic trigger is not the three no-handle paths but a saturated queue, which is the condition this sweep exists for. The comment now says so plainly, and also records that error_message is not cleared by a later terminal write, so a wrongly-claimed run can complete still showing the customer EXEC_NOT_STARTED. This is recorded, not fixed: making the sweep safe under queue saturation is a behaviour change (a positive dispatched_at, or a predicate that also checks for a live pg_queue_message) and is escalated rather than attempted here. OTHER COMMENT AND MESSAGE DEFECTS ---------------------------------- - The zero-match warning added last round fired the same alarming text in both directions. On adopt a missing Beat row just means there was nothing to disable - PG is the firer - and that path runs on every pipeline save. Warning on the benign direction is how the real one, on rollback, gets filtered out. Both sites now branch the whole message, not just the verb. - reaper.py still stated "only needs to outlast a callback - seconds" directly above the paragraph retracting it. - reconcile_pg_schedules' argparse help - the copy an operator actually sees - still said --mirror-only is "the mode automation runs". - Three docstrings still justified the worker/backend boundary with "the sweep needs the rate limiter and the API storage connector"; the storage half was removed last round. - The sweep's module docstring asserted the predicate as fact while the body declares it unsound. The header is what a reader hits first, and it is what would justify reinstating the delete or shortening the grace period. TYPE-DESIGN HARDENING ---------------------- Three independent agents confirmed the _fields-derived SELECT emits a column list identical to the previous hardcoded one. Two gaps remained: the docstrings still described the dependency in the direction that no longer applies (the field names ARE the query now, so a reorder is safe and a RENAME is the hazard), and nothing pinned the tuples in a lane without Postgres - every test that executes those queries skips without a live database. Added DB-free assertions pinning both _fields tuples, and quoted the derived identifiers so the derivation is total. Verified: workers suites 39 passed; rig unit-runner 5 passed; entrypoint.sh passes bash -n. The backend suite still cannot start in this checkout, so the Django-side fixes above are unverified locally and rely on CI. Co-Authored-By: Claude Opus 5 --- backend/entrypoint.sh | 17 +++-- .../commands/mirror_pg_periodic_tasks.py | 13 ++-- .../commands/reconcile_pg_schedules.py | 2 +- backend/scheduler/ownership.py | 33 ++++++---- .../tests/test_pg_finalization_fixes.py | 13 ++-- .../execution_log_internal_views.py | 2 +- .../tests/test_undispatched_sweep.py | 14 +++- .../workflow_v2/undispatched_sweep.py | 48 +++++++++++--- .../queue_backend/pg_queue/pg_scheduler.py | 25 ++++++-- workers/queue_backend/pg_queue/reaper.py | 5 +- workers/shared/api/internal_client.py | 2 +- workers/tests/test_pg_scheduler.py | 64 +++++++++++++++++-- 12 files changed, 187 insertions(+), 51 deletions(-) diff --git a/backend/entrypoint.sh b/backend/entrypoint.sh index b3236d7dda..9766b01add 100755 --- a/backend/entrypoint.sh +++ b/backend/entrypoint.sh @@ -36,7 +36,9 @@ if [ "$migrate" = true ]; then # 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. A rolling restart therefore can move schedules + # 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 @@ -45,14 +47,21 @@ if [ "$migrate" = true ]; then # 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, and a no-op for rows whose ownership already - # matches, so a restart that changes nothing writes nothing. + # 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. Whatever fired the schedules before keeps firing them in that case. + # 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" diff --git a/backend/pg_queue/management/commands/mirror_pg_periodic_tasks.py b/backend/pg_queue/management/commands/mirror_pg_periodic_tasks.py index 2db4958496..7073d98c3c 100644 --- a/backend/pg_queue/management/commands/mirror_pg_periodic_tasks.py +++ b/backend/pg_queue/management/commands/mirror_pg_periodic_tasks.py @@ -403,13 +403,16 @@ def _set_ownership( matched = PeriodicTask.objects.filter(name=row.name).update( **beat_updates ) - if not matched: - # See ownership.py: a bulk .update() returning 0 is - # indistinguishable from success. On release that means a - # periodic PG has let go of and Beat cannot pick up. + 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"{verb} — it may now have no firer" + 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 diff --git a/backend/pg_queue/management/commands/reconcile_pg_schedules.py b/backend/pg_queue/management/commands/reconcile_pg_schedules.py index 7f5a1e78e1..f5e9cd6f44 100644 --- a/backend/pg_queue/management/commands/reconcile_pg_schedules.py +++ b/backend/pg_queue/management/commands/reconcile_pg_schedules.py @@ -76,7 +76,7 @@ def add_arguments(self, parser: Any) -> None: 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 — see handle()." + "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( diff --git a/backend/scheduler/ownership.py b/backend/scheduler/ownership.py index 452a5e2581..c7c9c28173 100644 --- a/backend/scheduler/ownership.py +++ b/backend/scheduler/ownership.py @@ -262,17 +262,28 @@ def reconcile_ownership_for( beat_updates["last_run_at"] = timezone.now() matched = PeriodicTask.objects.filter(name=pipeline_id).update(**beat_updates) if not matched: - # 0 rows means Beat has no row for this pipeline — on the RELEASE - # direction that is a schedule with no firer at all: PG has let go - # and Beat has nothing to take over. Silent before this: a bulk - # .update() reports success by returning 0, so a rollback could - # orphan a schedule and log nothing. - logger.warning( - "reconcile_ownership_for: no Beat PeriodicTask named %s to " - "%s — the schedule may now have no firer", - pipeline_id, - "release to" if not pg_owned else "disable for adoption", - ) + # 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/workflow_manager/execution/tests/test_pg_finalization_fixes.py b/backend/workflow_manager/execution/tests/test_pg_finalization_fixes.py index 294b659515..1f8a4ec5f9 100644 --- a/backend/workflow_manager/execution/tests/test_pg_finalization_fixes.py +++ b/backend/workflow_manager/execution/tests/test_pg_finalization_fixes.py @@ -294,13 +294,18 @@ def test_execution_time_left_alone_when_no_file_timestamp(self): the selection filter, not the early return this test exists to pin. """ ex = self._exec(ExecutionStatus.EXECUTING, files=[ExecutionStatus.COMPLETED]) - WorkflowFileExecution.objects.filter(workflow_execution=ex).update( - modified_at=None - ) + # 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"]) - self.view._restamp_execution_time_from_files(ex) + 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 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 c7094e2d8e..0c07b04742 100644 --- a/backend/workflow_manager/workflow_v2/execution_log_internal_views.py +++ b/backend/workflow_manager/workflow_v2/execution_log_internal_views.py @@ -187,7 +187,7 @@ class SweepUndispatchedExecutionsAPIView(APIView): 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 and the API storage connector — both backend-side. + 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. diff --git a/backend/workflow_manager/workflow_v2/tests/test_undispatched_sweep.py b/backend/workflow_manager/workflow_v2/tests/test_undispatched_sweep.py index f0d7ec3c54..fb496e6473 100644 --- a/backend/workflow_manager/workflow_v2/tests/test_undispatched_sweep.py +++ b/backend/workflow_manager/workflow_v2/tests/test_undispatched_sweep.py @@ -27,11 +27,21 @@ def _execution(**overrides): - """A PENDING, never-dispatched execution old enough to sweep.""" - from workflow_manager.workflow_v2.models import WorkflowExecution + """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, diff --git a/backend/workflow_manager/workflow_v2/undispatched_sweep.py b/backend/workflow_manager/workflow_v2/undispatched_sweep.py index 9d98423aab..4824afc31f 100644 --- a/backend/workflow_manager/workflow_v2/undispatched_sweep.py +++ b/backend/workflow_manager/workflow_v2/undispatched_sweep.py @@ -28,7 +28,12 @@ **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 never happened — no JSONB probing, no cross-table joins, and +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. """ @@ -196,7 +201,11 @@ def _release_abandoned_resources(execution_id: str, workflow_id: str) -> None: **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. Each side is isolated so one failing does not skip the other. + 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 @@ -211,7 +220,15 @@ def _release_abandoned_resources(execution_id: str, workflow_id: str) -> None: .get(id=execution_id) .workflow.organization ) - APIDeploymentRateLimiter.release_slot(organization, execution_id) + # 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 " @@ -232,11 +249,26 @@ def _release_abandoned_resources(execution_id: str, workflow_id: str) -> None: # * 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. Marking it - # ERROR is wrong but self-corrects — the running worker's own terminal write - # supersedes it, and error→completed is explicitly permitted by the status guard. - # Deleting its input does not self-correct: the worker is still going to read it, - # and the row's user-facing message invites a re-run against data that is gone. + # 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 diff --git a/workers/queue_backend/pg_queue/pg_scheduler.py b/workers/queue_backend/pg_queue/pg_scheduler.py index caa08c56bd..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,7 +149,7 @@ def dispatch_due_schedules(conn: PgConnection) -> int: base = cur.fetchone()[0] cur.execute( f""" - SELECT {", ".join(_DueSchedule._fields)} + 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) @@ -234,8 +240,15 @@ def dispatch_due_schedules(conn: PgConnection) -> int: class _DuePeriodicTask(NamedTuple): """One row from the generic-periodic due scan (UN-3796). - Sibling of :class:`_DueSchedule`. Same reason for existing: the field names are - bound to the SELECT's column order at exactly one site. + 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 @@ -291,7 +304,7 @@ def dispatch_due_periodic_tasks(conn: PgConnection) -> int: base = cur.fetchone()[0] cur.execute( f""" - SELECT {", ".join(_DuePeriodicTask._fields)} + 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) diff --git a/workers/queue_backend/pg_queue/reaper.py b/workers/queue_backend/pg_queue/reaper.py index ffa10b8629..b43438c44f 100644 --- a/workers/queue_backend/pg_queue/reaper.py +++ b/workers/queue_backend/pg_queue/reaper.py @@ -181,7 +181,8 @@ def reaper_interval_from_env() -> float: #: 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 only needs to outlast a callback that is about to fire — seconds. +#: 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 @@ -1378,7 +1379,7 @@ def _sweep_undispatched_executions(self) -> None: 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 and the API storage connector. + 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. diff --git a/workers/shared/api/internal_client.py b/workers/shared/api/internal_client.py index 2342688c35..bba4af5296 100644 --- a/workers/shared/api/internal_client.py +++ b/workers/shared/api/internal_client.py @@ -1454,7 +1454,7 @@ def sweep_undispatched_executions(self) -> dict[str, Any]: find it and PENDING — a non-terminal state — is never resolved. Backend-side because the sweep touches ``WorkflowExecution``, the API-deployment - rate limiter and the API storage connector; the reaper only supplies leader + 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. """ diff --git a/workers/tests/test_pg_scheduler.py b/workers/tests/test_pg_scheduler.py index 91b4af04bd..e5ea53bc08 100644 --- a/workers/tests/test_pg_scheduler.py +++ b/workers/tests/test_pg_scheduler.py @@ -23,6 +23,41 @@ 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. @@ -82,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 @@ -113,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() @@ -348,8 +393,12 @@ def _messages_on(conn, queue): 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}%",)) + 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() @@ -360,7 +409,10 @@ def test_due_owned_row_fires_onto_its_own_queue_and_advances(self, clean_periodi 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, + conn, + pg_owned=True, + enabled=True, + next_run_at=past, task_kwargs={"retention_days": 30}, ) From 1b7a26f6548b71ac4723f4841f033df528de7e94 Mon Sep 17 00:00:00 2001 From: ali Date: Tue, 25 Aug 2026 15:34:43 +0530 Subject: [PATCH 30/33] =?UTF-8?q?UN-3796=20[FIX]=20Raise=20the=20undispatc?= =?UTF-8?q?hed-sweep=20grace=20to=20an=20hour=20=E2=80=94=20PG=20drops=20w?= =?UTF-8?q?hat=20it=20terminalises?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The sweep infers "never dispatched" from both handles being NULL. Three paths in workflow_helper reach that state AFTER the message is on its transport, so the predicate can claim a row whose message is still queued. On Celery that was survivable: the orchestrator ran on the row regardless and its own terminal write superseded 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 the message is acked and the work is DROPPED, with a worker WARNING as the only trace. The guard is explicitly PG-only; its own comment says "Celery has no redelivery ... the check is a no-op there". So this is a risk the transport migration introduces, in the phase we are entering. It needs two things at once: an unrecorded handle AND a queue wait longer than the grace. Raising 900 -> 3600 puts the grace clear of any realistic dequeue latency, which collapses the overlap without new machinery or a migration. Not the fix. The fix is to make dispatch a POSITIVE fact - stamp dispatched_at in the same call that records the handle, including on the three branches that currently return without writing one, and key the predicate and the 0026 partial index on it. That needs a migration and is tracked separately. Blast radius of the raise: a genuinely undispatched execution now shows PENDING for up to an hour before it errors, rather than 15 minutes. For ETL that is the whole cost - the sweep never writes file rows and file history is only recorded after a file is actually processed, so nothing is marked done and the next scheduled run re-discovers the file. For API deployments there is no next run, so the caller re-submits, which is what the error message already tells them. Still well under the barrier stuck-timeout, so the two sweeps stay disjoint. Overridable per environment via UNDISPATCHED_EXECUTION_GRACE_SECONDS. Tests pin the value and its upper bound, with the reason, so lowering it back toward the queue's p99 wait has to be a deliberate act with a latency figure in hand rather than a tidy-up. Co-Authored-By: Claude Opus 5 --- .../tests/test_undispatched_sweep.py | 34 +++++++++++++++++++ .../workflow_v2/undispatched_sweep.py | 23 ++++++++++++- 2 files changed, 56 insertions(+), 1 deletion(-) diff --git a/backend/workflow_manager/workflow_v2/tests/test_undispatched_sweep.py b/backend/workflow_manager/workflow_v2/tests/test_undispatched_sweep.py index fb496e6473..c78df28c58 100644 --- a/backend/workflow_manager/workflow_v2/tests/test_undispatched_sweep.py +++ b/backend/workflow_manager/workflow_v2/tests/test_undispatched_sweep.py @@ -169,6 +169,40 @@ def test_a_row_dispatched_mid_sweep_is_not_clobbered(self): assert ex.status == ExecutionStatus.PENDING.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. diff --git a/backend/workflow_manager/workflow_v2/undispatched_sweep.py b/backend/workflow_manager/workflow_v2/undispatched_sweep.py index 4824afc31f..7ef86aaafd 100644 --- a/backend/workflow_manager/workflow_v2/undispatched_sweep.py +++ b/backend/workflow_manager/workflow_v2/undispatched_sweep.py @@ -55,8 +55,29 @@ # 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 = 900 # 15 minutes +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. From 02ecc60a2efef4d152df71e07b083e88be8e1536 Mon Sep 17 00:00:00 2001 From: ali Date: Wed, 26 Aug 2026 10:59:39 +0530 Subject: [PATCH 31/33] =?UTF-8?q?UN-3796=20[FIX]=20Pass=20modified=5Fat=20?= =?UTF-8?q?explicitly=20=E2=80=94=20BaseModelManager.update()=20was=20un-a?= =?UTF-8?q?geing=20the=20row?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The last integration failure on this branch, and the cause is a manager convention rather than anything in the code under test. BaseModelManager.update() does `kwargs.setdefault("modified_at", timezone.now())` (backend/utils/models/base_model.py:39-41). So WorkflowExecution.objects.filter(pk=ex.pk).update(created_at=created) silently re-stamped modified_at to NOW, undoing the ageing _age() had just applied. The execution then failed `modified_at__lt=cutoff` at selection, the endpoint reported scanned=0, and the test failed on STATUS — with nothing in the failure output hinting that a timestamp had moved. I could not reproduce it locally (this checkout cannot start the backend suite) and it took a CI round trip plus reading the manager to find. Fixed by passing modified_at explicitly, which the manager's own docstring documents as the supported override. Checked the sibling suite for the same trap: test_undispatched_sweep.py has two updates without modified_at, but that sweep selects on created_at, so re-stamping modified_at cannot affect it. Left alone rather than changed defensively. This was the single remaining integration-backend failure. The previous CI run went from 12 failures to 2 after the earlier remediation commits; this closes the one that was mine. (The other is an e2e API-deployment timeout — see the PR thread; not attributed either way without evidence.) Co-Authored-By: Claude Opus 5 --- .../execution/tests/test_pg_finalization_fixes.py | 11 ++++++++++- 1 file changed, 10 insertions(+), 1 deletion(-) 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 1f8a4ec5f9..d9cc5ba4da 100644 --- a/backend/workflow_manager/execution/tests/test_pg_finalization_fixes.py +++ b/backend/workflow_manager/execution/tests/test_pg_finalization_fixes.py @@ -272,7 +272,16 @@ def test_execution_time_comes_from_the_last_file_not_from_now(self): # 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) - WorkflowExecution.objects.filter(pk=ex.pk).update(created_at=created) + # 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) ) From 1c4f1c785c56114eff9b6f46f9357a69731e8add Mon Sep 17 00:00:00 2001 From: ali Date: Wed, 26 Aug 2026 12:22:06 +0530 Subject: [PATCH 32/33] UN-3796 [FIX] Record dispatch as a positive fact so the sweep stops guessing Closes the root cause behind three independently-reported findings: the standardized review's C1 (Critical), its follow-up N5, and Greptile's P1 on PR #2254. All three are the same defect seen from different angles. The defect ---------- The undispatched sweep decides "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 on its transport: * _record_dispatch_handle raising and the caller swallowing it, under the comment "continuing - the orchestrator is already running" (:687) * the handle coming back empty and the recorder returning without writing (:551) * a PG handle that will not parse as a bigint, same early return (:564) So the sweep could claim a RUNNING execution and mark it ERROR. On Celery that self-corrected: the orchestrator ran regardless and its terminal write superseded the ERROR. On PG it does NOT - general/tasks.py returns skipped_terminal_execution and file_processing/tasks.py raises _TerminalExecutionSkip, so the message is acked and the work is silently DROPPED. Prior rounds removed the irreversible half (the staged-input delete) and widened the grace 900s -> 3600s. Both were mitigations. This removes the inference. The fix ------- `dispatched_at` is stamped by the dispatcher the instant dispatch returns, BEFORE any handle bookkeeping can fail - one write upstream of all three paths. The sweep then asks a question with a true answer. Single-phase rollout, deliberately. The predicate KEEPS both 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 During a rolling upgrade an old pod dispatches without stamping; that row matches `dispatched_at IS NULL` but is excluded by `task_id IS NOT NULL`. Drop the handle checks and this would need two releases, or it would sweep live work mid-deploy. Existing rows need no backfill for the same reason. A gap this opened, and closed ------------------------------ Making the sweep require `dispatched_at IS NULL` means a row that WAS dispatched but whose handle write failed stops matching it. It also did not match recover_stuck_pg_executions, which tested the handles alone - so it would have matched NEITHER sweep and could sit non-terminal forever, trading "wrongly terminalised" for "never terminalised". The recovery predicate is now three-way (dispatched_at OR either handle), which keeps the two sweeps disjoint AND jointly exhaustive: undispatched takes rows where all three are absent, recovery takes rows where any one is present. That is a cleaner split than before this change. Migrations ---------- 0027 adds the column - nullable, additive, transactional, no backfill. 0028 re-keys the partial index, mirroring 0026's pattern (atomic = False, CONCURRENTLY, IF NOT EXISTS, the INVALID-index guard, and the out-of-band build instructions). It creates the replacement BEFORE dropping the old one so the predicate stays served throughout; an interruption between the two leaves both, which costs a little write overhead and nothing else. Cost and residual risk, stated plainly --------------------------------------- One extra small UPDATE per dispatch. The residual hole is the stamp itself failing, which leaves that one execution in exactly the pre-existing ambiguity - three paths narrowed to one, not zero. The log line at that site says so. With the inference gone, the 3600s grace is no longer load-bearing and could be returned to 15 minutes, terminalising genuine orphans 4x sooner. Left as-is here: that is a separate, reversible tuning decision and this change is already the larger one. Tests: a dispatched row with NO handles is never swept (the exact shape the three paths produce); an unstamped row WITH a handle is not swept (the rolling-deploy guarantee); a genuinely undispatched row still IS swept (the fix must not disable the sweep); and the model/migration index-parity test now pins all three signals. NOT VERIFIED LOCALLY - the backend suite cannot start in this checkout. Compile, ruff and scoped pre-commit are clean, and lint is unchanged against HEAD. CI is what will verify the migrations and the tests. Co-Authored-By: Claude Opus 5 --- backend/workflow_manager/internal_views.py | 19 ++- .../workflow_manager/workflow_v2/execution.py | 31 ++++ .../0027_workflowexecution_dispatched_at.py | 68 +++++++++ .../0028_undispatched_idx_dispatched_at.py | 136 ++++++++++++++++++ .../workflow_v2/models/execution.py | 14 +- .../test_undispatched_execution_index.py | 13 +- .../tests/test_undispatched_sweep.py | 57 ++++++++ .../workflow_v2/undispatched_sweep.py | 2 + .../workflow_v2/workflow_helper.py | 22 +++ 9 files changed, 354 insertions(+), 8 deletions(-) create mode 100644 backend/workflow_manager/workflow_v2/migrations/0027_workflowexecution_dispatched_at.py create mode 100644 backend/workflow_manager/workflow_v2/migrations/0028_undispatched_idx_dispatched_at.py diff --git a/backend/workflow_manager/internal_views.py b/backend/workflow_manager/internal_views.py index ef817977e5..0aa0fb1f3a 100644 --- a/backend/workflow_manager/internal_views.py +++ b/backend/workflow_manager/internal_views.py @@ -710,9 +710,22 @@ def recover_stuck_pg_executions(self, request): files = WorkflowFileExecution.objects.filter(workflow_execution=OuterRef("pk")) stuck_ids = list( WorkflowExecution.objects.filter( - # Dispatched on EITHER transport. Both-NULL means never dispatched — - # undispatched_sweep.py's row, not ours. - Q(queue_message_id__isnull=False) | Q(task_id__isnull=False), + # 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, 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/migrations/0027_workflowexecution_dispatched_at.py b/backend/workflow_manager/workflow_v2/migrations/0027_workflowexecution_dispatched_at.py new file mode 100644 index 0000000000..f468666e32 --- /dev/null +++ b/backend/workflow_manager/workflow_v2/migrations/0027_workflowexecution_dispatched_at.py @@ -0,0 +1,68 @@ +"""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 successfully handed to its " + "transport. NULL means it was never dispatched — the positive fact " + "the undispatched sweep needs, because absent task_id/" + "queue_message_id also occurs for executions that ARE running." + ), + ), + ), + ] 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 a511904ba7..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) @@ -253,9 +264,10 @@ class Meta: # ExecutionStatus by tests/test_undispatched_execution_index.py. models.Index( fields=["created_at"], - name="we_undispatched_idx", + name="we_undispatched_dispatch_idx", condition=Q( status="PENDING", + dispatched_at__isnull=True, task_id__isnull=True, queue_message_id__isnull=True, ), 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 index 9708ffbf6a..4268a2f7fa 100644 --- a/backend/workflow_manager/workflow_v2/tests/test_undispatched_execution_index.py +++ b/backend/workflow_manager/workflow_v2/tests/test_undispatched_execution_index.py @@ -1,6 +1,6 @@ """Guard: the undispatched-execution partial index stays in sync with reality. -``we_undispatched_idx`` (``WorkflowExecution.Meta.indexes``) hardcodes the literal +``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 @@ -29,7 +29,7 @@ from unstract.core.data_models import ExecutionStatus # noqa: E402 -INDEX_NAME = "we_undispatched_idx" +INDEX_NAME = "we_undispatched_dispatch_idx" _MIGRATION = ( Path(__file__).resolve().parent.parent / "migrations" @@ -54,12 +54,17 @@ def test_the_status_literal_is_the_pending_enum_value(self): matching every row the sweep looks for.""" assert _condition_children()["status"] == ExecutionStatus.PENDING.value - def test_both_dispatch_handles_are_in_the_predicate(self): - """Dropping either makes the index non-matching for the sweep's WHERE clause + 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 diff --git a/backend/workflow_manager/workflow_v2/tests/test_undispatched_sweep.py b/backend/workflow_manager/workflow_v2/tests/test_undispatched_sweep.py index c78df28c58..e69c2d42d0 100644 --- a/backend/workflow_manager/workflow_v2/tests/test_undispatched_sweep.py +++ b/backend/workflow_manager/workflow_v2/tests/test_undispatched_sweep.py @@ -169,6 +169,63 @@ def test_a_row_dispatched_mid_sweep_is_not_clobbered(self): 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. diff --git a/backend/workflow_manager/workflow_v2/undispatched_sweep.py b/backend/workflow_manager/workflow_v2/undispatched_sweep.py index 7ef86aaafd..cc3d77f356 100644 --- a/backend/workflow_manager/workflow_v2/undispatched_sweep.py +++ b/backend/workflow_manager/workflow_v2/undispatched_sweep.py @@ -143,6 +143,7 @@ def _positive_int_from_env(name: str, default: int) -> int: 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 @@ -150,6 +151,7 @@ def _positive_int_from_env(name: str, default: int) -> int: 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 diff --git a/backend/workflow_manager/workflow_v2/workflow_helper.py b/backend/workflow_manager/workflow_v2/workflow_helper.py index 66b59ee715..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 ) From a713fb6e4a1fd60a49d328da79a6f7829843f446 Mon Sep 17 00:00:00 2001 From: ali Date: Wed, 26 Aug 2026 13:24:47 +0530 Subject: [PATCH 33/33] UN-3796 [FIX] Match 0027's field to the model so makemigrations is clean The migration and the model declared different db_comment text for dispatched_at, so Django saw a pending field alteration and `makemigrations --check` proposed a 0029. Caught by running the check against a live Postgres rather than by reading. Verified end to end against the running stack (DB_SCHEMA=public on a scratch database, since a fresh test DB has no `unstract` schema for search_path): * makemigrations --check --dry-run -> "No changes detected in app 'workflow_v2'" * migrate 0026 -> 0027 -> 0028 applies cleanly * dispatched_at: nullable, timestamptz * we_undispatched_idx dropped, we_undispatched_dispatch_idx created, and zero INVALID indexes (the CONCURRENTLY interruption case 0026's guard exists for) * EXPLAIN on the sweep's exact predicate -> "Index Scan using we_undispatched_dispatch_idx", so the partial index is matched, not bypassed * migrate back to 0026 unapplies both, restores we_undispatched_idx and drops the column - rollback is real, not assumed Tests: 139 passed across workflow_v2 and execution, including the three new positive-fact regressions and the recovery suite. One unrelated failure (test_workflow_author) reproduces identically with these changes stashed, so it is pre-existing and not attributed here. Co-Authored-By: Claude Opus 5 --- .../migrations/0027_workflowexecution_dispatched_at.py | 9 +++++---- 1 file changed, 5 insertions(+), 4 deletions(-) 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 index f468666e32..9d483d152d 100644 --- a/backend/workflow_manager/workflow_v2/migrations/0027_workflowexecution_dispatched_at.py +++ b/backend/workflow_manager/workflow_v2/migrations/0027_workflowexecution_dispatched_at.py @@ -58,10 +58,11 @@ class Migration(migrations.Migration): null=True, blank=True, db_comment=( - "Set when the orchestrator task was successfully handed to its " - "transport. NULL means it was never dispatched — the positive fact " - "the undispatched sweep needs, because absent task_id/" - "queue_message_id also occurs for executions that ARE running." + "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." ), ), ),