scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-11-13 08:52:02 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 174

feat(db): classify errors, recycle poisoned pools and take the rollup lock_timeout like Prisma

Browse source 
db::Error now reports a kind that matches PrismaDBExceptionHandler's buckets:
unreachable, capacity, deadlock, serialization failure, lock timeout, read-only
after failover, stale cached plan and rejected data. is_unavailable gates
fail-open and is_requeue_safe says whether a failed write provably never
applied. begin() takes the rollup lock_timeout from the caller instead of
hardcoding 5000ms, so SPEND_ROLLUP_LOCK_TIMEOUT_MS applies on both paths

Pools::recycle_after is the native counterpart of recreating the Prisma
client: after a 25006 (the writer endpoint now points at a read-only node) or
"cached plan must not change result type" (a migration changed a table under
a prepared statement), every pooled connection opened before the error is
dropped on its next acquire, so the pool reconnects and prepares fresh

begin, lock_team and add_team_member_spend take whether statements stay named and
run through Statements, so they send unnamed statements when the cache is off

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
This commit is contained in:
Yujong Lee 2026-11-03 14:59:52 +09:00
parent fc7a3029e2
commit 89900a9019
8 changed files with 315 additions and 51 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

8
litellm-rust/crates/db/.sqlx/query-5c74f47a2b84574133c520d235192bf33e16fea1dbb0c32cf89a376a874927f4.json → litellm-rust/crates/db/.sqlx/query-805611b4198a8f11e66e29b77afa021f3a07befaedd77d4949e9401b120a5322.json generated
View file

@ -1,6 +1,6 @@
{
"db_name": "PostgreSQL",
"query": "SELECT set_config('lock_timeout', '5000ms', true)",
"query": "SELECT set_config('lock_timeout', $1::text, true)",
"describe": {
"columns": [
{
@ -11,11 +11,13 @@
}
],
"parameters": {
"Left": []
"Left": [
"Text"
]
},
"nullable": [
null
]
},
"hash": "5c74f47a2b84574133c520d235192bf33e16fea1dbb0c32cf89a376a874927f4"
"hash": "805611b4198a8f11e66e29b77afa021f3a07befaedd77d4949e9401b120a5322"
}

3
litellm-rust/crates/db/AGENTS.md
View file

@ -53,7 +53,8 @@ Python and SQLx run side by side until every repository operation is native, so
| In-container PgBouncer | `pgbouncer.py` | Inherited: the pooled URL carries PgBouncer's credential |
| Startup connection budget | `postgres_connection_budget_message` | Counts the Rust writer and reader pools' shares |
| Read replica routing | `routing_prisma_wrapper.py`, `DATABASE_URL_READ_REPLICA` | `Pools`: callers pick `writer()`, `read()` or `read_after_write()`. `read()` prefers the reader and `read_after_write()` the writer (`WriterPinnedClient`); either falls back to the other side when its pool cannot connect and skips that side for 30s, so reads keep working on the reader while the writer is down. Writes have no fallback |
| Error classification, `lock_timeout` on rollups | `exception_handler.py`, `rollup_lock_timeout.py` | Planned |
| Error classification, `lock_timeout` on rollups | `exception_handler.py`, `rollup_lock_timeout.py` | `Error::kind`, `is_unavailable`, `is_requeue_safe`; `begin` takes the caller's `lock_timeout` |
| Recreate the client after 25006 (writer became read-only) or a stale cached plan | `exception_handler.py`, `_query_first_with_cached_plan_fallback` | `Pools::recycle_after` drops every pooled connection opened before the error, so the next acquire reconnects and prepares fresh. Callers retry a `StalePlan` once |
| `ServiceTypes.DB` events per round trip | `db_span.py`, `log_db_metrics.py`, `prometheus_services.py`, `otel/model/spans.py` | Python host only. Wrap each native operation in `observe` with the call type and table its Prisma version uses, and the bridge forwards events to `ServiceLoggingEvents`. Nesting under an outer Python `db_span` is not detected. A Rust host has no sink yet and must emit the names below |
| Auth lookup gate, fail open on DB outage, reader choice per query | `db_lookup_gate.py`, `allow_requests_on_db_unavailable` | With each ported auth read |
| Connect at boot, so a database that cannot be reached is reported before traffic | `PrismaClient.connect` | The pool stays lazy and never fails boot. `Database.open` pings it once in the background and logs a warning when it cannot connect, so a mistranslated endpoint is not first found by a ported query |

9
litellm-rust/crates/db/src/endpoint.rs
View file

@ -105,13 +105,18 @@ impl Endpoint {
}
}
pub fn pool(&self) -> Result<PgPool, Error> {
pub fn pool_options(&self) -> Result<PgPoolOptions, Error> {
let limits = self.pool;
Ok(PgPoolOptions::new()
.min_connections(0)
.max_connections(limits.max_connections)
.acquire_timeout(seconds(limits.acquire_timeout_seconds)?.unwrap_or(WAIT_FOREVER))
.idle_timeout(seconds(limits.idle_timeout_seconds)?)
.idle_timeout(seconds(limits.idle_timeout_seconds)?))
}
pub fn pool(&self) -> Result<PgPool, Error> {
Ok(self
.pool_options()?
.connect_lazy_with(self.connect_options()))
}

126
litellm-rust/crates/db/src/error.rs
View file

@ -10,21 +10,75 @@ pub enum Error {
InvalidDuration(#[from] std::time::TryFromFloatSecsError),
}
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum ErrorKind {
Unreachable,
Capacity,
Deadlock,
SerializationFailure,
LockTimeout,
ReadOnly,
StalePlan,
DataRejected,
Other,
}
impl Error {
pub fn is_retryable(&self) -> bool {
pub fn kind(&self) -> ErrorKind {
match self {
Self::Query(sqlx::Error::Io(_)) => ErrorKind::Unreachable,
Self::Query(sqlx::Error::PoolTimedOut) => ErrorKind::Capacity,
Self::Query(sqlx::Error::Database(error)) => {
matches!(
error.code().as_deref(),
Some("40P01" | "55P03" | "40001" | "53300")
)
sqlstate_kind(error.code().as_deref(), error.message())
}
Self::Query(sqlx::Error::Io(_) | sqlx::Error::PoolTimedOut) => true,
_ => ErrorKind::Other,
}
}
pub fn is_retryable(&self) -> bool {
matches!(
self.kind(),
ErrorKind::Unreachable
| ErrorKind::Capacity
| ErrorKind::Deadlock
| ErrorKind::SerializationFailure
| ErrorKind::LockTimeout
| ErrorKind::StalePlan
)
}
pub fn needs_fresh_connections(&self) -> bool {
matches!(self.kind(), ErrorKind::ReadOnly | ErrorKind::StalePlan)
}
pub fn is_unavailable(&self) -> bool {
matches!(self.kind(), ErrorKind::Unreachable | ErrorKind::Capacity)
}
pub fn is_requeue_safe(&self) -> bool {
match self {
Self::Query(sqlx::Error::PoolTimedOut | sqlx::Error::PoolClosed) => true,
Self::Query(sqlx::Error::Database(_)) => self.kind() != ErrorKind::DataRejected,
_ => false,
}
}
}
fn sqlstate_kind(code: Option<&str>, message: &str) -> ErrorKind {
match code {
Some("0A000") if message.contains("cached plan must not change result type") => {
ErrorKind::StalePlan
}
Some("53300") => ErrorKind::Capacity,
Some("40P01") => ErrorKind::Deadlock,
Some("40001") => ErrorKind::SerializationFailure,
Some("55P03") => ErrorKind::LockTimeout,
Some("25006") => ErrorKind::ReadOnly,
Some(code) if code.starts_with("22") || code.starts_with("23") => ErrorKind::DataRejected,
_ => ErrorKind::Other,
}
}
#[cfg(test)]
mod tests {
use std::error::Error as StdError;
@ -32,29 +86,59 @@ mod tests {
use rstest::rstest;
use sqlx::PgPool;
use super::Error;
use super::{Error, ErrorKind};
use crate::test_support::database_pool;
#[rstest]
#[case::closed_pool(sqlx::Error::PoolClosed, false)]
#[case::pool_timeout(sqlx::Error::PoolTimedOut, true)]
#[case::io(sqlx::Error::Io(std::io::ErrorKind::ConnectionReset.into()), true)]
#[case::missing_row(sqlx::Error::RowNotFound, false)]
fn classifies_client_errors(#[case] error: sqlx::Error, #[case] retryable: bool) {
assert_eq!(Error::from(error).is_retryable(), retryable);
#[case::closed_pool(sqlx::Error::PoolClosed, ErrorKind::Other, false, false, true)]
#[case::pool_timeout(sqlx::Error::PoolTimedOut, ErrorKind::Capacity, true, true, true)]
#[case::io(
sqlx::Error::Io(std::io::ErrorKind::ConnectionReset.into()),
ErrorKind::Unreachable,
true,
true,
false
)]
#[case::missing_row(sqlx::Error::RowNotFound, ErrorKind::Other, false, false, false)]
fn classifies_client_errors(
#[case] error: sqlx::Error,
#[case] kind: ErrorKind,
#[case] retryable: bool,
#[case] unavailable: bool,
#[case] requeue_safe: bool,
) {
let error = Error::from(error);
assert_eq!(error.kind(), kind);
assert_eq!(error.is_retryable(), retryable);
assert_eq!(error.is_unavailable(), unavailable);
assert_eq!(error.is_requeue_safe(), requeue_safe);
}
#[rstest]
#[case::connection_capacity("53300", true)]
#[case::deadlock("40P01", true)]
#[case::lock_unavailable("55P03", true)]
#[case::serialization_failure("40001", true)]
#[case::unique_violation("23505", false)]
#[case::connection_capacity("53300", ErrorKind::Capacity, true, true, true)]
#[case::deadlock("40P01", ErrorKind::Deadlock, true, false, true)]
#[case::lock_unavailable("55P03", ErrorKind::LockTimeout, true, false, true)]
#[case::serialization_failure("40001", ErrorKind::SerializationFailure, true, false, true)]
#[case::read_only_after_failover("25006", ErrorKind::ReadOnly, false, false, true)]
#[case::unique_violation("23505", ErrorKind::DataRejected, false, false, false)]
#[case::numeric_overflow("22003", ErrorKind::DataRejected, false, false, false)]
#[case::undefined_table("42P01", ErrorKind::Other, false, false, true)]
#[case::feature_not_supported_without_a_stale_plan(
"0A000",
ErrorKind::Other,
false,
false,
true
)]
#[tokio::test]
#[ignore = "requires DATABASE_URL pointing to a Prisma-migrated PostgreSQL database"]
async fn classifies_server_errors(
#[case] code: &str,
#[case] kind: ErrorKind,
#[case] retryable: bool,
#[case] unavailable: bool,
#[case] requeue_safe: bool,
#[future(awt)] database_pool: Result<PgPool, Box<dyn StdError>>,
) -> Result<(), Box<dyn StdError>> {
let pool = database_pool?;
@ -74,7 +158,11 @@ mod tests {
error.as_database_error().and_then(|error| error.code()),
Some(code.into())
);
assert_eq!(Error::from(error).is_retryable(), retryable);
let error = Error::from(error);
assert_eq!(error.kind(), kind);
assert_eq!(error.is_retryable(), retryable);
assert_eq!(error.is_unavailable(), unavailable);
assert_eq!(error.is_requeue_safe(), requeue_safe);
tx.rollback().await?;
Ok(())
}

2
litellm-rust/crates/db/src/lib.rs
View file

@ -24,7 +24,7 @@ pub use daily_spend::{DailyEndUserRow, DailyTeamRow, DailyUserRow};
pub use deadline::within;
pub use end_user::{EndUser, RateLimits, find_end_user};
pub use endpoint::{ClientIdentity, Endpoint, PoolLimits, SslMode};
pub use error::Error;
pub use error::{Error, ErrorKind};
pub use failure::DbFailure;
pub use health::health_check;
pub use observe::{Call, Observer, Outcome, QueryEvent, observe};

169
litellm-rust/crates/db/src/pools.rs
View file

@ -1,4 +1,7 @@
use std::time::{Duration, Instant};
use std::{
sync::{Arc, Mutex, PoisonError},
time::{Duration, Instant},
};
use litellm_circuit_breaker::{Breaker, Config, Failure};
use sqlx::{PgPool, Postgres, pool::PoolConnection};
@ -9,16 +12,46 @@ const RECOVERY: Duration = Duration::from_secs(30);
type Connection = PoolConnection<Postgres>;
#[derive(Default)]
struct Recycled(Mutex<Option<Instant>>);
impl Recycled {
fn mark(&self, now: Instant) {
*self.0.lock().unwrap_or_else(PoisonError::into_inner) = Some(now);
}
fn predates(&self, now: Instant, age: Duration) -> bool {
let Some(recycled_at) = *self.0.lock().unwrap_or_else(PoisonError::into_inner) else {
return false;
};
now.checked_sub(age)
.is_none_or(|opened_at| opened_at <= recycled_at)
}
}
#[derive(Clone)]
struct Side {
pool: PgPool,
breaker: Breaker,
recycled: Arc<Recycled>,
}
impl Side {
fn new(endpoint: &Endpoint) -> Result<Self, Error> {
let recycled = Arc::new(Recycled::default());
let pool = endpoint
.pool_options()?
.before_acquire({
let recycled = Arc::clone(&recycled);
move |_, metadata| {
let fresh = !recycled.predates(Instant::now(), metadata.age);
Box::pin(async move { Ok(fresh) })
}
})
.connect_lazy_with(endpoint.connect_options());
Ok(Self {
pool: endpoint.pool()?,
pool,
recycled,
breaker: Breaker::new(
Config {
failure_threshold: 1,
@ -99,6 +132,18 @@ impl Pools {
acquire_preferring(&self.writer, self.reader.as_ref()).await
}
pub fn recycle_after(&self, error: &Error) -> bool {
if !error.needs_fresh_connections() {
return false;
}
let now = Instant::now();
self.writer.recycled.mark(now);
if let Some(reader) = &self.reader {
reader.recycled.mark(now);
}
true
}
pub async fn close(&self) {
self.writer.pool.close().await;
if let Some(reader) = &self.reader {
@ -109,12 +154,124 @@ impl Pools {
#[cfg(test)]
mod tests {
use std::{error::Error as StdError, time::Instant};
use std::{
error::Error as StdError,
time::{Duration, Instant},
};
use sqlx::PgConnection;
use rstest::rstest;
use sqlx::{Connection as _, PgConnection};
use super::Pools;
use crate::{Endpoint, PoolLimits, test_support::database_endpoint};
use super::{Pools, Recycled};
use crate::{Endpoint, Error, ErrorKind, PoolLimits, test_support::database_endpoint};
#[rstest]
#[case::never_recycled(None, 5, false)]
#[case::opened_before_the_recycle(Some(10), 30, true)]
#[case::opened_after_the_recycle(Some(10), 5, false)]
fn discards_connections_opened_before_a_recycle(
#[case] recycled_seconds_ago: Option<u64>,
#[case] age_seconds: u64,
#[case] stale: bool,
) {
let now = Instant::now() + Duration::from_secs(60);
let recycled = Recycled::default();
if let Some(ago) = recycled_seconds_ago {
recycled.mark(now - Duration::from_secs(ago));
}
assert_eq!(
recycled.predates(now, Duration::from_secs(age_seconds)),
stale
);
}
async fn backend_pid(pools: &Pools) -> Result<i32, sqlx::Error> {
#[expect(clippy::disallowed_methods, reason = "A probe query needs no schema")]
sqlx::query_scalar("SELECT pg_backend_pid()")
.fetch_one(&mut *pools.writer().acquire().await?)
.await
}
async fn raised(pools: &Pools, code: &str) -> Result<Error, sqlx::Error> {
#[expect(
clippy::disallowed_methods,
reason = "Test query must raise the requested SQLSTATE"
)]
let error = sqlx::query(sqlx::AssertSqlSafe(format!(
"DO $$ BEGIN RAISE EXCEPTION USING ERRCODE = '{code}'; END $$"
)))
.execute(pools.writer())
.await
.expect_err("PostgreSQL must raise the requested error");
Ok(Error::from(error))
}
#[rstest]
#[case::read_only_writer("25006", true)]
#[case::lock_timeout("55P03", false)]
#[tokio::test]
#[ignore = "requires DATABASE_URL pointing to a Prisma-migrated PostgreSQL database"]
async fn recycles_pooled_connections_only_after_errors_that_poison_them(
#[case] code: &str,
#[case] recycled: bool,
) -> Result<(), Box<dyn StdError>> {
let pools = Pools::new(&database_endpoint()?, None)?;
let before = backend_pid(&pools).await?;
assert_eq!(pools.recycle_after(&raised(&pools, code).await?), recycled);
assert_eq!(backend_pid(&pools).await? != before, recycled);
Ok(())
}
#[tokio::test]
#[ignore = "requires DATABASE_URL pointing to a Prisma-migrated PostgreSQL database"]
async fn a_schema_change_under_a_cached_plan_is_recovered_by_a_recycle()
-> Result<(), Box<dyn StdError>> {
let endpoint = database_endpoint()?;
let pools = Pools::new(
&Endpoint {
pool: PoolLimits {
max_connections: 1,
..endpoint.pool
},
..endpoint.clone()
},
None,
)?;
let table = format!("native_stale_plan_probe_{}", std::process::id());
let select = format!("SELECT * FROM {table}");
#[expect(clippy::disallowed_methods, reason = "A scratch table needs no schema")]
let run = |sql: &str| sqlx::query(sqlx::AssertSqlSafe(sql.to_owned()));
let mut admin = PgConnection::connect_with(&endpoint.connect_options()).await?;
run(&format!("CREATE TABLE {table} (id int)"))
.execute(&mut admin)
.await?;
let outcome = async {
run(&select).fetch_all(pools.writer()).await?;
run(&format!("ALTER TABLE {table} ADD COLUMN extra int"))
.execute(&mut admin)
.await?;
let stale = Error::from(
run(&select)
.fetch_all(pools.writer())
.await
.expect_err("the cached plan must be stale"),
);
let recycled = pools.recycle_after(&stale);
run(&select).fetch_all(pools.writer()).await?;
Ok::<_, Box<dyn StdError>>((stale.kind(), recycled))
}
.await;
run(&format!("DROP TABLE {table}"))
.execute(&mut admin)
.await?;
assert_eq!(outcome?, (ErrorKind::StalePlan, true));
Ok(())
}
fn read_only(endpoint: &Endpoint) -> Endpoint {
Endpoint {

30
litellm-rust/crates/db/src/spend/team_lock.rs
View file

@ -1,28 +1,38 @@
use std::time::Duration;
use sqlx::{PgPool, Postgres, Transaction};
use crate::Error;
use crate::{Error, Statements};
pub async fn begin(pool: &PgPool) -> Result<Transaction<'static, Postgres>, Error> {
pub async fn begin(
pool: &PgPool,
lock_timeout: Duration,
named: bool,
) -> Result<Transaction<'static, Postgres>, Error> {
let mut transaction = pool.begin().await?;
sqlx::query!("SELECT set_config('lock_timeout', '5000ms', true)")
.fetch_one(&mut *transaction)
.await?;
sqlx::query!(
"SELECT set_config('lock_timeout', $1::text, true)",
format!("{}ms", lock_timeout.as_millis())
)
.fetch_one(Statements::new(&mut *transaction, named))
.await?;
Ok(transaction)
}
pub async fn lock_team(
transaction: &mut Transaction<'_, Postgres>,
team_id: &str,
named: bool,
) -> Result<(), Error> {
sqlx::query_file!("queries/spend/lock_team.sql", team_id)
.fetch_one(&mut **transaction)
.fetch_one(Statements::new(&mut **transaction, named))
.await?;
Ok(())
}
#[cfg(test)]
mod tests {
use std::error::Error as StdError;
use std::{error::Error as StdError, time::Duration};
use super::{begin, lock_team};
use crate::test_support::database_pool;
@ -42,7 +52,7 @@ mod tests {
.max_connections(2)
.connect_with(database_pool.connect_options().as_ref().clone())
.await?;
let mut first = begin(&pool).await?;
let mut first = begin(&pool, Duration::from_millis(1500), true).await?;
let mut second = pool.begin().await?;
#[expect(
clippy::disallowed_methods,
@ -51,7 +61,7 @@ mod tests {
let original_timeout: String = sqlx::query_scalar("SELECT current_setting('lock_timeout')")
.fetch_one(&mut *second)
.await?;
lock_team(&mut first, team_id).await?;
lock_team(&mut first, team_id, true).await?;
#[expect(
clippy::disallowed_methods,
reason = "Test assertions are outside the compile-checked query set"
@ -60,7 +70,7 @@ mod tests {
sqlx::query_as("SELECT pg_backend_pid(), current_setting('lock_timeout')")
.fetch_one(&mut *first)
.await?;
assert_eq!(timeout, "5s");
assert_eq!(timeout, "1500ms");
#[expect(
clippy::disallowed_methods,
reason = "Test assertions are outside the compile-checked query set"

19
litellm-rust/crates/db/src/spend/team_member_spend.rs
View file

@ -1,7 +1,7 @@
use sqlx::{Postgres, Transaction};
use super::lock_team;
use crate::Error;
use crate::{Error, Statements};
#[macro_rules_attribute::apply(crate::db_input)]
pub struct TeamMemberSpend {
@ -13,6 +13,7 @@ pub struct TeamMemberSpend {
pub async fn add_team_member_spend(
transaction: &mut Transaction<'_, Postgres>,
members: &[TeamMemberSpend],
named: bool,
) -> Result<(), Error> {
if members.is_empty() {
return Ok(());
@ -27,7 +28,7 @@ pub async fn add_team_member_spend(
.collect();
teams.dedup();
for team in teams {
lock_team(transaction, team).await?;
lock_team(transaction, team, named).await?;
}
let rows = serde_json::Value::Array(
members
@ -42,14 +43,14 @@ pub async fn add_team_member_spend(
.collect(),
);
sqlx::query_file!("queries/spend/upsert_team_member_spend.sql", rows)
.execute(&mut **transaction)
.execute(Statements::new(&mut **transaction, named))
.await?;
Ok(())
}
#[cfg(test)]
mod tests {
use std::error::Error as StdError;
use std::{error::Error as StdError, time::Duration};
use super::super::begin;
use super::{TeamMemberSpend, add_team_member_spend};
@ -90,7 +91,7 @@ mod tests {
#[future(awt)] database_pool: Result<PgPool, Box<dyn StdError>>,
) -> Result<(), Box<dyn StdError>> {
let pool = database_pool?;
let mut tx = begin(&pool).await?;
let mut tx = begin(&pool, Duration::from_secs(5), true).await?;
#[expect(
clippy::disallowed_methods,
reason = "Fixture writes are outside the compile-checked query set"
@ -113,10 +114,10 @@ mod tests {
)
.execute(&mut *tx)
.await?;
add_team_member_spend(&mut tx, &[]).await?;
add_team_member_spend(&mut tx, &members(0.0)).await?;
add_team_member_spend(&mut tx, &members(cost)).await?;
add_team_member_spend(&mut tx, &members(cost)).await?;
add_team_member_spend(&mut tx, &[], true).await?;
add_team_member_spend(&mut tx, &members(0.0), true).await?;
add_team_member_spend(&mut tx, &members(cost), true).await?;
add_team_member_spend(&mut tx, &members(cost), true).await?;
#[expect(
clippy::disallowed_methods,
reason = "Test assertions are outside the compile-checked query set"