6093eac7bf
Wave 64-65 cleanup: Proto regeneration and build system updates from Tonic 0.12→0.14 upgrade Files updated: - Cargo.lock: Dependency resolution for Tonic 0.14.2 - All build.rs: Updated for tonic-prost-build - Proto files: Regenerated with tonic-prost 0.14 - Examples/tests: Updated for new gRPC API 🤖 Generated with [Claude Code](https://claude.com/claude-code) Co-Authored-By: Claude <noreply@anthropic.com>
558 lines
18 KiB
Rust
558 lines
18 KiB
Rust
use crate::error::{DatabaseError, DatabaseResult};
|
|
use config::database::PoolConfig;
|
|
use sqlx::postgres::{PgPool, PgPoolOptions};
|
|
use std::sync::atomic::{AtomicU64, Ordering};
|
|
use std::sync::Arc;
|
|
use std::time::Duration;
|
|
use tokio::sync::RwLock;
|
|
use tracing::{debug, error, info, warn};
|
|
|
|
// PoolConfig is now imported from the config crate
|
|
|
|
/// Extension trait for PoolConfig validation
|
|
///
|
|
/// Provides validation methods for pool configuration to ensure
|
|
/// settings are valid before creating a connection pool.
|
|
trait PoolConfigValidation {
|
|
/// Validate the pool configuration settings
|
|
///
|
|
/// # Errors
|
|
///
|
|
/// Returns `DatabaseError::Configuration` if any validation fails:
|
|
/// - `min_connections` is greater than `max_connections`
|
|
/// - `max_connections` is 0
|
|
/// - `acquire_timeout_secs` is 0
|
|
/// - `database_url` is not a valid PostgreSQL connection string
|
|
fn validate(&self) -> DatabaseResult<()>;
|
|
}
|
|
|
|
impl PoolConfigValidation for PoolConfig {
|
|
fn validate(&self) -> DatabaseResult<()> {
|
|
if self.min_connections > self.max_connections {
|
|
return Err(DatabaseError::Configuration {
|
|
message: "min_connections cannot be greater than max_connections".to_string(),
|
|
});
|
|
}
|
|
|
|
if self.max_connections == 0 {
|
|
return Err(DatabaseError::Configuration {
|
|
message: "max_connections must be greater than 0".to_string(),
|
|
});
|
|
}
|
|
|
|
if self.acquire_timeout_secs == 0 {
|
|
return Err(DatabaseError::Configuration {
|
|
message: "acquire_timeout_secs must be greater than 0".to_string(),
|
|
});
|
|
}
|
|
|
|
if !self.database_url.starts_with("postgres://")
|
|
&& !self.database_url.starts_with("postgresql://")
|
|
{
|
|
return Err(DatabaseError::Configuration {
|
|
message: "database_url must be a valid PostgreSQL connection string".to_string(),
|
|
});
|
|
}
|
|
|
|
Ok(())
|
|
}
|
|
}
|
|
|
|
/// Connection pool statistics
|
|
///
|
|
/// Provides comprehensive metrics about the database connection pool
|
|
/// including usage patterns, health status, and performance indicators.
|
|
#[derive(Debug, Clone, Default)]
|
|
pub struct PoolStats {
|
|
/// Total number of connections created since pool initialization
|
|
pub total_connections_created: u64,
|
|
/// Number of currently active connections in use
|
|
pub active_connections: u32,
|
|
/// Number of idle connections available for use
|
|
pub idle_connections: u32,
|
|
/// Total number of connection acquisitions requested
|
|
pub total_acquisitions: u64,
|
|
/// Number of failed connection acquisition attempts
|
|
pub failed_acquisitions: u64,
|
|
/// Number of connections closed due to reaching maximum lifetime
|
|
pub connections_closed_max_lifetime: u64,
|
|
/// Number of connections closed due to idle timeout
|
|
pub connections_closed_idle_timeout: u64,
|
|
/// Number of failed health check attempts
|
|
pub failed_health_checks: u64,
|
|
}
|
|
|
|
/// Enhanced database connection pool with monitoring and health checks
|
|
///
|
|
/// Provides a high-level interface to PostgreSQL connection pooling with:
|
|
/// - Automatic connection lifecycle management
|
|
/// - Health monitoring and statistics collection
|
|
/// - Configurable timeouts and pool sizing
|
|
/// - Background health check tasks
|
|
#[derive(Debug)]
|
|
pub struct DatabasePool {
|
|
/// The underlying SQLx PostgreSQL pool
|
|
inner: PgPool,
|
|
/// Pool configuration settings
|
|
config: PoolConfig,
|
|
/// Pool statistics protected by async RwLock
|
|
stats: Arc<RwLock<PoolStats>>,
|
|
/// Atomic counter for total connection acquisitions
|
|
total_acquisitions: Arc<AtomicU64>,
|
|
/// Atomic counter for failed connection acquisitions
|
|
failed_acquisitions: Arc<AtomicU64>,
|
|
/// Atomic counter for total connections created
|
|
total_connections_created: Arc<AtomicU64>,
|
|
}
|
|
|
|
impl DatabasePool {
|
|
/// Create a new database pool with the given configuration
|
|
///
|
|
/// # Arguments
|
|
///
|
|
/// * `config` - Pool configuration including connection limits, timeouts, and database URL
|
|
///
|
|
/// # Returns
|
|
///
|
|
/// A new `DatabasePool` instance ready for use
|
|
///
|
|
/// # Errors
|
|
///
|
|
/// Returns `DatabaseError::Configuration` if the configuration is invalid
|
|
/// or `DatabaseError::ConnectionPool` if the pool cannot be created
|
|
///
|
|
/// # Examples
|
|
///
|
|
/// ```rust,no_run
|
|
/// use database::{DatabasePool, PoolConfig};
|
|
///
|
|
/// #[tokio::main]
|
|
/// async fn main() -> Result<(), Box<dyn std::error::Error>> {
|
|
/// let config = PoolConfig::default();
|
|
/// let pool = DatabasePool::new(config).await?;
|
|
/// Ok(())
|
|
/// }
|
|
/// ```
|
|
pub async fn new(config: PoolConfig) -> DatabaseResult<Self> {
|
|
// Validate configuration
|
|
config.validate()?;
|
|
|
|
info!(
|
|
"Creating database pool with {} min connections, {} max connections",
|
|
config.min_connections, config.max_connections
|
|
);
|
|
|
|
// Build the pool
|
|
let pool = PgPoolOptions::new()
|
|
.min_connections(config.min_connections)
|
|
.max_connections(config.max_connections)
|
|
.acquire_timeout(Duration::from_secs(config.acquire_timeout_secs))
|
|
.max_lifetime(Duration::from_secs(config.max_lifetime_secs))
|
|
.idle_timeout(Duration::from_secs(config.idle_timeout_secs))
|
|
.test_before_acquire(config.test_before_acquire)
|
|
.connect(&config.database_url)
|
|
.await
|
|
.map_err(|e| DatabaseError::ConnectionPool {
|
|
message: format!("Failed to create connection pool: {}", e),
|
|
})?;
|
|
|
|
let database_pool = Self {
|
|
inner: pool,
|
|
config: config.clone(),
|
|
stats: Arc::new(RwLock::new(PoolStats::default())),
|
|
total_acquisitions: Arc::new(AtomicU64::new(0)),
|
|
failed_acquisitions: Arc::new(AtomicU64::new(0)),
|
|
total_connections_created: Arc::new(AtomicU64::new(0)),
|
|
};
|
|
|
|
// Start health check task if enabled
|
|
if config.health_check_enabled {
|
|
database_pool.start_health_check_task().await;
|
|
}
|
|
|
|
info!("Database pool created successfully");
|
|
Ok(database_pool)
|
|
}
|
|
|
|
/// Get a connection from the pool
|
|
///
|
|
/// Acquires a connection from the pool, waiting up to the configured
|
|
/// acquire timeout if no connections are immediately available.
|
|
///
|
|
/// # Returns
|
|
///
|
|
/// A pooled connection that will be returned to the pool when dropped
|
|
///
|
|
/// # Errors
|
|
///
|
|
/// Returns `DatabaseError::Timeout` if the acquire timeout is exceeded
|
|
/// or other connection-related errors from the underlying pool
|
|
///
|
|
/// # Examples
|
|
///
|
|
/// ```rust,no_run
|
|
/// # use database::{DatabasePool, PoolConfig};
|
|
/// # #[tokio::main]
|
|
/// # async fn main() -> Result<(), Box<dyn std::error::Error>> {
|
|
/// # let pool = DatabasePool::new(PoolConfig::default()).await?;
|
|
/// let conn = pool.acquire().await?;
|
|
/// // Use connection...
|
|
/// // Connection automatically returned to pool when dropped
|
|
/// # Ok(())
|
|
/// # }
|
|
/// ```
|
|
pub async fn acquire(&self) -> DatabaseResult<sqlx::pool::PoolConnection<sqlx::Postgres>> {
|
|
self.total_acquisitions.fetch_add(1, Ordering::Relaxed);
|
|
|
|
debug!("Acquiring connection from pool");
|
|
|
|
match self.inner.acquire().await {
|
|
Ok(conn) => {
|
|
debug!("Successfully acquired connection from pool");
|
|
Ok(conn)
|
|
},
|
|
Err(e) => {
|
|
self.failed_acquisitions.fetch_add(1, Ordering::Relaxed);
|
|
error!("Failed to acquire connection from pool: {}", e);
|
|
Err(DatabaseError::from(e))
|
|
},
|
|
}
|
|
}
|
|
|
|
/// Get a reference to the underlying pool for direct use
|
|
///
|
|
/// Provides direct access to the SQLx pool for operations that
|
|
/// require the underlying pool interface.
|
|
///
|
|
/// # Returns
|
|
///
|
|
/// A reference to the underlying `PgPool`
|
|
///
|
|
/// # Usage
|
|
///
|
|
/// This method is typically used for executing queries directly
|
|
/// without acquiring a connection explicitly.
|
|
pub fn inner(&self) -> &PgPool {
|
|
&self.inner
|
|
}
|
|
|
|
/// Get current pool statistics
|
|
///
|
|
/// Returns a snapshot of the current pool state including
|
|
/// connection counts, acquisition metrics, and health status.
|
|
///
|
|
/// # Returns
|
|
///
|
|
/// Current pool statistics including active/idle connections,
|
|
/// total acquisitions, failed attempts, and health check results
|
|
///
|
|
/// # Examples
|
|
///
|
|
/// ```rust,no_run
|
|
/// # use database::{DatabasePool, PoolConfig};
|
|
/// # #[tokio::main]
|
|
/// # async fn main() -> Result<(), Box<dyn std::error::Error>> {
|
|
/// # let pool = DatabasePool::new(PoolConfig::default()).await?;
|
|
/// let stats = pool.stats().await;
|
|
/// println!("Active connections: {}", stats.active_connections);
|
|
/// println!("Failed acquisitions: {}", stats.failed_acquisitions);
|
|
/// # Ok(())
|
|
/// # }
|
|
/// ```
|
|
pub async fn stats(&self) -> PoolStats {
|
|
let mut stats = self.stats.read().await.clone();
|
|
|
|
// Update atomic counters
|
|
stats.total_acquisitions = self.total_acquisitions.load(Ordering::Relaxed);
|
|
stats.failed_acquisitions = self.failed_acquisitions.load(Ordering::Relaxed);
|
|
stats.total_connections_created = self.total_connections_created.load(Ordering::Relaxed);
|
|
|
|
// Get current pool state
|
|
stats.active_connections = self.inner.size();
|
|
stats.idle_connections = self.inner.num_idle() as u32;
|
|
|
|
stats
|
|
}
|
|
|
|
/// Check if the pool is healthy
|
|
///
|
|
/// Performs a simple query to verify the database connection
|
|
/// and pool functionality. Updates health check statistics.
|
|
///
|
|
/// # Returns
|
|
///
|
|
/// `true` if the health check passes, `false` otherwise
|
|
///
|
|
/// # Errors
|
|
///
|
|
/// This method does not return errors for failed health checks,
|
|
/// instead it returns `false` and increments the failed health
|
|
/// check counter in the statistics.
|
|
///
|
|
/// # Examples
|
|
///
|
|
/// ```rust,no_run
|
|
/// # use database::{DatabasePool, PoolConfig};
|
|
/// # #[tokio::main]
|
|
/// # async fn main() -> Result<(), Box<dyn std::error::Error>> {
|
|
/// # let pool = DatabasePool::new(PoolConfig::default()).await?;
|
|
/// if pool.health_check().await? {
|
|
/// println!("Database is healthy");
|
|
/// } else {
|
|
/// println!("Database health check failed");
|
|
/// }
|
|
/// # Ok(())
|
|
/// # }
|
|
/// ```
|
|
pub async fn health_check(&self) -> DatabaseResult<bool> {
|
|
debug!("Performing pool health check");
|
|
|
|
match sqlx::query("SELECT 1").fetch_one(&self.inner).await {
|
|
Ok(_) => {
|
|
debug!("Pool health check passed");
|
|
Ok(true)
|
|
},
|
|
Err(e) => {
|
|
warn!("Pool health check failed: {}", e);
|
|
let mut stats = self.stats.write().await;
|
|
stats.failed_health_checks += 1;
|
|
Ok(false)
|
|
},
|
|
}
|
|
}
|
|
|
|
/// Get the pool configuration
|
|
///
|
|
/// Returns a reference to the configuration used to create this pool.
|
|
///
|
|
/// # Returns
|
|
///
|
|
/// A reference to the `PoolConfig` used during pool creation
|
|
pub fn config(&self) -> &PoolConfig {
|
|
&self.config
|
|
}
|
|
|
|
/// Close the pool gracefully
|
|
///
|
|
/// Closes all connections in the pool and prevents new connections
|
|
/// from being created. This is an irreversible operation.
|
|
///
|
|
/// # Examples
|
|
///
|
|
/// ```rust,no_run
|
|
/// # use database::{DatabasePool, PoolConfig};
|
|
/// # #[tokio::main]
|
|
/// # async fn main() -> Result<(), Box<dyn std::error::Error>> {
|
|
/// # let pool = DatabasePool::new(PoolConfig::default()).await?;
|
|
/// // Use pool for operations...
|
|
/// pool.close().await;
|
|
/// assert!(pool.is_closed());
|
|
/// # Ok(())
|
|
/// # }
|
|
/// ```
|
|
pub async fn close(&self) {
|
|
info!("Closing database pool");
|
|
self.inner.close().await;
|
|
info!("Database pool closed");
|
|
}
|
|
|
|
/// Check if the pool is closed
|
|
///
|
|
/// Returns `true` if the pool has been closed and can no longer
|
|
/// create new connections.
|
|
///
|
|
/// # Returns
|
|
///
|
|
/// `true` if the pool is closed, `false` if it's still active
|
|
pub fn is_closed(&self) -> bool {
|
|
self.inner.is_closed()
|
|
}
|
|
|
|
/// Start the health check background task
|
|
///
|
|
/// Spawns a background task that periodically performs health checks
|
|
/// on the database connection pool. The task runs until the pool is closed.
|
|
///
|
|
/// The health check interval is configured via the pool configuration.
|
|
/// Failed health checks are recorded in the pool statistics.
|
|
async fn start_health_check_task(&self) {
|
|
if !self.config.health_check_enabled {
|
|
return;
|
|
}
|
|
|
|
let pool = self.inner.clone();
|
|
let stats = self.stats.clone();
|
|
let interval = Duration::from_secs(self.config.health_check_interval_secs);
|
|
|
|
tokio::spawn(async move {
|
|
let mut interval_timer = tokio::time::interval(interval);
|
|
|
|
loop {
|
|
interval_timer.tick().await;
|
|
|
|
if pool.is_closed() {
|
|
debug!("Pool is closed, stopping health check task");
|
|
break;
|
|
}
|
|
|
|
match sqlx::query("SELECT 1").fetch_one(&pool).await {
|
|
Ok(_) => {
|
|
debug!("Scheduled health check passed");
|
|
},
|
|
Err(e) => {
|
|
warn!("Scheduled health check failed: {}", e);
|
|
let mut stats = stats.write().await;
|
|
stats.failed_health_checks += 1;
|
|
},
|
|
}
|
|
}
|
|
});
|
|
|
|
debug!(
|
|
"Health check task started with {}s interval",
|
|
self.config.health_check_interval_secs
|
|
);
|
|
}
|
|
|
|
/// Execute a test query to validate the connection
|
|
///
|
|
/// Performs a simple SELECT 1 query to verify database connectivity.
|
|
/// This is similar to a health check but returns an error on failure.
|
|
///
|
|
/// # Returns
|
|
///
|
|
/// `Ok(())` if the ping succeeds
|
|
///
|
|
/// # Errors
|
|
///
|
|
/// Returns a `DatabaseError` if the ping query fails
|
|
///
|
|
/// # Examples
|
|
///
|
|
/// ```rust,no_run
|
|
/// # use database::{DatabasePool, PoolConfig};
|
|
/// # #[tokio::main]
|
|
/// # async fn main() -> Result<(), Box<dyn std::error::Error>> {
|
|
/// # let pool = DatabasePool::new(PoolConfig::default()).await?;
|
|
/// pool.ping().await?;
|
|
/// println!("Database connection is active");
|
|
/// # Ok(())
|
|
/// # }
|
|
/// ```
|
|
pub async fn ping(&self) -> DatabaseResult<()> {
|
|
sqlx::query("SELECT 1")
|
|
.fetch_one(&self.inner)
|
|
.await
|
|
.map_err(DatabaseError::from)?;
|
|
Ok(())
|
|
}
|
|
|
|
/// Get current pool size
|
|
///
|
|
/// Returns the total number of connections currently managed
|
|
/// by the pool (both active and idle).
|
|
///
|
|
/// # Returns
|
|
///
|
|
/// The current total number of connections in the pool
|
|
pub fn size(&self) -> u32 {
|
|
self.inner.size()
|
|
}
|
|
|
|
/// Get number of idle connections
|
|
///
|
|
/// Returns the number of connections that are currently idle
|
|
/// and available for use.
|
|
///
|
|
/// # Returns
|
|
///
|
|
/// The number of idle connections available in the pool
|
|
pub fn num_idle(&self) -> usize {
|
|
self.inner.num_idle()
|
|
}
|
|
|
|
/// Reset pool statistics
|
|
///
|
|
/// Resets all statistical counters to zero. This includes
|
|
/// acquisition counts, failure counts, and other metrics.
|
|
///
|
|
/// # Examples
|
|
///
|
|
/// ```rust,no_run
|
|
/// # use database::{DatabasePool, PoolConfig};
|
|
/// # #[tokio::main]
|
|
/// # async fn main() -> Result<(), Box<dyn std::error::Error>> {
|
|
/// # let pool = DatabasePool::new(PoolConfig::default()).await?;
|
|
/// // After some operations...
|
|
/// pool.reset_stats().await;
|
|
/// let stats = pool.stats().await;
|
|
/// assert_eq!(stats.total_acquisitions, 0);
|
|
/// # Ok(())
|
|
/// # }
|
|
/// ```
|
|
pub async fn reset_stats(&self) {
|
|
let mut stats = self.stats.write().await;
|
|
*stats = PoolStats::default();
|
|
self.total_acquisitions.store(0, Ordering::Relaxed);
|
|
self.failed_acquisitions.store(0, Ordering::Relaxed);
|
|
self.total_connections_created.store(0, Ordering::Relaxed);
|
|
info!("Pool statistics reset");
|
|
}
|
|
}
|
|
|
|
impl Clone for DatabasePool {
|
|
fn clone(&self) -> Self {
|
|
Self {
|
|
inner: self.inner.clone(),
|
|
config: self.config.clone(),
|
|
stats: self.stats.clone(),
|
|
total_acquisitions: self.total_acquisitions.clone(),
|
|
failed_acquisitions: self.failed_acquisitions.clone(),
|
|
total_connections_created: self.total_connections_created.clone(),
|
|
}
|
|
}
|
|
}
|
|
|
|
#[cfg(test)]
|
|
mod tests {
|
|
use super::*;
|
|
|
|
#[test]
|
|
fn test_pool_config_validation() {
|
|
let mut config = PoolConfig::default();
|
|
assert!(config.validate().is_ok());
|
|
|
|
// Test min > max
|
|
config.min_connections = 10;
|
|
config.max_connections = 5;
|
|
assert!(config.validate().is_err());
|
|
|
|
// Test max = 0
|
|
config.min_connections = 0;
|
|
config.max_connections = 0;
|
|
assert!(config.validate().is_err());
|
|
|
|
// Test invalid URL
|
|
config.max_connections = 10;
|
|
config.database_url = "invalid://url".to_string();
|
|
assert!(config.validate().is_err());
|
|
}
|
|
|
|
#[test]
|
|
fn test_pool_config_default() {
|
|
let config = PoolConfig::default();
|
|
assert_eq!(config.min_connections, 1);
|
|
assert_eq!(config.max_connections, 10);
|
|
assert!(config.test_before_acquire);
|
|
assert!(config.health_check_enabled);
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn test_pool_stats_default() {
|
|
let stats = PoolStats::default();
|
|
assert_eq!(stats.total_connections_created, 0);
|
|
assert_eq!(stats.active_connections, 0);
|
|
assert_eq!(stats.failed_acquisitions, 0);
|
|
}
|
|
}
|