Lines
78.87 %
Functions
92.86 %
Branches
100 %
use super::registry::{PoolRegistry, TENANT_MAX_CONNECTIONS, TenantPool};
use crate::config::system_config;
use crate::db::DBError;
use crate::user::connection::TenantConnection;
use sqlx::postgres::PgPoolOptions;
use sqlx::types::Uuid;
use std::sync::Arc;
use std::time::Duration;
use tokio::sync::OnceCell;
/// Default idle life of a tenant pool, overridable via the
/// `userregistrytimeout` config key.
const DEFAULT_IDLE_SECS: u64 = 300;
/// How stale a tenant's quarantine status may be before it is re-read.
///
/// The status lives in the global directory and is written by the fleet
/// migrator in *another process*, so a warm pool here would otherwise keep
/// serving a database that has just been declared unusable — the registry is
/// per-process and nothing can invalidate it from outside. Re-reading on every
/// acquire would double the load on the admin pool (whose whole budget is 5);
/// re-reading on a short clock bounds the staleness instead, at one query per
/// tenant per interval.
const QUARANTINE_RECHECK: Duration = Duration::from_secs(30);
async fn acquire_connection(
registry: &PoolRegistry,
user_id: Uuid,
) -> Result<TenantConnection, DBError> {
let slot = match registry.touch(user_id)? {
Some(slot) => slot,
None => registry.install(user_id)?,
};
if registry.claim_recheck(user_id, QUARANTINE_RECHECK)? {
match routing(user_id).await {
// A tenant quarantined mid-flight must stop being served, so this
// drops the pool as well as refusing the request.
Ok(route) if route.quarantined || route.disabled => {
registry.forget(user_id, &slot)?;
return Err(refusal(&route));
}
Ok(_) => {}
Err(e) => {
// The claim already advanced the clock; give it back so the
// next caller retries rather than trusting a stale answer for
// the rest of the interval.
registry.release_recheck(user_id, QUARANTINE_RECHECK)?;
return Err(e);
// Cloning the handle pins the pool for the length of this request: an
// eviction between here and `acquire` only un-registers it, so the caller
// never races a close.
let tenant = match slot.get_or_try_init(|| build(registry, user_id)).await {
Ok(tenant) => Arc::clone(tenant),
let conn = tenant.pool().acquire().await.map_err(DBError::Sqlx)?;
// The pool travels with the connection: releasing its permit here, while
// the connection is still open, is what let the resident bound be exceeded.
Ok(TenantConnection::tenant(conn, tenant))
/// The directory stores a database name; the DSN is composed here.
/// The residency permit is taken before connecting and moved into the
/// [`TenantPool`], so a build that fails or is cancelled returns it by dropping
/// it rather than leaving the slot to be reclaimed later.
async fn build(registry: &PoolRegistry, user_id: Uuid) -> Result<Arc<TenantPool>, DBError> {
let permit = registry.reserve().await?;
let route = routing(user_id).await?;
// Also here, not only on recheck: a *fresh* pool would otherwise be built
// for a disabled user and served until the first recheck.
if route.quarantined || route.disabled {
let url = crate::provision::dsn_for(&route.database_name)?;
let pool = PgPoolOptions::new()
.max_connections(TENANT_MAX_CONNECTIONS)
.idle_timeout(registry.idle_timeout())
.acquire_timeout(Duration::from_secs(10))
.connect(&url)
.await
.map_err(DBError::Sqlx)?;
Ok(Arc::new(TenantPool::new(pool, permit)))
/// Where a tenant lives and whether it may currently be served.
struct Route {
database_name: String,
quarantined: bool,
disabled: bool,
/// Not the same event, so not the same user-facing text.
fn refusal(route: &Route) -> DBError {
if route.quarantined {
DBError::TenantQuarantined
} else {
DBError::TenantDisabled
async fn routing(user_id: Uuid) -> Result<Route, DBError> {
let mut conn = crate::db::get_connection().await?;
let row = sqlx::query_file!("../server/sql/select/system/tenant_routing.sql", &user_id)
.fetch_one(&mut *conn)
Ok(Route {
database_name: row.database_name,
quarantined: row.quarantined_at.is_some(),
disabled: !row.verified,
})
static POOL_REGISTRY: OnceCell<PoolRegistry> = OnceCell::const_new();
async fn pool_registry() -> &'static PoolRegistry {
POOL_REGISTRY
.get_or_init(|| async move {
let timeout = match system_config("userregistrytimeout").await {
Ok(Some(value)) => value.to_string().parse().unwrap_or(DEFAULT_IDLE_SECS),
_ => DEFAULT_IDLE_SECS,
PoolRegistry::new(Duration::from_secs(timeout))
pub(in crate::user) async fn get_connection(user_id: Uuid) -> Result<TenantConnection, DBError> {
acquire_connection(pool_registry().await, user_id).await