ri/database/
pool.rs

1//! Copyright © 2025-2026 Wenze Wei. All Rights Reserved.
2//!
3//! This file is part of Ri.
4//! The Ri project belongs to the Dunimd Team.
5//!
6//! Licensed under the Apache License, Version 2.0 (the "License");
7//! You may not use this file except in compliance with the License.
8//! You may obtain a copy of the License at
9//!
10//!     http://www.apache.org/licenses/LICENSE-2.0
11//!
12//! Unless required by applicable law or agreed to in writing, software
13//! distributed under the License is distributed on an "AS IS" BASIS,
14//! WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
15//! See the License for the specific language governing permissions and
16//! limitations under the License.
17
18use crate::core::RiResult;
19use crate::database::{RiDatabase, RiDatabaseConfig, RiDBResult, RiDBRow};
20use dashmap::DashMap;
21use std::sync::atomic::{AtomicU64, AtomicBool, Ordering};
22use std::sync::Arc;
23use tokio::sync::Semaphore;
24use tokio::time::{Duration, Instant};
25use std::sync::RwLock;
26
27fn handle_poisoned_lock<T>(lock_result: Result<T, std::sync::PoisonError<T>>) -> T {
28    match lock_result {
29        Ok(guard) => guard,
30        Err(e) => {
31            log::warn!("[Ri.Database] Lock was poisoned, recovering...");
32            e.into_inner()
33        }
34    }
35}
36
37#[cfg_attr(feature = "pyo3", pyo3::prelude::pyclass)]
38#[derive(Debug, Clone, Default)]
39pub struct RiDatabaseMetrics {
40    pub active_connections: u64,
41    pub idle_connections: u64,
42    pub total_connections: u64,
43    pub queries_executed: u64,
44    pub query_duration_ms: f64,
45    pub errors: u64,
46    pub utilization_rate: f64,
47}
48
49#[cfg_attr(feature = "pyo3", pyo3::prelude::pyclass)]
50#[derive(Debug, Clone)]
51pub struct RiDynamicPoolConfig {
52    pub enable_dynamic_scaling: bool,
53    pub scale_up_threshold: f64,
54    pub scale_down_threshold: f64,
55    pub scale_down_cooldown_secs: u64,
56    pub min_connections: u32,
57    pub max_connections: u32,
58    pub scale_up_step: u32,
59    pub scale_down_step: u32,
60}
61
62impl Default for RiDynamicPoolConfig {
63    fn default() -> Self {
64        Self {
65            enable_dynamic_scaling: true,
66            scale_up_threshold: 0.8,
67            scale_down_threshold: 0.3,
68            scale_down_cooldown_secs: 300,
69            min_connections: 2,
70            max_connections: 50,
71            scale_up_step: 2,
72            scale_down_step: 1,
73        }
74    }
75}
76
77#[cfg(feature = "pyo3")]
78#[pyo3::prelude::pymethods]
79impl RiDynamicPoolConfig {
80    #[new]
81    fn py_new() -> Self {
82        Self::default()
83    }
84
85    #[staticmethod]
86    fn create() -> Self {
87        Self::default()
88    }
89
90    fn get_enable_dynamic_scaling(&self) -> bool {
91        self.enable_dynamic_scaling
92    }
93
94    fn set_enable_dynamic_scaling(&mut self, value: bool) {
95        self.enable_dynamic_scaling = value;
96    }
97
98    fn get_scale_up_threshold(&self) -> f64 {
99        self.scale_up_threshold
100    }
101
102    fn set_scale_up_threshold(&mut self, value: f64) {
103        self.scale_up_threshold = value;
104    }
105
106    fn get_scale_down_threshold(&self) -> f64 {
107        self.scale_down_threshold
108    }
109
110    fn set_scale_down_threshold(&mut self, value: f64) {
111        self.scale_down_threshold = value;
112    }
113
114    fn get_scale_down_cooldown_secs(&self) -> u64 {
115        self.scale_down_cooldown_secs
116    }
117
118    fn set_scale_down_cooldown_secs(&mut self, value: u64) {
119        self.scale_down_cooldown_secs = value;
120    }
121
122    fn get_min_connections(&self) -> u32 {
123        self.min_connections
124    }
125
126    fn set_min_connections(&mut self, value: u32) {
127        self.min_connections = value;
128    }
129
130    fn get_max_connections(&self) -> u32 {
131        self.max_connections
132    }
133
134    fn set_max_connections(&mut self, value: u32) {
135        self.max_connections = value;
136    }
137
138    fn get_scale_up_step(&self) -> u32 {
139        self.scale_up_step
140    }
141
142    fn set_scale_up_step(&mut self, value: u32) {
143        self.scale_up_step = value;
144    }
145
146    fn get_scale_down_step(&self) -> u32 {
147        self.scale_down_step
148    }
149
150    fn set_scale_down_step(&mut self, value: u32) {
151        self.scale_down_step = value;
152    }
153}
154
155#[derive(Clone)]
156pub struct PooledDatabase {
157    id: u32,
158    inner: Arc<dyn RiDatabase>,
159    pool: Arc<RiDatabasePool>,
160}
161
162impl Drop for PooledDatabase {
163    fn drop(&mut self) {
164        // Automatically release the connection back to the pool when dropped
165        // This prevents resource leaks if the user forgets to call release()
166        let pool = self.pool.clone();
167        let id = self.id;
168        let inner = self.inner.clone();
169        
170        // Use tokio::spawn to handle the async release operation
171        // since Drop cannot be async
172        tokio::spawn(async move {
173            pool.active_connections.fetch_sub(1, Ordering::SeqCst);
174            pool.idle_connections.fetch_add(1, Ordering::SeqCst);
175            
176            pool.available.insert(id, PoolConnection { 
177                db: inner,
178                acquired_at: Instant::now(),
179                created_at: Instant::now(),
180            });
181
182            let _ = pool.check_and_scale().await;
183        });
184    }
185}
186
187impl PooledDatabase {
188    pub fn new(id: u32, inner: Arc<dyn RiDatabase>, pool: Arc<RiDatabasePool>) -> Self {
189        Self { id, inner, pool }
190    }
191
192    pub fn id(&self) -> u32 {
193        self.id
194    }
195
196    pub async fn execute(&self, sql: &str) -> RiResult<u64> {
197        self.inner.execute(sql).await
198    }
199
200    pub async fn query(&self, sql: &str) -> RiResult<RiDBResult> {
201        self.inner.query(sql).await
202    }
203
204    pub async fn query_one(&self, sql: &str) -> RiResult<Option<RiDBRow>> {
205        self.inner.query_one(sql).await
206    }
207
208    pub async fn ping(&self) -> RiResult<bool> {
209        self.inner.ping().await
210    }
211
212    pub fn is_connected(&self) -> bool {
213        self.inner.is_connected()
214    }
215
216    pub fn pool_metrics(&self) -> RiDatabaseMetrics {
217        self.pool.metrics()
218    }
219}
220
221#[async_trait::async_trait]
222impl RiDatabase for PooledDatabase {
223    fn database_type(&self) -> crate::database::DatabaseType {
224        self.inner.database_type()
225    }
226
227    async fn execute(&self, sql: &str) -> RiResult<u64> {
228        self.inner.execute(sql).await
229    }
230
231    async fn query(&self, sql: &str) -> RiResult<RiDBResult> {
232        self.inner.query(sql).await
233    }
234
235    async fn query_one(&self, sql: &str) -> RiResult<Option<RiDBRow>> {
236        self.inner.query_one(sql).await
237    }
238
239    async fn ping(&self) -> RiResult<bool> {
240        self.inner.ping().await
241    }
242
243    fn is_connected(&self) -> bool {
244        self.inner.is_connected()
245    }
246
247    async fn close(&self) -> RiResult<()> {
248        self.pool.close().await
249    }
250
251    async fn batch_execute(&self, sql: &str, params: &[Vec<serde_json::Value>]) -> RiResult<Vec<u64>> {
252        self.inner.batch_execute(sql, params).await
253    }
254
255    async fn batch_query(&self, sql: &str, params: &[Vec<serde_json::Value>]) -> RiResult<Vec<RiDBResult>> {
256        self.inner.batch_query(sql, params).await
257    }
258
259    async fn execute_with_params(&self, sql: &str, params: &[serde_json::Value]) -> RiResult<u64> {
260        self.inner.execute_with_params(sql, params).await
261    }
262
263    async fn query_with_params(&self, sql: &str, params: &[serde_json::Value]) -> RiResult<RiDBResult> {
264        self.inner.query_with_params(sql, params).await
265    }
266
267    async fn transaction(&self) -> RiResult<Box<dyn crate::database::RiDatabaseTransaction>> {
268        self.inner.transaction().await
269    }
270}
271
272struct PoolConnection {
273    db: Arc<dyn RiDatabase>,
274    acquired_at: Instant,
275    created_at: Instant,
276}
277
278struct LowUtilizationTracker {
279    below_threshold_since: Option<Instant>,
280    was_below_threshold: bool,
281}
282
283impl LowUtilizationTracker {
284    fn new() -> Self {
285        Self {
286            below_threshold_since: None,
287            was_below_threshold: false,
288        }
289    }
290
291    fn update(&mut self, is_below_threshold: bool) {
292        if is_below_threshold && !self.was_below_threshold {
293            self.below_threshold_since = Some(Instant::now());
294            self.was_below_threshold = true;
295        } else if !is_below_threshold {
296            self.below_threshold_since = None;
297            self.was_below_threshold = false;
298        }
299    }
300
301    fn duration_below_threshold(&self) -> Option<Duration> {
302        self.below_threshold_since.map(|since| since.elapsed())
303    }
304}
305
306#[cfg_attr(feature = "pyo3", pyo3::prelude::pyclass)]
307pub struct RiDatabasePool {
308    config: RiDatabaseConfig,
309    connections: Arc<DashMap<u32, PoolConnection>>,
310    available: Arc<DashMap<u32, PoolConnection>>,
311    connection_ids: Arc<AtomicU64>,
312    semaphore: Arc<Semaphore>,
313    max_idle_time: Duration,
314    max_lifetime: Duration,
315    idle_connections: Arc<AtomicU64>,
316    active_connections: Arc<AtomicU64>,
317    total_connections: Arc<AtomicU64>,
318    queries_executed: Arc<AtomicU64>,
319    errors: Arc<AtomicU64>,
320    dynamic_config: Arc<RwLock<RiDynamicPoolConfig>>,
321    low_utilization_tracker: Arc<RwLock<LowUtilizationTracker>>,
322    scaling_in_progress: Arc<AtomicBool>,
323    current_max_connections: Arc<AtomicU64>,
324}
325
326/// Maximum allowed connections in a pool to prevent resource exhaustion
327const MAX_POOL_SIZE: u32 = 1000;
328
329/// Minimum allowed connections in a pool
330const MIN_POOL_SIZE: u32 = 1;
331
332impl RiDatabasePool {
333    pub async fn new(config: RiDatabaseConfig) -> RiResult<Self> {
334        // Security: Validate connection pool limits
335        let max_connections = config.max_connections.clamp(MIN_POOL_SIZE, MAX_POOL_SIZE);
336        let min_connections = config.min_idle_connections.clamp(0, max_connections);
337        
338        if config.max_connections > MAX_POOL_SIZE {
339            log::warn!(
340                "[Ri.Database] max_connections {} exceeds limit {}, clamping to {}",
341                config.max_connections, MAX_POOL_SIZE, max_connections
342            );
343        }
344
345        let dynamic_config = RiDynamicPoolConfig {
346            enable_dynamic_scaling: true,
347            scale_up_threshold: 0.8,
348            scale_down_threshold: 0.3,
349            scale_down_cooldown_secs: 300,
350            min_connections,
351            max_connections,
352            scale_up_step: 2,
353            scale_down_step: 1,
354        };
355
356        let pool = Self {
357            config: config.clone(),
358            connections: Arc::new(DashMap::new()),
359            available: Arc::new(DashMap::new()),
360            connection_ids: Arc::new(AtomicU64::new(0)),
361            semaphore: Arc::new(Semaphore::new(max_connections as usize)),
362            max_idle_time: Duration::from_secs(config.idle_timeout_secs),
363            max_lifetime: Duration::from_secs(config.max_lifetime_secs),
364            idle_connections: Arc::new(AtomicU64::new(0)),
365            active_connections: Arc::new(AtomicU64::new(0)),
366            total_connections: Arc::new(AtomicU64::new(0)),
367            queries_executed: Arc::new(AtomicU64::new(0)),
368            errors: Arc::new(AtomicU64::new(0)),
369            dynamic_config: Arc::new(RwLock::new(dynamic_config)),
370            low_utilization_tracker: Arc::new(RwLock::new(LowUtilizationTracker::new())),
371            scaling_in_progress: Arc::new(AtomicBool::new(false)),
372            current_max_connections: Arc::new(AtomicU64::new(max_connections as u64)),
373        };
374
375        for _ in 0..min_connections {
376            if let Ok(conn) = pool.create_connection().await {
377                let id = pool.connection_ids.fetch_add(1, Ordering::SeqCst) as u32;
378                let now = Instant::now();
379                pool.available.insert(id, PoolConnection { 
380                    db: conn, 
381                    acquired_at: now,
382                    created_at: now,
383                });
384                pool.idle_connections.fetch_add(1, Ordering::SeqCst);
385                pool.total_connections.fetch_add(1, Ordering::SeqCst);
386            }
387        }
388
389        Ok(pool)
390    }
391
392    async fn create_connection(&self) -> RiResult<Arc<dyn RiDatabase>> {
393        match self.config.database_type {
394            #[cfg(feature = "postgres")]
395            crate::database::DatabaseType::Postgres => {
396                let connection_string = self.config.connection_string();
397                let db = crate::database::postgres::PostgresDatabase::new(&connection_string, self.config.clone()).await
398                    .map_err(|e| crate::core::RiError::Config(e.to_string()))?;
399                Ok(Arc::new(db) as Arc<dyn RiDatabase>)
400            }
401            #[cfg(feature = "mysql")]
402            crate::database::DatabaseType::MySQL => {
403                let connection_string = self.config.connection_string();
404                let db = crate::database::mysql::MySQLDatabase::new(&connection_string, self.config.clone()).await
405                    .map_err(|e| crate::core::RiError::Config(e.to_string()))?;
406                Ok(Arc::new(db) as Arc<dyn RiDatabase>)
407            }
408            #[cfg(feature = "sqlite")]
409            crate::database::DatabaseType::SQLite => {
410                let url = format!("sqlite:{}", self.config.database);
411                let db = tokio::runtime::Handle::current().block_on(
412                    crate::database::sqlite::SQLiteDatabase::new(&url, self.config.clone())
413                );
414                match db {
415                    Ok(db) => Ok(Arc::new(db) as Arc<dyn RiDatabase>),
416                    Err(e) => Err(crate::core::RiError::Config(e.to_string())),
417                }
418            }
419            _ => Err(crate::core::RiError::Config("Unsupported database type".to_string())),
420        }
421    }
422
423    pub fn utilization_rate(&self) -> f64 {
424        let active = self.active_connections.load(Ordering::SeqCst);
425        let total = self.total_connections.load(Ordering::SeqCst);
426        
427        if total == 0 {
428            return 0.0;
429        }
430        
431        (active as f64) / (total as f64)
432    }
433
434    pub async fn check_and_scale(&self) -> RiResult<()> {
435        let dynamic_config = handle_poisoned_lock(self.dynamic_config.read()).clone();
436        
437        if !dynamic_config.enable_dynamic_scaling {
438            return Ok(());
439        }
440
441        if self.scaling_in_progress.compare_exchange(
442            false,
443            true,
444            Ordering::SeqCst,
445            Ordering::SeqCst,
446        ).is_err() {
447            return Ok(());
448        }
449
450        let result = self.do_scaling(&dynamic_config).await;
451        
452        self.scaling_in_progress.store(false, Ordering::SeqCst);
453        
454        result
455    }
456
457    async fn do_scaling(&self, dynamic_config: &RiDynamicPoolConfig) -> RiResult<()> {
458        let utilization = self.utilization_rate();
459        let total = self.total_connections.load(Ordering::SeqCst) as u32;
460        let _active = self.active_connections.load(Ordering::SeqCst) as u32;
461        
462        if utilization > dynamic_config.scale_up_threshold {
463            {
464                let mut tracker = handle_poisoned_lock(self.low_utilization_tracker.write());
465                tracker.update(false);
466            }
467            
468            if total < dynamic_config.max_connections {
469                let to_add = std::cmp::min(
470                    dynamic_config.scale_up_step,
471                    dynamic_config.max_connections - total,
472                );
473                
474                for _ in 0..to_add {
475                    match self.create_connection().await {
476                        Ok(conn) => {
477                            let id = self.connection_ids.fetch_add(1, Ordering::SeqCst) as u32;
478                            let now = Instant::now();
479                            self.available.insert(id, PoolConnection { 
480                                db: conn, 
481                                acquired_at: now,
482                                created_at: now,
483                            });
484                            self.idle_connections.fetch_add(1, Ordering::SeqCst);
485                            self.total_connections.fetch_add(1, Ordering::SeqCst);
486                        }
487                        Err(e) => {
488                            self.errors.fetch_add(1, Ordering::SeqCst);
489                            return Err(e);
490                        }
491                    }
492                }
493                
494                let new_total = self.total_connections.load(Ordering::SeqCst);
495                self.current_max_connections.store(new_total, Ordering::SeqCst);
496            }
497        } else if utilization < dynamic_config.scale_down_threshold {
498            let should_scale_down = {
499                let mut tracker = handle_poisoned_lock(self.low_utilization_tracker.write());
500                tracker.update(true);
501                
502                if let Some(duration) = tracker.duration_below_threshold() {
503                    duration >= Duration::from_secs(dynamic_config.scale_down_cooldown_secs)
504                } else {
505                    false
506                }
507            };
508            
509            if should_scale_down && total > dynamic_config.min_connections {
510                let idle = self.idle_connections.load(Ordering::SeqCst) as u32;
511                let to_remove = std::cmp::min(
512                    std::cmp::min(dynamic_config.scale_down_step, idle),
513                    total - dynamic_config.min_connections,
514                );
515                
516                if to_remove > 0 {
517                    self.remove_idle_connections(to_remove).await;
518                }
519            }
520        } else {
521            let mut tracker = handle_poisoned_lock(self.low_utilization_tracker.write());
522            tracker.update(false);
523        }
524        
525        Ok(())
526    }
527
528    async fn remove_idle_connections(&self, count: u32) {
529        let mut removed = 0u32;
530        let now = Instant::now();
531        
532        let to_remove: Vec<u32> = self.available
533            .iter()
534            .filter(|entry| {
535                let conn = entry.value();
536                now.duration_since(conn.acquired_at) > self.max_idle_time
537            })
538            .take(count as usize)
539            .map(|entry| *entry.key())
540            .collect();
541        
542        for id in to_remove {
543            if removed >= count {
544                break;
545            }
546            
547            if let Some((_, conn)) = self.available.remove(&id) {
548                let _ = conn.db.close().await;
549                self.idle_connections.fetch_sub(1, Ordering::SeqCst);
550                self.total_connections.fetch_sub(1, Ordering::SeqCst);
551                removed += 1;
552            }
553        }
554        
555        if removed < count {
556            let additional_remove: Vec<u32> = self.available
557                .iter()
558                .take((count - removed) as usize)
559                .map(|entry| *entry.key())
560                .collect();
561            
562            for id in additional_remove {
563                if removed >= count {
564                    break;
565                }
566                
567                if let Some((_, conn)) = self.available.remove(&id) {
568                    let _ = conn.db.close().await;
569                    self.idle_connections.fetch_sub(1, Ordering::SeqCst);
570                    self.total_connections.fetch_sub(1, Ordering::SeqCst);
571                    removed += 1;
572                }
573            }
574        }
575        
576        let new_total = self.total_connections.load(Ordering::SeqCst);
577        self.current_max_connections.store(new_total, Ordering::SeqCst);
578    }
579
580    pub async fn get(&self) -> RiResult<PooledDatabase> {
581        let _permit = self.semaphore.acquire().await.map_err(|e| crate::core::RiError::Config(e.to_string()))?;
582
583        let mut reused_db = None;
584        let mut reused_id = None;
585
586        let now = Instant::now();
587        
588        for entry in self.available.iter() {
589            let id = *entry.key();
590            let conn = entry.value();
591            if now.duration_since(conn.acquired_at) > self.max_idle_time || now.duration_since(conn.created_at) > self.max_lifetime {
592                self.available.remove(&id);
593                let _ = conn.db.close().await;
594                self.idle_connections.fetch_sub(1, Ordering::SeqCst);
595                self.total_connections.fetch_sub(1, Ordering::SeqCst);
596            } else {
597                reused_db = Some(conn.db.clone());
598                reused_id = Some(id);
599                self.available.remove(&id);
600                self.idle_connections.fetch_sub(1, Ordering::SeqCst);
601                self.active_connections.fetch_add(1, Ordering::SeqCst);
602                break;
603            }
604        }
605
606        let (db, id) = if let Some((existing_db, existing_id)) = reused_db.zip(reused_id) {
607            (existing_db, existing_id)
608        } else {
609            match self.create_connection().await {
610                Ok(new_conn) => {
611                    let id = self.connection_ids.fetch_add(1, Ordering::SeqCst) as u32;
612                    self.total_connections.fetch_add(1, Ordering::SeqCst);
613                    self.active_connections.fetch_add(1, Ordering::SeqCst);
614                    (new_conn, id)
615                }
616                Err(e) => {
617                    self.errors.fetch_add(1, Ordering::SeqCst);
618                    return Err(e);
619                }
620            }
621        };
622
623        let _ = self.check_and_scale().await;
624
625        Ok(PooledDatabase::new(id, db, Arc::new(self.clone())))
626    }
627
628    pub async fn release(&self, db: PooledDatabase) {
629        self.active_connections.fetch_sub(1, Ordering::SeqCst);
630        self.idle_connections.fetch_add(1, Ordering::SeqCst);
631        
632        self.available.insert(db.id(), PoolConnection { 
633            db: db.inner.clone(),
634            acquired_at: Instant::now(),
635            created_at: Instant::now(),
636        });
637
638        let _ = self.check_and_scale().await;
639    }
640
641    pub async fn close(&self) -> RiResult<()> {
642        self.semaphore.close();
643        for entry in self.connections.iter() {
644            let _ = entry.value().db.close().await;
645        }
646        self.connections.clear();
647        self.available.clear();
648        Ok(())
649    }
650
651    pub fn metrics(&self) -> RiDatabaseMetrics {
652        let active = self.active_connections.load(Ordering::SeqCst);
653        let total = self.total_connections.load(Ordering::SeqCst);
654        
655        RiDatabaseMetrics {
656            active_connections: active,
657            idle_connections: self.idle_connections.load(Ordering::SeqCst),
658            total_connections: total,
659            queries_executed: self.queries_executed.load(Ordering::SeqCst),
660            query_duration_ms: 0.0,
661            errors: self.errors.load(Ordering::SeqCst),
662            utilization_rate: if total > 0 { (active as f64) / (total as f64) } else { 0.0 },
663        }
664    }
665
666    pub fn get_dynamic_config(&self) -> RiDynamicPoolConfig {
667        handle_poisoned_lock(self.dynamic_config.read()).clone()
668    }
669
670    pub fn set_dynamic_config(&self, config: RiDynamicPoolConfig) {
671        let mut current = handle_poisoned_lock(self.dynamic_config.write());
672        *current = config;
673    }
674
675    pub fn update_dynamic_config<F>(&self, f: F) 
676    where
677        F: FnOnce(&mut RiDynamicPoolConfig),
678    {
679        let mut config = handle_poisoned_lock(self.dynamic_config.write());
680        f(&mut config);
681    }
682
683    pub fn is_scaling_in_progress(&self) -> bool {
684        self.scaling_in_progress.load(Ordering::SeqCst)
685    }
686
687    pub fn get_current_max_connections(&self) -> u32 {
688        self.current_max_connections.load(Ordering::SeqCst) as u32
689    }
690
691    pub async fn force_scale_up(&self, count: u32) -> RiResult<()> {
692        let dynamic_config = handle_poisoned_lock(self.dynamic_config.read()).clone();
693        let total = self.total_connections.load(Ordering::SeqCst) as u32;
694        
695        let to_add = std::cmp::min(count, dynamic_config.max_connections.saturating_sub(total));
696        
697        for _ in 0..to_add {
698            match self.create_connection().await {
699                Ok(conn) => {
700                    let id = self.connection_ids.fetch_add(1, Ordering::SeqCst) as u32;
701                    let now = Instant::now();
702                    self.available.insert(id, PoolConnection { 
703                        db: conn, 
704                        acquired_at: now,
705                        created_at: now,
706                    });
707                    self.idle_connections.fetch_add(1, Ordering::SeqCst);
708                    self.total_connections.fetch_add(1, Ordering::SeqCst);
709                }
710                Err(e) => {
711                    self.errors.fetch_add(1, Ordering::SeqCst);
712                    return Err(e);
713                }
714            }
715        }
716        
717        let new_total = self.total_connections.load(Ordering::SeqCst);
718        self.current_max_connections.store(new_total, Ordering::SeqCst);
719        
720        Ok(())
721    }
722
723    pub async fn force_scale_down(&self, count: u32) -> RiResult<()> {
724        let dynamic_config = handle_poisoned_lock(self.dynamic_config.read()).clone();
725        let total = self.total_connections.load(Ordering::SeqCst) as u32;
726        
727        let to_remove = std::cmp::min(
728            count,
729            total.saturating_sub(dynamic_config.min_connections)
730        );
731        
732        if to_remove > 0 {
733            self.remove_idle_connections(to_remove).await;
734        }
735        
736        Ok(())
737    }
738}
739
740#[cfg(feature = "pyo3")]
741#[pyo3::prelude::pymethods]
742impl RiDatabasePool {
743    #[new]
744    fn py_new(config: RiDatabaseConfig) -> Self {
745        let dynamic_config = RiDynamicPoolConfig {
746            enable_dynamic_scaling: true,
747            scale_up_threshold: 0.8,
748            scale_down_threshold: 0.3,
749            scale_down_cooldown_secs: 300,
750            min_connections: config.min_idle_connections,
751            max_connections: config.max_connections,
752            scale_up_step: 2,
753            scale_down_step: 1,
754        };
755
756        Self {
757            config: config.clone(),
758            connections: Arc::new(DashMap::new()),
759            available: Arc::new(DashMap::new()),
760            connection_ids: Arc::new(AtomicU64::new(0)),
761            semaphore: Arc::new(Semaphore::new(config.max_connections as usize)),
762            max_idle_time: Duration::from_secs(config.idle_timeout_secs),
763            max_lifetime: Duration::from_secs(config.max_lifetime_secs),
764            idle_connections: Arc::new(AtomicU64::new(0)),
765            active_connections: Arc::new(AtomicU64::new(0)),
766            total_connections: Arc::new(AtomicU64::new(0)),
767            queries_executed: Arc::new(AtomicU64::new(0)),
768            errors: Arc::new(AtomicU64::new(0)),
769            dynamic_config: Arc::new(RwLock::new(dynamic_config)),
770            low_utilization_tracker: Arc::new(RwLock::new(LowUtilizationTracker::new())),
771            scaling_in_progress: Arc::new(AtomicBool::new(false)),
772            current_max_connections: Arc::new(AtomicU64::new(config.max_connections as u64)),
773        }
774    }
775
776    fn status(&self) -> String {
777        format!(
778            "Pool status - Active: {}, Idle: {}, Total: {}, Queries: {}, Errors: {}, Utilization: {:.2}%",
779            self.active_connections.load(Ordering::SeqCst),
780            self.idle_connections.load(Ordering::SeqCst),
781            self.total_connections.load(Ordering::SeqCst),
782            self.queries_executed.load(Ordering::SeqCst),
783            self.errors.load(Ordering::SeqCst),
784            self.utilization_rate() * 100.0
785        )
786    }
787
788    fn get_config(&self) -> RiDatabaseConfig {
789        self.config.clone()
790    }
791
792    fn get_utilization_rate(&self) -> f64 {
793        self.utilization_rate()
794    }
795
796    fn get_metrics(&self) -> RiDatabaseMetrics {
797        self.metrics()
798    }
799
800    fn get_dynamic_pool_config(&self) -> RiDynamicPoolConfig {
801        self.get_dynamic_config()
802    }
803
804    fn set_dynamic_pool_config(&self, config: RiDynamicPoolConfig) {
805        self.set_dynamic_config(config);
806    }
807
808    fn set_enable_dynamic_scaling(&self, enable: bool) {
809        self.update_dynamic_config(|c| c.enable_dynamic_scaling = enable);
810    }
811
812    fn set_scale_up_threshold(&self, threshold: f64) {
813        self.update_dynamic_config(|c| c.scale_up_threshold = threshold);
814    }
815
816    fn set_scale_down_threshold(&self, threshold: f64) {
817        self.update_dynamic_config(|c| c.scale_down_threshold = threshold);
818    }
819
820    fn set_scale_down_cooldown_secs(&self, secs: u64) {
821        self.update_dynamic_config(|c| c.scale_down_cooldown_secs = secs);
822    }
823
824    fn set_min_connections(&self, min: u32) {
825        self.update_dynamic_config(|c| c.min_connections = min);
826    }
827
828    fn set_max_connections(&self, max: u32) {
829        self.update_dynamic_config(|c| c.max_connections = max);
830    }
831
832    fn set_scale_up_step(&self, step: u32) {
833        self.update_dynamic_config(|c| c.scale_up_step = step);
834    }
835
836    fn set_scale_down_step(&self, step: u32) {
837        self.update_dynamic_config(|c| c.scale_down_step = step);
838    }
839
840    fn is_scaling(&self) -> bool {
841        self.is_scaling_in_progress()
842    }
843
844    fn get_current_pool_size(&self) -> u32 {
845        self.get_current_max_connections()
846    }
847}
848
849impl Clone for RiDatabasePool {
850    fn clone(&self) -> Self {
851        Self {
852            config: self.config.clone(),
853            connections: self.connections.clone(),
854            available: self.available.clone(),
855            connection_ids: self.connection_ids.clone(),
856            semaphore: self.semaphore.clone(),
857            max_idle_time: self.max_idle_time,
858            max_lifetime: self.max_lifetime,
859            idle_connections: self.idle_connections.clone(),
860            active_connections: self.active_connections.clone(),
861            total_connections: self.total_connections.clone(),
862            queries_executed: self.queries_executed.clone(),
863            errors: self.errors.clone(),
864            dynamic_config: self.dynamic_config.clone(),
865            low_utilization_tracker: self.low_utilization_tracker.clone(),
866            scaling_in_progress: self.scaling_in_progress.clone(),
867            current_max_connections: self.current_max_connections.clone(),
868        }
869    }
870}
ri/database/pool.rs

ri/database/
pool.rs
