|
1 | 1 | use deadpool_postgres::{Manager, ManagerConfig, Pool, RecyclingMethod, Runtime}; |
2 | 2 | use std::str::FromStr; |
3 | 3 | use tokio_postgres::NoTls; |
| 4 | +use tokio::time::{sleep, Duration}; |
| 5 | +use crate::service::MetricsService; |
| 6 | +use std::sync::Arc; |
| 7 | +use std::cmp; |
4 | 8 |
|
5 | 9 | pub type DbPool = Pool; |
6 | 10 |
|
@@ -38,6 +42,104 @@ pub async fn run_migrations(database_url: &str) -> Result<(), Box<dyn std::error |
38 | 42 | Ok(()) |
39 | 43 | } |
40 | 44 |
|
| 45 | +/// Start a background task that monitors database connections and updates metrics. |
| 46 | +/// |
| 47 | +/// - `database_url`: Postgres connection string used for monitoring queries. |
| 48 | +/// - `configured_max_size`: the configured pool max size from configuration. |
| 49 | +/// - `check_interval_secs`: how often to poll Postgres for connection counts. |
| 50 | +/// |
| 51 | +/// Returns a JoinHandle for the spawned task. The task will run until cancelled. |
| 52 | +pub fn start_db_pool_monitoring( |
| 53 | + database_url: String, |
| 54 | + configured_max_size: usize, |
| 55 | + check_interval_secs: u64, |
| 56 | +) -> tokio::task::JoinHandle<()> { |
| 57 | + tokio::spawn(async move { |
| 58 | + loop { |
| 59 | + match tokio_postgres::Config::from_str(&database_url) |
| 60 | + .and_then(|cfg| cfg.connect(NoTls)) |
| 61 | + { |
| 62 | + Ok((client, connection)) => { |
| 63 | + // detach connection handling |
| 64 | + tokio::spawn(async move { |
| 65 | + if let Err(e) = connection.await { |
| 66 | + tracing::warn!(error = %e, "Postgres monitor connection error"); |
| 67 | + } |
| 68 | + }); |
| 69 | + |
| 70 | + // Query active connections for this database |
| 71 | + match client |
| 72 | + .query_one( |
| 73 | + "SELECT count(*) FROM pg_stat_activity WHERE datname = current_database()", |
| 74 | + &[], |
| 75 | + ) |
| 76 | + .await |
| 77 | + { |
| 78 | + Ok(row) => { |
| 79 | + let active: i64 = row.get(0); |
| 80 | + let active_usize = cmp::min(active as usize, 1_000_000); |
| 81 | + MetricsService::update_db_pool_status(configured_max_size, active_usize); |
| 82 | + } |
| 83 | + Err(e) => tracing::warn!(error = %e, "Failed to query pg_stat_activity"), |
| 84 | + } |
| 85 | + |
| 86 | + let _ = client.close().await; |
| 87 | + } |
| 88 | + Err(e) => tracing::error!(error = %e, "Failed to connect to Postgres for monitoring"), |
| 89 | + } |
| 90 | + |
| 91 | + sleep(Duration::from_secs(check_interval_secs)).await; |
| 92 | + } |
| 93 | + }) |
| 94 | +} |
| 95 | + |
| 96 | +/// Health check for database connectivity. |
| 97 | +pub async fn health_check_db(database_url: &str) -> bool { |
| 98 | + match tokio_postgres::Config::from_str(database_url) |
| 99 | + .and_then(|cfg| cfg.connect(NoTls)) |
| 100 | + .await |
| 101 | + { |
| 102 | + Ok((mut client, connection)) => { |
| 103 | + // drive connection |
| 104 | + tokio::spawn(async move { |
| 105 | + let _ = connection.await; |
| 106 | + }); |
| 107 | + |
| 108 | + let res = client.query_one("SELECT 1", &[]).await.is_ok(); |
| 109 | + let _ = client.close().await; |
| 110 | + res |
| 111 | + } |
| 112 | + Err(_) => false, |
| 113 | + } |
| 114 | +} |
| 115 | + |
| 116 | +/// Recommend a new pool size based on current utilization and configuration. |
| 117 | +/// This function does not mutate or rebuild the pool; it only suggests a size. |
| 118 | +pub fn recommend_pool_size( |
| 119 | + current_max: usize, |
| 120 | + active_connections: usize, |
| 121 | + min_pool_size: usize, |
| 122 | + resize_step: usize, |
| 123 | + high_threshold: f64, |
| 124 | + low_threshold: f64, |
| 125 | +) -> usize { |
| 126 | + if current_max == 0 { |
| 127 | + return current_max; |
| 128 | + } |
| 129 | + |
| 130 | + let utilization = (active_connections as f64 / current_max as f64) * 100.0; |
| 131 | + |
| 132 | + if utilization >= high_threshold { |
| 133 | + // scale up |
| 134 | + current_max.saturating_add(resize_step) |
| 135 | + } else if utilization <= low_threshold && current_max > min_pool_size { |
| 136 | + // scale down but not below min |
| 137 | + current_max.saturating_sub(resize_step).max(min_pool_size) |
| 138 | + } else { |
| 139 | + current_max |
| 140 | + } |
| 141 | +} |
| 142 | + |
41 | 143 | /// Reset migrations for testing purposes |
42 | 144 | /// This drops all tables, types, and the migration history to allow re-running migrations |
43 | 145 | /// WARNING: Only use this in test environments! This will destroy all data in the database. |
|
0 commit comments