use std::sync::atomic::{AtomicU64, AtomicUsize, Ordering};
use std::time::{Duration, Instant};
#[derive(Debug, Clone)]
pub struct LockOptimizeConfig {
pub enabled: bool,
pub acquire_timeout: Duration,
pub max_hold_time_warning: Duration,
pub adaptive_spin: bool,
pub max_spin_iterations: usize,
}
impl Default for LockOptimizeConfig {
fn default() -> Self {
Self {
enabled: true,
acquire_timeout: Duration::from_secs(5),
max_hold_time_warning: Duration::from_millis(100),
adaptive_spin: true,
max_spin_iterations: 1000,
}
}
}
#[derive(Debug, Default)]
pub struct LockStats {
pub locks_acquired: AtomicU64,
pub locks_released_early: AtomicU64,
pub total_hold_time_ns: AtomicU64,
pub max_hold_time_ns: AtomicU64,
pub contentions: AtomicU64,
pub spin_successes: AtomicU64,
pub spin_failures: AtomicU64,
}
impl LockStats {
pub fn new() -> Self {
Self::default()
}
pub fn record_acquire(&self) {
self.locks_acquired.fetch_add(1, Ordering::Relaxed);
}
pub fn record_release(&self, hold_time: Duration) {
let ns = hold_time.as_nanos() as u64;
self.total_hold_time_ns.fetch_add(ns, Ordering::Relaxed);
let mut current = self.max_hold_time_ns.load(Ordering::Relaxed);
while ns > current {
match self
.max_hold_time_ns
.compare_exchange_weak(current, ns, Ordering::Relaxed, Ordering::Relaxed)
{
Ok(_) => break,
Err(actual) => current = actual,
}
}
}
pub fn record_early_release(&self) {
self.locks_released_early.fetch_add(1, Ordering::Relaxed);
}
pub fn record_contention(&self) {
self.contentions.fetch_add(1, Ordering::Relaxed);
}
pub fn record_spin_success(&self) {
self.spin_successes.fetch_add(1, Ordering::Relaxed);
}
pub fn record_spin_failure(&self) {
self.spin_failures.fetch_add(1, Ordering::Relaxed);
}
pub fn total_acquired(&self) -> u64 {
self.locks_acquired.load(Ordering::Relaxed)
}
pub fn avg_hold_time(&self) -> Duration {
let total = self.total_hold_time_ns.load(Ordering::Relaxed);
let count = self.locks_acquired.load(Ordering::Relaxed);
total.checked_div(count).map(Duration::from_nanos).unwrap_or(Duration::ZERO)
}
pub fn max_hold_time(&self) -> Duration {
Duration::from_nanos(self.max_hold_time_ns.load(Ordering::Relaxed))
}
pub fn contention_rate(&self) -> f64 {
let acquired = self.locks_acquired.load(Ordering::Relaxed);
let contentions = self.contentions.load(Ordering::Relaxed);
if acquired == 0 {
0.0
} else {
contentions as f64 / acquired as f64
}
}
pub fn spin_success_rate(&self) -> f64 {
let successes = self.spin_successes.load(Ordering::Relaxed);
let failures = self.spin_failures.load(Ordering::Relaxed);
let total = successes + failures;
if total == 0 { 0.0 } else { successes as f64 / total as f64 }
}
pub fn reset(&self) {
self.locks_acquired.store(0, Ordering::Relaxed);
self.locks_released_early.store(0, Ordering::Relaxed);
self.total_hold_time_ns.store(0, Ordering::Relaxed);
self.max_hold_time_ns.store(0, Ordering::Relaxed);
self.contentions.store(0, Ordering::Relaxed);
self.spin_successes.store(0, Ordering::Relaxed);
self.spin_failures.store(0, Ordering::Relaxed);
}
}
pub struct LockOptimizer {
config: LockOptimizeConfig,
stats: LockStats,
current_spin: AtomicUsize,
}
impl LockOptimizer {
pub fn new(config: LockOptimizeConfig) -> Self {
Self {
config,
stats: LockStats::new(),
current_spin: AtomicUsize::new(100),
}
}
pub fn with_defaults() -> Self {
Self::new(LockOptimizeConfig::default())
}
pub fn config(&self) -> &LockOptimizeConfig {
&self.config
}
pub fn stats(&self) -> &LockStats {
&self.stats
}
pub fn on_acquire(&self) {
if !self.config.enabled {
return;
}
self.stats.record_acquire();
}
pub fn on_release(&self, hold_time: Duration) {
if !self.config.enabled {
return;
}
self.stats.record_release(hold_time);
if hold_time < self.config.acquire_timeout / 2 {
self.stats.record_early_release();
}
}
pub fn on_contention(&self) {
if !self.config.enabled {
return;
}
self.stats.record_contention();
}
pub fn try_spin<F>(&self, mut try_acquire: F) -> bool
where
F: FnMut() -> bool,
{
if !self.config.enabled || !self.config.adaptive_spin {
return false;
}
let spin_count = self.current_spin.load(Ordering::Relaxed).min(self.config.max_spin_iterations);
for _ in 0..spin_count {
if try_acquire() {
self.stats.record_spin_success();
self.adapt_spin(true);
return true;
}
std::hint::spin_loop();
}
self.stats.record_spin_failure();
self.adapt_spin(false);
false
}
fn adapt_spin(&self, success: bool) {
let current = self.current_spin.load(Ordering::Relaxed);
let new_count = if success {
(current * 2).min(self.config.max_spin_iterations)
} else {
(current / 2).max(10)
};
self.current_spin.store(new_count, Ordering::Relaxed);
}
pub fn current_spin_count(&self) -> usize {
self.current_spin.load(Ordering::Relaxed)
}
pub fn is_hold_time_excessive(&self, hold_time: Duration) -> bool {
hold_time > self.config.max_hold_time_warning
}
pub fn reset_stats(&self) {
self.stats.reset();
}
}
pub struct LockGuard<'a> {
optimizer: &'a LockOptimizer,
start: Instant,
}
impl<'a> LockGuard<'a> {
pub fn new(optimizer: &'a LockOptimizer) -> Self {
optimizer.on_acquire();
Self {
optimizer,
start: Instant::now(),
}
}
}
impl Drop for LockGuard<'_> {
fn drop(&mut self) {
self.optimizer.on_release(self.start.elapsed());
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_lock_stats() {
let stats = LockStats::new();
stats.record_acquire();
stats.record_acquire();
stats.record_release(Duration::from_millis(10));
stats.record_release(Duration::from_millis(20));
assert_eq!(stats.total_acquired(), 2);
assert!(stats.avg_hold_time() >= Duration::from_millis(15));
assert_eq!(stats.max_hold_time(), Duration::from_millis(20));
}
#[test]
fn test_contention_rate() {
let stats = LockStats::new();
stats.record_acquire();
stats.record_acquire();
stats.record_acquire();
stats.record_contention();
assert!((stats.contention_rate() - 0.3333333333333333).abs() < 0.01);
}
#[test]
fn test_spin_stats() {
let stats = LockStats::new();
stats.record_spin_success();
stats.record_spin_success();
stats.record_spin_failure();
assert!((stats.spin_success_rate() - 0.6666666666666666).abs() < 0.01);
}
#[test]
fn test_lock_optimizer() {
let optimizer = LockOptimizer::with_defaults();
{
let _guard = LockGuard::new(&optimizer);
std::thread::sleep(Duration::from_millis(10));
}
assert_eq!(optimizer.stats().total_acquired(), 1);
assert!(optimizer.stats().avg_hold_time() >= Duration::from_millis(10));
}
#[test]
fn test_adaptive_spin() {
let optimizer = LockOptimizer::with_defaults();
let acquired = optimizer.try_spin(|| true);
assert!(acquired);
assert!(optimizer.current_spin_count() > 100);
let acquired = optimizer.try_spin(|| false);
assert!(!acquired);
assert!(optimizer.current_spin_count() < 200); }
#[test]
fn test_disabled_optimizer() {
let config = LockOptimizeConfig {
enabled: false,
..Default::default()
};
let optimizer = LockOptimizer::new(config);
optimizer.on_acquire();
optimizer.on_release(Duration::from_millis(10));
assert_eq!(optimizer.stats().total_acquired(), 0);
}
}