1use std::collections::BTreeMap;
4use std::error::Error;
5use std::sync::Mutex;
6
7use crate::core::connection::{CloseReason, PooledConnection};
8use crate::observability::events::PoolListener;
9use crate::resilience::breaker::BreakerState;
10
11const ACQUIRE_WAIT_WINDOW: usize = 1024;
15
16#[derive(Debug, Clone, PartialEq)]
18pub enum StatValue {
19 Str(String),
21 Int(i64),
23 Float(f64),
25 Map(BTreeMap<String, i64>),
27 Null,
29}
30
31fn percentile(sorted: &[f64], q: f64) -> f64 {
33 if sorted.is_empty() {
34 return 0.0;
35 }
36 let k = (sorted.len() - 1) as f64 * q;
37 let lo = k as usize;
38 let hi = (lo + 1).min(sorted.len() - 1);
39 sorted[lo] + (sorted[hi] - sorted[lo]) * (k - lo as f64)
40}
41
42fn average(values: &[f64]) -> f64 {
43 if values.is_empty() {
44 return 0.0;
45 }
46 values.iter().sum::<f64>() / values.len() as f64
47}
48
49#[derive(Debug, Clone, PartialEq)]
52pub struct PoolStats {
53 pub name: String,
55 pub max_size: usize,
57 pub size: usize,
59 pub idle: usize,
61 pub in_use: usize,
63 pub waiting: usize,
65
66 pub created: u64,
68 pub create_failures: u64,
70 pub acquired: u64,
72 pub released: u64,
74 pub timeouts: u64,
76 pub closed: u64,
78 pub leaks: u64,
80 pub closed_by_reason: BTreeMap<String, i64>,
82
83 pub avg_acquire_wait: f64,
85 pub p50_acquire_wait: f64,
87 pub p95_acquire_wait: f64,
89 pub p99_acquire_wait: f64,
91
92 pub avg_create_latency: f64,
94 pub p50_create_latency: f64,
96 pub p95_create_latency: f64,
98 pub p99_create_latency: f64,
100
101 pub breaker_state: Option<String>,
103}
104
105impl PoolStats {
106 pub fn to_dictionary(&self) -> BTreeMap<&'static str, StatValue> {
108 let mut m = BTreeMap::new();
109 m.insert("name", StatValue::Str(self.name.clone()));
110 m.insert("max_size", StatValue::Int(self.max_size as i64));
111 m.insert("size", StatValue::Int(self.size as i64));
112 m.insert("idle", StatValue::Int(self.idle as i64));
113 m.insert("in_use", StatValue::Int(self.in_use as i64));
114 m.insert("waiting", StatValue::Int(self.waiting as i64));
115 m.insert("created", StatValue::Int(self.created as i64));
116 m.insert(
117 "create_failures",
118 StatValue::Int(self.create_failures as i64),
119 );
120 m.insert("acquired", StatValue::Int(self.acquired as i64));
121 m.insert("released", StatValue::Int(self.released as i64));
122 m.insert("timeouts", StatValue::Int(self.timeouts as i64));
123 m.insert("closed", StatValue::Int(self.closed as i64));
124 m.insert("leaks", StatValue::Int(self.leaks as i64));
125 m.insert(
126 "closed_by_reason",
127 StatValue::Map(self.closed_by_reason.clone()),
128 );
129 m.insert("avg_acquire_wait", StatValue::Float(self.avg_acquire_wait));
130 m.insert("p50_acquire_wait", StatValue::Float(self.p50_acquire_wait));
131 m.insert("p95_acquire_wait", StatValue::Float(self.p95_acquire_wait));
132 m.insert("p99_acquire_wait", StatValue::Float(self.p99_acquire_wait));
133 m.insert(
134 "avg_create_latency",
135 StatValue::Float(self.avg_create_latency),
136 );
137 m.insert(
138 "p50_create_latency",
139 StatValue::Float(self.p50_create_latency),
140 );
141 m.insert(
142 "p95_create_latency",
143 StatValue::Float(self.p95_create_latency),
144 );
145 m.insert(
146 "p99_create_latency",
147 StatValue::Float(self.p99_create_latency),
148 );
149 m.insert(
150 "breaker_state",
151 match &self.breaker_state {
152 Some(s) => StatValue::Str(s.clone()),
153 None => StatValue::Null,
154 },
155 );
156 m
157 }
158}
159
160pub(crate) struct StatsGauges {
162 pub name: String,
163 pub max_size: usize,
164 pub size: usize,
165 pub idle: usize,
166 pub in_use: usize,
167 pub waiting: usize,
168 pub breaker_state: Option<String>,
169}
170
171struct Counters {
172 created: u64,
173 create_failures: u64,
174 acquired: u64,
175 released: u64,
176 timeouts: u64,
177 closed: u64,
178 leaks: u64,
179 by_reason: BTreeMap<String, i64>,
180 waits: Vec<f64>,
181 creates: Vec<f64>,
182}
183
184pub(crate) struct StatsCollector<T> {
187 inner: Mutex<Counters>,
188 _marker: std::marker::PhantomData<fn(T)>,
189}
190
191fn push_window(window: &mut Vec<f64>, value: f64) {
192 window.push(value);
193 if window.len() > ACQUIRE_WAIT_WINDOW {
194 let overflow = window.len() - ACQUIRE_WAIT_WINDOW;
195 window.drain(0..overflow);
196 }
197}
198
199impl<T> StatsCollector<T> {
200 pub(crate) fn new() -> Self {
201 Self {
202 inner: Mutex::new(Counters {
203 created: 0,
204 create_failures: 0,
205 acquired: 0,
206 released: 0,
207 timeouts: 0,
208 closed: 0,
209 leaks: 0,
210 by_reason: BTreeMap::new(),
211 waits: Vec::new(),
212 creates: Vec::new(),
213 }),
214 _marker: std::marker::PhantomData,
215 }
216 }
217
218 pub(crate) fn snapshot(&self, g: StatsGauges) -> PoolStats {
219 let c = self.inner.lock().unwrap();
220 let mut waits = c.waits.clone();
221 waits.sort_by(|a, b| a.partial_cmp(b).unwrap());
222 let mut creates = c.creates.clone();
223 creates.sort_by(|a, b| a.partial_cmp(b).unwrap());
224 PoolStats {
225 name: g.name,
226 max_size: g.max_size,
227 size: g.size,
228 idle: g.idle,
229 in_use: g.in_use,
230 waiting: g.waiting,
231 created: c.created,
232 create_failures: c.create_failures,
233 acquired: c.acquired,
234 released: c.released,
235 timeouts: c.timeouts,
236 closed: c.closed,
237 leaks: c.leaks,
238 closed_by_reason: c.by_reason.clone(),
239 avg_acquire_wait: average(&waits),
240 p50_acquire_wait: percentile(&waits, 0.50),
241 p95_acquire_wait: percentile(&waits, 0.95),
242 p99_acquire_wait: percentile(&waits, 0.99),
243 avg_create_latency: average(&creates),
244 p50_create_latency: percentile(&creates, 0.50),
245 p95_create_latency: percentile(&creates, 0.95),
246 p99_create_latency: percentile(&creates, 0.99),
247 breaker_state: g.breaker_state,
248 }
249 }
250}
251
252impl<T> PoolListener<T> for StatsCollector<T> {
253 fn on_connection_created(&self, _conn: &PooledConnection<T>, duration: f64) {
254 let mut c = self.inner.lock().unwrap();
255 c.created += 1;
256 push_window(&mut c.creates, duration);
257 }
258
259 fn on_connection_closed(&self, _conn: &PooledConnection<T>, reason: CloseReason) {
260 let mut c = self.inner.lock().unwrap();
261 c.closed += 1;
262 *c.by_reason.entry(reason.as_str().to_string()).or_insert(0) += 1;
263 }
264
265 fn on_create_failed(&self, _error: &(dyn Error + 'static)) {
266 self.inner.lock().unwrap().create_failures += 1;
267 }
268
269 fn on_acquired(&self, _conn: &PooledConnection<T>, wait_time: f64) {
270 let mut c = self.inner.lock().unwrap();
271 c.acquired += 1;
272 push_window(&mut c.waits, wait_time);
273 }
274
275 fn on_released(&self, _conn: &PooledConnection<T>, _hold_time: f64) {
276 self.inner.lock().unwrap().released += 1;
277 }
278
279 fn on_acquire_timeout(&self, _waited: f64) {
280 self.inner.lock().unwrap().timeouts += 1;
281 }
282
283 fn on_connection_leak(&self, _conn: &PooledConnection<T>, _held_for: f64) {
284 self.inner.lock().unwrap().leaks += 1;
285 }
286
287 fn on_breaker_state_change(&self, _old: BreakerState, _new: BreakerState) {}
288}
289
290#[cfg(test)]
291mod tests {
292 use super::*;
293
294 fn gauges() -> StatsGauges {
295 StatsGauges {
296 name: "p".into(),
297 max_size: 10,
298 size: 2,
299 idle: 1,
300 in_use: 1,
301 waiting: 0,
302 breaker_state: None,
303 }
304 }
305
306 #[test]
307 fn percentile_and_average_edge_cases() {
308 assert_eq!(percentile(&[], 0.5), 0.0);
309 assert_eq!(average(&[]), 0.0);
310 let v = vec![1.0, 2.0, 3.0, 4.0];
311 assert_eq!(percentile(&v, 0.0), 1.0);
312 assert_eq!(percentile(&v, 1.0), 4.0);
313 assert_eq!(percentile(&v, 0.5), 2.5); assert_eq!(average(&v), 2.5);
315 }
316
317 #[test]
318 fn collector_counts_and_snapshots() {
319 let c = StatsCollector::<u32>::new();
320 let pc = PooledConnection::new(1u32, 0.0, None);
321 c.on_connection_created(&pc, 0.2);
322 c.on_connection_created(&pc, 0.4);
323 c.on_acquired(&pc, 0.1);
324 c.on_released(&pc, 1.0);
325 c.on_acquire_timeout(2.0);
326 c.on_connection_closed(&pc, CloseReason::Idle);
327 c.on_connection_closed(&pc, CloseReason::Idle);
328 c.on_connection_leak(&pc, 5.0);
329 let e: Box<dyn Error + Send + Sync> = "x".into();
330 c.on_create_failed(&*e);
331 c.on_breaker_state_change(BreakerState::Closed, BreakerState::Open);
332
333 let s = c.snapshot(gauges());
334 assert_eq!(s.created, 2);
335 assert_eq!(s.acquired, 1);
336 assert_eq!(s.released, 1);
337 assert_eq!(s.timeouts, 1);
338 assert_eq!(s.closed, 2);
339 assert_eq!(s.leaks, 1);
340 assert_eq!(s.create_failures, 1);
341 assert_eq!(s.closed_by_reason.get("idle"), Some(&2));
342 assert!((s.avg_create_latency - 0.3).abs() < 1e-9);
343 assert_eq!(s.name, "p");
344 assert_eq!(s.max_size, 10);
345 }
346
347 #[test]
348 fn to_dictionary_has_snakecase_keys() {
349 let c = StatsCollector::<u32>::new();
350 let mut g = gauges();
351 g.breaker_state = Some("closed".into());
352 let d = c.snapshot(g).to_dictionary();
353 for key in [
354 "name",
355 "max_size",
356 "size",
357 "idle",
358 "in_use",
359 "waiting",
360 "created",
361 "create_failures",
362 "acquired",
363 "released",
364 "timeouts",
365 "closed",
366 "leaks",
367 "closed_by_reason",
368 "avg_acquire_wait",
369 "p50_acquire_wait",
370 "p95_acquire_wait",
371 "p99_acquire_wait",
372 "avg_create_latency",
373 "p50_create_latency",
374 "p95_create_latency",
375 "p99_create_latency",
376 "breaker_state",
377 ] {
378 assert!(d.contains_key(key), "missing key {key}");
379 }
380 assert_eq!(d["breaker_state"], StatValue::Str("closed".into()));
381 assert_eq!(d["name"], StatValue::Str("p".into()));
382 assert_eq!(d["max_size"], StatValue::Int(10));
383 let d2 = c.snapshot(gauges()).to_dictionary();
385 assert_eq!(d2["breaker_state"], StatValue::Null);
386 }
387
388 #[test]
389 fn window_is_bounded() {
390 let c = StatsCollector::<u32>::new();
391 let pc = PooledConnection::new(1u32, 0.0, None);
392 for i in 0..(ACQUIRE_WAIT_WINDOW + 50) {
393 c.on_acquired(&pc, i as f64);
394 }
395 let inner = c.inner.lock().unwrap();
396 assert_eq!(inner.waits.len(), ACQUIRE_WAIT_WINDOW);
397 assert_eq!(inner.waits[0], 50.0);
399 }
400}