perf(database): reuse prepared statements on the inference hot path

[PierreLeGuen]

Reuses prepared statements for inference-path repository queries to avoid repeating the Parse/Describe round trip on warm connections. This branch includes #1045 so the new batched signature write uses the cache too. Merge #1045 first. The PR continues to target main so the repository's required CI workflows run.

Stale result shapes and missing server-side statement handles invalidate caches throughout the pool. Nontransactional queries re-prepare and retry on the same connection; billing transactions roll back and restart through the existing bounded repository retry.

Validation includes PostgreSQL 16 regressions for:

• Workspace reads after a column is added beneath warmed statements.
• Real billing inserts after schema invalidation on every pool connection, verifying one usage record and one balance increment.
• A missing cached balance-update statement after the usage insert succeeds, verifying that retry rolls back the partial transaction and charges exactly once.
• Duplicate billing and a subsequent fresh write on the same connection after rollback.
• Repeated signature batches retaining the same prepared statement and updating both stored signatures.

Validation: 224 relevant PostgreSQL 16 E2E tests and 49 credit/usage unit tests passed with #1044, #1045, and #1046 combined, including all four cache regressions. Focused mutation checks confirm the billing regressions fail without pool-wide eviction, and the batch-reuse test fails with raw execute. Repository unit tests, formatting, and strict Clippy passed; #1045's included service changes passed 168 attestation/completion tests.

Deploy separately after #1044 and #1045, without an accompanying schema migration. Performance gains still need representative warm-traffic validation: PostgreSQL can select generic plans for reused statements, so monitor query/request latency, pool waits, and database load as well as cache-invalidation warnings.

[ironloopai]

Review · Status

🟩 Completed

IronLoop completed the review and posted it to GitHub.

Result

Open submitted review →

Run details

• Run: 8b85a34a-167e-4e3c-84a5-9a02593be041
• Base: main at 7c852f2
• Head: perf/cached-prepared-statements at e064000
• Created: 2026-09-11 00:11 UTC
• Updated: 2026-09-11 00:15 UTC

_{Automatic trigger · attempt 1 of 3 · completed in 4m 40s}

[claude]

Review — prepared statement reuse on the inference hot path

Solid, well-motivated change, and the rolling-deploy failure mode is the right thing to have thought about. Two issues I would want resolved before merge, both about the other ways a cached statement handle goes bad. (No prior review threads on this PR, so nothing to build on.)

⚠️ 1. Critical — cached_* inside a transaction that rolls back leaves a dangling handle (billing path)

crates/database/src/repositories/organization_usage.rs:96 prepares the usage-log INSERT inside the transaction, and :184 explicitly rolls that transaction back on the duplicate (ON CONFLICT DO NOTHING → no row) path.

PostgreSQL’s prepared-statement catalog is transactional: a named statement Parsed inside a transaction block is deallocated when that transaction rolls back. deadpool’s per-connection StatementCache does not know that, so it keeps handing back a Statement for a server-side statement that no longer exists.

Failure scenario:

1. Fresh pooled connection; the first record_usage on it is a duplicate (retry, stream restart, resumed response).
2. prepare_cached Parses the INSERT inside the transaction → rollback() at :184 → server drops it.
3. Every later record_usage served by that connection Binds a dead statement → SQLSTATE 26000 "prepared statement ... does not exist".
4. 26000 is not handled in map_db_error (utils.rs:15-59) → DatabaseError, not retryable → retry_db! gives up, and nothing ever evicts the cache entry. The connection is poisoned for its whole lifetime, and the poisoning accumulates across connections.

The same shape applies to any abort of that transaction (e.g. the organization_balance upsert at :149 failing).

Suggested fix — prepare before opening the transaction, so the Parse is not rolled back:

// prepare_cached takes &client; do it before client.transaction() takes &mut
let insert_usage = client.prepare_cached(INSERT_USAGE_SQL).await.map_err(map_db_error)?;
let upsert_balance = client.prepare_cached(UPSERT_BALANCE_SQL).await.map_err(map_db_error)?;
let transaction = client.transaction().await.map_err(map_db_error)?;
let maybe_row = transaction.query_opt(&insert_usage, &[...]).await.map_err(map_db_error)?;

and, defensively, treat SqlState::INVALID_SQL_STATEMENT_NAME exactly like the stale-plan case in both forget_if_stale and map_db_error (evict + classify retryable) — same class of "handle outlived the server object" bug, three lines.

⚠️ 2. Recovery is per-(connection, statement), but retry_db! only gets 3 attempts

forget_if_stale (statement_cache.rs:106-116) evicts one SQL key from one connection’s cache, and retry_db! (retry.rs:7) allows MAX_ATTEMPTS = 3 with a fresh pool.get() on each attempt. Prod runs DATABASE_MAX_CONNECTIONS=16 (crates/config/src/types.rs:757).

After a migration adds a column to workspaces, all warm connections hold a stale SELECT w.*. Under concurrency, attempt 1 lands on conn A (evict, fail), attempt 2 on conn B (evict, fail), attempt 3 on conn C (evict, fail) → 500. So the first several requests per invalidated statement do surface as errors; the description’s "requests keep returning 200" holds for a single connection, not for a 16-connection pool under load.

Cheap improvement: on a stale plan, clear that connection’s whole cache rather than one key — one error then repairs all 11 hot statements on the connection instead of one: self.statement_cache.clear();. Dropping the connection from the pool so it is rebuilt clean would be even more decisive, but clear() alone shrinks the failure burst substantially.

Minor

• statement_cache.rs:28 — locale-dependent detection. db_err.message().contains("cached plan must not change result type") breaks if the server runs with a non-English lc_messages, silently disabling exactly the recovery this PR is built around. Consider keying primarily on SqlState::FEATURE_NOT_SUPPORTED for statements you executed from the cache (you already know it came from prepare_cached), with the message as a refinement rather than a requirement.
• prepared_statement_cache.rs:29-33 — the test warms one connection, not several. The 8 probes are sequential and each returns its connection before the next pool.get(), so deadpool never opens a second one; the comment "Several requests seed several connections" does not hold. The test therefore only covers single-connection recovery (which needs just 2 of the 3 attempts) and would not catch issue 2. Firing the warm-up and the post-ALTER probes concurrently (futures::future::join_all) over the 4-connection test pool would exercise the real path.
• prepared_statement_cache.rs:41 — uuid::Uuid::new_v4().simple().to_string()[..8].to_string() has a redundant inner to_string().

Logging looks clean — the warn at statement_cache.rs:110 deliberately omits the SQL, which is right.

⚠️ Issues found — 1 and 2 above are the blockers.

[chatgpt-codex-connector]

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: e06400008e

ℹ️ 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".

[chatgpt-codex-connector]

Codex Review Summary

This comment shows the latest Codex review activity on this pull request.

Review	Status	Commit	Review trigger
📝 Code Review	✅ Completed 2026-09-11T00:11:29.268594Z	e064000	PR opened

ℹ️ 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" or "@codex security review".

Codex reacts with 👀 while any review is running, comments if it has suggestions, and reacts with 👍 once all reviews finish with no findings.

[ironloopai]

Review · Summary

Found two actionable issues in the prepared-statement cache change.

Findings: 🟠 Medium 1 · 🟡 Low 1

_{Code-specific findings are attached to the diff.}

Validation

• ✅ Captured CI evidence — The release build, unit tests, and integration tests passed in the captured pull-request evidence.

Review details

• Run: 8b85a34a-167e-4e3c-84a5-9a02593be041
• Attempts: 1

[PierreLeGuen]

Addressed in the follow-up commit:

1. Prepared statement inside a rolled-back transaction — I checked this against PostgreSQL 16 before changing anything: BEGIN; PREPARE p AS SELECT 1; ROLLBACK; EXECUTE p; succeeds, so prepared statements are session-scoped and a rollback does not deallocate them (the same catalog serves protocol-level Parse). The billing path is therefore not poisoned by the duplicate-record rollback. I still added SQLSTATE 26000 to the stale-handle detection as a defensive measure, handled exactly like the stale-plan case, since it costs three lines.
2. Recovery per (connection, statement) vs three retries — eviction is now pool-wide (StatementCaches::clear on the pool manager), and non-transactional statements are re-prepared and re-run once on the same connection, so a schema change is repaired by the first request that hits it rather than one connection per retry attempt. The transactional billing write evicts pool-wide too and is restarted by retry_db!.
3. Locale-dependent detection — keyed on the RevalidateCachedQuery routine, message as fallback.
4. Test warmed one connection — warm-up and post-ALTER probes now run concurrently with join_all over the test pool. Redundant to_string() removed.

Full e2e set rerun locally (203 passed), clippy and fmt clean.

[github-actions]

🔍 OpenCodeReview found 1 issue(s) in this PR.

• ✅ 1 posted as inline comment(s)
• 📝 0 posted as summary

[lloydmak99]

The prepared-statement caching and its recovery paths hold up: stale-handle errors (SQLSTATE 0A000/26000) are classified as retryable, evict the cache, and self-heal on both the non-transactional (run_cached! on Client) and transactional (record_usage → map_tx_error → retry_db!) paths. Recovery logic and error classification are consistent across all converted call sites. No merge-blocking issues found.

Optional follow-up:

• crates/database/src/repositories/statement_cache.rs:60,:97 — a transient per-connection 26000 triggers a pool-wide StatementCaches::clear(), re-Parsing all cached statements on every connection. It's self-correcting, but a per-connection clear (reserving the pool-wide clear for 0A000) would avoid the extra round trips. Worth a comment on the trade-off; not required to merge.

Checks: git diff --check, cargo fmt --check, and GitHub lint/unit/integration/E2E/security all passed; cargo compilation couldn't run locally (no C linker), compensated by verifying the deadpool 0.14.1/0.12.3 API surface against vendored sources.