Handle CockroachDB serialization failures

Miroslav Vadkertirequested to merge
mvadkert/nucleus:handle-serialization-failures into main

Summary

CockroachDB runs all transactions at SERIALIZABLE isolation and aborts conflicting transactions with a serialization failure (SQLSTATE 40001), expecting the client to retry them. The API had no retry handling, so these conflicts surfaced as 500 Internal Server Error to users and workers.

Root cause

Commit 6314eda4 (Add last_used_at timestamp to tokens) made every authenticated request potentially issue UPDATE tokens SET last_used_at=.... Shared worker tokens are used by many workers concurrently, so once the 60s throttle window expires, several in-flight requests race to update the same token row — CockroachDB aborts all but one with WriteTooOldError, failing even read-only GETs. This is TESTING-FARM-API-43 (11k+ events since the deploy).

The same error class also affects concurrent updates of the same request row in update_test_request() / delete_test_request(): TESTING-FARM-API-33 and TESTING-FARM-API-2V.

Changes

  • core/database.py: new is_serialization_failure() and retry_on_serialization_failure() helpers — rollback between attempts, full-jitter exponential backoff: each sleep is drawn uniformly from zero up to an exponentially growing, capped bound (colliding transactions are near-simultaneous, so a deterministic backoff would make the losers wake up together and re-collide).
  • crud/token.py: _touch_last_used_at() swallows serialization failures after a rollback — losing the race means a concurrent request has just recorded the usage, so there is nothing left to do, and recording usage must never fail the authenticated request.
  • crud/test_request.py: update_test_request() and delete_test_request() retry their transaction via the helper. Because a serialization failure means a concurrent transaction changed the row, the state transition is re-validated against the freshly read state inside the retried closure — a request that meanwhile reached a terminal state is no longer canceled/updated, and a losing writer no longer clobbers the race winner's state. Both functions return the state they actually committed from.
  • middleware/test_request.py: update and cancellation state-transition metrics are recorded after the CRUD call commits, using the returned prior state — so a rejected retry records nothing and the from_state label reflects the true committed-from state rather than a stale pre-conflict read.
  • settings.yaml: retry behavior configurable via DATABASE_SERIALIZATION_FAILURE_RETRIES (default 3), DATABASE_SERIALIZATION_FAILURE_BACKOFF_SECONDS (default 0.1) and DATABASE_SERIALIZATION_FAILURE_BACKOFF_MAX_SECONDS (default 2.0), overridable with TF_API_-prefixed environment variables.
  • pyproject.toml: add types-psycopg2 stubs as a development dependency (pinned below the releases that require Python >= 3.10), so mypy type-checks the psycopg2.errors usage.
  • New regression tests, including the serialization-race scenarios for update/delete state re-validation and committed-prior-state metrics; verified they fail with the exact production OperationalError (psycopg2.errors.SerializationFailure) when run against the code without the fix.

Fixes TESTING-FARM-API-43 Fixes TESTING-FARM-API-33 Fixes TESTING-FARM-API-2V in Sentry

Generated-by: Claude Code

Edited by Miroslav Vadkerti

Merge request reports