feat: backfill historical spend counters with readiness checks

[henrypark133]

Existing spend counters only contain charges recorded after #1116. Add historical reconciliation so later readers can replace lifetime usage scans without undercounting old spending.

This PR is stacked on #1116. It adds a per-organization readiness marker and a backfill-spend-counters operator shipped in the application image. The operator captures raw-minus-counter corrections in a repeatable-read snapshot, releases that snapshot, then adds corrections under the existing organization accounting lock. Counter updates and readiness commit together; competing workers and reruns skip completed organizations. Historical scans have a configurable deadline; apply statements and lock waits are capped at five seconds. Legacy total_spent remains unchanged.

Rollout order is required:

1. Deploy feat: maintain transactional inference and service spend counters #1116 to every usage writer and drain all older processes and in-flight transactions.
2. Deploy this migration/operator. Do not rewrite or delete raw history during reconciliation.
3. Run /app/backfill-spend-counters with the normal database environment (optionally --organization UUID or --statement-timeout-seconds N). The command does not run migrations. All-organization mode verifies readiness before reporting completion.
4. Deploy the stacked reader PR only after reconciliation completes. Returning to pre-counter writers invalidates readiness and requires fresh reconciliation before counter readers can resume; existing markers cannot detect old writers.

Readers stay unchanged in this prerequisite PR; #1120 switches key lists, API-key admission checks, and billing summaries together. No production backfill was run by this change.

Validation: three PostgreSQL integration tests passed, covering real inference/service postings between snapshot and apply, competing workers, reruns, service-only and deleted keys, counter-only negative drift, both key and organization overflow rollback, empty/inactive/missing balances, and migration defaults. The tests fail against a no-op reconciliation. The actual CLI reconciled a historical fixture and reran without changing its totals. The new test target is included in PostgreSQL CI. No Markdown plans are included.

Checks: cargo nextest run -p database --test spend_counter_backfill (3 passed), cargo clippy -p database --all-targets --all-features -- -D warnings, and formatting passed.

[ironloopai]

Review · Status

🟩 Completed

IronLoop completed the review and posted it to GitHub.

Result

Open submitted review →

Run details

• Run: a718bb8a-139f-43ff-9d96-19b98670706b
• Base: codex/analytics-spend-counters at a9789f2
• Head: codex/analytics-counter-backfill at 4cd053e
• Created: 2026-09-22 17:55 UTC
• Updated: 2026-09-22 18:01 UTC

_{Automatic trigger · attempt 1 of 3 · completed in 6m 15s}

[ironloopai]

Review · Summary

🟢 No actionable findings

No actionable issues found in the reviewed change.

Validation

• ✅ Static review — Reviewed the migration, reconciliation flow, CLI packaging, and CI coverage wiring against the current transaction and accounting-lock behavior.

Review details

• Run: a718bb8a-139f-43ff-9d96-19b98670706b
• Attempts: 1

[claude]

Review — feat: backfill historical spend counters with readiness checks

I traced the reconciliation algebra against the actual C1 writers before commenting. The core approach is sound, and I want to state that explicitly since it is the part most likely to be wrong:

• organization_usage.rs:198-236 and organization_service_usage.rs:365-387 both write the raw log row, the organization_balance increment, and increment_api_key_spend inside one transaction, under lock_organization_accounting. So raw and counted advance atomically.
• prepare() reads the readiness flag, the balance counters, and the raw/counted join from a single REPEATABLE READ READ ONLY snapshot (the SELECT set_config(...) acquires it), so both sides of the delta come from the same point in time.
• Therefore counted_now + (raw_snapshot − counted_snapshot) == raw_now at apply time, and the marker plus corrections commit together under the same org lock the writers take. The recheck at spend_counters_backfill.rs:445-459 closes the double-apply window.

Also confirmed: api_key_id is NOT NULL on both usage log tables (no NULL-key row falls out of the FULL OUTER JOIN), organization_id is indexed on both, Database::pool() returns the write/leader pool (no replica-lag read), the V0004 trigger guarantees a balance row per org, ADD COLUMN plus a separate SET DEFAULT avoids a table rewrite, and nothing reads organization_balance via SELECT * / RETURNING * — so V0082 is rolling-update safe.

The issues below are not correctness bugs in the delta math.

---

⚠️ 1. --statement-timeout-seconds does not cover the apply phase (spend_counters_backfill.rs:239, 439)

APPLY_TIMEOUT is a hard-coded 5s const, and apply() ignores the operator-supplied timeout entirely — the flag only reaches prepare(). For an organization with a large drifted-key set, the single UNNEST upsert plus the balance update must finish within 5s while holding the org accounting lock. If it cannot, that org is permanently unreconcilable and the operator has no knob, which in turn blocks ensure_spend_counters_ready and therefore the #1120 reader rollout.

The help text (backfill-spend-counters.rs:176-183) reinforces the wrong mental model by presenting the flag as the operation timeout.

Either thread the option through, or split the flag so the cap is reachable:

pub async fn apply(self, apply_timeout: Duration) -> Result<SpendBackfillOutcome> {
    validate_timeout(apply_timeout)?;
    ...
    set_timeout(&transaction, "statement_timeout", apply_timeout).await?;

Keeping lock_timeout pinned at 5s is correct and should stay — it is the statement cap that needs to be operator-tunable.

⚠️ 2. One anomalous organization aborts the entire multi-org run (backfill-spend-counters.rs:100-111)

reconcile_one(...).await? propagates straight out of main, so the first failure kills the process. Two reachable, non-transient failures do this:

• prepare() bails with organization {id} has no organization_balance row (spend_counters_backfill.rs:307)
• checked_correction bails on counted > raw (spend_counters_backfill.rs:625)

Both are anomalies worth surfacing, but neither should stop the other N−1 organizations. Since incomplete_organizations orders by id, a single bad org near the start leaves everything after it unreconciled, and a rerun just re-hits the same org. Suggest accumulating per-org failures, continuing, and exiting non-zero with a count at the end — ensure_spend_counters_ready already provides the hard gate, so a partial run followed by a targeted --organization fix is the natural operator workflow:

let mut failures = 0usize;
for organization_id in ids {
    if let Err(error) = reconcile_one(&database, organization_id, options.statement_timeout).await {
        tracing::error!(%organization_id, %error, "organization reconciliation failed");
        failures += 1;
    }
}

(IDs and counts only, consistent with the privacy rules the existing logging already follows.)

⚠️ 3. The production entry point has no test coverage

incomplete_organizations is referenced only by spend_counters_backfill.rs and the binary — never by tests/spend_counter_backfill.rs. The three tests drive prepare / apply / ensure_spend_counters_ready directly, so the keyset pagination, the after = ids.last() advance, the BATCH_SIZE boundary, and the missing-balance-row path are all unexercised. That is the code that actually runs in production. A test with more than BATCH_SIZE incomplete orgs (or a temporarily lowered batch size) would cover the loop cheaply.

⚠️ 4. Test databases leak on failure, and DROP DATABASE can race (tests/spend_counter_backfill.rs, cleanup)

cleanup() is only reached on the success path — every failing or panicking test permanently leaves a spend_counter_backfill_<uuid> database behind on the shared self-hosted runner.

Separately, drop(self.pool) does not synchronously terminate the backends: dropping a tokio-postgres Client signals its connection task, which closes asynchronously. DROP DATABASE can therefore intermittently fail with 55006 (objects_in_use). On PG 13+ the fix is one clause:

.batch_execute(&format!("DROP DATABASE IF EXISTS {} WITH (FORCE)", self.database_name))

Combined with cleanup in a guard (or Drop), this removes both the flake and the leak.

Minor, related: each test creates a fresh database and runs all 82 migrations under the e2e-db group 4-way concurrency, against slow-timeout = { period = "60s", terminate-after = 2 }. Worth watching for timeout flakes on a busy runner.

---

Notes, not blockers

• Database::from_config starts Patroni discovery and background refresh tasks; the CLI never calls database.shutdown(). Harmless at process exit, but an explicit shutdown would be tidier.
• The CLI constructs the full Database (all ~10 repositories) solely to obtain pool(). Incidental.
• ORDER BY api_key_id on the FULL OUTER JOIN resolves to the COALESCE output alias in PostgreSQL (output names take precedence), so it is not ambiguous — but it is also not needed.
• No privacy concerns: every log statement emits IDs, counts, and timeouts only, consistent with CLAUDE.md.
• The prior automated review (ironloopai, pullrequestreview-5281776874) reported no actionable findings; the items above are additive rather than contradictory.

⚠️ Issues found — #1 and #2 are the ones worth resolving before this runs against production, since both can strand the rollout that #1120 depends on.

[henrypark133]

Code Review (multi-agent)

Intent: Backfill historical spend counters with a repeatable-read reconciliation operator and per-organization readiness markers so future counter-based readers don't undercount pre-existing spend.

Stats: 5 findings (from 7 raw, 7 after filter, 6 after anchor gate, 5 after dedup) across 3 files. Reviewers run: correctness, security, performance, design, coverage. Reviewers failed: none. Body-only: 0.

Correctness / Regression-escape

1. High No test exercises prepare() for an org with no organization_balance row via the real batch path (crates/database/src/spend_counters_backfill.rs:60-79, confidence 92) — anchor: crates/database/src/spend_counters_backfill.rs:78
   incomplete_organizations() includes orgs with no organization_balance row (LEFT JOIN). prepare() then bails hard with "has no organization_balance row", which propagates via ? through reconcile_one and aborts the entire all-organization batch loop, skipping ensure_spend_counters_ready for every remaining org — not just the one with the missing row. No test exercises this path (only ensure_spend_counters_ready is tested for the missing-balance case, never prepare()/the batch path). Verified independently in the source.
   Also flagged by: correctness/Medium (same root cause, framed as an error-handling design question: should reconcile_one skip+report instead of hard-aborting the batch?)

Coverage

2. High incomplete_organizations() has zero test coverage (crates/database/src/spend_counters_backfill.rs:334-358, confidence 93) — anchor: crates/database/src/spend_counters_backfill.rs:334
   No test calls incomplete_organizations directly or indirectly (not even through the binary — nothing in the test suite, CI workflow, or an e2e harness invokes backfill-spend-counters). Untested: the LEFT JOIN including orgs with no balance row, excluding ready orgs, the after-cursor pagination, and ordering.
3. Medium parse_args has no unit tests (crates/database/src/bin/backfill-spend-counters.rs:83-174, confidence 58) — anchor: crates/database/src/bin/backfill-spend-counters.rs:83 (candidate — validate claim)
   No #[cfg(test)] module exists for this pure function. Untested: --organization UUID parsing, --statement-timeout-seconds parsing and its zero-timeout bail, --help, and the unknown-flag bail.
4. Medium LOCK_TIMEOUT/APPLY_TIMEOUT 5s caps and validate_timeout bounds are untested (crates/database/src/spend_counters_backfill.rs:22-23, confidence 52) — anchor: crates/database/src/spend_counters_backfill.rs:22 (candidate — validate claim)
   No test holds the accounting lock on one connection while calling apply() on another to confirm lock_timeout actually fires a bounded error, and no test calls prepare() with a zero or > i32::MAX timeout to hit validate_timeout's bail. The PR body explicitly calls out the 5s cap as a safety property.

Performance / Concurrency

5. High Backfill apply() holds the hot-path org accounting lock during live traffic (crates/database/src/spend_counters_backfill.rs:200-260, confidence 78) — anchor: crates/database/src/spend_counters_backfill.rs:212
   apply() takes the same FOR UPDATE organization row lock as live usage-recording writers (organization_usage.rs/organization_service_usage.rs via lock_organization_accounting), then runs a multi-row UPSERT into api_key_spend sized by the org's full historical distinct API-key count plus the organization_balance UPDATE, all in one transaction before commit. Per the PR's rollout plan this operator runs against a database already receiving live inference/service traffic. statement_timeout (5s) applies per-statement, not per-transaction, so total lock-hold time across the 3 sequential statements isn't capped at 5s — worst case is closer to 10-15s for a large organization. Verified independently: mechanism and magnitude hold up; only the word "unbounded" was an overstatement (it's bounded by API-key count and per-statement timeout, not literally unbounded).
   Fix: Chunk the api_key_spend UPSERT into small batches committed incrementally under short-lived lock acquisitions, or throttle per-organization apply() calls, so lock hold time is bounded independent of historical API-key count.

---

🤖 Generated with Claude Code multi-agent code review.

[github-actions]

✅ OpenCodeReview: Review partially complete: 0 finding(s); 5 of 8 selected item(s) failed.

[henrypark133]

Review pass: added production-entry-point, pagination, missing-row, argument, and lock-timeout coverage. CLI help now says --statement-timeout-seconds controls snapshot statements; apply statements and lock waits retain their fixed 5-second caps. Oversized values fail before connecting.

Deliberate choices retained after a strict design review:

• Atomic per-organization apply: incremental commits need a separate durable progress/replay protocol; the high-priority lock-duration thread remains open for that rollout tradeoff. Each statement is bounded, but total transaction/lock duration is not a single 5-second wall-clock guarantee.
• Fail fast on missing balances or counters exceeding raw history. Continuing past an accounting anomaly would change the operator policy. The actual CLI now has a test asserting nonzero exit and no processing of later organizations.
• No force/timeout escape mode. A larger organization that exceeds the apply budget needs measured follow-up work before reader rollout.

Test-file changes: spend_counter_backfill.rs runs each test body in a task and drops its owned UUID database with WITH (FORCE) before propagating errors/panics; existing assertions remain. New spend_counter_backfill/coverage.rs covers keyset ordering, missing balances, timeout bounds, 101-organization CLI traversal, and fail-fast CLI behavior, with subprocess deadlines. The first combined run exposed that map_db_error converts SQLSTATE 57014 into RepositoryError::QueryTimeout; the new timeout test was corrected to assert that existing typed contract or exact 55P03, while retaining the elapsed-time and readiness assertions. Binary-local tests cover argument parsing. Cleanup covers test-body failures; a database creation/migration failure still needs operator cleanup of that failed setup.

Local aggregate validation (all four PRs plus review fixes):

cargo nextest run --test e2e_all --test spend_counter_backfill --test counter_readers --bin backfill-spend-counters --no-fail-fast
Summary: 819 passed, 0 failed, 12 existing skips

Formatting and diff checks passed. The final run used DATABASE_* configuration without PG* overrides. This local aggregate validation is not a production backfill or deployment.

GitHub CI verified on 0fb7a71a: Test Suite, lint, unit, integration, E2E, script tests, cargo audit, and cargo deny all succeeded. Build release was skipped by the workflow. No CI reruns were needed.

CI follow-up after retargeting to main

Head 1cf55eb7: CI remains red after one failed-jobs rerun. Both E2E attempts ran on self-hosted runner gpu03.

• Attempt 1: backfill_reconciles_snapshot_delta_and_is_idempotent lost its PostgreSQL connection during migration V12, before the test body. 800 tests passed; one failed; five did not run.
• Attempt 2: organization_priority::priority_updates_bypass_warm_real_auth_caches_on_both_instances first logged a PostgreSQL communication error and Connection closed while recording usage. Its retry then hit SQLSTATE 55P03 locking the same organization. 636 tests passed; one failed; the remaining tests did not run. Backfill tests were not reached.
• Logs support a recurring connection/runner problem; the precise disconnect cause and blocker identity are not established. The delayed server-side transaction cleanup is a hypothesis, not proven.
• Lint, unit tests, integration tests, scripts, and security checks passed. No source or test files changed; no assertions weakened.

Failed run and rerun. Investigate gpu03 PostgreSQL/container networking before another run. This pass used its single evidence-based rerun; no further unchanged reruns were triggered.

[lloydmak99]

Pure-addition change: a new V0082 migration, a reconciliation module + operator binary, tests, and CI/Docker wiring, with no modifications to existing reader/writer paths.

The reconciliation algebra checks out: prepare() captures raw_snapshot − counted_snapshot under a single REPEATABLE READ READ ONLY snapshot, and apply() adds that delta to current counters under the same org accounting lock the C1 writers use, so counted_final = raw_now. Double-apply is prevented by re-reading spend_counters_ready_at under the lock and gating the balance UPDATE on IS NULL (confirmed by the concurrent-worker test). The migration is rolling-update safe: ADD COLUMN then a separate SET DEFAULT NOW() leaves existing rows NULL for backfill with no table rewrite, and no SELECT */RETURNING * on organization_balance breaks old pods. Usage-log FKs are RESTRICT, so the api_key_spend upsert can't orphan-fail, and the keyset loop terminates cleanly.

No blocking correctness, data-loss, security, or API-contract issues found.

Checks: cargo fmt --all -- --check and git diff --check passed; CI cargo audit / cargo deny advisories passed. Compilation and DB-backed tests couldn't run locally (no cc linker / no local PostgreSQL) — those run in CI via cargo nextest.