pub use linkrs_metrics::utils::CacheStats;
use parking_lot::RwLock;
use std::sync::atomic::{AtomicUsize, Ordering};
use std::sync::Arc;
#[derive(Debug, Default)]
pub struct MemoryStats {
current_bytes: AtomicUsize,
max_bytes: AtomicUsize,
entry_count: AtomicUsize,
}
impl MemoryStats {
pub fn new(max_bytes: usize) -> Self {
Self {
current_bytes: AtomicUsize::new(0),
max_bytes: AtomicUsize::new(max_bytes),
entry_count: AtomicUsize::new(0),
}
}
pub fn update(&self, current_bytes: usize, entry_count: usize) {
self.current_bytes.store(current_bytes, Ordering::Relaxed);
self.entry_count.store(entry_count, Ordering::Relaxed);
}
pub fn set_max_bytes(&self, max_bytes: usize) {
self.max_bytes.store(max_bytes, Ordering::Relaxed);
}
pub fn current_bytes(&self) -> usize {
self.current_bytes.load(Ordering::Relaxed)
}
pub fn max_bytes(&self) -> usize {
self.max_bytes.load(Ordering::Relaxed)
}
pub fn entry_count(&self) -> usize {
self.entry_count.load(Ordering::Relaxed)
}
pub fn usage_ratio(&self) -> f64 {
let max = self.max_bytes();
if max == 0 {
0.0
} else {
self.current_bytes() as f64 / max as f64
}
}
pub fn reset(&self) {
self.current_bytes.store(0, Ordering::Relaxed);
self.entry_count.store(0, Ordering::Relaxed);
}
}
impl Clone for MemoryStats {
fn clone(&self) -> Self {
Self {
current_bytes: AtomicUsize::new(self.current_bytes()),
max_bytes: AtomicUsize::new(self.max_bytes()),
entry_count: AtomicUsize::new(self.entry_count()),
}
}
}
#[derive(Debug)]
pub struct PlanCacheStats {
pub counters: CacheStats,
pub memory: MemoryStats,
pub avg_query_size: Arc<RwLock<usize>>,
pub total_query_size: AtomicUsize,
}
impl Clone for PlanCacheStats {
fn clone(&self) -> Self {
Self {
counters: self.counters.clone(),
memory: self.memory.clone(),
avg_query_size: Arc::new(RwLock::new(*self.avg_query_size.read())),
total_query_size: AtomicUsize::new(self.total_query_size.load(Ordering::Relaxed)),
}
}
}
impl PlanCacheStats {
pub fn new(memory_budget: usize) -> Self {
Self {
counters: CacheStats::new(),
memory: MemoryStats::new(memory_budget),
avg_query_size: Arc::new(RwLock::new(0)),
total_query_size: AtomicUsize::new(0),
}
}
pub fn record_query_size(&self, size: usize) {
self.total_query_size.fetch_add(size, Ordering::Relaxed);
let total = self.total_query_size.load(Ordering::Relaxed);
let count = self.memory.entry_count();
if let Some(avg) = total.checked_div(count) {
*self.avg_query_size.write() = avg;
}
}
pub fn avg_query_size(&self) -> usize {
*self.avg_query_size.read()
}
pub fn estimated_memory(&self) -> usize {
const PER_ENTRY_OVERHEAD: usize = 1024;
let total_query = self.total_query_size.load(Ordering::Relaxed);
let entries = self.memory.entry_count();
total_query + (entries * PER_ENTRY_OVERHEAD)
}
pub fn hit_rate(&self) -> f64 {
self.counters.hit_rate()
}
pub fn reset(&self) {
self.counters.reset();
self.memory.reset();
self.total_query_size.store(0, Ordering::Relaxed);
*self.avg_query_size.write() = 0;
}
pub fn snapshot(&self) -> PlanCacheStatsSnapshot {
PlanCacheStatsSnapshot {
hits: self.counters.hits(),
misses: self.counters.misses(),
evictions: self.counters.evictions(),
expirations: self.counters.expirations(),
entry_count: self.memory.entry_count(),
current_memory: self.memory.current_bytes(),
max_memory: self.memory.max_bytes(),
hit_rate: self.hit_rate(),
avg_query_size: self.avg_query_size(),
}
}
}
impl Default for PlanCacheStats {
fn default() -> Self {
Self::new(50 * 1024 * 1024)
}
}
#[derive(Debug, Clone, Default)]
pub struct PlanCacheStatsSnapshot {
pub hits: u64,
pub misses: u64,
pub evictions: u64,
pub expirations: u64,
pub entry_count: usize,
pub current_memory: usize,
pub max_memory: usize,
pub hit_rate: f64,
pub avg_query_size: usize,
}
#[derive(Debug, Clone)]
pub struct CteCacheStats {
pub counters: CacheStats,
pub memory: MemoryStats,
}
impl CteCacheStats {
pub fn new(max_size: usize) -> Self {
Self {
counters: CacheStats::new(),
memory: MemoryStats::new(max_size),
}
}
pub fn hit_rate(&self) -> f64 {
self.counters.hit_rate()
}
pub fn memory_usage_ratio(&self) -> f64 {
self.memory.usage_ratio()
}
pub fn reset(&self) {
self.counters.reset();
self.memory.reset();
}
pub fn snapshot(&self) -> CteCacheStatsSnapshot {
CteCacheStatsSnapshot {
hits: self.counters.hits(),
misses: self.counters.misses(),
evictions: self.counters.evictions(),
rejections: self.counters.rejections(),
entry_count: self.memory.entry_count(),
current_memory: self.memory.current_bytes(),
max_memory: self.memory.max_bytes(),
hit_rate: self.hit_rate(),
memory_usage_ratio: self.memory_usage_ratio(),
}
}
}
impl Default for CteCacheStats {
fn default() -> Self {
Self::new(64 * 1024 * 1024)
}
}
#[derive(Debug, Clone, Default)]
pub struct CteCacheStatsSnapshot {
pub hits: u64,
pub misses: u64,
pub evictions: u64,
pub rejections: u64,
pub entry_count: usize,
pub current_memory: usize,
pub max_memory: usize,
pub hit_rate: f64,
pub memory_usage_ratio: f64,
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_cache_stats() {
let stats = CacheStats::new();
stats.record_hit();
stats.record_hit();
stats.record_miss();
assert_eq!(stats.hits(), 2);
assert_eq!(stats.misses(), 1);
assert_eq!(stats.total_requests(), 3);
assert!((stats.hit_rate() - 0.6666666666666666).abs() < 0.01);
}
#[test]
fn test_cache_stats_reset() {
let stats = CacheStats::new();
stats.record_hit();
stats.record_miss();
stats.reset();
assert_eq!(stats.hits(), 0);
assert_eq!(stats.misses(), 0);
}
#[test]
fn test_memory_stats() {
let stats = MemoryStats::new(1000);
stats.update(500, 10);
assert_eq!(stats.current_bytes(), 500);
assert_eq!(stats.entry_count(), 10);
assert_eq!(stats.usage_ratio(), 0.5);
}
#[test]
fn test_plan_cache_stats() {
let stats = PlanCacheStats::new(1024 * 1024);
stats.counters.record_hit();
stats.counters.record_miss();
stats.memory.update(512 * 1024, 5);
let snapshot = stats.snapshot();
assert_eq!(snapshot.hits, 1);
assert_eq!(snapshot.misses, 1);
assert_eq!(snapshot.entry_count, 5);
}
#[test]
fn test_cte_cache_stats() {
let stats = CteCacheStats::new(1024 * 1024);
stats.counters.record_hit();
stats.counters.record_hit();
stats.memory.update(256 * 1024, 3);
let snapshot = stats.snapshot();
assert_eq!(snapshot.hits, 2);
assert_eq!(snapshot.entry_count, 3);
assert!((snapshot.hit_rate - 1.0).abs() < 0.01);
}
}