fix(storage): converge canonical drain for perpetually snoozing jobs (#456)

[hardbyte]

Fixes #456.

The bug

canonical_live_backlog counts all of awa.scheduled_jobs, and the drain phase assumes that backlog converges as jobs come due and finish. For workloads whose handlers snooze on every run (recurring per-entity heartbeats implemented as JobResult::Snooze) it never does: the canonical executor's snooze path writes the job straight back into canonical scheduled_jobs, so the backlog self-replenishes and awa storage finalize --wait blocks forever.

Reproduced on a restored production-shaped snapshot with tens of thousands of jobs on a ~24h snooze cycle. After enter-mixed-transition, new inserts routed to queue storage correctly, but every canonical job that came due was promoted to canonical jobs_hot, executed, and re-snoozed back into canonical scheduled_jobs. Queue-storage deferred_jobs stayed empty; canonical_live_backlog stayed flat.

This affects 0.7 as well as 0.6 — the ADR-037 gate blocks awa migrate on an unfinalized cluster, but it does not stop a 0.7 worker from executing canonical work (canonical still works in 0.7, per docs/upgrade-0.6-to-0.7.md, and is removed in 0.8).

The fix

awa_model::reschedule owns the snooze / retry-backoff / retry-after write and branches on awa.active_queue_storage_schema():

• canonical (NULL) — the same jobs_hot -> scheduled_jobs move as before, so pre-flip behaviour is byte-for-byte unchanged.
• mixed_transition / active — takes the canonical row and re-inserts the attempt into the active queue-storage schema's deferred_jobs, carrying over attempt, errors, progress, run_lease, timestamps, and the unique claim.

Every canonical job therefore leaves the canonical plane at its next completion, and the existing drain/finalize flow converges for perpetual snoozers with no operator intervention.

No migration, deliberately

This ships as runtime SQL against the existing v040 schema rather than a v044 function. A migration would be unreachable by the clusters that need it: the ADR-037 gate refuses awa migrate on exactly the unfinalized clusters where this bug manifests. It also keeps the fix backportable to the 0.6 line, whose migration numbering 0.7 already owns past v040. (A queue-storage-only variant of this fix was prototyped as v044 first and discarded for these reasons.)

run_lease semantics

Unchanged, and strictly stronger. The take of the canonical row carries the same state = 'running' AND run_lease = $n guard as the historical UPDATE. The successor takes a fresh id from the queue-storage sequence (the two planes have different id sequences, so reusing the canonical id could collide with a later queue-storage insert), which means a late completion from a rescued attempt matches no canonical row and reports stale. Both stale paths are covered by tests.

Adjacent fix: leadership no longer gates drain

Two related stalls, both observed in the same rehearsal:

• a queue-storage leader did not promote due canonical scheduled_jobs, and
• a canonical-resolved (pre-flip) leader did not promote queue-storage deferred_jobs — so re-scheduled jobs that had already migrated sat unpromoted until leadership moved.

Each maintenance leader now promotes the other plane's deferred backlog for the duration of the transition; the pre-flip case goes through a new standby_queue_storage handle, which is the configured-but-inactive queue-storage runtime a canonical-resolved worker already carries. Both paths are guarded no-ops on finalized and fresh installs (active_queue_storage_schema() / empty canonical tables), so steady-state cost is one cheap query per promote tick.

Adjacent fix: rescue and terminal writes survive the routing flip

Review of the issue's "related sharp edge" (mid-flight jobs stuck running at the flip, rescue not covering the canonical plane) found it was two gaps, both fixed here:

• Rescue never covered the canonical plane post-flip, under either leader. The leadership half mirrors the promote fix: the three rescue sweeps (stale heartbeat, expired deadline, callback timeout) now cover both planes from either leader, with the pre-flip case going through the same standby_queue_storage handle. The deeper half: the canonical rescue statements wrote through the awa.jobs compatibility view, which exposes only the queue-storage plane once routing flips — so even a canonical-resolved leader's rescue silently matched nothing. The canonical rescue (batched and the Canonical rescue sweep wedges permanently on a single unique-key conflict (idx_awa_jobs_unique) #388 per-row fallback) now performs the physical jobs_hot -> scheduled_jobs move directly, the same shape as the reschedule module; the table triggers preserve the unique-claim conflict semantics. Cross-plane rows skip local cancellation signalling and DLQ moves, since the planes' id sequences are independent.
• The executor's fast-path terminal completions (fail, cancel, callback park) also wrote through the view, so post-flip they no-opped as "already rescued/cancelled" and the job wedged in running. Callback-carrying jobs were the worst hit, since the reschedule migration deliberately keeps them canonical and relies on them reaching a terminal or parked state on their own. These now write awa.jobs_hot directly, matching the _with_followups variants that already did.

A new maintenance test drives a queue-storage-resolved leader against seeded stuck canonical rows (stale heartbeat + expired callback) in mixed_transition and asserts both rescue into scheduled_jobs; it runs on its own database because it deliberately holds the transition unfinalized.

Tests

rolling_transition_rehearsal_test.rs gains a snooze-forever workload on its own queue. This is the gap that let #456 through: the existing snooze_once handler snoozes 100ms once and then completes, so the drain converges and the perpetual case is unrepresented. canonical_backlog_count now also counts scheduled_jobs, as its doc comment already claimed it did. The rehearsal asserts the snoozer bank has left the canonical plane and keeps executing on queue storage after finalize — it fails on the pre-fix code (at_finalize=763 after=763) and passes with it.

A focused test in sql_only_storage_upgrade_test.rs covers canonical routing (id preserved, snooze does not count the attempt), a stale wrong-run_lease call, migrated routing (fresh id, retryable state, payload metadata/tags/errors/progress preserved, unique claim re-pointed at the successor), the stale post-migration re-call, and a canonical backlog of zero with nothing having completed.

migration_test.rs also swaps two hardcoded 43s for CURRENT_VERSION (with the version-floor rewinds pinned to 43 deliberately, so the pending range keeps crossing v043). Behaviour-neutral today; it stops the next migration bump from failing those tests spuriously.

Verification

• cargo test --workspace — 75 suites, 713 tests, clean run from a fresh schema
• cargo test -p awa --test rolling_transition_rehearsal_test -- --ignored — passes end-to-end
• pytest — 315 passed, 1 skipped
• cargo fmt --all and cargo clippy --all-targets --all-features -- -D warnings clean
• The rescue/terminal-write commit additionally re-ran locally: the full awa-worker lib suite plus the rescue_unique_conflict, lifecycle_hook, external_wait, enqueue_spec, and executor_guard integration suites, and sql_only_storage_upgrade_test — all green.

Two environment notes for anyone reproducing locally, neither a code issue: the test Postgres needs max_locks_per_transaction=2048 (the substrate installs many partitions), and a stale awa_health_test database carrying a leftover future-version row fails health_endpoint_test until dropped.

A 0.6.x backport of the fix (without the migration_test.rs hardening) follows separately against release/0.6.3.

Summary by CodeRabbit

• Bug Fixes

  • Improved rescheduling during mixed storage transitions so perpetually snoozing jobs no longer stall canonical backlog draining.
  • Rescheduled jobs now move to the deferred backlog while preserving timing, progress, errors, and unique-claim behavior.
  • Improved handling of stale completions, duplicate claims, callback-bearing jobs, and cross-plane deferred work promotion.

• Tests

  • Expanded transition, migration, locking, and rolling-rehearsal coverage, including perpetual snooze scenarios.

• Documentation

  • Updated upgrade and operational guidance for deferred-job behavior, backlog monitoring, and troubleshooting.

[coderabbitai]

📝 Walkthrough

Walkthrough

Canonical job rescheduling now uses shared model logic. During mixed transitions, successors move to queue-storage deferred_jobs. Executors, maintenance promotion, tests, correctness models, and operational documentation cover perpetual snoozing and cross-plane backlog draining.

Changes

Canonical drain convergence

Layer / File(s)	Summary
Reschedule model and storage routing awa-model/src/lib.rs, awa-model/src/reschedule.rs, awa-model/src/queue_storage.rs	Adds guarded rescheduling, transition routing, queue-storage migration, duplicate cancellation, and preserved job metadata.
Executor completion integration awa-worker/src/executor.rs, awa/tests/lifecycle_hook_test.rs	Routes retry, backoff, snooze, and follow-up paths through rescheduling and updates retry event snapshots.
Maintenance routing-flip promotion awa-worker/src/client.rs, awa-worker/src/maintenance.rs	Promotes due deferred work across both storage planes during transitions.
Transition validation and rehearsal awa-model/tests/sql_only_storage_upgrade_test.rs, awa/tests/rolling_transition_rehearsal_test.rs, awa/tests/migration_test.rs, correctness/run-tlc-suite.sh	Adds coverage for routing, locking, claims, delays, perpetual snoozing, migration, and TLC scenarios.
Transition model and operational contract correctness/storage/*, CHANGELOG.md, docs/*.md, skills/awa-operations/SKILL.md	Models snoozing backlog migration and documents cross-plane promotion, drain convergence, rollout checks, and older-build workarounds.

Estimated code review effort: 4 (Complex) | ~60 minutes

Sequence Diagram(s)

sequenceDiagram
  participant Worker
  participant Executor
  participant Reschedule
  participant Canonical
  participant QueueStorage
  participant Maintenance
  Worker->>Executor: complete canonical job
  Executor->>Reschedule: reschedule with run lease
  Reschedule->>Canonical: remove guarded running row
  Reschedule->>QueueStorage: insert successor into deferred_jobs
  Maintenance->>QueueStorage: promote due deferred job
  QueueStorage-->>Worker: execute successor after routing transition

Loading

Possibly related PRs

• hardbyte/awa#362: Shares queue-storage test execution and runtime rescheduling paths.
• hardbyte/awa#390: Shares maintenance and unique-claim conflict handling.
• hardbyte/awa#425: Shares storage-transition machinery and migration tests.

Suggested labels: full-ci

Poem

A rabbit watched the snoozing jobs move through,
From canonical rows to deferred queue.
Claims and attempts stayed aligned,
Stale completions were left behind.
The drain reached zero when the transition grew true.

🚥 Pre-merge checks | ✅ 5

✅ Passed checks (5 passed)

Check name	Status	Explanation
Linked Issues check	✅ Passed	The changes implement cross-plane rescheduling, preserve lease and retry data, promote both deferred backlogs, and add perpetual-snooze coverage for issue #456.
Out of Scope Changes check	✅ Passed	The code, tests, models, and documentation changes directly support issue #456 and its storage-transition requirements.
Docstring Coverage	✅ Passed	No functions found in the changed files to evaluate docstring coverage. Skipping docstring coverage check.
Description Check	✅ Passed	Check skipped - CodeRabbit’s high-level summary is enabled.
Title check	✅ Passed	The title clearly and concisely describes the main fix: making canonical drain converge for perpetually snoozing jobs during storage transitions.

✨ Finishing Touches

📝 Generate docstrings

• Create stacked PR
• Commit on current branch

---

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

_{Comment @coderabbitai help to get the list of available commands.}

[chatgpt-codex-connector]

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: c1334ad6f2

ℹ️ About Codex in GitHub

Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you

• Open a pull request for review
• Mark a draft as ready
• Comment "@codex review".

If Codex has suggestions, it will comment; otherwise it will react with 👍.

Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".

[coderabbitai]

Actionable comments posted: 4

🧹 Nitpick comments (4)

awa-worker/src/maintenance.rs (1)

2177-2201: 🚀 Performance & Scalability | 🔵 Trivial | ⚡ Quick win

Cross-plane promotion runs unconditionally on every promote tick. Both new blocks add steady-state DB work to every leader at the promote cadence (default 250ms, twice per tick) with no transition-state gate, so clusters that finalized long ago keep paying for mixed-transition machinery.

• awa-worker/src/maintenance.rs#L2177-L2201: gate the canonical mirror-drain on the transition state (or a cached "canonical plane drained" flag) instead of running promote_due_batch unconditionally.
• awa-worker/src/maintenance.rs#L2208-L2215: resolve/cache awa.active_queue_storage_schema() once per pass (or with a coarse refresh) rather than per state per tick.

🤖 Prompt for AI Agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@awa-worker/src/maintenance.rs` around lines 2177 - 2201, The promote tick
performs unnecessary steady-state database work. In the canonical promotion
block around promote_due_batch, gate mirror draining on the active
storage-transition state or a cached canonical-plane-drained flag; also update
the per-state logic at awa-worker/src/maintenance.rs lines 2208-2215 to resolve
or cache awa.active_queue_storage_schema() once per pass or via a coarse refresh
instead of once per state per tick.

awa-model/tests/sql_only_storage_upgrade_test.rs (3)

447-456: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low value

UNION ALL … ORDER BY 1 DESC LIMIT 1 reads as a max-of-two-counts trick.

It works (max is 0 only when both are 0), but a single summed scalar is clearer and asserts the same thing.

♻️ Clearer query

-    let canonical_left: i64 = sqlx::query_scalar(
-        "SELECT count(*)::bigint FROM awa.jobs_hot WHERE id = $1 \
-         UNION ALL SELECT count(*)::bigint FROM awa.scheduled_jobs WHERE id = $1 \
-         ORDER BY 1 DESC LIMIT 1",
-    )
+    let canonical_left: i64 = sqlx::query_scalar(
+        "SELECT (SELECT count(*)::bigint FROM awa.jobs_hot WHERE id = $1) \
+              + (SELECT count(*)::bigint FROM awa.scheduled_jobs WHERE id = $1)",
+    )

🤖 Prompt for AI Agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@awa-model/tests/sql_only_storage_upgrade_test.rs` around lines 447 - 456,
Update the canonical_left query in the SQL-only storage upgrade test to use one
scalar that sums the counts from awa.jobs_hot and awa.scheduled_jobs for the
bound job2, instead of UNION ALL with ORDER BY and LIMIT. Preserve the existing
fetch and assert_eq! behavior.

---

401-408: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low value

The count(*) OVER () assertion is tautological.

WHERE id = $1 on a primary key can never yield more than one row, so count(*) OVER () is always 1 when fetch_one succeeds. If the intent is to assert the canonical plane holds exactly one reschedule_test row, count without the id filter; otherwise just select state.

♻️ Simplify

-    let (state, count): (String, i64) = sqlx::query_as(
-        "SELECT state::text, count(*) OVER () FROM awa.scheduled_jobs WHERE id = $1",
-    )
-    .bind(job1)
-    .fetch_one(&pool)
-    .await
-    .expect("scheduled row after canonical snooze");
-    assert_eq!((state.as_str(), count), ("scheduled", 1));
+    let state: String =
+        sqlx::query_scalar("SELECT state::text FROM awa.scheduled_jobs WHERE id = $1")
+            .bind(job1)
+            .fetch_one(&pool)
+            .await
+            .expect("scheduled row after canonical snooze");
+    assert_eq!(state, "scheduled");

🤖 Prompt for AI Agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@awa-model/tests/sql_only_storage_upgrade_test.rs` around lines 401 - 408,
Update the query and assertion around the canonical snooze check to avoid
counting rows filtered by the primary-key id. Count the canonical
reschedule_test rows without the id predicate if verifying exactly one row, or
select only state if row-count validation is not intended; keep the existing
scheduled-state assertion.

---

528-544: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick win

Cleanup/finalize only runs on the success path.

Every assertion above uses assert!/panic!, so any failure leaves the shared test database in mixed_transition with leftover deferred rows and claims — which, per your own comment, makes a later migrations::run hit the ADR-037 unfinalized-cluster gate and turns one failure into cascading failures in unrelated tests. Consider running the body via AssertUnwindSafe(...).catch_unwind() (or a drop guard that finalizes) and re-raising after cleanup.

🤖 Prompt for AI Agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@awa-model/tests/sql_only_storage_upgrade_test.rs` around lines 528 - 544, The
cleanup and storage finalization after the reschedule test must also run when
assertions panic. Wrap the test body containing the assertions in an
unwind-catching mechanism such as AssertUnwindSafe(...).catch_unwind(), perform
the existing deferred-job, claims, canonical-work, and storage_finalize cleanup
afterward, then re-raise the original panic.

🤖 Prompt for all review comments with AI agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

Inline comments:
In `@awa-model/src/reschedule.rs`:
- Around line 279-288: Update promote_due around the payload construction and
migrated DELETE ... RETURNING flow to preserve canonical lease callback state,
including callback_id, callback_timeout_at, callback_filter,
callback_on_complete, callback_on_fail, and callback_transform, on a compatible
successor claim; if that cannot be migrated, explicitly skip jobs with callback
wiring instead of promoting them as plain queued jobs.
- Around line 336-340: Update the rescheduling query around the run_at-producing
SELECT and its finalized_at bind so the Retryable timestamp is generated by
PostgreSQL using the same DB clock, rather than Utc::now() in the worker. Fold
the finalized_at value into that SELECT and bind the returned value, preserving
None for non-Retryable states.

In `@awa-worker/src/executor.rs`:
- Around line 1501-1521: Update rescheduled_event_row to accept the newly
recorded error entry and the appropriate finalized_at value, then assign both to
the cloned JobRow so the Retried event matches the previously re-read row.
Ensure the synthesized row includes the appended errors entry and clears or
updates finalized_at for the rescheduled state.

In `@awa/tests/rolling_transition_rehearsal_test.rs`:
- Around line 639-650: Give the queue_storage_target client in the
SnoozeForeverWorker registration its own Arc<AtomicI64> counter instead of
reusing snooze_performs, mirroring the separate rust_handled and
qs_target_handled counters. Update the related phase-7 assertion to track this
queue-storage-specific counter so post-finalize progress is attributed to queue
storage.

---

Nitpick comments:
In `@awa-model/tests/sql_only_storage_upgrade_test.rs`:
- Around line 447-456: Update the canonical_left query in the SQL-only storage
upgrade test to use one scalar that sums the counts from awa.jobs_hot and
awa.scheduled_jobs for the bound job2, instead of UNION ALL with ORDER BY and
LIMIT. Preserve the existing fetch and assert_eq! behavior.
- Around line 401-408: Update the query and assertion around the canonical
snooze check to avoid counting rows filtered by the primary-key id. Count the
canonical reschedule_test rows without the id predicate if verifying exactly one
row, or select only state if row-count validation is not intended; keep the
existing scheduled-state assertion.
- Around line 528-544: The cleanup and storage finalization after the reschedule
test must also run when assertions panic. Wrap the test body containing the
assertions in an unwind-catching mechanism such as
AssertUnwindSafe(...).catch_unwind(), perform the existing deferred-job, claims,
canonical-work, and storage_finalize cleanup afterward, then re-raise the
original panic.

In `@awa-worker/src/maintenance.rs`:
- Around line 2177-2201: The promote tick performs unnecessary steady-state
database work. In the canonical promotion block around promote_due_batch, gate
mirror draining on the active storage-transition state or a cached
canonical-plane-drained flag; also update the per-state logic at
awa-worker/src/maintenance.rs lines 2208-2215 to resolve or cache
awa.active_queue_storage_schema() once per pass or via a coarse refresh instead
of once per state per tick.

🪄 Autofix (Beta)

Fix all unresolved CodeRabbit comments on this PR:

• Push a commit to this branch (recommended)
• Create a new PR with the fixes

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: Organization UI

Review profile: CHILL

Plan: Pro Plus

Run ID: f283bba9-e73e-4bc2-8526-9b488b77fe3a

📥 Commits

Reviewing files that changed from the base of the PR and between 2ddf310 and c1334ad.

📒 Files selected for processing (9)

• CHANGELOG.md
• awa-model/src/lib.rs
• awa-model/src/reschedule.rs
• awa-model/tests/sql_only_storage_upgrade_test.rs
• awa-worker/src/client.rs
• awa-worker/src/executor.rs
• awa-worker/src/maintenance.rs
• awa/tests/migration_test.rs
• awa/tests/rolling_transition_rehearsal_test.rs

[coderabbitai]

Actionable comments posted: 1

🤖 Prompt for all review comments with AI agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

Inline comments:
In `@docs/upgrade-0.5-to-0.6.md`:
- Line 174: The documented manual backlog migration must be atomic. Update the
workaround to require one transaction, lock each scheduled job and its
corresponding awa.job_unique_claims row before deletion, and then reinsert
through awa.insert_job_compat(...), preserving the original job payload and
uniqueness metadata; alternatively explicitly require all producers to be
stopped.

🪄 Autofix (Beta)

Fix all unresolved CodeRabbit comments on this PR:

• Push a commit to this branch (recommended)
• Create a new PR with the fixes

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: Organization UI

Review profile: CHILL

Plan: Pro Plus

Run ID: 600ca76a-07cd-4d83-b56e-2dbcfebc938c

📥 Commits

Reviewing files that changed from the base of the PR and between c1334ad and c8d69d4.

📒 Files selected for processing (17)

• awa-model/src/reschedule.rs
• awa-model/tests/sql_only_storage_upgrade_test.rs
• awa-worker/src/executor.rs
• awa/tests/lifecycle_hook_test.rs
• correctness/run-tlc-suite.sh
• correctness/storage/AwaStorageTransition.cfg
• correctness/storage/AwaStorageTransition.tla
• correctness/storage/AwaStorageTransitionCurrentGate.cfg
• correctness/storage/AwaStorageTransitionMigrate07Ungated.cfg
• correctness/storage/AwaStorageTransitionRescheduleStaysCanonical.cfg
• correctness/storage/MAPPING.md
• docs/architecture.md
• docs/configuration.md
• docs/queue-storage-substrate.md
• docs/upgrade-0.5-to-0.6.md
• docs/upgrade-0.6-to-0.7.md
• skills/awa-operations/SKILL.md