use std::collections::HashMap;
use std::fmt;
use std::ops::Deref;
use std::sync::{
Arc, Mutex, RwLock,
atomic::{AtomicU64, Ordering},
};
use std::time::{SystemTime, UNIX_EPOCH};
use serde::Serialize;
use sqlx::PgPool;
#[cfg(feature = "clickhouse")]
use crate::runtime::executors::clickhouse::ClickHouseExecutor;
#[cfg(feature = "mongodb")]
use crate::runtime::executors::mongodb::MongoDbExecutor;
#[cfg(feature = "neo4j")]
use crate::runtime::executors::neo4j::Neo4jExecutor;
#[cfg(feature = "qdrant")]
use crate::runtime::executors::qdrant::QdrantHttpClient;
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize)]
pub enum ClientState {
Warming,
Ready,
Draining,
Closed,
Unhealthy,
}
impl ClientState {
pub fn as_str(self) -> &'static str {
match self {
Self::Warming => "warming",
Self::Ready => "ready",
Self::Draining => "draining",
Self::Closed => "closed",
Self::Unhealthy => "unhealthy",
}
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize)]
pub enum PoolingMode {
Session,
Transaction,
Statement,
}
impl PoolingMode {
pub fn from_labels(labels: &HashMap<String, String>) -> Self {
labels
.get("pooling_mode")
.or_else(|| labels.get("pooling"))
.or_else(|| labels.get("pool_mode"))
.map(|value| Self::from_value(value))
.unwrap_or(Self::Session)
}
pub fn from_value(value: &str) -> Self {
match value.trim().to_ascii_lowercase().as_str() {
"transaction" | "transaction_pool" | "transaction_pooling" => Self::Transaction,
"statement" | "statement_pool" | "statement_pooling" | "pipelined" => Self::Statement,
_ => Self::Session,
}
}
pub fn as_str(self) -> &'static str {
match self {
Self::Session => "session",
Self::Transaction => "transaction",
Self::Statement => "statement",
}
}
}
#[derive(Debug, Clone, PartialEq, Eq, Hash, Serialize)]
pub struct ClientKey {
pub project_id: String,
pub backend: String,
pub instance: String,
pub role: String,
}
impl ClientKey {
pub fn new(project_id: &str, backend: &str, instance: &str, role: &str) -> Self {
Self {
project_id: normalize_part(project_id, "default"),
backend: normalize_part(backend, "unknown"),
instance: normalize_part(instance, "default"),
role: normalize_part(role, "read_write"),
}
}
}
#[derive(Debug, Clone, Serialize)]
pub struct ClientSnapshot {
pub project_id: String,
pub backend: String,
pub instance: String,
pub role: String,
pub state: String,
pub pooling_mode: String,
pub created_at_unix_ms: u64,
pub last_health_probe_unix_ms: u64,
pub active_leases: u64,
pub total_leases: u64,
pub error_count: u64,
pub last_error: Option<String>,
pub lease_budget: Option<u64>,
pub labels: HashMap<String, String>,
pub pg_active_connections: Option<u32>,
pub pg_idle_connections: Option<u32>,
}
#[derive(Clone)]
pub struct ClientLease<T> {
value: T,
active: Arc<AtomicU64>,
}
impl<T> ClientLease<T> {
fn new(value: T, active: Arc<AtomicU64>, total: Arc<AtomicU64>) -> Self {
active.fetch_add(1, Ordering::Relaxed);
total.fetch_add(1, Ordering::Relaxed);
Self { value, active }
}
pub fn into_inner(self) -> T
where
T: Clone,
{
self.value.clone()
}
}
impl<T> Deref for ClientLease<T> {
type Target = T;
fn deref(&self) -> &Self::Target {
&self.value
}
}
impl<T> Drop for ClientLease<T> {
fn drop(&mut self) {
self.active.fetch_sub(1, Ordering::Relaxed);
}
}
#[derive(Clone)]
enum ManagedClientKind {
Postgres(PgPool),
#[cfg(feature = "redis")]
Redis(redis::Client),
#[cfg(feature = "qdrant")]
Qdrant(QdrantHttpClient),
#[cfg(feature = "s3")]
S3(aws_sdk_s3::Client),
#[cfg(feature = "mongodb")]
Mongodb(MongoDbExecutor),
#[cfg(feature = "neo4j")]
Neo4j(Neo4jExecutor),
#[cfg(feature = "clickhouse")]
ClickHouse(ClickHouseExecutor),
}
#[derive(Clone)]
struct ManagedClientEntry {
key: ClientKey,
kind: ManagedClientKind,
state: ClientState,
pooling_mode: PoolingMode,
labels: HashMap<String, String>,
created_at_unix_ms: u64,
last_health_probe_unix_ms: Arc<AtomicU64>,
active_leases: Arc<AtomicU64>,
total_leases: Arc<AtomicU64>,
error_count: Arc<AtomicU64>,
last_error: Arc<Mutex<Option<String>>>,
lease_budget: Option<u64>,
}
impl ManagedClientEntry {
fn new(
key: ClientKey,
kind: ManagedClientKind,
role: &str,
labels: HashMap<String, String>,
) -> Self {
let pooling_mode = PoolingMode::from_labels(&labels);
let key = ClientKey {
role: normalize_part(role, "read_write"),
..key
};
Self {
key,
kind,
state: ClientState::Ready,
pooling_mode,
lease_budget: lease_budget_from_labels(&labels),
labels,
created_at_unix_ms: unix_millis(),
last_health_probe_unix_ms: Arc::new(AtomicU64::new(0)),
active_leases: Arc::new(AtomicU64::new(0)),
total_leases: Arc::new(AtomicU64::new(0)),
error_count: Arc::new(AtomicU64::new(0)),
last_error: Arc::new(Mutex::new(None)),
}
}
fn snapshot(&self) -> ClientSnapshot {
let _kind = self.kind_label();
let (pg_active_connections, pg_idle_connections) = match &self.kind {
ManagedClientKind::Postgres(pool) => {
let size = pool.size();
let idle = pool.num_idle() as u32;
(Some(size.saturating_sub(idle)), Some(idle))
}
#[cfg(any(
feature = "redis",
feature = "qdrant",
feature = "s3",
feature = "mongodb",
feature = "neo4j",
feature = "clickhouse"
))]
_ => (None, None),
};
ClientSnapshot {
project_id: self.key.project_id.clone(),
backend: self.key.backend.clone(),
instance: self.key.instance.clone(),
role: self.key.role.clone(),
state: self.state.as_str().to_string(),
pooling_mode: self.pooling_mode.as_str().to_string(),
created_at_unix_ms: self.created_at_unix_ms,
last_health_probe_unix_ms: self.last_health_probe_unix_ms.load(Ordering::Relaxed),
active_leases: self.active_leases.load(Ordering::Relaxed),
total_leases: self.total_leases.load(Ordering::Relaxed),
error_count: self.error_count.load(Ordering::Relaxed),
last_error: self.last_error.lock().ok().and_then(|guard| guard.clone()),
lease_budget: self.lease_budget,
labels: self.labels.clone(),
pg_active_connections,
pg_idle_connections,
}
}
fn kind_label(&self) -> &'static str {
match &self.kind {
ManagedClientKind::Postgres(pool) => {
let _ = pool;
"postgres"
}
#[cfg(feature = "redis")]
ManagedClientKind::Redis(client) => {
let _ = client;
"redis"
}
#[cfg(feature = "qdrant")]
ManagedClientKind::Qdrant(client) => {
let _ = client;
"qdrant"
}
#[cfg(feature = "s3")]
ManagedClientKind::S3(client) => {
let _ = client;
"s3"
}
#[cfg(feature = "mongodb")]
ManagedClientKind::Mongodb(executor) => {
let _ = executor;
"mongodb"
}
#[cfg(feature = "neo4j")]
ManagedClientKind::Neo4j(executor) => {
let _ = executor;
"neo4j"
}
#[cfg(feature = "clickhouse")]
ManagedClientKind::ClickHouse(executor) => {
let _ = executor;
"clickhouse"
}
}
}
}
#[derive(Clone, Default)]
pub struct ConnectionManager {
inner: Arc<RwLock<HashMap<ClientKey, ManagedClientEntry>>>,
}
impl fmt::Debug for ConnectionManager {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
let len = self
.inner
.read()
.map(|guard| guard.len())
.unwrap_or_default();
f.debug_struct("ConnectionManager")
.field("clients", &len)
.finish()
}
}
impl ConnectionManager {
pub(crate) fn register_postgres(
&self,
instance: &str,
role: &str,
pool: PgPool,
labels: HashMap<String, String>,
) {
self.register(
ClientKey::new("default", "postgres", instance, role),
ManagedClientKind::Postgres(pool),
role,
labels,
);
}
#[cfg(feature = "redis")]
pub(crate) fn register_redis(
&self,
instance: &str,
role: &str,
client: redis::Client,
labels: HashMap<String, String>,
) {
self.register(
ClientKey::new("default", "redis", instance, role),
ManagedClientKind::Redis(client),
role,
labels,
);
}
#[cfg(feature = "qdrant")]
pub(crate) fn register_qdrant(
&self,
instance: &str,
role: &str,
client: QdrantHttpClient,
labels: HashMap<String, String>,
) {
self.register(
ClientKey::new("default", "qdrant", instance, role),
ManagedClientKind::Qdrant(client),
role,
labels,
);
}
#[cfg(feature = "s3")]
pub(crate) fn register_s3(
&self,
backend: &str,
instance: &str,
role: &str,
client: aws_sdk_s3::Client,
labels: HashMap<String, String>,
) {
self.register(
ClientKey::new("default", backend, instance, role),
ManagedClientKind::S3(client),
role,
labels,
);
}
#[cfg(feature = "mongodb")]
pub(crate) fn register_mongodb(
&self,
instance: &str,
role: &str,
executor: MongoDbExecutor,
labels: HashMap<String, String>,
) {
self.register(
ClientKey::new("default", "mongodb", instance, role),
ManagedClientKind::Mongodb(executor),
role,
labels,
);
}
#[cfg(feature = "neo4j")]
pub(crate) fn register_neo4j(
&self,
instance: &str,
role: &str,
executor: Neo4jExecutor,
labels: HashMap<String, String>,
) {
self.register(
ClientKey::new("default", "neo4j", instance, role),
ManagedClientKind::Neo4j(executor),
role,
labels,
);
}
#[cfg(feature = "clickhouse")]
pub(crate) fn register_clickhouse(
&self,
instance: &str,
role: &str,
executor: ClickHouseExecutor,
labels: HashMap<String, String>,
) {
self.register(
ClientKey::new("default", "clickhouse", instance, role),
ManagedClientKind::ClickHouse(executor),
role,
labels,
);
}
fn register(
&self,
key: ClientKey,
kind: ManagedClientKind,
role: &str,
labels: HashMap<String, String>,
) {
if let Ok(mut guard) = self.inner.write() {
guard.insert(
key.clone(),
ManagedClientEntry::new(key, kind, role, labels),
);
}
}
pub(crate) fn lease_postgres(&self, instance: &str) -> Option<ClientLease<PgPool>> {
let entry = self.entry("postgres", instance)?;
if !entry.lease_budget_allows() {
return None;
}
match entry.kind {
ManagedClientKind::Postgres(pool) => Some(ClientLease::new(
pool,
entry.active_leases,
entry.total_leases,
)),
#[cfg(any(
feature = "redis",
feature = "qdrant",
feature = "s3",
feature = "mongodb",
feature = "neo4j",
feature = "clickhouse"
))]
_ => None,
}
}
pub fn mark_state(&self, backend: &str, instance: &str, state: ClientState) -> bool {
let Some(key) = self.find_key(backend, instance) else {
return false;
};
self.inner
.write()
.ok()
.and_then(|mut guard| {
guard.get_mut(&key).map(|entry| {
entry.state = state;
})
})
.is_some()
}
pub fn mark_draining(&self, backend: &str, instance: &str) -> bool {
self.mark_state(backend, instance, ClientState::Draining)
}
pub fn mark_closed(&self, backend: &str, instance: &str) -> bool {
self.mark_state(backend, instance, ClientState::Closed)
}
pub fn record_health_result(
&self,
backend: &str,
instance: &str,
ok: bool,
error: Option<String>,
) -> bool {
let Some(key) = self.find_key(backend, instance) else {
return false;
};
self.inner
.write()
.ok()
.and_then(|mut guard| {
guard.get_mut(&key).map(|entry| {
entry
.last_health_probe_unix_ms
.store(unix_millis(), Ordering::Relaxed);
if ok {
entry.state = ClientState::Ready;
if let Ok(mut last_error) = entry.last_error.lock() {
*last_error = None;
}
} else {
entry.state = ClientState::Unhealthy;
entry.error_count.fetch_add(1, Ordering::Relaxed);
if let Ok(mut last_error) = entry.last_error.lock() {
*last_error = error;
}
}
})
})
.is_some()
}
pub fn replace_all_from(&self, shadow: &ConnectionManager) {
let replacement = shadow.inner.read().map(|guard| guard.clone());
if let (Ok(replacement), Ok(mut current)) = (replacement, self.inner.write()) {
let mut next = replacement;
for entry in current.values_mut() {
if entry.active_leases.load(Ordering::Relaxed) > 0 {
entry.state = ClientState::Draining;
let mut retained = entry.clone();
retained.key.instance = format!(
"{}__draining_{}",
retained.key.instance, retained.created_at_unix_ms
);
retained.key.role = format!("{}_draining", retained.key.role);
next.insert(retained.key.clone(), retained);
} else {
entry.state = ClientState::Closed;
}
}
*current = next;
}
}
pub fn active_leases(&self, backend: &str, instance: &str) -> Option<u64> {
self.entry_any_state(backend, instance)
.map(|entry| entry.active_leases.load(Ordering::Relaxed))
}
pub fn snapshots(&self) -> Vec<ClientSnapshot> {
let mut snapshots = self
.inner
.read()
.map(|guard| {
guard
.values()
.map(ManagedClientEntry::snapshot)
.collect::<Vec<_>>()
})
.unwrap_or_default();
snapshots.sort_by(|a, b| {
(
a.project_id.as_str(),
a.backend.as_str(),
a.instance.as_str(),
a.role.as_str(),
)
.cmp(&(
b.project_id.as_str(),
b.backend.as_str(),
b.instance.as_str(),
b.role.as_str(),
))
});
snapshots
}
pub fn metrics_text(&self) -> String {
let snapshots = self.snapshots();
if snapshots.is_empty() {
return String::new();
}
let mut out = String::new();
out.push_str("# TYPE udb_connection_client_state gauge\n");
out.push_str("# TYPE udb_connection_client_active_leases gauge\n");
out.push_str("# TYPE udb_connection_client_total_leases counter\n");
out.push_str("# TYPE udb_connection_client_errors_total counter\n");
out.push_str("# TYPE udb_connection_client_lease_budget gauge\n");
out.push_str("# TYPE udb_connection_client_last_health_probe_unix_ms gauge\n");
out.push_str("# TYPE udb_connection_pg_active_connections gauge\n");
out.push_str("# TYPE udb_connection_pg_idle_connections gauge\n");
for snapshot in snapshots {
let labels = format!(
"project=\"{}\",backend=\"{}\",instance=\"{}\",role=\"{}\",state=\"{}\",pooling_mode=\"{}\"",
metric_escape(&snapshot.project_id),
metric_escape(&snapshot.backend),
metric_escape(&snapshot.instance),
metric_escape(&snapshot.role),
metric_escape(&snapshot.state),
metric_escape(&snapshot.pooling_mode),
);
out.push_str(&format!("udb_connection_client_state{{{labels}}} 1\n"));
out.push_str(&format!(
"udb_connection_client_active_leases{{{labels}}} {}\n",
snapshot.active_leases
));
out.push_str(&format!(
"udb_connection_client_total_leases{{{labels}}} {}\n",
snapshot.total_leases
));
out.push_str(&format!(
"udb_connection_client_errors_total{{{labels}}} {}\n",
snapshot.error_count
));
out.push_str(&format!(
"udb_connection_client_last_health_probe_unix_ms{{{labels}}} {}\n",
snapshot.last_health_probe_unix_ms
));
if let Some(budget) = snapshot.lease_budget {
out.push_str(&format!(
"udb_connection_client_lease_budget{{{labels}}} {budget}\n"
));
}
if let Some(active) = snapshot.pg_active_connections {
out.push_str(&format!(
"udb_connection_pg_active_connections{{{labels}}} {active}\n"
));
}
if let Some(idle) = snapshot.pg_idle_connections {
out.push_str(&format!(
"udb_connection_pg_idle_connections{{{labels}}} {idle}\n"
));
}
}
out
}
fn entry(&self, backend: &str, instance: &str) -> Option<ManagedClientEntry> {
self.entry_any_state(backend, instance)
.filter(|entry| entry.state == ClientState::Ready)
}
fn entry_any_state(&self, backend: &str, instance: &str) -> Option<ManagedClientEntry> {
let backend = normalize_part(backend, "unknown");
let instance = normalize_part(instance, "default");
self.inner.read().ok().and_then(|guard| {
guard
.values()
.find(|entry| entry.key.backend == backend && entry.key.instance == instance)
.cloned()
})
}
fn find_key(&self, backend: &str, instance: &str) -> Option<ClientKey> {
let backend = normalize_part(backend, "unknown");
let instance = normalize_part(instance, "default");
self.inner.read().ok().and_then(|guard| {
guard
.keys()
.find(|key| key.backend == backend && key.instance == instance)
.cloned()
})
}
}
impl ManagedClientEntry {
fn lease_budget_allows(&self) -> bool {
self.lease_budget
.map(|budget| self.active_leases.load(Ordering::Relaxed) < budget)
.unwrap_or(true)
}
}
fn normalize_part(value: &str, default: &str) -> String {
let value = value.trim();
if value.is_empty() {
default.to_string()
} else {
value.to_ascii_lowercase()
}
}
fn unix_millis() -> u64 {
SystemTime::now()
.duration_since(UNIX_EPOCH)
.map(|duration| duration.as_millis() as u64)
.unwrap_or_default()
}
fn metric_escape(value: &str) -> String {
value
.replace('\\', "\\\\")
.replace('"', "\\\"")
.replace('\n', "\\n")
}
fn lease_budget_from_labels(labels: &HashMap<String, String>) -> Option<u64> {
let cap = labels
.get("connection_cap")
.or_else(|| labels.get("max_connections"))
.or_else(|| labels.get("upstream_max_connections"))
.and_then(|value| value.trim().parse::<u64>().ok())?;
if cap == 0 {
return None;
}
let safety_factor = labels
.get("safety_factor")
.or_else(|| labels.get("connection_safety_factor"))
.and_then(|value| value.trim().parse::<f64>().ok())
.filter(|value| *value > 0.0 && *value <= 1.0)
.unwrap_or(1.0);
Some(((cap as f64) * safety_factor).floor().max(1.0) as u64)
}
#[cfg(test)]
mod tests {
use super::*;
use sqlx::postgres::PgPoolOptions;
#[test]
fn pooling_mode_reads_backend_labels() {
assert_eq!(
PoolingMode::from_labels(&HashMap::from([(
"pooling_mode".to_string(),
"transaction".to_string()
)])),
PoolingMode::Transaction
);
assert_eq!(PoolingMode::from_value("pipelined"), PoolingMode::Statement);
assert_eq!(PoolingMode::from_value("unknown"), PoolingMode::Session);
}
#[tokio::test]
async fn postgres_registration_produces_snapshot_and_metrics() {
let pool = PgPoolOptions::new()
.connect_lazy("postgres://user:pass@localhost/db")
.expect("lazy pool should not connect");
let manager = ConnectionManager::default();
manager.register_postgres(
"primary",
"read_write",
pool,
HashMap::from([("pooling_mode".to_string(), "session".to_string())]),
);
let snapshots = manager.snapshots();
assert_eq!(snapshots.len(), 1);
assert_eq!(snapshots[0].backend, "postgres");
assert_eq!(snapshots[0].instance, "primary");
assert_eq!(snapshots[0].pooling_mode, "session");
assert_eq!(snapshots[0].state, "ready");
assert!(
manager
.metrics_text()
.contains("udb_connection_client_state")
);
}
#[tokio::test]
async fn postgres_lease_accounting_drops_back_to_zero() {
let pool = PgPoolOptions::new()
.connect_lazy("postgres://user:pass@localhost/db")
.expect("lazy pool should not connect");
let manager = ConnectionManager::default();
manager.register_postgres("primary", "read_write", pool, HashMap::new());
{
let _lease = manager.lease_postgres("primary").unwrap();
let snapshot = manager.snapshots().pop().unwrap();
assert_eq!(snapshot.active_leases, 1);
assert_eq!(snapshot.total_leases, 1);
}
let snapshot = manager.snapshots().pop().unwrap();
assert_eq!(snapshot.active_leases, 0);
assert_eq!(snapshot.total_leases, 1);
}
#[tokio::test]
async fn draining_clients_stop_new_leases_but_keep_accounting() {
let pool = PgPoolOptions::new()
.connect_lazy("postgres://user:pass@localhost/db")
.expect("lazy pool should not connect");
let manager = ConnectionManager::default();
manager.register_postgres("primary", "read_write", pool, HashMap::new());
let lease = manager.lease_postgres("primary").unwrap();
assert!(manager.mark_draining("postgres", "primary"));
assert!(manager.lease_postgres("primary").is_none());
assert_eq!(manager.active_leases("postgres", "primary"), Some(1));
drop(lease);
assert_eq!(manager.active_leases("postgres", "primary"), Some(0));
}
#[tokio::test]
async fn health_errors_and_budget_are_reported() {
let pool = PgPoolOptions::new()
.connect_lazy("postgres://user:pass@localhost/db")
.expect("lazy pool should not connect");
let manager = ConnectionManager::default();
manager.register_postgres(
"primary",
"read_write",
pool,
HashMap::from([
("connection_cap".to_string(), "2".to_string()),
("safety_factor".to_string(), "0.5".to_string()),
]),
);
assert_eq!(manager.snapshots()[0].lease_budget, Some(1));
let lease = manager.lease_postgres("primary").unwrap();
assert!(manager.lease_postgres("primary").is_none());
drop(lease);
assert!(manager.record_health_result(
"postgres",
"primary",
false,
Some("probe failed".to_string())
));
let snapshot = manager.snapshots().pop().unwrap();
assert_eq!(snapshot.state, "unhealthy");
assert_eq!(snapshot.error_count, 1);
assert_eq!(snapshot.last_error.as_deref(), Some("probe failed"));
assert!(snapshot.last_health_probe_unix_ms > 0);
}
#[tokio::test]
async fn shadow_manager_can_replace_current_registry() {
let old_pool = PgPoolOptions::new()
.connect_lazy("postgres://user:pass@localhost/old")
.expect("lazy pool should not connect");
let new_pool = PgPoolOptions::new()
.connect_lazy("postgres://user:pass@localhost/new")
.expect("lazy pool should not connect");
let current = ConnectionManager::default();
current.register_postgres("primary", "read_write", old_pool, HashMap::new());
let shadow = ConnectionManager::default();
shadow.register_postgres(
"primary",
"read_write",
new_pool,
HashMap::from([("generation".to_string(), "2".to_string())]),
);
current.replace_all_from(&shadow);
let snapshots = current.snapshots();
assert_eq!(snapshots.len(), 1);
assert_eq!(
snapshots[0].labels.get("generation").map(String::as_str),
Some("2")
);
}
#[tokio::test]
async fn shadow_replacement_retains_draining_clients_with_active_leases() {
let old_pool = PgPoolOptions::new()
.connect_lazy("postgres://user:pass@localhost/old")
.expect("lazy pool should not connect");
let new_pool = PgPoolOptions::new()
.connect_lazy("postgres://user:pass@localhost/new")
.expect("lazy pool should not connect");
let current = ConnectionManager::default();
current.register_postgres("primary", "read_write", old_pool, HashMap::new());
let lease = current.lease_postgres("primary").unwrap();
let shadow = ConnectionManager::default();
shadow.register_postgres(
"primary",
"read_write",
new_pool,
HashMap::from([("generation".to_string(), "2".to_string())]),
);
current.replace_all_from(&shadow);
let snapshots = current.snapshots();
assert_eq!(snapshots.len(), 2);
assert!(snapshots.iter().any(|snapshot| {
snapshot.instance == "primary"
&& snapshot.labels.get("generation").map(String::as_str) == Some("2")
}));
assert!(snapshots.iter().any(|snapshot| {
snapshot.state == "draining" && snapshot.instance.starts_with("primary__draining_")
}));
drop(lease);
}
#[test]
fn hot_path_handlers_do_not_open_backend_connections() {
let manifest_dir = std::path::PathBuf::from(env!("CARGO_MANIFEST_DIR"));
let mut files = vec![manifest_dir.join("src/runtime/core/tx_object.rs")];
let handlers_dir = manifest_dir.join("src/runtime/service");
for entry in std::fs::read_dir(&handlers_dir).expect("read service dir") {
let path = entry.expect("read service entry").path();
let Some(name) = path.file_name().and_then(|name| name.to_str()) else {
continue;
};
if name.starts_with("handlers_") && name.ends_with(".rs") {
files.push(path);
}
}
let forbidden = [
"connect(",
"Client::open",
"from_env(",
"std::env::var",
"env::var",
];
let mut violations = Vec::new();
for file in files {
let content = std::fs::read_to_string(&file).expect("read source file");
for pattern in forbidden {
if content.contains(pattern) {
violations.push(format!("{} contains {}", file.display(), pattern));
}
}
}
assert!(
violations.is_empty(),
"hot path connection/env access violations:\n{}",
violations.join("\n")
);
}
#[test]
fn runtime_env_reads_are_confined_to_startup_config_or_compat_boundaries() {
let manifest_dir = std::path::PathBuf::from(env!("CARGO_MANIFEST_DIR"));
let runtime_dir = manifest_dir.join("src/runtime");
let allowed = [
"src/runtime/authn/",
"src/runtime/authz/",
"src/runtime/canonical_store/dialect.rs",
"src/runtime/cdc/mod.rs",
"src/runtime/cdc/source.rs",
"src/runtime/cdc/encryption.rs",
"src/runtime/cdc/engine_dlq.rs",
"src/runtime/cdc/engine_tail.rs",
"src/runtime/channels.rs",
"src/runtime/config/",
"src/runtime/consistency_fence.rs",
"src/runtime/xa_recovery.rs",
"src/runtime/core/helpers.rs",
"src/runtime/core/setup_data.rs",
"src/runtime/core/native_store.rs",
"src/runtime/executor_utils.rs",
"src/runtime/executors/clickhouse.rs",
"src/runtime/executors/http.rs",
"src/runtime/executors/mssql.rs",
"src/runtime/executors/mongodb.rs",
"src/runtime/executors/neo4j.rs",
"src/runtime/native_catalog.rs",
"src/runtime/observability.rs",
"src/runtime/otel.rs",
"src/runtime/projection/mod.rs",
"src/runtime/replica.rs",
"src/runtime/security.rs",
"src/runtime/signalling/",
"src/runtime/singleton.rs",
"src/runtime/service/mod.rs",
"src/runtime/service/asset_service/",
"src/runtime/service/auth_service/",
"src/runtime/service/method_security.rs",
"src/runtime/service/storage_service/",
"src/runtime/service/webrtc_service/",
];
let mut stack = vec![runtime_dir];
let mut violations = Vec::new();
while let Some(dir) = stack.pop() {
for entry in std::fs::read_dir(&dir).expect("read runtime dir") {
let path = entry.expect("read runtime entry").path();
if path.is_dir() {
stack.push(path);
continue;
}
if path.extension().and_then(|ext| ext.to_str()) != Some("rs") {
continue;
}
let rel = path
.strip_prefix(&manifest_dir)
.expect("runtime source under manifest")
.to_string_lossy()
.replace('\\', "/");
if rel == "src/runtime/connection_manager.rs"
|| rel.contains("/tests.rs")
|| rel.contains("_tests.rs")
|| rel.contains("tests/")
|| allowed.iter().any(|prefix| rel.starts_with(prefix))
{
continue;
}
let content = std::fs::read_to_string(&path).expect("read runtime source");
for pattern in ["std::env::var", "env::var", "from_env("] {
if content.contains(pattern) {
violations.push(format!("{rel} contains {pattern}"));
}
}
}
}
assert!(
violations.is_empty(),
"runtime env access outside allowlisted startup/config boundaries:\n{}",
violations.join("\n")
);
}
}