1use std::sync::atomic::{AtomicU64, Ordering};
11use std::sync::{Arc, Mutex};
12use std::time::{Duration, Instant};
13
14use rushdb_types::ids::{OwnerKind, TenantId};
15use rushdb_types::stats::Histogram;
16
17use crate::cache::{CacheStats, LogicalCache};
18use crate::resources::{ByteBudget, StorageResources};
19use crate::StorageError;
20
21pub type PlannerStatsKey = (TenantId, OwnerKind, Option<String>, u64);
23
24#[derive(Clone, Debug)]
27pub struct PlannerStatsEntry {
28 pub histogram: Option<Arc<Histogram>>,
30 pub computed_at_ci: u64,
32}
33
34struct Slot {
35 entry: PlannerStatsEntry,
36 last_used: AtomicU64,
38}
39
40pub type Clock = Arc<dyn Fn() -> Instant + Send + Sync>;
42
43pub struct PlannerStats {
45 cache: Mutex<Option<LogicalCache<PlannerStatsKey, Slot>>>,
46 idle: Duration,
47 max_age_commits: u64,
48 clock: Clock,
49 epoch: Instant,
50}
51
52impl PlannerStats {
53 fn stamp(&self) -> u64 {
54 u64::try_from(
55 (self.clock)()
56 .saturating_duration_since(self.epoch)
57 .as_nanos(),
58 )
59 .unwrap_or(u64::MAX)
60 }
61
62 pub fn new(
69 resources: &StorageResources,
70 parent: Option<&ByteBudget>,
71 clock: Clock,
72 ) -> Result<Self, StorageError> {
73 let invalid =
74 |e: crate::resources::ResourceError| StorageError::InvalidConfiguration(e.to_string());
75 let cache = if resources.planner_stats_bytes == 0 {
76 None
77 } else {
78 let budget = match parent {
79 Some(parent) => parent
80 .child(
81 "planner-stats",
82 resources.planner_stats_bytes.min(parent.snapshot().limit),
83 )
84 .map_err(invalid)?,
85 None => ByteBudget::new("planner-stats", resources.planner_stats_bytes)
86 .map_err(invalid)?,
87 };
88 Some(LogicalCache::new(budget).map_err(invalid)?)
89 };
90 Ok(Self {
91 cache: Mutex::new(cache),
92 idle: Duration::from_secs(resources.planner_stats_idle_seconds),
93 max_age_commits: resources.planner_stats_max_age_commits,
94 epoch: clock(),
95 clock,
96 })
97 }
98
99 #[must_use]
102 pub fn get(&self, key: &PlannerStatsKey, ci: u64) -> Option<PlannerStatsEntry> {
103 let now = self.stamp();
104 let mut guard = self.cache.lock().ok()?;
105 let cache = guard.as_mut()?;
106 let slot = cache.get(key)?;
107 if ci.saturating_sub(slot.entry.computed_at_ci) > self.max_age_commits {
108 return None;
109 }
110 slot.last_used.store(now, Ordering::Relaxed);
111 Some(slot.entry.clone())
112 }
113
114 pub fn insert(&self, key: PlannerStatsKey, entry: PlannerStatsEntry) -> bool {
117 let bytes = std::mem::size_of::<Slot>()
118 + std::mem::size_of::<PlannerStatsKey>()
119 + key.2.as_ref().map_or(0, String::capacity)
120 + entry.histogram.as_ref().map_or(0, |h| h.resident_bytes());
121 let Ok(mut guard) = self.cache.lock() else {
122 return false;
123 };
124 let Some(cache) = guard.as_mut() else {
125 return false;
126 };
127 let slot = Slot {
128 entry,
129 last_used: AtomicU64::new(self.stamp()),
130 };
131 cache.insert(key, slot, bytes)
132 }
133
134 pub fn evict_idle(&self) -> usize {
137 let now = self.stamp();
138 let idle = u64::try_from(self.idle.as_nanos()).unwrap_or(u64::MAX);
139 let Ok(mut guard) = self.cache.lock() else {
140 return 0;
141 };
142 guard.as_mut().map_or(0, |cache| {
143 cache
144 .retain(|_, slot| now.saturating_sub(slot.last_used.load(Ordering::Relaxed)) < idle)
145 })
146 }
147
148 #[must_use]
150 pub fn stats(&self) -> CacheStats {
151 self.cache
152 .lock()
153 .ok()
154 .and_then(|guard| guard.as_ref().map(LogicalCache::stats))
155 .unwrap_or_default()
156 }
157}
158
159#[cfg(test)]
160mod tests {
161 use super::*;
162
163 fn mock_clock() -> (Clock, Arc<AtomicU64>) {
165 let base = Instant::now();
166 let offset = Arc::new(AtomicU64::new(0));
167 let read = Arc::clone(&offset);
168 let clock: Clock =
169 Arc::new(move || base + Duration::from_secs(read.load(Ordering::Relaxed)));
170 (clock, offset)
171 }
172
173 fn cache(bytes: usize) -> (PlannerStats, Arc<AtomicU64>) {
174 let (clock, offset) = mock_clock();
175 let resources = StorageResources {
176 planner_stats_bytes: bytes,
177 ..StorageResources::default()
178 };
179 (PlannerStats::new(&resources, None, clock).unwrap(), offset)
180 }
181
182 fn entry(ci: u64) -> PlannerStatsEntry {
183 let mut histogram = Histogram::default();
184 histogram.add(&[1, 2, 3], 7);
185 PlannerStatsEntry {
186 histogram: Some(Arc::new(histogram)),
187 computed_at_ci: ci,
188 }
189 }
190
191 fn key(pid: u64) -> PlannerStatsKey {
192 (TenantId::from_ordinal(2), OwnerKind::Record, None, pid)
193 }
194
195 #[test]
196 fn stats_cache_misses_load_from_disk_and_evict_after_idle() {
197 let (stats, clock) = cache(1 << 20);
198 assert!(stats.get(&key(1), 5).is_none(), "cold cache misses");
199 assert!(stats.insert(key(1), entry(5)));
200 assert_eq!(stats.get(&key(1), 6).unwrap().computed_at_ci, 5);
201 clock.store(599, Ordering::Relaxed);
202 assert_eq!(stats.evict_idle(), 0, "used 599 s ago: still warm");
203 assert!(stats.get(&key(1), 6).is_some());
205 clock.store(599 + 600, Ordering::Relaxed);
206 assert_eq!(stats.evict_idle(), 1);
207 assert_eq!(stats.stats().bytes, 0, "an idle cache holds nothing");
208 assert!(stats.get(&key(1), 6).is_none());
209 }
210
211 #[test]
212 fn stats_cache_is_budget_bounded() {
213 let budget = 16 * 1024;
214 let (stats, _) = cache(budget);
215 for pid in 0..500 {
216 stats.insert(key(pid), entry(1));
217 assert!(stats.stats().bytes <= budget as u64);
218 }
219 let occupancy = stats.stats();
220 assert!(occupancy.evictions > 0, "LRU made room: {occupancy:?}");
221 assert!(stats.get(&key(499), 1).is_some(), "the newest entry stays");
222 }
223
224 #[test]
225 fn estimates_refresh_after_max_age_commits() {
226 let (stats, _) = cache(1 << 20);
227 stats.insert(key(1), entry(10));
228 assert!(stats.get(&key(1), 10 + 10_000).is_some());
229 assert!(
230 stats.get(&key(1), 10 + 10_001).is_none(),
231 "older than planner-stats-max-age-commits: reload"
232 );
233 }
234
235 #[test]
236 fn zero_budget_disables_caching() {
237 let (stats, _) = cache(0);
238 assert!(!stats.insert(key(1), entry(1)));
239 assert!(stats.get(&key(1), 1).is_none());
240 }
241}