#[derive(Clone, Copy)]
pub(crate) enum Site {
Unlabeled,
PlanCache,
ParseCache,
ResultCache,
ArtIndexRegistry,
ArtPkRegistry,
StatementRegistry,
}
#[allow(dead_code)]
pub(crate) struct SiteStat {
pub name: &'static str,
pub acquisitions: u64,
pub contended: u64,
pub contended_wait_nanos: u64,
}
#[cfg(feature = "lock-census")]
use std::sync::atomic::{AtomicBool, AtomicU64, Ordering};
#[cfg(feature = "lock-census")]
use std::time::{Duration, Instant};
#[cfg(feature = "lock-census")]
const SITE_COUNT: usize = 6;
#[cfg(feature = "lock-census")]
const SITE_NAMES: [&str; SITE_COUNT] = [
"plan_cache_shard",
"parse_cache_shard",
"result_cache_shard",
"art_index_registry",
"art_pk_registry",
"statement_registry",
];
#[cfg(feature = "lock-census")]
struct SiteCounters {
acquisitions: AtomicU64,
contended: AtomicU64,
contended_wait_nanos: AtomicU64,
}
#[cfg(feature = "lock-census")]
impl SiteCounters {
const fn new() -> Self {
Self {
acquisitions: AtomicU64::new(0),
contended: AtomicU64::new(0),
contended_wait_nanos: AtomicU64::new(0),
}
}
}
#[cfg(feature = "lock-census")]
static ENABLED: AtomicBool = AtomicBool::new(false);
#[cfg(feature = "lock-census")]
static SITES: [SiteCounters; SITE_COUNT] = [
SiteCounters::new(),
SiteCounters::new(),
SiteCounters::new(),
SiteCounters::new(),
SiteCounters::new(),
SiteCounters::new(),
];
#[cfg(feature = "lock-census")]
fn site_index(site: Site) -> Option<usize> {
match site {
Site::Unlabeled => None,
Site::PlanCache => Some(0),
Site::ParseCache => Some(1),
Site::ResultCache => Some(2),
Site::ArtIndexRegistry => Some(3),
Site::ArtPkRegistry => Some(4),
Site::StatementRegistry => Some(5),
}
}
#[cfg(feature = "lock-census")]
#[inline]
fn is_enabled() -> bool {
ENABLED.load(Ordering::Relaxed)
}
#[cfg(feature = "lock-census")]
#[inline]
fn record_acquire(site: Site) {
if let Some(c) = site_index(site).and_then(|i| SITES.get(i)) {
c.acquisitions.fetch_add(1, Ordering::Relaxed);
}
}
#[cfg(feature = "lock-census")]
#[inline]
fn record_contended(site: Site, waited: Duration) {
if let Some(c) = site_index(site).and_then(|i| SITES.get(i)) {
c.acquisitions.fetch_add(1, Ordering::Relaxed);
c.contended.fetch_add(1, Ordering::Relaxed);
c.contended_wait_nanos.fetch_add(waited.as_nanos() as u64, Ordering::Relaxed);
}
}
#[cfg(feature = "lock-census")]
pub(crate) fn set_enabled(on: bool) {
ENABLED.store(on, Ordering::Relaxed);
}
#[cfg(not(feature = "lock-census"))]
#[inline]
pub(crate) fn set_enabled(_on: bool) {}
#[cfg(feature = "lock-census")]
#[inline]
pub(crate) fn mutex_lock<T>(site: Site, m: &std::sync::Mutex<T>) -> std::sync::MutexGuard<'_, T> {
if !is_enabled() {
return m.lock().unwrap_or_else(|e| e.into_inner());
}
match m.try_lock() {
Ok(guard) => {
record_acquire(site);
guard
}
Err(std::sync::TryLockError::WouldBlock) => {
let start = Instant::now();
let guard = m.lock().unwrap_or_else(|e| e.into_inner());
record_contended(site, start.elapsed());
guard
}
Err(std::sync::TryLockError::Poisoned(p)) => {
record_acquire(site);
p.into_inner()
}
}
}
#[cfg(not(feature = "lock-census"))]
#[inline]
pub(crate) fn mutex_lock<T>(_site: Site, m: &std::sync::Mutex<T>) -> std::sync::MutexGuard<'_, T> {
m.lock().unwrap_or_else(|e| e.into_inner())
}
#[cfg(feature = "lock-census")]
#[inline]
pub(crate) fn rwlock_read<T>(site: Site, l: &std::sync::RwLock<T>) -> std::sync::RwLockReadGuard<'_, T> {
if !is_enabled() {
return l.read().unwrap_or_else(|e| e.into_inner());
}
match l.try_read() {
Ok(guard) => {
record_acquire(site);
guard
}
Err(std::sync::TryLockError::WouldBlock) => {
let start = Instant::now();
let guard = l.read().unwrap_or_else(|e| e.into_inner());
record_contended(site, start.elapsed());
guard
}
Err(std::sync::TryLockError::Poisoned(p)) => {
record_acquire(site);
p.into_inner()
}
}
}
#[cfg(not(feature = "lock-census"))]
#[inline]
pub(crate) fn rwlock_read<T>(_site: Site, l: &std::sync::RwLock<T>) -> std::sync::RwLockReadGuard<'_, T> {
l.read().unwrap_or_else(|e| e.into_inner())
}
#[cfg(feature = "lock-census")]
#[inline]
pub(crate) fn pl_rwlock_read<T>(
site: Site,
l: &parking_lot::RwLock<T>,
) -> parking_lot::RwLockReadGuard<'_, T> {
if !is_enabled() {
return l.read();
}
match l.try_read() {
Some(guard) => {
record_acquire(site);
guard
}
None => {
let start = Instant::now();
let guard = l.read();
record_contended(site, start.elapsed());
guard
}
}
}
#[cfg(not(feature = "lock-census"))]
#[inline]
pub(crate) fn pl_rwlock_read<T>(
_site: Site,
l: &parking_lot::RwLock<T>,
) -> parking_lot::RwLockReadGuard<'_, T> {
l.read()
}
#[cfg(feature = "lock-census")]
pub(crate) fn snapshot() -> Vec<SiteStat> {
SITES
.iter()
.zip(SITE_NAMES)
.map(|(counters, name)| SiteStat {
name,
acquisitions: counters.acquisitions.load(Ordering::Relaxed),
contended: counters.contended.load(Ordering::Relaxed),
contended_wait_nanos: counters.contended_wait_nanos.load(Ordering::Relaxed),
})
.collect()
}
#[cfg(not(feature = "lock-census"))]
#[inline]
pub(crate) fn snapshot() -> Vec<SiteStat> {
Vec::new()
}
#[cfg(all(test, feature = "lock-census"))]
#[allow(clippy::indexing_slicing)]
mod tests {
use super::{mutex_lock, set_enabled, snapshot, Site};
use std::sync::Mutex;
#[test]
fn census_records_and_snapshots() {
let names: Vec<&str> = snapshot().iter().map(|s| s.name).collect();
assert_eq!(
names,
vec![
"plan_cache_shard",
"parse_cache_shard",
"result_cache_shard",
"art_index_registry",
"art_pk_registry",
"statement_registry",
]
);
let m = Mutex::new(0u32);
set_enabled(false);
let before_off = snapshot()[0].acquisitions;
{
let _g = mutex_lock(Site::PlanCache, &m);
}
assert_eq!(snapshot()[0].acquisitions, before_off, "disabled census must not record");
set_enabled(true);
let before = snapshot()[0].acquisitions;
let contended_before = snapshot()[0].contended;
{
let _g = mutex_lock(Site::PlanCache, &m);
}
let after = snapshot();
assert!(after[0].acquisitions > before, "enabled census must record acquisition");
assert_eq!(after[0].contended, contended_before, "uncontended lock must not count as contended");
let unlabeled_acq = snapshot()[0].acquisitions;
{
let _g = mutex_lock(Site::Unlabeled, &m);
}
assert_eq!(snapshot()[0].acquisitions, unlabeled_acq, "Unlabeled must not record");
set_enabled(false);
}
}