use async_trait::async_trait;
use chrono::{DateTime, Utc};
use parking_lot::RwLock;
use serde::{Deserialize, Serialize};
use std::collections::HashMap;
use std::sync::Arc;
use std::time::Duration;
use tracing::{debug, error, info, warn};
use uuid::Uuid;
#[derive(Debug, Clone, PartialEq, Eq, Hash, Serialize, Deserialize)]
pub enum PoolEventType {
ConnectionAcquired,
ConnectionReleased,
AcquisitionTimeout,
PoolExhausted,
ConnectionCreated,
ConnectionClosed,
ValidationFailed,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct PoolEvent {
pub event_id: Uuid,
pub event_type: PoolEventType,
pub timestamp: DateTime<Utc>,
pub pool_size: u32,
pub idle_connections: u32,
pub wait_time_ms: Option<u64>,
pub error: Option<String>,
pub metadata: HashMap<String, String>,
}
impl PoolEvent {
pub fn new(event_type: PoolEventType, pool_size: u32, idle_connections: u32) -> Self {
Self {
event_id: Uuid::new_v4(),
event_type,
timestamp: Utc::now(),
pool_size,
idle_connections,
wait_time_ms: None,
error: None,
metadata: HashMap::new(),
}
}
pub fn with_wait_time(mut self, wait_time: Duration) -> Self {
self.wait_time_ms = Some(wait_time.as_millis() as u64);
self
}
pub fn with_error(mut self, error: String) -> Self {
self.error = Some(error);
self
}
pub fn add_metadata(mut self, key: String, value: String) -> Self {
self.metadata.insert(key, value);
self
}
}
#[async_trait]
pub trait PoolEventHandler: Send + Sync {
async fn handle_event(&self, event: PoolEvent);
}
#[derive(Debug, Clone)]
pub struct LoggingEventHandler {
pub log_levels: HashMap<PoolEventType, LogLevel>,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum LogLevel {
Debug,
Info,
Warn,
Error,
}
impl Default for LoggingEventHandler {
fn default() -> Self {
let mut log_levels = HashMap::new();
log_levels.insert(PoolEventType::ConnectionAcquired, LogLevel::Debug);
log_levels.insert(PoolEventType::ConnectionReleased, LogLevel::Debug);
log_levels.insert(PoolEventType::AcquisitionTimeout, LogLevel::Warn);
log_levels.insert(PoolEventType::PoolExhausted, LogLevel::Error);
log_levels.insert(PoolEventType::ConnectionCreated, LogLevel::Info);
log_levels.insert(PoolEventType::ConnectionClosed, LogLevel::Info);
log_levels.insert(PoolEventType::ValidationFailed, LogLevel::Warn);
Self { log_levels }
}
}
#[async_trait]
impl PoolEventHandler for LoggingEventHandler {
async fn handle_event(&self, event: PoolEvent) {
let level = self
.log_levels
.get(&event.event_type)
.copied()
.unwrap_or(LogLevel::Info);
let message = format!(
"Pool event: {:?}, size: {}, idle: {}",
event.event_type, event.pool_size, event.idle_connections
);
match level {
LogLevel::Debug => debug!(
event_id = %event.event_id,
event_type = ?event.event_type,
pool_size = event.pool_size,
idle = event.idle_connections,
wait_time_ms = ?event.wait_time_ms,
"{}", message
),
LogLevel::Info => info!(
event_id = %event.event_id,
event_type = ?event.event_type,
pool_size = event.pool_size,
idle = event.idle_connections,
wait_time_ms = ?event.wait_time_ms,
"{}", message
),
LogLevel::Warn => warn!(
event_id = %event.event_id,
event_type = ?event.event_type,
pool_size = event.pool_size,
idle = event.idle_connections,
wait_time_ms = ?event.wait_time_ms,
error = ?event.error,
"{}", message
),
LogLevel::Error => error!(
event_id = %event.event_id,
event_type = ?event.event_type,
pool_size = event.pool_size,
idle = event.idle_connections,
wait_time_ms = ?event.wait_time_ms,
error = ?event.error,
"{}", message
),
}
}
}
#[derive(Debug, Clone)]
pub struct MetricsEventHandler {
counts: Arc<RwLock<HashMap<PoolEventType, u64>>>,
total_wait_time_ms: Arc<RwLock<u64>>,
wait_time_samples: Arc<RwLock<u64>>,
}
impl Default for MetricsEventHandler {
fn default() -> Self {
Self {
counts: Arc::new(RwLock::new(HashMap::new())),
total_wait_time_ms: Arc::new(RwLock::new(0)),
wait_time_samples: Arc::new(RwLock::new(0)),
}
}
}
impl MetricsEventHandler {
pub fn get_count(&self, event_type: &PoolEventType) -> u64 {
self.counts.read().get(event_type).copied().unwrap_or(0)
}
pub fn get_avg_wait_time_ms(&self) -> f64 {
let total = *self.total_wait_time_ms.read();
let samples = *self.wait_time_samples.read();
if samples == 0 {
0.0
} else {
total as f64 / samples as f64
}
}
pub fn reset(&self) {
self.counts.write().clear();
*self.total_wait_time_ms.write() = 0;
*self.wait_time_samples.write() = 0;
}
}
#[async_trait]
impl PoolEventHandler for MetricsEventHandler {
async fn handle_event(&self, event: PoolEvent) {
let mut counts = self.counts.write();
*counts.entry(event.event_type.clone()).or_insert(0) += 1;
drop(counts);
if let Some(wait_time_ms) = event.wait_time_ms {
*self.total_wait_time_ms.write() += wait_time_ms;
*self.wait_time_samples.write() += 1;
}
}
}
#[derive(Clone)]
pub struct AlertEventHandler {
alert_callback: Arc<dyn Fn(PoolEvent) + Send + Sync>,
}
impl AlertEventHandler {
pub fn new<F>(callback: F) -> Self
where
F: Fn(PoolEvent) + Send + Sync + 'static,
{
Self {
alert_callback: Arc::new(callback),
}
}
}
#[async_trait]
impl PoolEventHandler for AlertEventHandler {
async fn handle_event(&self, event: PoolEvent) {
match event.event_type {
PoolEventType::AcquisitionTimeout
| PoolEventType::PoolExhausted
| PoolEventType::ValidationFailed => {
(self.alert_callback)(event);
}
_ => {}
}
}
}
#[derive(Clone)]
pub struct PoolEventManager {
handlers: Arc<RwLock<Vec<Arc<dyn PoolEventHandler>>>>,
}
impl Default for PoolEventManager {
fn default() -> Self {
Self::new()
}
}
impl PoolEventManager {
pub fn new() -> Self {
Self {
handlers: Arc::new(RwLock::new(Vec::new())),
}
}
pub fn add_handler(&self, handler: Arc<dyn PoolEventHandler>) {
self.handlers.write().push(handler);
}
pub async fn emit(&self, event: PoolEvent) {
let handlers = self.handlers.read().clone();
for handler in handlers {
handler.handle_event(event.clone()).await;
}
}
pub async fn emit_acquired(&self, pool_size: u32, idle: u32, wait_time: Duration) {
let event = PoolEvent::new(PoolEventType::ConnectionAcquired, pool_size, idle)
.with_wait_time(wait_time);
self.emit(event).await;
}
pub async fn emit_released(&self, pool_size: u32, idle: u32) {
let event = PoolEvent::new(PoolEventType::ConnectionReleased, pool_size, idle);
self.emit(event).await;
}
pub async fn emit_timeout(&self, pool_size: u32, idle: u32, timeout: Duration) {
let event = PoolEvent::new(PoolEventType::AcquisitionTimeout, pool_size, idle)
.with_wait_time(timeout)
.with_error(format!(
"Connection acquisition timed out after {}ms",
timeout.as_millis()
));
self.emit(event).await;
}
pub async fn emit_exhausted(&self, pool_size: u32) {
let event = PoolEvent::new(PoolEventType::PoolExhausted, pool_size, 0)
.with_error("Pool has no available connections".to_string());
self.emit(event).await;
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_pool_event_creation() {
let event = PoolEvent::new(PoolEventType::ConnectionAcquired, 10, 5);
assert_eq!(event.event_type, PoolEventType::ConnectionAcquired);
assert_eq!(event.pool_size, 10);
assert_eq!(event.idle_connections, 5);
assert!(event.wait_time_ms.is_none());
assert!(event.error.is_none());
}
#[test]
fn test_pool_event_with_wait_time() {
let event = PoolEvent::new(PoolEventType::ConnectionAcquired, 10, 5)
.with_wait_time(Duration::from_millis(150));
assert_eq!(event.wait_time_ms, Some(150));
}
#[test]
fn test_pool_event_with_error() {
let event = PoolEvent::new(PoolEventType::AcquisitionTimeout, 10, 0)
.with_error("Timeout occurred".to_string());
assert_eq!(event.error, Some("Timeout occurred".to_string()));
}
#[tokio::test]
async fn test_logging_event_handler() {
let handler = LoggingEventHandler::default();
let event = PoolEvent::new(PoolEventType::ConnectionAcquired, 10, 5);
handler.handle_event(event).await;
}
#[tokio::test]
async fn test_metrics_event_handler() {
let handler = MetricsEventHandler::default();
for _ in 0..5 {
let event = PoolEvent::new(PoolEventType::ConnectionAcquired, 10, 5)
.with_wait_time(Duration::from_millis(100));
handler.handle_event(event).await;
}
assert_eq!(handler.get_count(&PoolEventType::ConnectionAcquired), 5);
assert!((handler.get_avg_wait_time_ms() - 100.0).abs() < 0.1);
}
#[tokio::test]
async fn test_metrics_event_handler_reset() {
let handler = MetricsEventHandler::default();
let event = PoolEvent::new(PoolEventType::ConnectionAcquired, 10, 5)
.with_wait_time(Duration::from_millis(100));
handler.handle_event(event).await;
handler.reset();
assert_eq!(handler.get_count(&PoolEventType::ConnectionAcquired), 0);
assert_eq!(handler.get_avg_wait_time_ms(), 0.0);
}
#[tokio::test]
async fn test_alert_event_handler() {
let alert_triggered = Arc::new(RwLock::new(false));
let alert_triggered_clone = alert_triggered.clone();
let handler = AlertEventHandler::new(move |_event| {
*alert_triggered_clone.write() = true;
});
let event = PoolEvent::new(PoolEventType::ConnectionAcquired, 10, 5);
handler.handle_event(event).await;
assert!(!*alert_triggered.read());
let event = PoolEvent::new(PoolEventType::PoolExhausted, 10, 0);
handler.handle_event(event).await;
assert!(*alert_triggered.read());
}
#[tokio::test]
async fn test_pool_event_manager() {
let manager = PoolEventManager::new();
let metrics = Arc::new(MetricsEventHandler::default());
manager.add_handler(metrics.clone());
manager
.emit_acquired(10, 5, Duration::from_millis(100))
.await;
manager.emit_released(10, 6).await;
assert_eq!(metrics.get_count(&PoolEventType::ConnectionAcquired), 1);
assert_eq!(metrics.get_count(&PoolEventType::ConnectionReleased), 1);
}
#[tokio::test]
async fn test_pool_event_manager_multiple_handlers() {
let manager = PoolEventManager::new();
let metrics1 = Arc::new(MetricsEventHandler::default());
let metrics2 = Arc::new(MetricsEventHandler::default());
manager.add_handler(metrics1.clone());
manager.add_handler(metrics2.clone());
let event = PoolEvent::new(PoolEventType::ConnectionAcquired, 10, 5);
manager.emit(event).await;
assert_eq!(metrics1.get_count(&PoolEventType::ConnectionAcquired), 1);
assert_eq!(metrics2.get_count(&PoolEventType::ConnectionAcquired), 1);
}
#[tokio::test]
async fn test_emit_timeout() {
let manager = PoolEventManager::new();
let metrics = Arc::new(MetricsEventHandler::default());
manager.add_handler(metrics.clone());
manager.emit_timeout(10, 0, Duration::from_secs(5)).await;
assert_eq!(metrics.get_count(&PoolEventType::AcquisitionTimeout), 1);
}
#[tokio::test]
async fn test_emit_exhausted() {
let manager = PoolEventManager::new();
let metrics = Arc::new(MetricsEventHandler::default());
manager.add_handler(metrics.clone());
manager.emit_exhausted(10).await;
assert_eq!(metrics.get_count(&PoolEventType::PoolExhausted), 1);
}
}