use std::sync::atomic::{AtomicU64, Ordering};
use std::sync::{Arc, Mutex, OnceLock};
use std::time::{Duration, Instant};
use nametrics_core::atomic::{add, AtomicHistogram, LATENCY_SECONDS};
use nametrics_core::{
LegacyMetricsSource, MetricDescriptor, MetricKind, MetricSample, MetricValue,
};
use nanotify_core::{AlertRoute, EventKind, Notification, NotificationFields, Severity};
use crate::{DatabaseDriver, MAX_MANAGED_DATASOURCES};
#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
pub enum StatementLogging {
#[default]
Disabled,
Debug,
Trace,
}
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum ConnectionPurpose {
Mapper,
Migration,
Probe,
Direct,
}
impl ConnectionPurpose {
pub const fn as_str(self) -> &'static str {
match self {
Self::Mapper => "mapper",
Self::Migration => "migration",
Self::Probe => "probe",
Self::Direct => "direct",
}
}
}
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum AcquireOutcome {
Ok,
Timeout,
Closed,
Worker,
Cancelled,
Other,
}
impl AcquireOutcome {
const fn status(self) -> usize {
match self {
Self::Ok => 0,
Self::Cancelled => 2,
_ => 1,
}
}
}
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum ConnectionRejection {
TxRequired,
TxForbidden,
CrossDatasource,
CrossDriver,
UnknownDatasource,
RegistryUnavailable,
TxConnectionBusy,
}
#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
pub enum WaitLogLevel {
Info,
#[default]
Warn,
}
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub struct WaitObservationPolicy {
pub threshold: Duration,
pub log_enabled: bool,
pub log_level: WaitLogLevel,
pub log_cooldown: Duration,
}
pub struct AcquireAlertPolicy {
pub route: AlertRoute,
pub severity: Severity,
pub cooldown: Duration,
pub purposes: Vec<ConnectionPurpose>,
pub service: String,
pub instance: String,
pub environment: Option<String>,
pub cluster: Option<String>,
}
impl Default for WaitObservationPolicy {
fn default() -> Self {
Self {
threshold: Duration::from_millis(250),
log_enabled: false,
log_level: WaitLogLevel::Warn,
log_cooldown: Duration::from_secs(60),
}
}
}
struct WaitCells {
outcomes: [AtomicU64; 6],
durations: [AtomicHistogram; 3],
in_flight: AtomicU64,
}
impl WaitCells {
const fn new() -> Self {
Self {
outcomes: [const { AtomicU64::new(0) }; 6],
durations: [const { AtomicHistogram::new(LATENCY_SECONDS) }; 3],
in_flight: AtomicU64::new(0),
}
}
}
pub struct ConnectionSlot {
driver: DatabaseDriver,
datasource: Box<str>,
acquire: [WaitCells; 4],
transaction: WaitCells,
rejections: [AtomicU64; 7],
policies: OnceLock<(WaitObservationPolicy, WaitObservationPolicy)>,
log_next: [AtomicU64; 2],
acquire_alert: OnceLock<AcquireAlertPolicy>,
alert_next: AtomicU64,
alert_sequence: AtomicU64,
}
impl ConnectionSlot {
fn new(driver: DatabaseDriver, datasource: &str) -> Self {
Self {
driver,
datasource: datasource.into(),
acquire: [const { WaitCells::new() }; 4],
transaction: WaitCells::new(),
rejections: [const { AtomicU64::new(0) }; 7],
policies: OnceLock::new(),
log_next: [const { AtomicU64::new(0) }; 2],
acquire_alert: OnceLock::new(),
alert_next: AtomicU64::new(0),
alert_sequence: AtomicU64::new(0),
}
}
pub fn reject(&self, reason: ConnectionRejection) {
add(&self.rejections[reason as usize], 1);
}
pub fn acquire(&'static self, purpose: ConnectionPurpose) -> WaitGuard {
WaitGuard::new(self, purpose, false)
}
pub fn transaction_slot(&'static self) -> WaitGuard {
WaitGuard::new(self, ConnectionPurpose::Direct, true)
}
}
static SLOTS: [OnceLock<ConnectionSlot>; MAX_MANAGED_DATASOURCES * 2] =
[const { OnceLock::new() }; MAX_MANAGED_DATASOURCES * 2];
static REGISTRATION: Mutex<()> = Mutex::new(());
static UNKNOWN_MYSQL: OnceLock<ConnectionSlot> = OnceLock::new();
static UNKNOWN_PGSQL: OnceLock<ConnectionSlot> = OnceLock::new();
static UNKNOWN_SQLITE: OnceLock<ConnectionSlot> = OnceLock::new();
static CLOCK: OnceLock<Instant> = OnceLock::new();
pub fn register_datasource(
driver: DatabaseDriver,
datasource: &str,
) -> anyhow::Result<&'static ConnectionSlot> {
crate::validate_datasource_name(datasource)?;
let _registration = REGISTRATION
.lock()
.unwrap_or_else(|error| error.into_inner());
CLOCK.get_or_init(Instant::now);
fallback(driver);
if let Some(slot) = known_slot(driver, datasource) {
return Ok(slot);
}
let cell = SLOTS
.iter()
.find(|cell| cell.get().is_none())
.ok_or_else(|| anyhow::anyhow!("connection observation datasource capacity exceeded"))?;
Ok(cell.get_or_init(|| ConnectionSlot::new(driver, datasource)))
}
pub fn configure_datasource(
driver: DatabaseDriver,
datasource: &str,
acquire: WaitObservationPolicy,
slot: WaitObservationPolicy,
) -> anyhow::Result<()> {
let cells = register_datasource(driver, datasource)?;
let policy = (acquire, slot);
anyhow::ensure!(
*cells.policies.get_or_init(|| policy) == policy,
"connection observation policy is frozen; restart is required"
);
Ok(())
}
pub fn configure_acquire_alert(
driver: DatabaseDriver,
datasource: &str,
policy: AcquireAlertPolicy,
) -> anyhow::Result<()> {
anyhow::ensure!(
!policy.purposes.is_empty() && policy.purposes.len() <= 4,
"acquire alert purposes must be nonempty and bounded"
);
let mut seen = [false; 4];
for purpose in &policy.purposes {
anyhow::ensure!(!seen[*purpose as usize], "duplicate acquire alert purpose");
seen[*purpose as usize] = true;
}
register_datasource(driver, datasource)?
.acquire_alert
.set(policy)
.map_err(|_| anyhow::anyhow!("acquire alert policy is frozen; restart is required"))
}
pub fn connection_slot(driver: DatabaseDriver, datasource: &str) -> &'static ConnectionSlot {
known_slot(driver, datasource).unwrap_or_else(|| fallback(driver))
}
fn known_slot(driver: DatabaseDriver, datasource: &str) -> Option<&'static ConnectionSlot> {
SLOTS
.iter()
.filter_map(OnceLock::get)
.find(|slot| slot.driver == driver && slot.datasource.as_ref() == datasource)
}
fn fallback(driver: DatabaseDriver) -> &'static ConnectionSlot {
match driver {
DatabaseDriver::MySql => &UNKNOWN_MYSQL,
DatabaseDriver::PostgreSql => &UNKNOWN_PGSQL,
DatabaseDriver::Sqlite => &UNKNOWN_SQLITE,
}
.get_or_init(|| ConnectionSlot::new(driver, "__unknown__"))
}
pub struct WaitGuard {
slot: &'static ConnectionSlot,
purpose: ConnectionPurpose,
transaction: bool,
started: Instant,
finished: bool,
}
impl WaitGuard {
fn new(slot: &'static ConnectionSlot, purpose: ConnectionPurpose, transaction: bool) -> Self {
let guard = Self {
slot,
purpose,
transaction,
started: Instant::now(),
finished: false,
};
guard.cells().in_flight.fetch_add(1, Ordering::Relaxed);
guard
}
fn cells(&self) -> &WaitCells {
if self.transaction {
&self.slot.transaction
} else {
&self.slot.acquire[self.purpose as usize]
}
}
pub fn finish(&mut self, outcome: AcquireOutcome) {
if self.finished {
return;
}
self.finished = true;
let elapsed = self.started.elapsed();
let cells = self.cells();
add(&cells.outcomes[outcome as usize], 1);
cells.durations[outcome.status()].observe(elapsed);
cells.in_flight.fetch_sub(1, Ordering::Relaxed);
if outcome == AcquireOutcome::Ok {
self.log_wait(elapsed);
}
if outcome == AcquireOutcome::Timeout && !self.transaction {
self.notify_timeout(elapsed);
}
}
fn notify_timeout(&self, elapsed: Duration) {
let Some(policy) = self.slot.acquire_alert.get() else {
return;
};
if !policy.purposes.contains(&self.purpose) {
return;
}
let now = CLOCK
.get_or_init(Instant::now)
.elapsed()
.as_millis()
.min(u64::MAX as u128) as u64;
if !policy.cooldown.is_zero()
&& self
.slot
.alert_next
.fetch_update(Ordering::Relaxed, Ordering::Relaxed, |value| {
(value <= now).then(|| {
now.saturating_add(policy.cooldown.as_millis().min(u64::MAX as u128) as u64)
})
})
.is_err()
{
return;
}
let sequence = self.slot.alert_sequence.fetch_add(1, Ordering::Relaxed);
let occurred_at_unix_ms = std::time::SystemTime::now()
.duration_since(std::time::UNIX_EPOCH)
.unwrap_or_default()
.as_millis()
.min(u64::MAX as u128) as u64;
use std::hash::{Hash, Hasher};
let mut identity = std::collections::hash_map::DefaultHasher::new();
(
policy.instance.as_str(),
driver_name(self.slot.driver),
self.slot.datasource.as_ref(),
)
.hash(&mut identity);
let notification = Notification::new(NotificationFields {
id: format!(
"acquire-{:x}-{occurred_at_unix_ms}-{sequence}",
identity.finish()
),
event: EventKind::AcquireTimeout,
severity: policy.severity,
occurred_at_unix_ms,
service: policy.service.clone(),
instance: policy.instance.clone(),
environment: policy.environment.clone(),
cluster: policy.cluster.clone(),
driver: driver_name(self.slot.driver).to_owned(),
datasource: self.slot.datasource.to_string(),
operation: "acquire".to_owned(),
purpose: Some(self.purpose.as_str().to_owned()),
duration: elapsed,
outcome: "timeout".to_owned(),
error_kind: Some("timeout".to_owned()),
..NotificationFields::default()
});
let _ = policy.route.try_send(notification);
}
fn log_wait(&self, elapsed: Duration) {
let Some(policies) = self.slot.policies.get() else {
return;
};
let policy = if self.transaction {
policies.1
} else {
policies.0
};
if !policy.log_enabled || elapsed < policy.threshold {
return;
}
let next = &self.slot.log_next[usize::from(self.transaction)];
let now = CLOCK
.get_or_init(Instant::now)
.elapsed()
.as_millis()
.min(u64::MAX as u128) as u64;
if policy.log_cooldown != Duration::ZERO
&& next
.fetch_update(Ordering::Relaxed, Ordering::Relaxed, |value| {
(value <= now).then(|| {
now.saturating_add(
policy.log_cooldown.as_millis().min(u64::MAX as u128) as u64
)
})
})
.is_err()
{
return;
}
let event = if self.transaction {
"transaction_slot_wait"
} else {
"connection_acquire_wait"
};
let result = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| {
match policy.log_level {
WaitLogLevel::Info => {
tracing::info!(target: "natx::connection", event, datasource = self.slot.datasource.as_ref(), driver = driver_name(self.slot.driver), purpose = self.purpose.as_str(), duration_seconds = elapsed.as_secs_f64(), "数据库连接等待超过阈值")
}
WaitLogLevel::Warn => {
tracing::warn!(target: "natx::connection", event, datasource = self.slot.datasource.as_ref(), driver = driver_name(self.slot.driver), purpose = self.purpose.as_str(), duration_seconds = elapsed.as_secs_f64(), "数据库连接等待超过阈值")
}
}
}));
if let Err(payload) = result {
if let Err(nested) =
std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| drop(payload)))
{
std::mem::forget(nested);
}
}
}
}
impl Drop for WaitGuard {
fn drop(&mut self) {
self.finish(AcquireOutcome::Cancelled);
}
}
fn driver_name(driver: DatabaseDriver) -> &'static str {
match driver {
DatabaseDriver::MySql => "mysql",
DatabaseDriver::PostgreSql => "postgresql",
DatabaseDriver::Sqlite => "sqlite",
}
}
const OUTCOMES: [&str; 6] = ["ok", "timeout", "closed", "worker", "cancelled", "other"];
const STATUSES: [&str; 3] = ["success", "failure", "cancelled"];
const PURPOSES: [&str; 4] = ["mapper", "migration", "probe", "direct"];
const REJECTIONS: [&str; 7] = [
"tx_required",
"tx_forbidden",
"cross_datasource",
"cross_driver",
"unknown_datasource",
"registry_unavailable",
"tx_connection_busy",
];
macro_rules! descriptor {
($name:ident, $text:literal, $kind:ident, $labels:expr, $bounds:expr) => {
static $name: MetricDescriptor = MetricDescriptor {
name: $text,
help: "数据库连接与等待状态。",
unit: if matches!(MetricKind::$kind, MetricKind::Histogram) {
"seconds"
} else {
""
},
kind: MetricKind::$kind,
label_names: $labels,
histogram_bounds: $bounds,
};
};
}
descriptor!(
ACQUIRES,
"natx_connection_acquires_total",
Counter,
&["datasource", "driver", "purpose", "outcome"],
&[]
);
descriptor!(
ACQUIRE_DURATION,
"natx_connection_acquire_duration_seconds",
Histogram,
&["datasource", "driver", "purpose", "status"],
LATENCY_SECONDS
);
descriptor!(
ACQUIRE_IN_FLIGHT,
"natx_connection_acquire_in_flight",
Gauge,
&["datasource", "driver", "purpose"],
&[]
);
descriptor!(
SLOT_WAITS,
"natx_transaction_slot_waits_total",
Counter,
&["datasource", "driver", "outcome"],
&[]
);
descriptor!(
SLOT_DURATION,
"natx_transaction_slot_wait_duration_seconds",
Histogram,
&["datasource", "driver", "status"],
LATENCY_SECONDS
);
descriptor!(
SLOT_IN_FLIGHT,
"natx_transaction_slot_wait_in_flight",
Gauge,
&["datasource", "driver"],
&[]
);
descriptor!(
REJECTION,
"natx_connection_rejections_total",
Counter,
&["datasource", "driver", "reason"],
&[]
);
descriptor!(
POOL_CONNECTIONS,
"natx_pool_connections",
Gauge,
&["datasource", "driver", "state"],
&[]
);
descriptor!(
POOL_MAX,
"natx_pool_max_connections",
Gauge,
&["datasource", "driver"],
&[]
);
static CONNECTION_DESCRIPTORS: &[&MetricDescriptor] = &[
&ACQUIRES,
&ACQUIRE_DURATION,
&ACQUIRE_IN_FLIGHT,
&SLOT_WAITS,
&SLOT_DURATION,
&SLOT_IN_FLIGHT,
&REJECTION,
];
static POOL_DESCRIPTORS: &[&MetricDescriptor] = &[&POOL_CONNECTIONS, &POOL_MAX];
pub struct ConnectionMetricsSource {
slots: Box<[&'static ConnectionSlot]>,
}
impl ConnectionMetricsSource {
pub fn new(driver: DatabaseDriver) -> Self {
let _registration = REGISTRATION
.lock()
.unwrap_or_else(|error| error.into_inner());
let slots = SLOTS
.iter()
.filter_map(OnceLock::get)
.filter(|slot| slot.driver == driver)
.chain(std::iter::once(fallback(driver)))
.collect::<Vec<_>>()
.into_boxed_slice();
Self {
slots,
}
}
pub fn series_budget(&self) -> usize {
self.slots.len() * (5 * (6 + 3 * (LATENCY_SECONDS.len() + 3) + 1) + 7)
}
}
impl LegacyMetricsSource for ConnectionMetricsSource {
fn descriptors(&self) -> &'static [&'static MetricDescriptor] {
CONNECTION_DESCRIPTORS
}
fn snapshot(&self) -> Option<Vec<MetricSample>> {
let mut samples = Vec::new();
for slot in &self.slots {
let base = vec![
("datasource", slot.datasource.to_string()),
("driver", driver_name(slot.driver).to_owned()),
];
for (purpose, cells) in PURPOSES.iter().zip(&slot.acquire) {
let mut labels = base.clone();
labels.push(("purpose", (*purpose).to_owned()));
snapshot_wait(
&mut samples,
cells,
labels,
&ACQUIRES,
&ACQUIRE_DURATION,
&ACQUIRE_IN_FLIGHT,
);
}
snapshot_wait(
&mut samples,
&slot.transaction,
base.clone(),
&SLOT_WAITS,
&SLOT_DURATION,
&SLOT_IN_FLIGHT,
);
for (reason, cell) in REJECTIONS.iter().zip(&slot.rejections) {
let mut labels = base.clone();
labels.push(("reason", (*reason).to_owned()));
samples.push(MetricSample {
name: REJECTION.name,
labels,
value: MetricValue::Counter(cell.load(Ordering::Relaxed)),
});
}
}
Some(samples)
}
fn render_prometheus(&self, _output: &mut String) {}
}
fn snapshot_wait(
samples: &mut Vec<MetricSample>,
cells: &WaitCells,
labels: Vec<(&'static str, String)>,
counter: &'static MetricDescriptor,
histogram: &'static MetricDescriptor,
gauge: &'static MetricDescriptor,
) {
for (outcome, cell) in OUTCOMES.iter().zip(&cells.outcomes) {
let mut labels = labels.clone();
labels.push(("outcome", (*outcome).to_owned()));
samples.push(MetricSample {
name: counter.name,
labels,
value: MetricValue::Counter(cell.load(Ordering::Relaxed)),
});
}
for (status, cell) in STATUSES.iter().zip(&cells.durations) {
let mut labels = labels.clone();
labels.push(("status", (*status).to_owned()));
samples.push(MetricSample {
name: histogram.name,
labels,
value: cell.snapshot(),
});
}
samples.push(MetricSample {
name: gauge.name,
labels,
value: MetricValue::Gauge(cells.in_flight.load(Ordering::Relaxed) as f64),
});
}
pub struct PoolState {
pub total: u32,
pub idle: u32,
pub max: u32,
}
pub trait PoolSnapshot: Send + Sync {
fn snapshot(&self) -> PoolState;
}
pub struct PoolMetricsSource {
driver: DatabaseDriver,
pools: Vec<(String, Arc<dyn PoolSnapshot>)>,
}
impl PoolMetricsSource {
pub fn new(
driver: DatabaseDriver,
pools: Vec<(String, Arc<dyn PoolSnapshot>)>,
) -> anyhow::Result<Self> {
anyhow::ensure!(
pools.len() <= MAX_MANAGED_DATASOURCES,
"pool observation capacity exceeded"
);
let mut names = std::collections::BTreeSet::new();
for (name, _) in &pools {
crate::validate_datasource_name(name)?;
anyhow::ensure!(names.insert(name), "duplicate pool observation datasource");
}
Ok(Self {
driver,
pools,
})
}
pub fn series_budget(&self) -> usize {
self.pools.len() * 4
}
}
impl LegacyMetricsSource for PoolMetricsSource {
fn descriptors(&self) -> &'static [&'static MetricDescriptor] {
POOL_DESCRIPTORS
}
fn snapshot(&self) -> Option<Vec<MetricSample>> {
let mut samples = Vec::with_capacity(self.series_budget());
for (name, pool) in &self.pools {
let state = pool.snapshot();
let base = vec![
("datasource", name.clone()),
("driver", driver_name(self.driver).to_owned()),
];
for (label, value) in [
("total", state.total),
("idle", state.idle),
("in_use", state.total.saturating_sub(state.idle)),
] {
let mut labels = base.clone();
labels.push(("state", label.to_owned()));
samples.push(MetricSample {
name: POOL_CONNECTIONS.name,
labels,
value: MetricValue::Gauge(value as f64),
});
}
samples.push(MetricSample {
name: POOL_MAX.name,
labels: base,
value: MetricValue::Gauge(state.max as f64),
});
}
Some(samples)
}
fn render_prometheus(&self, _output: &mut String) {}
}