Lines
100 %
Functions
81.82 %
Branches
use crate::db::DBError;
use sqlx::types::Uuid;
use sqlx::{Pool, Postgres};
use std::collections::HashMap;
use std::sync::{Arc, Mutex, MutexGuard};
use std::time::{Duration, Instant};
use tokio::sync::{OnceCell, OwnedSemaphorePermit, Semaphore};
#[cfg(test)]
mod tests;
/// Connections a single tenant pool may open.
pub(super) const TENANT_MAX_CONNECTIONS: u32 = 2;
/// Tenant pools resident in one process.
pub(super) const MAX_RESIDENT_POOLS: usize = 8;
/// How long a newcomer waits for a residency permit before giving up. Matches
/// the pool's own `acquire_timeout`, so a caller cannot wait appreciably longer
/// for a slot than for a connection once it has one.
const RESERVE_DEADLINE: Duration = Duration::from_secs(10);
/// How long each wait attempt runs before looking for a victim again.
const RESERVE_RETRY: Duration = Duration::from_millis(250);
/// A tenant pool and the residency permit that pays for it.
///
/// The permit lives *here*, not in the registry entry, and that placement is
/// what makes the budget hold under failure. The pool and its permit are
/// created and dropped as one value, so every path that discards a pool —
/// eviction, expiry, a failed build, a cancelled request — returns the permit
/// automatically. Nothing has to remember to.
/// It also decouples eviction from use: a caller holds an `Arc` of this while
/// working, so evicting the entry merely un-registers the pool. It stays alive,
/// and its permit stays spent, until the last user lets go.
pub(crate) struct TenantPool {
pool: Pool<Postgres>,
_permit: OwnedSemaphorePermit,
}
impl TenantPool {
pub(crate) const fn new(pool: Pool<Postgres>, permit: OwnedSemaphorePermit) -> Self {
Self {
pool,
_permit: permit,
pub(crate) const fn pool(&self) -> &Pool<Postgres> {
&self.pool
/// One user's slot. The cell is empty while a build is in flight; `OnceCell`
/// serializes creation per user, so concurrent first requests await one
/// initialization instead of each building a pool and overwriting the rest.
pub(super) type Slot = Arc<OnceCell<Arc<TenantPool>>>;
struct Entry {
cell: Slot,
last_used: Instant,
/// When this tenant's quarantine status was last read from the directory.
checked_at: Instant,
pub(super) struct PoolRegistry {
pools: Arc<Mutex<HashMap<Uuid, Entry>>>,
residency: Arc<Semaphore>,
expiration: Duration,
reserve_deadline: Duration,
/// Removes pools idle for longer than `expiration`. Dropping the entry drops
/// the pool and returns its permit; nothing else is required.
fn sweep_expired(pools: &mut HashMap<Uuid, Entry>, expiration: Duration) {
// `last_used` is stamped on acquire, so a connection held longer than the
// idle timeout looks idle to the clock while it is very much in use: only
// drop what nobody holds. An ABANDONED slot — never built, nobody
// building — is exactly such a thing, and requiring an initialised pool to
// sweep left one behind for every request cancelled mid-build, forever.
pools.retain(|_, e| e.last_used.elapsed() < expiration || !unheld(e));
/// Whether nobody holds this entry: no caller has the slot, and no connection
/// is out on the pool if one was ever built.
fn unheld(entry: &Entry) -> bool {
Arc::strong_count(&entry.cell) == 1
&& entry
.cell
.get()
.is_none_or(|pool| Arc::strong_count(pool) == 1)
/// Removes the least-recently-used slot so a newcomer can take its permit.
/// Returns whether anything was actually evicted — a registry holding only
/// in-flight builds has no permit to give back yet.
fn evict_lru(pools: &mut HashMap<Uuid, Entry>) -> bool {
let Some(victim) = pools
.iter()
.filter(|(_, e)| is_idle(e))
.min_by_key(|(_, e)| e.last_used)
.map(|(id, _)| *id)
else {
return false;
};
pools.remove(&victim).is_some()
/// Whether a pool can be dropped right now: nobody is using it, and nobody is
/// on their way to using it.
/// The permit travels with the pool and comes back only when its last
/// [`TenantConnection`] drops, so dropping a BUSY pool frees nothing — the
/// newcomer waits out the whole reserve deadline and fails, while the victim
/// rebuilds a second pool beside its still-live one and takes the tenant over
/// its own connection cap.
/// A slot still building is never idle: its permit is
/// taken inside the build, so dropping the slot frees nothing and only strands
/// the builder.
/// BOTH counts matter. A connection anchors the pool itself, but a caller that
/// has taken the slot and not yet reached the pool — it may be in `routing()`,
/// an admin-database round trip — holds only the slot, and evicting under it
/// puts that caller on an unregistered pool while its replacement is built.
fn is_idle(entry: &Entry) -> bool {
entry.cell.get().is_some() && unheld(entry)
impl PoolRegistry {
pub(super) fn new(expiration: Duration) -> Self {
pools: Arc::new(Mutex::new(HashMap::new())),
residency: Arc::new(Semaphore::new(MAX_RESIDENT_POOLS)),
expiration,
reserve_deadline: RESERVE_DEADLINE,
pub(super) const fn idle_timeout(&self) -> Duration {
self.expiration
fn lock(&self) -> Result<MutexGuard<'_, HashMap<Uuid, Entry>>, DBError> {
self.pools
.lock()
.map_err(|_| DBError::Config(crate::config::ConfigError::DB))
/// Marks the user's slot as used now and hands it back, if resident.
pub(super) fn touch(&self, user_id: Uuid) -> Result<Option<Slot>, DBError> {
let mut pools = self.lock()?;
sweep_expired(&mut pools, self.expiration);
Ok(pools.get_mut(&user_id).map(|entry| {
entry.last_used = Instant::now();
Arc::clone(&entry.cell)
}))
/// Reserves a residency permit, evicting the least-recently-used pool when
/// the process is already at its ceiling.
/// Retried on every pass, which is safe only because `is_idle` refuses to
/// take a pool anybody holds: a newcomer can therefore reclaim a resident
/// that fell idle while it waited, instead of giving up on the one victim
/// it found at the start and timing out beside an idle pool.
pub(super) async fn reserve(&self) -> Result<OwnedSemaphorePermit, DBError> {
let deadline = Instant::now() + self.reserve_deadline;
loop {
if let Ok(permit) = Arc::clone(&self.residency).try_acquire_owned() {
return Ok(permit);
{
evict_lru(&mut pools);
let remaining = deadline.saturating_duration_since(Instant::now());
if remaining.is_zero() {
return Err(DBError::PoolBudgetExhausted);
let wait = Arc::clone(&self.residency).acquire_owned();
if let Ok(Ok(permit)) = tokio::time::timeout(remaining.min(RESERVE_RETRY), wait).await {
/// Inserts an empty slot for `user_id`, or yields the existing one if
/// another task won the race.
pub(super) fn install(&self, user_id: Uuid) -> Result<Slot, DBError> {
if let Some(existing) = pools.get_mut(&user_id) {
existing.last_used = Instant::now();
return Ok(Arc::clone(&existing.cell));
let cell = Slot::default();
pools.insert(
user_id,
Entry {
cell: Arc::clone(&cell),
last_used: Instant::now(),
checked_at: Instant::now(),
},
);
Ok(cell)
/// Claims the right to re-read this tenant's quarantine status.
/// Returns true to exactly one caller per interval. A plain "is it stale?"
/// check would let every concurrent request for a tenant observe the same
/// stale timestamp and each issue its own directory query — a stampede
/// against an admin pool whose entire budget is 5. Claiming advances the
/// clock under the lock, so the losers proceed on the previous answer.
/// An absent entry answers true: that tenant is about to be built, and the
/// build reads the directory anyway.
pub(super) fn claim_recheck(&self, user_id: Uuid, interval: Duration) -> Result<bool, DBError> {
let Some(entry) = pools.get_mut(&user_id) else {
return Ok(true);
if entry.checked_at.elapsed() < interval {
return Ok(false);
entry.checked_at = Instant::now();
Ok(true)
/// Rolls back a claimed recheck whose read failed, so the next caller
/// retries instead of waiting out the interval on an unknown status.
pub(super) fn release_recheck(&self, user_id: Uuid, interval: Duration) -> Result<(), DBError> {
if let Some(entry) = self.lock()?.get_mut(&user_id) {
entry.checked_at = Instant::now()
.checked_sub(interval)
.unwrap_or_else(Instant::now);
Ok(())
/// Drops a slot whose build failed — but only if it is still *this* slot.
/// Without the identity check a caller whose build failed would evict
/// whatever entry happens to sit under that user id, including a fresh one
/// another task installed and populated in the meantime.
pub(super) fn forget(&self, user_id: Uuid, slot: &Slot) -> Result<(), DBError> {
if pools
.get(&user_id)
.is_some_and(|entry| Arc::ptr_eq(&entry.cell, slot))
pools.remove(&user_id);