use std::sync::{OnceLock, RwLock};
use sqlx::PgPool;
use tonic::Status;
use super::calc::wall_now_unix;
use super::config::DEFAULT_UNIT;
static ADMISSION_METERING_POOL: OnceLock<RwLock<Option<PgPool>>> = OnceLock::new();
fn admission_metering_cell() -> &'static RwLock<Option<PgPool>> {
ADMISSION_METERING_POOL.get_or_init(|| RwLock::new(None))
}
pub(crate) fn install_admission_metering_pool(pool: Option<PgPool>) {
if let Ok(mut guard) = admission_metering_cell().write() {
*guard = pool;
}
}
pub(crate) fn admission_metering_pool() -> Option<PgPool> {
admission_metering_cell().read().ok()?.clone()
}
pub(crate) fn admission_metering_method(service_label: &str, op_label: &str) -> String {
let service = service_label.trim();
let op = op_label.trim();
match (service.is_empty(), op.is_empty()) {
(true, true) => "native.unknown".to_string(),
(true, false) => format!("native.{op}"),
(false, true) => service.to_string(),
(false, false) => format!("{service}.{op}"),
}
}
pub(crate) async fn record_usage(
pool: &PgPool,
tenant_id: &str,
principal_id: &str,
method: &str,
unit: &str,
quantity: i64,
now_unix: i64,
) -> Result<(), Status> {
let tenant_id = tenant_id.trim();
let method = method.trim();
if tenant_id.is_empty() || method.is_empty() {
return Ok(());
}
let unit = {
let u = unit.trim();
if u.is_empty() { DEFAULT_UNIT } else { u }
};
let quantity = quantity.max(0);
let occurred = if now_unix > 0 {
now_unix
} else {
wall_now_unix()
};
let res = sqlx::query(
"INSERT INTO udb_metering.usage_events \
(tenant_id, principal_id, method, unit, quantity, occurred_at, occurred_at_unix) \
SELECT $1, $2, $3, $4, $5, to_timestamp($6), $7 \
WHERE set_config('app.current_tenant_id', $1, true) IS NOT NULL",
)
.bind(tenant_id)
.bind(principal_id.trim())
.bind(method)
.bind(unit)
.bind(quantity)
.bind(occurred as f64)
.bind(occurred)
.execute(pool)
.await;
if let Err(err) = res {
tracing::warn!(
target: "udb::metering",
error = %err,
tenant_id = %tenant_id,
method = %method,
"usage metering insert failed; swallowing so metering never fails the metered request",
);
}
Ok(())
}