use std::collections::HashMap;
use std::sync::atomic::{AtomicU64, Ordering};
use std::sync::{Arc, PoisonError, Weak};
#[cfg(feature = "postgres")]
use crate::sql::sqlx::postgres::{PgPool, PgPoolOptions};
use crate::sql::sqlx::{self, Database};
use tokio::sync::RwLock;
use tracing::Instrument as _;
use super::error::TenancyError;
use super::org::{Org, StorageMode};
use super::secrets::{ChainSecretsResolver, SecretsResolver, TenantSecrets};
#[derive(Debug, Clone)]
pub struct TenantPoolsConfig {
pub max_cached_database_pools: usize,
pub database_pool_max_connections: u32,
pub database_pool_min_connections: u32,
pub database_pool_acquire_timeout: std::time::Duration,
pub database_pool_idle_timeout: Option<std::time::Duration>,
pub database_pool_max_lifetime: Option<std::time::Duration>,
pub prewarm_active_tenants: bool,
pub max_cached_scoped_pools: usize,
pub scoped_pool_max_connections: u32,
pub media_storage: Option<crate::storage::StorageRegistry>,
}
impl Default for TenantPoolsConfig {
fn default() -> Self {
Self {
max_cached_database_pools: 64,
database_pool_max_connections: 16,
database_pool_min_connections: 0,
database_pool_acquire_timeout: std::time::Duration::from_secs(30),
database_pool_idle_timeout: Some(std::time::Duration::from_secs(10 * 60)),
database_pool_max_lifetime: Some(std::time::Duration::from_secs(30 * 60)),
prewarm_active_tenants: false,
max_cached_scoped_pools: 64,
scoped_pool_max_connections: 8,
media_storage: None,
}
}
}
#[derive(Debug, Clone, Copy, Default)]
pub struct PrewarmReport {
pub total_active: usize,
pub warmed: usize,
pub failed: usize,
pub skipped_cap: usize,
}
#[cfg(feature = "postgres")]
pub type DefaultTenantDb = sqlx::Postgres;
#[cfg(all(not(feature = "postgres"), feature = "sqlite"))]
pub type DefaultTenantDb = sqlx::Sqlite;
#[cfg(all(not(feature = "postgres"), not(feature = "sqlite"), feature = "mysql"))]
pub type DefaultTenantDb = sqlx::MySql;
#[cfg_attr(not(any(feature = "postgres", feature = "manage")), allow(dead_code))]
pub(crate) trait RegistryConnect: Database {
fn connect_registry(
url: &str,
) -> impl std::future::Future<Output = Result<sqlx::Pool<Self>, crate::sql::PoolError>> + Send;
}
#[cfg(feature = "postgres")]
impl RegistryConnect for sqlx::Postgres {
fn connect_registry(
url: &str,
) -> impl std::future::Future<Output = Result<sqlx::Pool<Self>, crate::sql::PoolError>> + Send
{
crate::sql::Pool::connect_postgres(url)
}
}
#[cfg(feature = "sqlite")]
impl RegistryConnect for sqlx::Sqlite {
fn connect_registry(
url: &str,
) -> impl std::future::Future<Output = Result<sqlx::Pool<Self>, crate::sql::PoolError>> + Send
{
crate::sql::Pool::connect_sqlite(url)
}
}
#[cfg(feature = "mysql")]
impl RegistryConnect for sqlx::MySql {
fn connect_registry(
url: &str,
) -> impl std::future::Future<Output = Result<sqlx::Pool<Self>, crate::sql::PoolError>> + Send
{
crate::sql::Pool::connect_mysql(url)
}
}
#[derive(Debug)]
pub enum TenantPool<DB: Database = DefaultTenantDb> {
#[cfg(feature = "postgres")]
Schema { schema: String, registry: PgPool },
Database { pool: Arc<sqlx::Pool<DB>> },
}
impl<DB: Database> Clone for TenantPool<DB> {
fn clone(&self) -> Self {
match self {
#[cfg(feature = "postgres")]
Self::Schema { schema, registry } => Self::Schema {
schema: schema.clone(),
registry: registry.clone(),
},
Self::Database { pool } => Self::Database {
pool: Arc::clone(pool),
},
}
}
}
impl<DB: Database> TenantPool<DB> {
#[must_use]
pub fn is_schema(&self) -> bool {
#[cfg(feature = "postgres")]
{
matches!(self, Self::Schema { .. })
}
#[cfg(not(feature = "postgres"))]
{
false
}
}
}
#[cfg(feature = "postgres")]
impl TenantPool<sqlx::Postgres> {
#[must_use]
pub fn pool(&self) -> &PgPool {
match self {
Self::Schema { registry, .. } => registry,
Self::Database { pool } => pool,
}
}
}
pub(super) struct CachedPool<DB: Database> {
pub(super) pool: Arc<sqlx::Pool<DB>>,
source: PoolSource,
last_used: AtomicU64,
}
#[derive(Clone, PartialEq, Eq, Hash)]
pub(super) enum PoolSource {
Database(Option<String>),
#[cfg_attr(not(feature = "postgres"), allow(dead_code))]
Schema(String),
}
impl PoolSource {
pub(super) fn database(org: &Org) -> Self {
Self::Database(org.database_url.clone())
}
}
impl<DB: Database> CachedPool<DB> {
pub(super) fn new(pool: Arc<sqlx::Pool<DB>>, source: PoolSource, tick: u64) -> Self {
Self {
pool,
source,
last_used: AtomicU64::new(tick),
}
}
pub(super) fn lookup<'a>(
cache: &'a HashMap<String, Self>,
slug: &str,
source: &PoolSource,
) -> Option<&'a Self> {
cache.get(slug).filter(|entry| entry.source == *source)
}
pub(super) fn touch(&self, tick: u64) {
self.last_used.store(tick, Ordering::Relaxed);
}
}
pub(super) struct SingleFlight<T>(std::sync::Mutex<HashMap<FlightKey, Weak<Flight<T>>>>);
#[derive(Clone, PartialEq, Eq, Hash)]
struct FlightKey {
slug: String,
source: PoolSource,
}
type Flight<T> = tokio::sync::Mutex<Option<Result<T, TenancyError>>>;
impl<T> Default for SingleFlight<T> {
fn default() -> Self {
Self(std::sync::Mutex::default())
}
}
impl<T: Clone> SingleFlight<T> {
pub(super) async fn run(
&self,
slug: &str,
source: &PoolSource,
build: impl std::future::Future<Output = Result<T, TenancyError>>,
) -> Result<T, TenancyError> {
let key = FlightKey {
slug: slug.to_owned(),
source: source.clone(),
};
let flight = {
let mut map = self.map();
map.retain(|_, f| f.strong_count() > 0);
if let Some(f) = map.get(&key).and_then(Weak::upgrade) {
f
} else {
let f = Arc::new(tokio::sync::Mutex::new(None));
map.insert(key.clone(), Arc::downgrade(&f));
f
}
};
let mut done = flight.lock().await;
let copy = |o: &Result<T, TenancyError>| match o {
Ok(v) => Ok(v.clone()),
Err(e) => Err(shared_error(e)),
};
if let Some(outcome) = done.as_ref() {
return copy(outcome);
}
let out = build.await;
*done = Some(copy(&out));
let mut map = self.map();
if map
.get(&key)
.is_some_and(|f| f.as_ptr() == Arc::as_ptr(&flight))
{
map.remove(&key);
}
out
}
fn map(&self) -> std::sync::MutexGuard<'_, HashMap<FlightKey, Weak<Flight<T>>>> {
self.0.lock().unwrap_or_else(PoisonError::into_inner)
}
#[cfg(test)]
fn len(&self) -> usize {
self.map().len()
}
}
fn shared_error(e: &TenancyError) -> TenancyError {
use super::secrets::SecretsError;
match e {
TenancyError::Resolution(s) => TenancyError::Resolution(s.clone()),
TenancyError::Validation(s) => TenancyError::Validation(s.clone()),
TenancyError::Busy => TenancyError::Busy,
TenancyError::Secrets(SecretsError::NotFound(s)) => {
TenancyError::Secrets(SecretsError::NotFound(s.clone()))
}
TenancyError::Secrets(SecretsError::Invalid(s)) => {
TenancyError::Secrets(SecretsError::Invalid(s.clone()))
}
TenancyError::Secrets(SecretsError::Backend(b)) => {
TenancyError::Secrets(SecretsError::Backend(b.to_string().into()))
}
TenancyError::Driver(sqlx::Error::PoolTimedOut) => {
TenancyError::Driver(sqlx::Error::PoolTimedOut)
}
other => TenancyError::Driver(sqlx::Error::Io(std::io::Error::other(other.to_string()))),
}
}
pub(super) fn evict_to_fit<DB: Database>(
cache: &mut HashMap<String, CachedPool<DB>>,
cap: usize,
) -> Vec<String> {
let mut evicted = Vec::new();
while cap > 0 && cache.len() >= cap {
let Some(victim) = cache
.iter()
.min_by_key(|(_, entry)| entry.last_used.load(Ordering::Relaxed))
.map(|(slug, _)| slug.clone())
else {
break;
};
cache.remove(&victim);
evicted.push(victim);
}
evicted
}
pub struct TenantPools<DB: Database = DefaultTenantDb> {
registry: sqlx::Pool<DB>,
config: TenantPoolsConfig,
secrets: TenantSecrets,
lru_tick: AtomicU64,
warned_near_cap: std::sync::atomic::AtomicBool,
cache: RwLock<HashMap<String, CachedPool<DB>>>,
scoped_cache: RwLock<HashMap<String, CachedPool<DB>>>,
held_cache: RwLock<HashMap<String, CachedPool<DB>>>,
held_database_cache: RwLock<HashMap<String, CachedPool<DB>>>,
building: SingleFlight<Arc<sqlx::Pool<DB>>>,
#[cfg_attr(not(feature = "postgres"), allow(dead_code))]
building_held_database: SingleFlight<Arc<sqlx::Pool<DB>>>,
#[cfg_attr(not(feature = "postgres"), allow(dead_code))]
building_scoped: SingleFlight<sqlx::Pool<DB>>,
#[cfg_attr(not(feature = "postgres"), allow(dead_code))]
building_held: SingleFlight<sqlx::Pool<DB>>,
}
#[derive(Clone, Copy, PartialEq, Eq)]
enum PoolUse {
Queries,
#[cfg_attr(not(feature = "postgres"), allow(dead_code))]
Held,
}
impl<DB: Database> TenantPools<DB> {
#[must_use]
pub fn new(registry: sqlx::Pool<DB>) -> Self {
Self::with_secrets(registry, ChainSecretsResolver::standard())
}
#[must_use]
pub fn with_secrets<R: SecretsResolver>(registry: sqlx::Pool<DB>, secrets: R) -> Self {
Self {
registry,
config: TenantPoolsConfig::default(),
secrets: TenantSecrets::new(secrets),
lru_tick: AtomicU64::new(0),
warned_near_cap: std::sync::atomic::AtomicBool::new(false),
cache: RwLock::new(HashMap::new()),
scoped_cache: RwLock::new(HashMap::new()),
held_cache: RwLock::new(HashMap::new()),
held_database_cache: RwLock::new(HashMap::new()),
building: SingleFlight::default(),
building_held_database: SingleFlight::default(),
building_scoped: SingleFlight::default(),
building_held: SingleFlight::default(),
}
}
fn next_tick(&self) -> u64 {
self.lru_tick.fetch_add(1, Ordering::Relaxed)
}
fn warn_if_near_cap(&self, len: usize, cap: usize, mode: &'static str) {
if self.claim_near_cap_warning(len, cap) {
tracing::warn!(
target: "rustango::tenancy::pools",
mode,
cached = len,
cap,
"tenant pool cache is near its cap; past it, the least \
recently used tenant is evicted and reconnects on its \
next request",
);
}
}
fn claim_near_cap_warning(&self, len: usize, cap: usize) -> bool {
if cap == 0 {
return false;
}
if len * 5 >= cap * 4 {
!self.warned_near_cap.swap(true, Ordering::Relaxed)
} else {
self.warned_near_cap.store(false, Ordering::Relaxed);
false
}
}
#[must_use]
pub fn config(mut self, config: TenantPoolsConfig) -> Self {
self.config = config;
self
}
#[must_use]
pub fn pool_config(&self) -> &TenantPoolsConfig {
&self.config
}
#[must_use]
pub fn registry_inner(&self) -> &sqlx::Pool<DB> {
&self.registry
}
pub(crate) fn secrets(&self) -> &TenantSecrets {
&self.secrets
}
#[must_use]
pub fn require_sealed_secrets(mut self) -> Self {
self.secrets = self.secrets.require_sealed();
self
}
pub(crate) fn fresh(&self) -> Self {
let mut pools = Self::new(self.registry.clone()).config(self.config.clone());
pools.secrets = self.secrets.clone();
pools
}
}
#[cfg(feature = "postgres")]
impl TenantPools<sqlx::Postgres> {
#[must_use]
pub fn registry(&self) -> &PgPool {
&self.registry
}
}
impl<DB: Database> TenantPools<DB>
where
crate::sql::Pool: From<sqlx::Pool<DB>>,
{
#[must_use]
pub fn registry_pool(&self) -> crate::sql::Pool {
crate::sql::Pool::from(self.registry.clone())
}
}
impl<DB: Database> TenantPools<DB> {
pub async fn database_pool_for_org(&self, org: &Org) -> Result<TenantPool<DB>, TenancyError> {
let mode = StorageMode::parse(&org.storage_mode).map_err(|got| {
TenancyError::Validation(format!(
"org `{}` has unknown storage_mode `{got}` (expected `schema` or `database`)",
org.slug
))
})?;
match mode {
StorageMode::Schema => Err(TenancyError::Validation(format!(
"org `{}` has `storage_mode = 'schema'` but TenantPools<{dbname}> is non-Postgres. \
Schema-mode is a Postgres-only optimization (uses `SET search_path` — no \
equivalent on MySQL/SQLite). Switch this org to `storage_mode = 'database'` and \
set `database_url` to its dedicated database / file; isolation semantics are \
equivalent.",
org.slug,
dbname = std::any::type_name::<DB>(),
))),
StorageMode::Database => {
let pool = self.pool_for_database_mode(org).await?;
Ok(TenantPool::Database { pool })
}
}
}
pub async fn database_acquire(&self, org: &Org) -> Result<TenantConn<DB>, TenancyError> {
let pool = self.database_pool_for_org(org).await?;
#[cfg_attr(not(feature = "postgres"), allow(irrefutable_let_patterns))]
let TenantPool::Database { pool } = pool
else {
unreachable!("database_pool_for_org rejects schema-mode")
};
Ok(TenantConn::database(pool.acquire().await?))
}
pub async fn invalidate(&self, slug: &str) {
self.cache.write().await.remove(slug);
self.scoped_cache.write().await.remove(slug);
self.held_cache.write().await.remove(slug);
self.held_database_cache.write().await.remove(slug);
}
}
pub trait TenantPoolInvalidator: Send + Sync {
fn invalidate<'a>(
&'a self,
slug: &'a str,
) -> std::pin::Pin<Box<dyn std::future::Future<Output = ()> + Send + 'a>>;
fn resolved_database_url<'a>(
&'a self,
org: &'a super::org::Org,
) -> std::pin::Pin<
Box<
dyn std::future::Future<Output = Result<String, super::error::TenancyError>>
+ Send
+ 'a,
>,
>;
fn tenant_secrets(&self) -> TenantSecrets;
fn decommission<'a>(
&'a self,
slug: &'a str,
action: super::decommission::Action,
) -> std::pin::Pin<
Box<
dyn std::future::Future<
Output = Result<super::decommission::Report, super::error::TenancyError>,
> + Send
+ 'a,
>,
>;
fn prewarm<'a>(
&'a self,
) -> std::pin::Pin<
Box<
dyn std::future::Future<Output = Result<PrewarmReport, super::error::TenancyError>>
+ Send
+ 'a,
>,
>;
}
impl<DB: Database> TenantPoolInvalidator for TenantPools<DB>
where
crate::sql::Pool: From<sqlx::Pool<DB>>,
{
fn invalidate<'a>(
&'a self,
slug: &'a str,
) -> std::pin::Pin<Box<dyn std::future::Future<Output = ()> + Send + 'a>> {
Box::pin(async move { TenantPools::<DB>::invalidate(self, slug).await })
}
fn resolved_database_url<'a>(
&'a self,
org: &'a super::org::Org,
) -> std::pin::Pin<
Box<
dyn std::future::Future<Output = Result<String, super::error::TenancyError>>
+ Send
+ 'a,
>,
> {
Box::pin(async move { TenantPools::<DB>::resolved_database_url(self, org).await })
}
fn tenant_secrets(&self) -> TenantSecrets {
self.secrets.clone()
}
fn decommission<'a>(
&'a self,
slug: &'a str,
action: super::decommission::Action,
) -> std::pin::Pin<
Box<
dyn std::future::Future<
Output = Result<super::decommission::Report, super::error::TenancyError>,
> + Send
+ 'a,
>,
> {
Box::pin(async move { super::decommission::decommission(self, slug, action).await })
}
fn prewarm<'a>(
&'a self,
) -> std::pin::Pin<
Box<
dyn std::future::Future<Output = Result<PrewarmReport, super::error::TenancyError>>
+ Send
+ 'a,
>,
> {
Box::pin(async move { TenantPools::<DB>::prewarm_database_tenants(self).await })
}
}
#[allow(dead_code)]
impl<DB: Database> TenantPools<DB> {
#[must_use]
pub fn into_invalidator(self: Arc<Self>) -> Arc<dyn TenantPoolInvalidator>
where
crate::sql::Pool: From<sqlx::Pool<DB>>,
{
self
}
pub async fn resolved_database_url(&self, org: &Org) -> Result<String, TenancyError> {
let reference = org.database_url.as_deref().ok_or_else(|| {
TenancyError::Validation(format!(
"org `{}` has no `database_url` to resolve (schema mode?)",
org.slug
))
})?;
self.resolve_secret(reference).await
}
pub(crate) async fn resolve_secret(&self, reference: &str) -> Result<String, TenancyError> {
Ok(self.secrets.resolve(reference).await?)
}
pub async fn cached_database_pool_count(&self) -> usize {
self.cache.read().await.len()
}
pub async fn cached_scoped_pool_count(&self) -> usize {
self.scoped_cache.read().await.len()
}
fn database_cache_for(&self, of: PoolUse) -> &RwLock<HashMap<String, CachedPool<DB>>> {
match of {
PoolUse::Queries => &self.cache,
PoolUse::Held => &self.held_database_cache,
}
}
async fn cached_database_pool(
&self,
slug: &str,
source: &PoolSource,
of: PoolUse,
) -> Option<Arc<sqlx::Pool<DB>>> {
let cache = self.database_cache_for(of).read().await;
let entry = CachedPool::lookup(&cache, slug, source)?;
entry.touch(self.next_tick());
Some(Arc::clone(&entry.pool))
}
async fn pool_for_database_mode(&self, org: &Org) -> Result<Arc<sqlx::Pool<DB>>, TenancyError> {
self.database_mode_pool(org, PoolUse::Queries).await
}
async fn database_mode_pool(
&self,
org: &Org,
of: PoolUse,
) -> Result<Arc<sqlx::Pool<DB>>, TenancyError> {
let source = PoolSource::database(org);
if let Some(pool) = self.cached_database_pool(&org.slug, &source, of).await {
return Ok(pool);
}
let build = self.build_database_mode(org, source.clone(), of);
if self.config.max_cached_database_pools == 0 {
return build.await;
}
let flight = match of {
PoolUse::Queries => &self.building,
PoolUse::Held => &self.building_held_database,
};
flight.run(&org.slug, &source, build).await
}
async fn build_database_mode(
&self,
org: &Org,
source: PoolSource,
of: PoolUse,
) -> Result<Arc<sqlx::Pool<DB>>, TenancyError> {
if let Some(pool) = self.cached_database_pool(&org.slug, &source, of).await {
return Ok(pool);
}
let span = tracing::info_span!("tenant_pool_init", slug = %org.slug, mode = "database");
async {
let resolve_start = std::time::Instant::now();
let reference = org.database_url.as_deref().ok_or_else(|| {
TenancyError::Validation(format!(
"org `{}` is `storage_mode = database` but has no `database_url`",
org.slug
))
})?;
let url = self.secrets.resolve(reference).await?;
tracing::debug!(
target: "rustango::tenancy::pools",
slug = %org.slug,
elapsed_ms = resolve_start.elapsed().as_millis() as u64,
"secrets resolver resolved tenant URL",
);
let connect_start = std::time::Instant::now();
let tag = org.id.get().map(|id| tenant_session_tag(*id));
let url = tag_postgres_url(&url, tag.as_deref());
let min = match of {
PoolUse::Queries => self.config.database_pool_min_connections,
PoolUse::Held => 0,
};
let pool = build_database_pool::<DB>(&url, &self.config, min).await?;
tracing::info!(
target: "rustango::tenancy::pools",
slug = %org.slug,
elapsed_ms = connect_start.elapsed().as_millis() as u64,
min_conn = self.config.database_pool_min_connections,
max_conn = self.config.database_pool_max_connections,
"tenant pool connected (database mode)",
);
let pool = Arc::new(pool);
let mut cache = self.database_cache_for(of).write().await;
if let Some(existing) = CachedPool::lookup(&cache, &org.slug, &source) {
existing.touch(self.next_tick());
return Ok(Arc::clone(&existing.pool));
}
cache.remove(&org.slug);
for slug in evict_to_fit(&mut cache, self.config.max_cached_database_pools) {
tracing::info!(
target: "rustango::tenancy::pools",
evicted = %slug,
for_slug = %org.slug,
cap = self.config.max_cached_database_pools,
"evicted the least recently used tenant pool to make room",
);
}
if self.config.max_cached_database_pools == 0 {
return Ok(pool);
}
cache.insert(
org.slug.clone(),
CachedPool::new(Arc::clone(&pool), source, self.next_tick()),
);
if of == PoolUse::Queries {
self.warn_if_near_cap(
cache.len(),
self.config.max_cached_database_pools,
"database",
);
}
Ok(pool)
}
.instrument(span)
.await
}
}
impl<DB: Database> TenantPools<DB>
where
crate::sql::Pool: From<sqlx::Pool<DB>>,
{
pub async fn scoped_pool_dyn(&self, org: &Org) -> Result<crate::sql::Pool, TenancyError> {
let mode = StorageMode::parse(&org.storage_mode).map_err(|got| {
TenancyError::Validation(format!(
"org `{}` has unknown storage_mode `{got}` (expected `schema` or `database`)",
org.slug
))
})?;
match mode {
StorageMode::Schema => {
#[cfg(feature = "postgres")]
{
if let Some(pg_pools) =
(self as &dyn std::any::Any).downcast_ref::<TenantPools<sqlx::Postgres>>()
{
let scoped = pg_pools.scoped_pool(org).await?;
return Ok(crate::sql::Pool::Postgres(scoped));
}
}
Err(TenancyError::Validation(format!(
"org `{}` has `storage_mode = 'schema'` but TenantPools<{dbname}> is \
non-Postgres. Schema-mode is a Postgres-only optimization (uses \
`SET search_path` — no equivalent on MySQL/SQLite). Switch this org to \
`storage_mode = 'database'` and set `database_url` to its dedicated \
database / file; isolation semantics are equivalent.",
org.slug,
dbname = std::any::type_name::<DB>(),
)))
}
StorageMode::Database => {
let pool = self.pool_for_database_mode(org).await?;
Ok(crate::sql::Pool::from((*pool).clone()))
}
}
}
pub(crate) async fn existing_pool_dyn(
&self,
org: &Org,
) -> Result<(crate::sql::Pool, bool), TenancyError> {
if org.storage_mode == "database" && org.database_url.is_some() {
let url = self.resolved_database_url(org).await?;
if url.starts_with("sqlite:") {
let pool = crate::sql::Pool::connect(&sqlite_without_create(&url))
.await
.map_err(|e| TenancyError::Validation(e.to_string()))?;
return Ok((pool, true));
}
}
Ok((self.scoped_pool_dyn(org).await?, false))
}
pub async fn prewarm_database_tenants(&self) -> Result<PrewarmReport, TenancyError> {
use crate::core::Column as _;
use crate::sql::FetcherPool as _;
let span = tracing::info_span!("tenant_pools_prewarm");
async {
let started = std::time::Instant::now();
let registry_pool = self.registry_pool();
let orgs: Vec<Org> = Org::objects()
.where_(Org::storage_mode.eq("database".to_owned()))
.where_(Org::active.eq(true))
.fetch(®istry_pool)
.await?;
let total = orgs.len();
let mut report = PrewarmReport {
total_active: total,
warmed: 0,
failed: 0,
skipped_cap: 0,
};
for org in orgs {
if self.cache.read().await.len() >= self.config.max_cached_database_pools {
tracing::warn!(
target: "rustango::tenancy::pools",
slug = %org.slug,
cap = self.config.max_cached_database_pools,
"skipping pre-warm: cache cap reached",
);
report.skipped_cap += 1;
continue;
}
match self.pool_for_database_mode(&org).await {
Ok(_) => report.warmed += 1,
Err(e) => {
tracing::warn!(
target: "rustango::tenancy::pools",
slug = %org.slug,
error = %e,
"pre-warm failed for tenant",
);
report.failed += 1;
}
}
}
tracing::info!(
target: "rustango::tenancy::pools",
elapsed_ms = started.elapsed().as_millis() as u64,
total = report.total_active,
warmed = report.warmed,
failed = report.failed,
skipped_cap = report.skipped_cap,
"prewarm complete",
);
Ok(report)
}
.instrument(span)
.await
}
}
#[cfg(feature = "postgres")]
impl TenantPools<sqlx::Postgres> {
pub async fn pool_for_org(
&self,
org: &Org,
) -> Result<TenantPool<sqlx::Postgres>, TenancyError> {
let mode = StorageMode::parse(&org.storage_mode).map_err(|got| {
TenancyError::Validation(format!(
"org `{}` has unknown storage_mode `{got}` (expected `schema` or `database`)",
org.slug
))
})?;
match mode {
StorageMode::Schema => {
let schema = org.effective_schema().to_owned();
Ok(TenantPool::Schema {
schema,
registry: self.registry.clone(),
})
}
StorageMode::Database => {
let pool = self.pool_for_database_mode(org).await?;
Ok(TenantPool::Database { pool })
}
}
}
pub async fn acquire(&self, org: &Org) -> Result<TenantConn<sqlx::Postgres>, TenancyError> {
match self.pool_for_org(org).await? {
TenantPool::Schema { schema, registry } if self.config.max_cached_scoped_pools == 0 => {
let conn = registry.acquire().await?;
let mut tc = TenantConn {
inner: Some(conn),
schema: Some(schema.clone()),
reset: Some(reset_pg_registry_conn),
};
sqlx::Connection::clear_cached_statements(&mut **tc).await?;
let stmt = format!("SET search_path TO {}", search_path(&schema));
rustango::sql::sqlx::query(&stmt).execute(&mut **tc).await?;
Ok(tc)
}
TenantPool::Schema { schema, registry } => {
let pool = self
.schema_mode_pool(org, schema.clone(), registry, PoolUse::Held)
.await?;
Ok(TenantConn {
inner: Some(pool.acquire().await?),
schema: Some(schema),
reset: Some(reset_pg_search_path),
})
}
TenantPool::Database { pool } if self.config.max_cached_database_pools == 0 => {
Ok(TenantConn::database(pool.acquire().await?))
}
TenantPool::Database { .. } => {
let held = self.database_mode_pool(org, PoolUse::Held).await?;
Ok(TenantConn::database(held.acquire().await?))
}
}
}
pub async fn scoped_pool(&self, org: &Org) -> Result<PgPool, TenancyError> {
match self.pool_for_org(org).await? {
TenantPool::Schema { schema, registry } => {
self.schema_mode_pool(org, schema, registry, PoolUse::Queries)
.await
}
TenantPool::Database { pool } => Ok((*pool).clone()),
}
}
fn scoped_cache_for(
&self,
of: PoolUse,
) -> &RwLock<HashMap<String, CachedPool<sqlx::Postgres>>> {
match of {
PoolUse::Queries => &self.scoped_cache,
PoolUse::Held => &self.held_cache,
}
}
async fn schema_mode_pool(
&self,
org: &Org,
schema: String,
registry: PgPool,
of: PoolUse,
) -> Result<PgPool, TenancyError> {
let source = PoolSource::Schema(schema.clone());
if let Some(pool) = self.cached_scoped_pool(&org.slug, &source, of).await {
return Ok(pool);
}
let build = self.build_schema_mode(org, schema, registry, source.clone(), of);
if self.config.max_cached_scoped_pools == 0 {
return build.await;
}
let flight = match of {
PoolUse::Queries => &self.building_scoped,
PoolUse::Held => &self.building_held,
};
flight.run(&org.slug, &source, build).await
}
async fn build_schema_mode(
&self,
org: &Org,
schema: String,
registry: PgPool,
source: PoolSource,
of: PoolUse,
) -> Result<PgPool, TenancyError> {
if let Some(pool) = self.cached_scoped_pool(&org.slug, &source, of).await {
return Ok(pool);
}
let span = tracing::info_span!("tenant_pool_init", slug = %org.slug, mode = "schema");
async {
let connect_start = std::time::Instant::now();
let scoped = self.build_scoped_pool(&schema, ®istry).await?;
if of == PoolUse::Queries && self.config.max_cached_scoped_pools > 0 {
warn_on_public_fallthrough(®istry, &org.slug, &schema).await;
}
tracing::info!(
target: "rustango::tenancy::pools",
slug = %org.slug,
schema = %schema,
elapsed_ms = connect_start.elapsed().as_millis() as u64,
max_conn = self.config.scoped_pool_max_connections,
"tenant pool connected (schema mode)",
);
let mut cache = self.scoped_cache_for(of).write().await;
if let Some(existing) = CachedPool::lookup(&cache, &org.slug, &source) {
existing.touch(self.next_tick());
return Ok((*existing.pool).clone());
}
cache.remove(&org.slug);
for slug in evict_to_fit(&mut cache, self.config.max_cached_scoped_pools) {
tracing::info!(
target: "rustango::tenancy::pools",
evicted = %slug,
for_slug = %org.slug,
cap = self.config.max_cached_scoped_pools,
"evicted the least recently used scoped pool to make room",
);
}
if self.config.max_cached_scoped_pools == 0 {
return Ok(scoped);
}
cache.insert(
org.slug.clone(),
CachedPool::new(Arc::new(scoped.clone()), source, self.next_tick()),
);
self.warn_if_near_cap(cache.len(), self.config.max_cached_scoped_pools, "schema");
Ok(scoped)
}
.instrument(span)
.await
}
async fn cached_scoped_pool(
&self,
slug: &str,
source: &PoolSource,
of: PoolUse,
) -> Option<PgPool> {
let cache = self.scoped_cache_for(of).read().await;
let entry = CachedPool::lookup(&cache, slug, source)?;
entry.touch(self.next_tick());
Some((*entry.pool).clone())
}
async fn build_scoped_pool(
&self,
schema: &str,
registry: &PgPool,
) -> Result<PgPool, TenancyError> {
let mut opts = (*registry.connect_options()).clone();
opts = opts.options([("search_path", search_path_option(schema))]);
Ok(scoped_pool_options(&self.config).connect_with(opts).await?)
}
}
#[cfg(feature = "postgres")]
async fn warn_on_public_fallthrough(registry: &PgPool, slug: &str, schema: &str) {
match tables_read_from_public(registry, schema).await {
Ok(tables) if tables.is_empty() => {}
Ok(tables) => tracing::warn!(
target: "rustango::tenancy::pools",
slug,
schema,
tables = %tables.join(", "),
"the tenant's schema lacks tables that `public` has, so its queries read and \
write those shared copies; run `migrate-tenants` for it or deactivate it. \
A later release stops falling through to `public`",
),
Err(e) => tracing::debug!(
target: "rustango::tenancy::pools",
slug,
error = %e,
"could not compare the tenant's tables with `public`",
),
}
}
#[cfg(feature = "postgres")]
async fn tables_read_from_public(
registry: &PgPool,
schema: &str,
) -> Result<Vec<String>, sqlx::Error> {
let tenant: Vec<String> = inventory::iter::<crate::core::ModelEntry>
.into_iter()
.map(|e| e.schema)
.filter(|s| s.scope == crate::core::ModelScope::Tenant && s.managed && !s.is_view)
.map(|s| s.table.to_owned())
.collect();
sqlx::query_scalar(
"SELECT c.relname::text FROM pg_catalog.pg_class c \
JOIN pg_catalog.pg_namespace n ON n.oid = c.relnamespace \
WHERE n.nspname = 'public' AND c.relname = ANY($2) AND c.relkind IN ('r', 'p') \
AND NOT EXISTS (SELECT 1 FROM pg_catalog.pg_class t \
JOIN pg_catalog.pg_namespace tn ON tn.oid = t.relnamespace \
WHERE tn.nspname = $1 AND t.relname = c.relname) \
ORDER BY 1",
)
.bind(schema)
.bind(tenant)
.fetch_all(registry)
.await
}
#[cfg(feature = "postgres")]
pub(crate) async fn connect_schema_scoped(
registry_url: &str,
schema: &str,
max: u32,
) -> Result<PgPool, TenancyError> {
let stmt: Arc<str> = Arc::from(set_search_path_sql(schema));
let pool = PgPoolOptions::new()
.max_connections(max)
.after_connect(move |conn, _meta| {
let stmt = Arc::clone(&stmt);
Box::pin(async move {
sqlx::query(&stmt).execute(conn).await?;
Ok(())
})
})
.connect(registry_url)
.await?;
crate::sql::tag_schema_scope(&pool, schema);
Ok(pool)
}
#[cfg(feature = "postgres")]
fn set_search_path_sql(schema: &str) -> String {
format!(
"SET search_path TO {}, public",
crate::sql::quote_with(schema, '"')
)
}
#[cfg(feature = "postgres")]
pub(crate) fn scoped_pool_options(config: &TenantPoolsConfig) -> PgPoolOptions {
tenant_pool_options(config, config.scoped_pool_max_connections, 0)
}
pub(crate) fn tenant_pool_options<DB: Database>(
config: &TenantPoolsConfig,
max: u32,
min: u32,
) -> sqlx::pool::PoolOptions<DB> {
let mut opts = sqlx::pool::PoolOptions::<DB>::new()
.max_connections(max)
.min_connections(min)
.acquire_timeout(config.database_pool_acquire_timeout);
if let Some(idle) = config.database_pool_idle_timeout {
opts = opts.idle_timeout(idle);
}
if let Some(lifetime) = config.database_pool_max_lifetime {
opts = opts.max_lifetime(lifetime);
}
crate::sql::checkout_check(opts, &crate::sql::pool_tuning())
}
pub(crate) async fn build_database_pool<DB: Database>(
url: &str,
config: &TenantPoolsConfig,
min: u32,
) -> Result<sqlx::Pool<DB>, TenancyError> {
let opts = tenant_pool_options::<DB>(config, config.database_pool_max_connections, min);
let url = ensure_sqlite_creates(url);
Ok(opts.connect(&url).await?)
}
fn sqlite_without_create(url: &str) -> String {
let (base, query) = url.split_once('?').unwrap_or((url, ""));
let mut params: Vec<&str> = query
.split('&')
.filter(|p| !p.is_empty() && !p.starts_with("mode="))
.collect();
params.push("mode=rw");
format!("{base}?{}", params.join("&"))
}
#[cfg(test)]
mod sqlite_without_create_tests {
use super::sqlite_without_create;
#[test]
fn any_mode_becomes_rw() {
assert_eq!(
sqlite_without_create("sqlite://a.db"),
"sqlite://a.db?mode=rw"
);
assert_eq!(
sqlite_without_create("sqlite://a.db?mode=rwc&journal_mode=wal"),
"sqlite://a.db?journal_mode=wal&mode=rw"
);
}
}
pub(crate) fn tenant_session_tag(org_id: i64) -> String {
format!("rustango-tenant:{org_id}")
}
fn tag_postgres_url<'a>(url: &'a str, tag: Option<&str>) -> std::borrow::Cow<'a, str> {
let pg = url.starts_with("postgres:") || url.starts_with("postgresql:");
match tag {
Some(tag) if pg => {
let sep = if url.contains('?') { '&' } else { '?' };
std::borrow::Cow::Owned(format!("{url}{sep}application_name={tag}"))
}
_ => std::borrow::Cow::Borrowed(url),
}
}
fn ensure_sqlite_creates(url: &str) -> std::borrow::Cow<'_, str> {
if !(url.starts_with("sqlite:") || url.starts_with("sqlite://")) {
return std::borrow::Cow::Borrowed(url);
}
if url.contains(":memory:") {
return std::borrow::Cow::Borrowed(url);
}
if url.contains("?mode=") || url.contains("&mode=") {
return std::borrow::Cow::Borrowed(url);
}
let sep = if url.contains('?') { '&' } else { '?' };
std::borrow::Cow::Owned(format!("{url}{sep}mode=rwc"))
}
#[cfg(test)]
mod tag_postgres_url_tests {
use super::tag_postgres_url;
#[test]
fn tags_only_postgres_urls() {
assert_eq!(
tag_postgres_url("postgres://h/db?sslmode=require", Some("t")),
"postgres://h/db?sslmode=require&application_name=t"
);
assert_eq!(
tag_postgres_url("postgresql://h/db", Some("t")),
"postgresql://h/db?application_name=t"
);
assert_eq!(tag_postgres_url("mysql://h/db", Some("t")), "mysql://h/db");
assert_eq!(tag_postgres_url("postgres://h/db", None), "postgres://h/db");
}
}
#[cfg(test)]
mod ensure_sqlite_creates_tests {
use super::ensure_sqlite_creates;
#[test]
fn appends_mode_rwc_to_bare_sqlite_path() {
let out = ensure_sqlite_creates("sqlite:///var/data/x.db");
assert_eq!(&*out, "sqlite:///var/data/x.db?mode=rwc");
}
#[test]
fn appends_with_ampersand_when_other_params_present() {
let out = ensure_sqlite_creates("sqlite:///x.db?journal_mode=wal");
assert_eq!(&*out, "sqlite:///x.db?journal_mode=wal&mode=rwc");
}
#[test]
fn passes_explicit_mode_unchanged() {
let out = ensure_sqlite_creates("sqlite:///x.db?mode=ro");
assert_eq!(&*out, "sqlite:///x.db?mode=ro");
}
#[test]
fn passes_memory_url_unchanged() {
let out = ensure_sqlite_creates("sqlite::memory:");
assert_eq!(&*out, "sqlite::memory:");
}
#[test]
fn passes_non_sqlite_url_unchanged() {
let out = ensure_sqlite_creates("postgres://u@h/db");
assert_eq!(&*out, "postgres://u@h/db");
let out = ensure_sqlite_creates("mysql://u@h/db");
assert_eq!(&*out, "mysql://u@h/db");
}
}
pub struct TenantConn<DB: Database = DefaultTenantDb> {
inner: Option<sqlx::pool::PoolConnection<DB>>,
schema: Option<String>,
reset: Option<ResetFn<DB>>,
}
type ResetFn<DB> = fn(
sqlx::pool::PoolConnection<DB>,
) -> std::pin::Pin<Box<dyn std::future::Future<Output = ()> + Send>>;
impl<DB: Database> TenantConn<DB> {
pub(crate) fn database(conn: sqlx::pool::PoolConnection<DB>) -> Self {
Self {
inner: Some(conn),
schema: None,
reset: None,
}
}
#[must_use]
pub fn schema(&self) -> Option<&str> {
self.schema.as_deref()
}
}
impl<DB: Database> std::ops::Deref for TenantConn<DB> {
type Target = sqlx::pool::PoolConnection<DB>;
fn deref(&self) -> &Self::Target {
self.inner.as_ref().expect("connection taken only on drop")
}
}
impl<DB: Database> std::ops::DerefMut for TenantConn<DB> {
fn deref_mut(&mut self) -> &mut Self::Target {
self.inner.as_mut().expect("connection taken only on drop")
}
}
impl<DB: Database> Drop for TenantConn<DB> {
fn drop(&mut self) {
let (Some(reset), Some(mut conn)) = (self.reset, self.inner.take()) else {
return;
};
let Ok(handle) = tokio::runtime::Handle::try_current() else {
conn.close_on_drop();
return;
};
handle.spawn(reset(conn));
}
}
#[cfg(feature = "postgres")]
fn reset_pg_search_path(
mut conn: sqlx::pool::PoolConnection<sqlx::Postgres>,
) -> std::pin::Pin<Box<dyn std::future::Future<Output = ()> + Send>> {
Box::pin(async move {
match sqlx::query("RESET search_path").execute(&mut *conn).await {
Ok(_) => drop(conn),
Err(error) => {
tracing::warn!(
target: "rustango::tenancy::pools",
%error,
"could not reset search_path on release; closing the connection \
rather than returning it to its pool",
);
let _ = conn.close().await;
}
}
})
}
#[cfg(feature = "postgres")]
fn reset_pg_registry_conn(
mut conn: sqlx::pool::PoolConnection<sqlx::Postgres>,
) -> std::pin::Pin<Box<dyn std::future::Future<Output = ()> + Send>> {
Box::pin(async move {
match sqlx::Connection::clear_cached_statements(&mut *conn).await {
Ok(()) => reset_pg_search_path(conn).await,
Err(error) => {
tracing::warn!(
target: "rustango::tenancy::pools",
%error,
"could not drop a tenant's prepared statements on release; closing the \
connection rather than returning it to its pool",
);
let _ = conn.close().await;
}
}
})
}
#[cfg(feature = "postgres")]
fn search_path(schema: &str) -> String {
format!("{},public", crate::sql::quote_with(schema, '"'))
}
#[cfg(feature = "postgres")]
fn search_path_option(schema: &str) -> String {
let mut out = String::new();
for c in search_path(schema).chars() {
if c == '\\' || c.is_ascii_whitespace() {
out.push('\\');
}
out.push(c);
}
out
}
#[cfg(test)]
mod tests {
use super::*;
#[cfg(feature = "sqlite")]
#[tokio::test]
async fn a_sqlite_registry_gets_the_framework_setup() {
let dir = tempfile::tempdir().unwrap();
let file = dir.path().join("registry.db");
let url = format!("sqlite:{}", file.display());
let mut tuning = crate::sql::PoolTuning::default();
tuning.max_connections = Some(3);
crate::sql::configure_pools(tuning);
let pool = <sqlx::Sqlite as RegistryConnect>::connect_registry(&url)
.await
.expect("a missing registry file is created");
assert!(file.exists());
let (mode,): (String,) = sqlx::query_as("PRAGMA journal_mode")
.fetch_one(&pool)
.await
.unwrap();
assert_eq!(mode.to_ascii_lowercase(), "wal");
let want = crate::sql::pool_tuning().max_connections.expect("tuned");
assert_eq!(
pool.options().get_max_connections(),
want,
"[database] pool size"
);
}
#[cfg(feature = "sqlite")]
struct CountingSecrets(Arc<std::sync::atomic::AtomicUsize>);
#[cfg(feature = "sqlite")]
#[async_trait::async_trait]
impl SecretsResolver for CountingSecrets {
async fn resolve(&self, r: &str) -> Result<String, super::super::secrets::SecretsError> {
self.0.fetch_add(1, Ordering::SeqCst);
tokio::time::sleep(std::time::Duration::from_millis(20)).await;
if r == "fail" {
return Err(super::super::secrets::SecretsError::NotFound(r.to_owned()));
}
Ok(r.to_owned())
}
}
#[cfg(feature = "sqlite")]
struct BarrierSecrets(Arc<tokio::sync::Barrier>);
#[cfg(feature = "sqlite")]
#[async_trait::async_trait]
impl SecretsResolver for BarrierSecrets {
async fn resolve(&self, r: &str) -> Result<String, super::super::secrets::SecretsError> {
self.0.wait().await;
Ok(r.to_owned())
}
}
#[cfg(feature = "sqlite")]
fn failing_sqlite_org() -> Org {
Org {
database_url: Some("fail".into()),
..cold_sqlite_org()
}
}
#[cfg(feature = "sqlite")]
#[tokio::test]
async fn a_failed_cold_build_fails_every_waiter_once() {
let n = Arc::new(std::sync::atomic::AtomicUsize::new(0));
let pools: TenantPools<sqlx::Sqlite> = TenantPools::with_secrets(
sqlx::SqlitePool::connect_lazy("sqlite::memory:").expect("lazy"),
CountingSecrets(n.clone()),
);
let org = failing_sqlite_org();
let got =
futures_util::future::join_all((0..32).map(|_| pools.pool_for_database_mode(&org)))
.await;
assert!(got.iter().all(Result::is_err));
assert_eq!(
n.load(Ordering::SeqCst),
1,
"one build for 32 failing requests"
);
assert!(pools.pool_for_database_mode(&org).await.is_err());
assert_eq!(n.load(Ordering::SeqCst), 2);
}
#[cfg(feature = "sqlite")]
#[tokio::test]
async fn a_failed_cold_build_fails_every_waiter_once_in_database_pools() {
use super::super::{BackendKind, DatabasePools};
let n = Arc::new(std::sync::atomic::AtomicUsize::new(0));
let pools: DatabasePools<sqlx::Sqlite> =
DatabasePools::with_secrets(BackendKind::Sqlite, CountingSecrets(n.clone()));
let org = failing_sqlite_org();
let got = futures_util::future::join_all((0..32).map(|_| pools.pool_for_org(&org))).await;
assert!(got.iter().all(Result::is_err));
assert_eq!(
n.load(Ordering::SeqCst),
1,
"one build for 32 failing requests"
);
}
#[cfg(feature = "sqlite")]
#[tokio::test]
async fn with_caching_off_cold_builds_run_in_parallel() {
use super::super::{BackendKind, DatabasePools};
let config = TenantPoolsConfig {
max_cached_database_pools: 0,
..TenantPoolsConfig::default()
};
let org = cold_sqlite_org();
let tenant = TenantPools::<sqlx::Sqlite>::with_secrets(
sqlx::SqlitePool::connect_lazy("sqlite::memory:").expect("lazy"),
BarrierSecrets(Arc::new(tokio::sync::Barrier::new(4))),
)
.config(config.clone());
let db = DatabasePools::<sqlx::Sqlite>::with_secrets(
BackendKind::Sqlite,
BarrierSecrets(Arc::new(tokio::sync::Barrier::new(4))),
)
.config(config);
let both = async {
futures_util::future::join_all((0..4).map(|_| tenant.pool_for_database_mode(&org)))
.await;
futures_util::future::join_all((0..4).map(|_| db.pool_for_org(&org))).await;
};
tokio::time::timeout(std::time::Duration::from_secs(5), both)
.await
.expect("cap 0 builds must not wait for each other");
}
#[tokio::test]
async fn abandoned_flights_are_pruned() {
let flights = SingleFlight::<u32>::default();
let source = PoolSource::Database(None);
for slug in ["a", "b", "c", "d"] {
let pending = flights.run(slug, &source, std::future::pending());
assert!(
tokio::time::timeout(std::time::Duration::from_millis(5), pending)
.await
.is_err()
);
}
assert_eq!(
flights.run("e", &source, async { Ok(1) }).await.ok(),
Some(1)
);
assert_eq!(flights.len(), 0, "dead and finished flights are dropped");
}
#[cfg(feature = "sqlite")]
fn cold_sqlite_org() -> Org {
Org {
slug: "cold".into(),
storage_mode: "database".into(),
backend_kind: "sqlite".into(),
database_url: Some("sqlite::memory:".into()),
..crate::testkit::org()
}
}
#[cfg(feature = "sqlite")]
#[tokio::test]
async fn a_cold_tenant_builds_its_pool_once() {
let n = Arc::new(std::sync::atomic::AtomicUsize::new(0));
let pools: TenantPools<sqlx::Sqlite> = TenantPools::with_secrets(
sqlx::SqlitePool::connect_lazy("sqlite::memory:").expect("lazy"),
CountingSecrets(n.clone()),
);
let org = cold_sqlite_org();
let got =
futures_util::future::join_all((0..32).map(|_| pools.pool_for_database_mode(&org)))
.await;
assert!(got.iter().all(Result::is_ok));
assert_eq!(
n.load(Ordering::SeqCst),
1,
"one build for 32 cold requests"
);
}
#[cfg(feature = "sqlite")]
#[tokio::test]
async fn a_cold_tenant_builds_its_pool_once_in_database_pools() {
use super::super::{BackendKind, DatabasePools};
let n = Arc::new(std::sync::atomic::AtomicUsize::new(0));
let pools: DatabasePools<sqlx::Sqlite> =
DatabasePools::with_secrets(BackendKind::Sqlite, CountingSecrets(n.clone()));
let org = cold_sqlite_org();
let got = futures_util::future::join_all((0..32).map(|_| pools.pool_for_org(&org))).await;
assert!(got.iter().all(Result::is_ok));
assert_eq!(
n.load(Ordering::SeqCst),
1,
"one build for 32 cold requests"
);
}
#[cfg(feature = "sqlite")]
struct SlowOldSecrets;
#[cfg(feature = "sqlite")]
#[async_trait::async_trait]
impl SecretsResolver for SlowOldSecrets {
async fn resolve(&self, r: &str) -> Result<String, super::super::secrets::SecretsError> {
if r.contains("old") {
tokio::time::sleep(std::time::Duration::from_millis(300)).await;
}
Ok(r.to_owned())
}
}
#[cfg(feature = "sqlite")]
fn moved_sqlite_orgs(dir: &std::path::Path) -> (Org, Org, String, String) {
let file = |name: &str| dir.join(name).to_string_lossy().into_owned();
let (old, new) = (file("old.db"), file("new.db"));
let org = |path: &str| Org {
slug: "moved".into(),
database_url: Some(format!("sqlite:{path}?mode=rwc")),
..cold_sqlite_org()
};
(org(&old), org(&new), old, new)
}
#[cfg(feature = "sqlite")]
fn sqlite_file(pool: &sqlx::SqlitePool) -> String {
pool.connect_options()
.get_filename()
.to_string_lossy()
.into_owned()
}
#[cfg(feature = "sqlite")]
#[tokio::test]
async fn a_build_for_another_source_is_not_shared() {
let dir = tempfile::tempdir().unwrap();
let (old_org, new_org, old, new) = moved_sqlite_orgs(dir.path());
let pools: TenantPools<sqlx::Sqlite> = TenantPools::with_secrets(
sqlx::SqlitePool::connect_lazy("sqlite::memory:").expect("lazy"),
SlowOldSecrets,
);
let late_new = async {
tokio::time::sleep(std::time::Duration::from_millis(50)).await;
pools.pool_for_database_mode(&new_org).await
};
let (a, b) = tokio::join!(pools.pool_for_database_mode(&old_org), late_new);
assert_eq!(sqlite_file(&a.unwrap()), old);
assert_eq!(sqlite_file(&b.unwrap()), new, "new org got the old pool");
}
#[cfg(feature = "sqlite")]
#[tokio::test]
async fn a_build_for_another_source_is_not_shared_in_database_pools() {
use super::super::{BackendKind, DatabasePools};
let dir = tempfile::tempdir().unwrap();
let (old_org, new_org, old, new) = moved_sqlite_orgs(dir.path());
let pools: DatabasePools<sqlx::Sqlite> =
DatabasePools::with_secrets(BackendKind::Sqlite, SlowOldSecrets);
let late_new = async {
tokio::time::sleep(std::time::Duration::from_millis(50)).await;
pools.pool_for_org(&new_org).await
};
let (a, b) = tokio::join!(pools.pool_for_org(&old_org), late_new);
assert_eq!(sqlite_file(a.unwrap().pool()), old);
assert_eq!(
sqlite_file(b.unwrap().pool()),
new,
"new org got the old pool"
);
}
#[cfg(feature = "postgres")]
#[tokio::test]
async fn a_schema_build_for_another_schema_is_not_shared() {
let (addr, accepted, server) = silent_pg_server().await;
let pools = Arc::new(TenantPools::new(lazy_pg(addr)));
let calls: Vec<_> = ["cold", "cold_v2"]
.into_iter()
.map(|schema| {
let org = Org {
schema_name: Some(schema.into()),
..cold_schema_org()
};
let p = pools.clone();
tokio::spawn(async move { p.scoped_pool(&org).await })
})
.collect();
let deadline = std::time::Instant::now() + std::time::Duration::from_secs(10);
while accepted.load(Ordering::SeqCst) < 2 && std::time::Instant::now() < deadline {
tokio::time::sleep(std::time::Duration::from_millis(10)).await;
}
let seen = accepted.load(Ordering::SeqCst);
for c in calls {
c.abort();
}
server.abort();
assert_eq!(seen, 2, "two schemas opened {seen} connections");
}
#[cfg(feature = "postgres")]
#[tokio::test]
async fn a_cold_schema_tenant_connects_once() {
let (addr, accepted, server) = silent_pg_server().await;
let pools = Arc::new(TenantPools::new(lazy_pg(addr)));
let org = cold_schema_org();
let calls: Vec<_> = (0..16)
.map(|_| {
let (p, o) = (pools.clone(), org.clone());
tokio::spawn(async move { p.scoped_pool(&o).await })
})
.collect();
let deadline = std::time::Instant::now() + std::time::Duration::from_secs(10);
while accepted.load(Ordering::SeqCst) == 0 && std::time::Instant::now() < deadline {
tokio::time::sleep(std::time::Duration::from_millis(10)).await;
}
tokio::time::sleep(std::time::Duration::from_millis(200)).await;
let seen = accepted.load(Ordering::SeqCst);
for c in calls {
c.abort();
}
server.abort();
assert_eq!(seen, 1, "16 cold requests opened {seen} connections");
}
#[cfg(feature = "postgres")]
#[tokio::test]
async fn a_timed_out_schema_build_fails_every_waiter_together() {
let (addr, _accepted, server) = silent_pg_server().await;
let timeout = std::time::Duration::from_millis(300);
let pools = Arc::new(TenantPools::new(lazy_pg(addr)).config(TenantPoolsConfig {
database_pool_acquire_timeout: timeout,
..TenantPoolsConfig::default()
}));
let org = cold_schema_org();
let start = std::time::Instant::now();
let calls = (0..8).map(|_| {
let (p, o) = (pools.clone(), org.clone());
tokio::spawn(async move { p.scoped_pool(&o).await })
});
let got = futures_util::future::join_all(calls).await;
let took = start.elapsed();
server.abort();
assert!(got.into_iter().all(|r| r.expect("task").is_err()));
assert!(took < timeout * 3, "8 waiters took {took:?}");
}
#[cfg(feature = "postgres")]
async fn silent_pg_server() -> (
std::net::SocketAddr,
Arc<std::sync::atomic::AtomicUsize>,
tokio::task::JoinHandle<()>,
) {
let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
let addr = listener.local_addr().unwrap();
let accepted = Arc::new(std::sync::atomic::AtomicUsize::new(0));
let n = accepted.clone();
let server = tokio::spawn(async move {
let mut held = Vec::new();
while let Ok((sock, _)) = listener.accept().await {
n.fetch_add(1, Ordering::SeqCst);
held.push(sock);
}
});
(addr, accepted, server)
}
#[cfg(feature = "postgres")]
fn lazy_pg(addr: std::net::SocketAddr) -> PgPool {
PgPoolOptions::new()
.connect_lazy(&format!("postgres://u:p@{addr}/db"))
.unwrap()
}
#[cfg(feature = "postgres")]
fn cold_schema_org() -> Org {
Org {
slug: "cold".into(),
storage_mode: "schema".into(),
schema_name: Some("cold".into()),
..crate::testkit::org()
}
}
#[cfg(feature = "sqlite")]
#[tokio::test]
async fn new_resolves_env_references() {
let pools: TenantPools<sqlx::Sqlite> =
TenantPools::new(sqlx::SqlitePool::connect_lazy("sqlite::memory:").expect("lazy"));
let path = std::env::var("PATH").expect("PATH");
assert_eq!(pools.resolve_secret("env://PATH").await.unwrap(), path);
}
#[cfg(feature = "sqlite")]
#[tokio::test]
async fn the_near_cap_warning_latches_and_resets() {
use std::sync::atomic::Ordering;
let pools: TenantPools<sqlx::Sqlite> =
TenantPools::new(sqlx::SqlitePool::connect_lazy("sqlite::memory:").expect("lazy"));
assert!(!pools.claim_near_cap_warning(7, 10), "70% is not near");
assert!(pools.claim_near_cap_warning(8, 10), "80% warns");
assert!(!pools.claim_near_cap_warning(9, 10), "no repeat at 90%");
assert!(!pools.claim_near_cap_warning(10, 10), "no repeat at cap");
assert!(!pools.claim_near_cap_warning(10, 10), "still no repeat");
assert!(!pools.claim_near_cap_warning(3, 10), "back under: silent");
assert!(pools.claim_near_cap_warning(9, 10), "second crossing warns");
pools.warned_near_cap.store(false, Ordering::Relaxed);
assert!(!pools.claim_near_cap_warning(0, 0), "cap 0 must not warn");
}
#[cfg(feature = "sqlite")]
#[tokio::test]
async fn the_pool_init_span_does_not_leak_onto_other_tasks() {
struct Parked(Arc<tokio::sync::Notify>);
#[async_trait::async_trait]
impl SecretsResolver for Parked {
async fn resolve(
&self,
r: &str,
) -> Result<String, super::super::secrets::SecretsError> {
self.0.notified().await;
Err(super::super::secrets::SecretsError::NotFound(r.to_owned()))
}
}
let buf = crate::testkit::CaptureWriter::default();
let subscriber = tracing_subscriber::fmt()
.with_writer(buf.clone())
.with_ansi(false)
.finish();
let _g = tracing::subscriber::set_default(subscriber);
let gate = Arc::new(tokio::sync::Notify::new());
let pools: TenantPools<sqlx::Sqlite> = TenantPools::with_secrets(
sqlx::SqlitePool::connect_lazy("sqlite::memory:").expect("lazy"),
Parked(gate.clone()),
);
let org = Org {
slug: "leaky".into(),
storage_mode: "database".into(),
database_url: Some("ref".into()),
..crate::testkit::org()
};
let (res, ()) = tokio::join!(pools.pool_for_database_mode(&org), async {
tokio::task::yield_now().await;
tracing::info!("other-task-event");
gate.notify_one();
});
assert!(res.is_err());
let out = buf.contents();
let line = out
.lines()
.find(|l| l.contains("other-task-event"))
.expect("logged");
assert!(!line.contains("tenant_pool_init"), "span leaked: {line}");
}
#[cfg(feature = "postgres")]
#[test]
fn set_search_path_quotes_odd_schema_names() {
assert_eq!(
set_search_path_sql("a\"b c\\d\x0Be"),
"SET search_path TO \"a\"\"b c\\d\x0Be\", public"
);
}
#[test]
fn config_defaults_are_sane() {
let c = TenantPoolsConfig::default();
assert!(c.max_cached_database_pools >= 16);
assert!(c.database_pool_max_connections >= 1);
}
#[test]
fn config_pool_timeout_defaults_preserve_pre_0_27_7_behavior() {
let c = TenantPoolsConfig::default();
assert!(!c.prewarm_active_tenants);
assert_eq!(c.database_pool_min_connections, 0);
assert!(c.database_pool_acquire_timeout >= std::time::Duration::from_secs(5));
assert!(c.database_pool_idle_timeout.is_some());
assert!(c.database_pool_max_lifetime.is_some());
}
#[cfg(feature = "postgres")]
#[test]
fn prewarm_report_zeroed_default() {
let r = PrewarmReport::default();
assert_eq!(r.total_active, 0);
assert_eq!(r.warmed, 0);
assert_eq!(r.failed, 0);
assert_eq!(r.skipped_cap, 0);
}
}