1use async_trait::async_trait;
6use crossbeam_queue::ArrayQueue;
7use futures::StreamExt;
8use serde::{Deserialize, Serialize};
10#[cfg(any(feature = "circuit-breaker", feature = "tenant-quota-rls-enhanced"))]
15use parking_lot::Mutex as PlMutex;
16#[cfg(feature = "rate-limit")]
17use parking_lot::RwLock as PlRwLock;
18use std::future::Future;
19use std::ops::{Deref, DerefMut};
20use std::pin::Pin;
21use std::sync::atomic::{AtomicBool, AtomicU32, AtomicU64, Ordering};
22use std::sync::Arc;
23use std::time::{Duration, Instant};
24use tokio::sync::Notify;
25
26#[cfg(feature = "circuit-breaker")]
30use crate::circuit_breaker::{CircuitBreaker, CircuitState, DefaultCircuitBreaker};
31use crate::error::PoolError;
32#[cfg(feature = "rate-limit")]
33use crate::rate_limiter::RateLimiter;
34#[cfg(feature = "tenant-quota-rls-enhanced")]
35use crate::tenant_quota_rls::{QuotaEnforcer, QuotaResource};
36
37pub type QueryRows = Vec<std::collections::HashMap<String, crate::value::Value>>;
39
40pub type QueryStreamItem =
42 Result<std::collections::HashMap<String, crate::value::Value>, crate::DbError>;
43
44pub trait Connection: Send + Sync {
51 fn execute<'a>(
53 &'a mut self,
54 sql: &'a str,
55 ) -> Pin<Box<dyn Future<Output = Result<u64, crate::DbError>> + Send + 'a>>;
56 fn query<'a>(
58 &'a mut self,
59 sql: &'a str,
60 ) -> Pin<Box<dyn Future<Output = Result<QueryRows, crate::DbError>> + Send + 'a>>;
61 fn begin_transaction<'a>(
63 &'a mut self,
64 ) -> Pin<Box<dyn Future<Output = Result<(), crate::DbError>> + Send + 'a>>;
65 fn commit<'a>(
67 &'a mut self,
68 ) -> Pin<Box<dyn Future<Output = Result<(), crate::DbError>> + Send + 'a>>;
69 fn rollback<'a>(
71 &'a mut self,
72 ) -> Pin<Box<dyn Future<Output = Result<(), crate::DbError>> + Send + 'a>>;
73 fn is_connected(&self) -> bool;
75 fn ping<'a>(&'a mut self) -> Pin<Box<dyn Future<Output = bool> + Send + 'a>>;
77 fn close<'a>(
79 &'a mut self,
80 ) -> Pin<Box<dyn Future<Output = Result<(), crate::DbError>> + Send + 'a>>;
81
82 fn execute_with_params<'a>(
88 &'a mut self,
89 sql: &'a str,
90 params: &'a [crate::value::Value],
91 ) -> Pin<Box<dyn Future<Output = Result<u64, crate::DbError>> + Send + 'a>> {
92 let _ = (sql, params);
93 Box::pin(async move {
94 Err(crate::DbError::Internal(
95 "execute_with_params not implemented for this adapter".to_string(),
96 ))
97 })
98 }
99
100 fn query_with_params<'a>(
106 &'a mut self,
107 sql: &'a str,
108 params: &'a [crate::value::Value],
109 ) -> Pin<Box<dyn Future<Output = Result<QueryRows, crate::DbError>> + Send + 'a>> {
110 let _ = (sql, params);
111 Box::pin(async move {
112 Err(crate::DbError::Internal(
113 "query_with_params not implemented for this adapter".to_string(),
114 ))
115 })
116 }
117
118 fn query_values<'a>(
123 &'a mut self,
124 sql: &'a str,
125 ) -> Pin<Box<dyn Future<Output = Result<crate::value::QueryValues, crate::DbError>> + Send + 'a>>
126 {
127 let _ = sql;
128 Box::pin(async move {
129 Err(crate::DbError::Internal(
130 "query_values not implemented for this adapter".to_string(),
131 ))
132 })
133 }
134
135 fn query_values_with_params<'a>(
139 &'a mut self,
140 sql: &'a str,
141 params: &'a [crate::value::Value],
142 ) -> Pin<Box<dyn Future<Output = Result<crate::value::QueryValues, crate::DbError>> + Send + 'a>>
143 {
144 let _ = (sql, params);
145 Box::pin(async move {
146 Err(crate::DbError::Internal(
147 "query_values_with_params not implemented for this adapter".to_string(),
148 ))
149 })
150 }
151
152 fn query_stream<'a>(
164 &'a mut self,
165 sql: &'a str,
166 ) -> Pin<Box<dyn futures::Stream<Item = QueryStreamItem> + Send + 'a>> {
167 let sql_owned = sql.to_string();
169 let stream = futures::stream::once(async move { self.query(&sql_owned).await })
171 .map(|result| {
173 let items: Vec<QueryStreamItem> = match result {
174 Ok(rows) => rows.into_iter().map(Ok).collect(),
175 Err(e) => vec![Err(e)],
176 };
177 futures::stream::iter(items)
178 })
179 .flatten();
180 Box::pin(stream)
181 }
182
183 fn query_stream_cursor<'a>(
194 &'a mut self,
195 sql: &'a str,
196 _batch_size: usize,
197 ) -> Pin<Box<dyn futures::Stream<Item = QueryStreamItem> + Send + 'a>> {
198 self.query_stream(sql)
199 }
200
201 fn execute_batch<'a>(
206 &'a mut self,
207 sqls: &'a [String],
208 ) -> Pin<Box<dyn Future<Output = Result<u64, crate::DbError>> + Send + 'a>> {
209 Box::pin(async move {
210 let mut total = 0u64;
211 for sql in sqls {
212 total += self.execute(sql).await?;
213 }
214 Ok(total)
215 })
216 }
217
218 fn execute_batch_params<'a>(
223 &'a mut self,
224 sql: &'a str,
225 params_batch: &'a [Vec<crate::value::Value>],
226 ) -> Pin<Box<dyn Future<Output = Result<u64, crate::DbError>> + Send + 'a>> {
227 Box::pin(async move {
228 let mut total = 0u64;
229 for params in params_batch {
230 total += self.execute_with_params(sql, params).await?;
231 }
232 Ok(total)
233 })
234 }
235}
236
237pub struct PooledConnection {
245 conn: Box<dyn Connection>,
246 created_at: Instant,
247 last_used_at: Instant,
248 pool: Option<Pool>,
249}
250
251impl PooledConnection {
252 fn new(conn: Box<dyn Connection>, pool: Pool) -> Self {
253 let now = Instant::now();
254 Self {
255 conn,
256 created_at: now,
257 last_used_at: now,
258 pool: Some(pool),
259 }
260 }
261
262 fn is_expired(&self, max_lifetime: Duration) -> bool {
263 self.created_at.elapsed() >= max_lifetime
264 }
265
266 fn is_idle_too_long(&self, idle_timeout: Duration) -> bool {
267 self.last_used_at.elapsed() >= idle_timeout
268 }
269
270 pub fn created_at(&self) -> Instant {
272 self.created_at
273 }
274
275 pub fn into_inner(mut self) -> Box<dyn Connection> {
280 self.pool = None; std::mem::replace(&mut self.conn, Box::new(ClosedConnection))
284 }
285}
286
287impl Drop for PooledConnection {
299 fn drop(&mut self) {
300 if let Some(pool) = self.pool.take() {
301 let conn = std::mem::replace(&mut self.conn, Box::new(ClosedConnection));
303 let pooled = PooledConnection {
304 conn,
305 created_at: self.created_at,
306 last_used_at: self.last_used_at,
307 pool: None,
308 };
309 if let Ok(handle) = tokio::runtime::Handle::try_current() {
311 handle.spawn(async move {
312 pool.release(pooled).await;
313 });
314 } else {
315 drop(pooled);
319 pool.total_count.fetch_sub(1, Ordering::SeqCst);
320 }
321 }
322 }
323}
324
325struct ClosedConnection;
329
330impl Connection for ClosedConnection {
331 fn execute<'a>(
332 &'a mut self,
333 _sql: &'a str,
334 ) -> Pin<Box<dyn Future<Output = Result<u64, crate::DbError>> + Send + 'a>> {
335 Box::pin(async {
336 Err(crate::DbError::ConnectionError(
337 "connection already returned to pool".to_string(),
338 ))
339 })
340 }
341
342 fn query<'a>(
343 &'a mut self,
344 _sql: &'a str,
345 ) -> Pin<Box<dyn Future<Output = Result<QueryRows, crate::DbError>> + Send + 'a>> {
346 Box::pin(async {
347 Err(crate::DbError::ConnectionError(
348 "connection already returned to pool".to_string(),
349 ))
350 })
351 }
352
353 fn begin_transaction<'a>(
354 &'a mut self,
355 ) -> Pin<Box<dyn Future<Output = Result<(), crate::DbError>> + Send + 'a>> {
356 Box::pin(async {
357 Err(crate::DbError::ConnectionError(
358 "connection already returned to pool".to_string(),
359 ))
360 })
361 }
362
363 fn commit<'a>(
364 &'a mut self,
365 ) -> Pin<Box<dyn Future<Output = Result<(), crate::DbError>> + Send + 'a>> {
366 Box::pin(async { Ok(()) })
367 }
368
369 fn rollback<'a>(
370 &'a mut self,
371 ) -> Pin<Box<dyn Future<Output = Result<(), crate::DbError>> + Send + 'a>> {
372 Box::pin(async { Ok(()) })
373 }
374
375 fn is_connected(&self) -> bool {
376 false
377 }
378
379 fn ping<'a>(&'a mut self) -> Pin<Box<dyn Future<Output = bool> + Send + 'a>> {
380 Box::pin(async { false })
381 }
382
383 fn close<'a>(
384 &'a mut self,
385 ) -> Pin<Box<dyn Future<Output = Result<(), crate::DbError>> + Send + 'a>> {
386 Box::pin(async { Ok(()) })
387 }
388}
389
390impl Deref for PooledConnection {
391 type Target = dyn Connection;
392
393 fn deref(&self) -> &Self::Target {
394 self.conn.as_ref()
395 }
396}
397
398impl DerefMut for PooledConnection {
399 fn deref_mut(&mut self) -> &mut Self::Target {
400 self.conn.as_mut()
401 }
402}
403
404#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
406pub enum TlsVersion {
407 #[default]
409 Tls12,
410 Tls13,
412}
413
414#[derive(Debug, Clone, Default)]
416pub struct TlsConfig {
417 pub enabled: bool,
419 pub ca_cert_path: Option<String>,
421 pub client_cert_path: Option<String>,
423 pub client_key_path: Option<String>,
425 pub min_version: TlsVersion,
427}
428
429#[derive(Debug, Clone)]
431pub enum PoolEvent {
432 ConnectionCreated,
434 ConnectionClosed,
436 ConnectionAcquired,
438 ConnectionReleased,
440 AcquireTimeout,
442}
443
444pub type PoolEventCallback = Arc<dyn Fn(PoolEvent) + Send + Sync>;
446
447pub struct PoolConfig {
449 pub max_size: u32,
451 pub min_idle: u32,
453 pub acquire_timeout: Duration,
455 pub idle_timeout: Duration,
457 pub max_lifetime: Duration,
459 pub connection_timeout: Duration,
461 pub tls: Option<TlsConfig>,
463 pub query_timeout: Option<Duration>,
465 pub max_rows: Option<usize>,
467 pub memory_limit: Option<usize>,
469 pub on_event: Option<PoolEventCallback>,
471 pub test_before_acquire: bool,
479 pub prewarm: bool,
491}
492
493impl Default for PoolConfig {
494 fn default() -> Self {
495 Self {
496 max_size: 100,
497 min_idle: 0,
498 acquire_timeout: Duration::from_secs(30),
499 idle_timeout: Duration::from_secs(600),
500 max_lifetime: Duration::from_secs(1800),
501 connection_timeout: Duration::from_secs(10),
502 tls: None,
503 query_timeout: Some(Duration::from_secs(30)),
504 max_rows: None,
505 memory_limit: None,
506 on_event: None,
507 test_before_acquire: false,
508 prewarm: false,
509 }
510 }
511}
512
513impl Clone for PoolConfig {
514 fn clone(&self) -> Self {
515 Self {
516 max_size: self.max_size,
517 min_idle: self.min_idle,
518 acquire_timeout: self.acquire_timeout,
519 idle_timeout: self.idle_timeout,
520 max_lifetime: self.max_lifetime,
521 connection_timeout: self.connection_timeout,
522 tls: self.tls.clone(),
523 query_timeout: self.query_timeout,
524 max_rows: self.max_rows,
525 memory_limit: self.memory_limit,
526 on_event: self.on_event.clone(),
527 test_before_acquire: self.test_before_acquire,
528 prewarm: self.prewarm,
529 }
530 }
531}
532
533impl PoolConfig {
534 pub fn validate(&self) -> Result<(), PoolError> {
536 if self.max_size == 0 {
537 return Err(PoolError::InvalidConfig("max_size cannot be 0".to_string()));
538 }
539 if self.min_idle > self.max_size {
540 return Err(PoolError::InvalidConfig(
541 "min_idle cannot exceed max_size".to_string(),
542 ));
543 }
544 const MAX_DURATION_SECS: u64 = u32::MAX as u64; for (name, dur) in [
550 ("acquire_timeout", self.acquire_timeout),
551 ("idle_timeout", self.idle_timeout),
552 ("max_lifetime", self.max_lifetime),
553 ("connection_timeout", self.connection_timeout),
554 ] {
555 if dur.as_secs() > MAX_DURATION_SECS {
556 return Err(PoolError::InvalidConfig(format!(
557 "{name} ({:?}) exceeds maximum allowed duration ({} seconds)",
558 dur, MAX_DURATION_SECS
559 )));
560 }
561 }
562 Ok(())
563 }
564
565 #[must_use]
567 pub fn with_prewarm(mut self, prewarm: bool) -> Self {
568 self.prewarm = prewarm;
569 self
570 }
571}
572
573pub struct PoolStatus {
575 pub idle: u32,
577 pub active: u32,
579 pub max: u32,
581 pub min: u32,
583 pub waiters: u32,
585}
586
587impl std::fmt::Debug for PoolStatus {
588 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
589 f.debug_struct("PoolStatus")
590 .field("idle", &self.idle)
591 .field("active", &self.active)
592 .field("max", &self.max)
593 .field("min", &self.min)
594 .field("waiters", &self.waiters)
595 .finish()
596 }
597}
598
599#[derive(Debug, Clone, Copy, Default, PartialEq, Eq, Serialize, Deserialize)]
606pub struct PoolMetrics {
607 pub acquire_count: u64,
609 pub acquire_failed_count: u64,
611 pub acquire_wait_time: Duration,
613 pub release_count: u64,
615 pub connection_created_count: u64,
617 pub connection_closed_count: u64,
619}
620
621impl PoolMetrics {
622 #[must_use]
624 pub fn average_acquire_wait_time(&self) -> Duration {
625 if self.acquire_count == 0 {
626 Duration::ZERO
627 } else {
628 self.acquire_wait_time / self.acquire_count as u32
629 }
630 }
631
632 #[must_use]
638 pub fn connection_reuse_rate(&self) -> f64 {
639 if self.acquire_count == 0 || self.connection_created_count > self.acquire_count {
640 0.0
641 } else {
642 let reused = self.acquire_count - self.connection_created_count;
643 reused as f64 / self.acquire_count as f64
644 }
645 }
646}
647
648#[derive(Debug, Clone)]
652pub struct PoolTuningAdvice {
653 pub suggested_max_size: Option<u32>,
655 pub suggested_min_idle: Option<u32>,
657 pub suggested_idle_timeout: Option<Duration>,
659 pub reason: String,
661}
662
663impl PoolTuningAdvice {
664 #[must_use]
666 pub fn is_optimal(&self) -> bool {
667 self.suggested_max_size.is_none()
668 && self.suggested_min_idle.is_none()
669 && self.suggested_idle_timeout.is_none()
670 }
671}
672
673pub struct PoolConfigBuilder {
675 config: PoolConfig,
676}
677
678impl PoolConfigBuilder {
679 pub fn new() -> Self {
681 Self {
682 config: PoolConfig::default(),
683 }
684 }
685
686 pub fn max_size(mut self, size: u32) -> Self {
688 self.config.max_size = size;
689 self
690 }
691
692 pub fn min_idle(mut self, count: u32) -> Self {
694 self.config.min_idle = count;
695 self
696 }
697
698 pub fn acquire_timeout(mut self, timeout_secs: u64) -> Self {
700 self.config.acquire_timeout = Duration::from_secs(timeout_secs);
701 self
702 }
703
704 pub fn idle_timeout(mut self, timeout_secs: u64) -> Self {
706 self.config.idle_timeout = Duration::from_secs(timeout_secs);
707 self
708 }
709
710 pub fn max_lifetime(mut self, lifetime_secs: u64) -> Self {
712 self.config.max_lifetime = Duration::from_secs(lifetime_secs);
713 self
714 }
715
716 pub fn tls(mut self, tls: TlsConfig) -> Self {
718 self.config.tls = Some(tls);
719 self
720 }
721
722 pub fn query_timeout(mut self, timeout: Duration) -> Self {
724 self.config.query_timeout = Some(timeout);
725 self
726 }
727
728 pub fn max_rows(mut self, max_rows: usize) -> Self {
730 self.config.max_rows = Some(max_rows);
731 self
732 }
733
734 pub fn memory_limit(mut self, memory_limit: usize) -> Self {
736 self.config.memory_limit = Some(memory_limit);
737 self
738 }
739
740 pub fn with_adaptive_tuning(
750 mut self,
751 capacity: usize,
752 idle_timeout_secs: u64,
753 acquire_timeout_ms: u64,
754 ) -> Self {
755 if capacity > 0 {
756 self.config.max_size = capacity as u32;
757 }
758 self.config.idle_timeout = Duration::from_secs(idle_timeout_secs);
759 self.config.acquire_timeout = Duration::from_millis(acquire_timeout_ms);
760 self
761 }
762
763 pub fn on_event(mut self, callback: PoolEventCallback) -> Self {
765 self.config.on_event = Some(callback);
766 self
767 }
768
769 pub fn test_before_acquire(mut self, enabled: bool) -> Self {
774 self.config.test_before_acquire = enabled;
775 self
776 }
777
778 pub fn prewarm(mut self, enabled: bool) -> Self {
783 self.config.prewarm = enabled;
784 self
785 }
786
787 pub fn build(self) -> Result<PoolConfig, PoolError> {
789 self.config.validate()?;
790 Ok(self.config)
791 }
792}
793
794impl Default for PoolConfigBuilder {
795 fn default() -> Self {
796 Self::new()
797 }
798}
799
800pub struct PoolCircuitBreakerLink {
817 original_capacity: u32,
819 current_capacity: u32,
821 shrink_count: std::sync::atomic::AtomicU64,
823 expand_count: std::sync::atomic::AtomicU64,
825 is_shrunk: std::sync::atomic::AtomicBool,
827}
828
829impl std::fmt::Debug for PoolCircuitBreakerLink {
830 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
831 f.debug_struct("PoolCircuitBreakerLink")
832 .field("original_capacity", &self.original_capacity)
833 .field("current_capacity", &self.current_capacity)
834 .field("is_shrunk", &self.is_shrunk.load(std::sync::atomic::Ordering::Relaxed))
835 .finish()
836 }
837}
838
839impl PoolCircuitBreakerLink {
840 pub fn new(original_capacity: u32) -> Self {
844 Self {
845 original_capacity,
846 current_capacity: original_capacity,
847 shrink_count: std::sync::atomic::AtomicU64::new(0),
848 expand_count: std::sync::atomic::AtomicU64::new(0),
849 is_shrunk: std::sync::atomic::AtomicBool::new(false),
850 }
851 }
852
853 pub fn shrink_pool(&mut self, factor: f64) -> u32 {
858 let factor = factor.clamp(0.1, 1.0);
859 let new_capacity = ((self.original_capacity as f64) * factor).round() as u32;
860 let new_capacity = new_capacity.max(1);
861 self.current_capacity = new_capacity;
862 self.shrink_count
863 .fetch_add(1, std::sync::atomic::Ordering::Relaxed);
864 self.is_shrunk
865 .store(true, std::sync::atomic::Ordering::Relaxed);
866 new_capacity
867 }
868
869 pub fn expand_pool(&mut self) -> u32 {
871 self.current_capacity = self.original_capacity;
872 self.expand_count
873 .fetch_add(1, std::sync::atomic::Ordering::Relaxed);
874 self.is_shrunk
875 .store(false, std::sync::atomic::Ordering::Relaxed);
876 self.current_capacity
877 }
878
879 pub fn current_capacity(&self) -> u32 {
881 self.current_capacity
882 }
883
884 pub fn original_capacity(&self) -> u32 {
886 self.original_capacity
887 }
888
889 pub fn shrink_count(&self) -> u64 {
891 self.shrink_count.load(std::sync::atomic::Ordering::Relaxed)
892 }
893
894 pub fn expand_count(&self) -> u64 {
896 self.expand_count.load(std::sync::atomic::Ordering::Relaxed)
897 }
898
899 pub fn is_shrunk(&self) -> bool {
901 self.is_shrunk.load(std::sync::atomic::Ordering::Relaxed)
902 }
903
904 pub fn on_circuit_state_change(&mut self, is_open: bool) -> u32 {
909 if is_open {
910 self.shrink_pool(0.5)
911 } else {
912 self.expand_pool()
913 }
914 }
915}
916
917#[cfg(test)]
918mod pool_circuit_breaker_link_tests {
919 use super::*;
920
921 #[test]
922 fn test_pool_circuit_breaker_link_new() {
923 let link = PoolCircuitBreakerLink::new(100);
924 assert_eq!(link.current_capacity(), 100);
925 assert_eq!(link.original_capacity(), 100);
926 assert!(!link.is_shrunk());
927 }
928
929 #[test]
930 fn test_shrink_pool_half() {
931 let mut link = PoolCircuitBreakerLink::new(100);
932 let new_cap = link.shrink_pool(0.5);
933 assert_eq!(new_cap, 50);
934 assert_eq!(link.current_capacity(), 50);
935 assert!(link.is_shrunk());
936 assert_eq!(link.shrink_count(), 1);
937 }
938
939 #[test]
940 fn test_shrink_pool_minimum_one() {
941 let mut link = PoolCircuitBreakerLink::new(2);
942 let new_cap = link.shrink_pool(0.1);
943 assert_eq!(new_cap, 1);
944 }
945
946 #[test]
947 fn test_expand_pool_restores_original() {
948 let mut link = PoolCircuitBreakerLink::new(100);
949 link.shrink_pool(0.3);
950 assert_eq!(link.current_capacity(), 30);
951 let restored = link.expand_pool();
952 assert_eq!(restored, 100);
953 assert!(!link.is_shrunk());
954 assert_eq!(link.expand_count(), 1);
955 }
956
957 #[test]
958 fn test_on_circuit_state_change_open() {
959 let mut link = PoolCircuitBreakerLink::new(100);
960 let cap = link.on_circuit_state_change(true);
961 assert_eq!(cap, 50);
962 assert!(link.is_shrunk());
963 }
964
965 #[test]
966 fn test_on_circuit_state_change_closed() {
967 let mut link = PoolCircuitBreakerLink::new(100);
968 link.on_circuit_state_change(true);
969 let cap = link.on_circuit_state_change(false);
970 assert_eq!(cap, 100);
971 assert!(!link.is_shrunk());
972 }
973
974 #[test]
975 fn test_shrink_factor_clamped() {
976 let mut link = PoolCircuitBreakerLink::new(100);
977 let cap = link.shrink_pool(0.0);
978 assert!(cap >= 10);
979 let cap2 = link.shrink_pool(2.0);
980 assert!(cap2 <= 100);
981 }
982
983 #[test]
984 fn test_multiple_shrink_expand_cycles() {
985 let mut link = PoolCircuitBreakerLink::new(100);
986 for _ in 0..3 {
987 link.on_circuit_state_change(true);
988 link.on_circuit_state_change(false);
989 }
990 assert_eq!(link.shrink_count(), 3);
991 assert_eq!(link.expand_count(), 3);
992 assert_eq!(link.current_capacity(), 100);
993 }
994
995 #[test]
996 fn test_debug_format() {
997 let link = PoolCircuitBreakerLink::new(50);
998 let s = format!("{:?}", link);
999 assert!(s.contains("PoolCircuitBreakerLink"));
1000 assert!(s.contains("50"));
1001 }
1002}
1003
1004#[async_trait]
1006pub trait ConnectionFactory: Send + Sync {
1007 async fn create(&self) -> Result<Box<dyn Connection>, crate::DbError>;
1009}
1010
1011pub struct Pool {
1017 config: PoolConfig,
1018 factory: Arc<dyn ConnectionFactory>,
1019 idle: Arc<ArrayQueue<PooledConnection>>,
1025 total_count: Arc<AtomicU32>,
1035 closed: Arc<AtomicBool>,
1037 notify: Arc<Notify>,
1038 waiters_count: Arc<AtomicU32>,
1040 dynamic_max_size: Arc<AtomicU32>,
1042 #[cfg(feature = "circuit-breaker")]
1048 circuit_breaker: Arc<PlMutex<DefaultCircuitBreaker>>,
1049 #[cfg(feature = "rate-limit")]
1058 rate_limiter: Arc<PlRwLock<Option<Arc<dyn RateLimiter>>>>,
1059 #[cfg(feature = "rate-limit")]
1061 rate_limit_key: String,
1062 #[cfg(feature = "tenant-quota-rls-enhanced")]
1067 quota_enforcer: Arc<PlMutex<Option<Arc<QuotaEnforcer>>>>,
1068 acquire_count: Arc<AtomicU64>,
1070 acquire_failed_count: Arc<AtomicU64>,
1072 acquire_wait_time_ns: Arc<AtomicU64>,
1074 release_count: Arc<AtomicU64>,
1076 connection_created_count: Arc<AtomicU64>,
1078 connection_closed_count: Arc<AtomicU64>,
1080}
1081
1082impl Clone for Pool {
1086 fn clone(&self) -> Self {
1087 Self {
1088 config: self.config.clone(),
1089 factory: self.factory.clone(),
1090 idle: self.idle.clone(),
1091 total_count: self.total_count.clone(),
1092 closed: self.closed.clone(),
1093 notify: Arc::clone(&self.notify),
1094 waiters_count: self.waiters_count.clone(),
1095 dynamic_max_size: self.dynamic_max_size.clone(),
1096 #[cfg(feature = "circuit-breaker")]
1097 circuit_breaker: Arc::clone(&self.circuit_breaker),
1098 #[cfg(feature = "rate-limit")]
1099 rate_limiter: Arc::clone(&self.rate_limiter),
1100 #[cfg(feature = "rate-limit")]
1101 rate_limit_key: self.rate_limit_key.clone(),
1102 #[cfg(feature = "tenant-quota-rls-enhanced")]
1103 quota_enforcer: Arc::clone(&self.quota_enforcer),
1104 acquire_count: self.acquire_count.clone(),
1105 acquire_failed_count: self.acquire_failed_count.clone(),
1106 acquire_wait_time_ns: self.acquire_wait_time_ns.clone(),
1107 release_count: self.release_count.clone(),
1108 connection_created_count: self.connection_created_count.clone(),
1109 connection_closed_count: self.connection_closed_count.clone(),
1110 }
1111 }
1112}
1113
1114impl Pool {
1115 pub fn new(config: PoolConfig, factory: Arc<dyn ConnectionFactory>) -> Result<Self, PoolError> {
1139 config.validate()?;
1140 let max_size = config.max_size as usize;
1143 let dynamic_max = config.max_size;
1144 Ok(Self {
1145 config,
1146 factory,
1147 idle: Arc::new(ArrayQueue::new(max_size)),
1148 total_count: Arc::new(AtomicU32::new(0)),
1149 closed: Arc::new(AtomicBool::new(false)),
1150 notify: Arc::new(Notify::new()),
1151 waiters_count: Arc::new(AtomicU32::new(0)),
1152 dynamic_max_size: Arc::new(AtomicU32::new(dynamic_max)),
1153 #[cfg(feature = "circuit-breaker")]
1156 circuit_breaker: Arc::new(PlMutex::new(DefaultCircuitBreaker::new(
1157 5,
1158 std::time::Duration::from_secs(30),
1159 ))),
1160 #[cfg(feature = "rate-limit")]
1163 rate_limiter: Arc::new(PlRwLock::new(None)),
1164 #[cfg(feature = "rate-limit")]
1165 rate_limit_key: "pool".to_string(),
1166 #[cfg(feature = "tenant-quota-rls-enhanced")]
1167 quota_enforcer: Arc::new(PlMutex::new(None)),
1168 acquire_count: Arc::new(AtomicU64::new(0)),
1169 acquire_failed_count: Arc::new(AtomicU64::new(0)),
1170 acquire_wait_time_ns: Arc::new(AtomicU64::new(0)),
1171 release_count: Arc::new(AtomicU64::new(0)),
1172 connection_created_count: Arc::new(AtomicU64::new(0)),
1173 connection_closed_count: Arc::new(AtomicU64::new(0)),
1174 })
1175 }
1176
1177 pub async fn new_async(
1184 config: PoolConfig,
1185 factory: Arc<dyn ConnectionFactory>,
1186 ) -> Result<Self, PoolError> {
1187 let pool = Self::new(config, factory)?;
1188 if pool.config.prewarm {
1189 pool.prewarm().await;
1190 }
1191 Ok(pool)
1192 }
1193
1194 pub async fn prewarm(&self) {
1211 if !self.config.prewarm {
1212 return;
1213 }
1214
1215 let min_idle = self.config.min_idle as usize;
1216 let mut warmed = 0;
1217
1218 for i in 0..min_idle {
1219 if self.closed.load(Ordering::Acquire) {
1221 break;
1222 }
1223
1224 let current_max = self.dynamic_max_size.load(Ordering::Acquire);
1226 let current = self.total_count.load(Ordering::Acquire);
1227 if current >= current_max {
1228 break;
1229 }
1230
1231 let created = loop {
1233 let current = self.total_count.load(Ordering::Acquire);
1234 if current >= current_max {
1235 break None;
1236 }
1237 match self.total_count.compare_exchange(
1238 current,
1239 current + 1,
1240 Ordering::SeqCst,
1241 Ordering::Acquire,
1242 ) {
1243 Ok(_) => break Some(()),
1244 Err(_) => continue,
1245 }
1246 };
1247
1248 if created.is_some() {
1249 match tokio::time::timeout(self.config.connection_timeout, self.factory.create())
1250 .await
1251 {
1252 Ok(Ok(conn)) => {
1253 #[cfg(feature = "circuit-breaker")]
1254 {
1255 self.circuit_breaker.lock().record_success();
1256 }
1257 self.emit_event(PoolEvent::ConnectionCreated);
1258 let pooled = PooledConnection::new(conn, self.clone());
1259 if self.idle.push(pooled).is_err() {
1261 let _ = self.total_count.fetch_sub(1, Ordering::SeqCst);
1263 tracing::warn!(
1264 target: "sz_orm::pool::prewarm",
1265 "prewarm connection {} failed: idle queue full",
1266 i
1267 );
1268 } else {
1269 warmed += 1;
1270 self.notify.notify_one();
1271 }
1272 }
1273 Ok(Err(e)) => {
1274 let _ = self.total_count.fetch_sub(1, Ordering::SeqCst);
1275 #[cfg(feature = "circuit-breaker")]
1276 {
1277 self.circuit_breaker.lock().record_failure();
1278 }
1279 tracing::warn!(
1280 target: "sz_orm::pool::prewarm",
1281 "prewarm connection {} failed: {}",
1282 i,
1283 e
1284 );
1285 }
1286 Err(_) => {
1287 let _ = self.total_count.fetch_sub(1, Ordering::SeqCst);
1288 #[cfg(feature = "circuit-breaker")]
1289 {
1290 self.circuit_breaker.lock().record_failure();
1291 }
1292 tracing::warn!(
1293 target: "sz_orm::pool::prewarm",
1294 "prewarm connection {} timeout",
1295 i
1296 );
1297 }
1298 }
1299 }
1300 }
1301
1302 if warmed > 0 {
1303 tracing::info!(
1304 target: "sz_orm::pool::prewarm",
1305 "pool prewarm completed: {}/{} connections established",
1306 warmed,
1307 min_idle
1308 );
1309 }
1310 }
1311
1312 #[cfg(feature = "auto-prewarm")]
1317 pub async fn progressive_prewarm(
1318 &self,
1319 batch_size: u32,
1320 interval: std::time::Duration,
1321 total_timeout: std::time::Duration,
1322 progress: &crate::prewarm::PrewarmProgress,
1323 ) {
1324 use std::time::Instant;
1325
1326 let min_idle = self.config.min_idle;
1327 if min_idle == 0 || !self.config.prewarm {
1328 progress.mark_completed();
1329 return;
1330 }
1331
1332 let start = Instant::now();
1333 let batch = batch_size.max(1);
1334 let mut warmed_total: u32 = 0;
1335
1336 while warmed_total < min_idle {
1337 if start.elapsed() >= total_timeout {
1338 tracing::warn!(
1339 target: "sz_orm::pool::prewarm",
1340 "progressive prewarm timeout: {}/{} connections established",
1341 warmed_total,
1342 min_idle
1343 );
1344 break;
1345 }
1346
1347 if self.closed.load(Ordering::Acquire) {
1348 break;
1349 }
1350
1351 let remaining = min_idle - warmed_total;
1352 let this_batch = batch.min(remaining);
1353
1354 for _ in 0..this_batch {
1355 let current_max = self.dynamic_max_size.load(Ordering::Acquire);
1356 let current = self.total_count.load(Ordering::Acquire);
1357 if current >= current_max {
1358 break;
1359 }
1360
1361 let created = loop {
1362 let current = self.total_count.load(Ordering::Acquire);
1363 if current >= current_max {
1364 break None;
1365 }
1366 match self.total_count.compare_exchange(
1367 current,
1368 current + 1,
1369 Ordering::SeqCst,
1370 Ordering::Acquire,
1371 ) {
1372 Ok(_) => break Some(()),
1373 Err(_) => continue,
1374 }
1375 };
1376
1377 if created.is_some() {
1378 match tokio::time::timeout(
1379 self.config.connection_timeout,
1380 self.factory.create(),
1381 )
1382 .await
1383 {
1384 Ok(Ok(conn)) => {
1385 #[cfg(feature = "circuit-breaker")]
1386 {
1387 self.circuit_breaker.lock().record_success();
1388 }
1389 self.emit_event(PoolEvent::ConnectionCreated);
1390 let pooled = PooledConnection::new(conn, self.clone());
1391 if self.idle.push(pooled).is_err() {
1392 let _ = self.total_count.fetch_sub(1, Ordering::SeqCst);
1393 progress.record_failure();
1394 } else {
1395 progress.record_success();
1396 warmed_total += 1;
1397 self.notify.notify_one();
1398 }
1399 }
1400 Ok(Err(_)) => {
1401 let _ = self.total_count.fetch_sub(1, Ordering::SeqCst);
1402 progress.record_failure();
1403 #[cfg(feature = "circuit-breaker")]
1404 {
1405 self.circuit_breaker.lock().record_failure();
1406 }
1407 }
1408 Err(_) => {
1409 let _ = self.total_count.fetch_sub(1, Ordering::SeqCst);
1410 progress.record_failure();
1411 #[cfg(feature = "circuit-breaker")]
1412 {
1413 self.circuit_breaker.lock().record_failure();
1414 }
1415 }
1416 }
1417 }
1418 }
1419
1420 if warmed_total < min_idle && interval > std::time::Duration::ZERO {
1421 tokio::time::sleep(interval).await;
1422 }
1423 }
1424
1425 progress.set_elapsed(start.elapsed());
1426 progress.mark_completed();
1427
1428 tracing::info!(
1429 target: "sz_orm::pool::prewarm",
1430 "progressive prewarm completed: {} warmed, {} failed, elapsed {:?}",
1431 progress.snapshot().warmed,
1432 progress.snapshot().failed,
1433 start.elapsed()
1434 );
1435 }
1436
1437 pub fn config(&self) -> &PoolConfig {
1439 &self.config
1440 }
1441
1442 #[cfg(feature = "circuit-breaker")]
1456 pub fn configure_circuit_breaker(
1457 &self,
1458 failure_threshold: usize,
1459 reset_timeout: std::time::Duration,
1460 ) {
1461 let new_cb = DefaultCircuitBreaker::new(failure_threshold, reset_timeout);
1462 let mut guard = self.circuit_breaker.lock();
1464 *guard = new_cb;
1465 }
1466
1467 #[cfg(feature = "circuit-breaker")]
1472 pub fn reset_circuit_breaker(&self) -> bool {
1473 let mut guard = self.circuit_breaker.lock();
1475 guard.reset()
1476 }
1477
1478 #[cfg(feature = "circuit-breaker")]
1480 pub fn circuit_state(&self) -> CircuitState {
1481 let guard = self.circuit_breaker.lock();
1483 guard.state()
1484 }
1485
1486 #[cfg(feature = "rate-limit")]
1495 pub fn set_rate_limiter(&self, limiter: Option<Arc<dyn RateLimiter>>) {
1496 let mut guard = self.rate_limiter.write();
1498 *guard = limiter;
1499 }
1500
1501 #[cfg(feature = "rate-limit")]
1503 pub fn with_rate_limit_key(mut self, key: impl Into<String>) -> Self {
1504 self.rate_limit_key = key.into();
1505 self
1506 }
1507
1508 #[cfg(feature = "tenant-quota-rls-enhanced")]
1513 pub fn set_quota_enforcer(&self, enforcer: Option<Arc<QuotaEnforcer>>) {
1514 let mut guard = self.quota_enforcer.lock();
1515 *guard = enforcer;
1516 }
1517
1518 #[cfg(feature = "tenant-quota-rls-enhanced")]
1527 pub async fn acquire_with_tenant(
1528 &self,
1529 tenant_id: &str,
1530 ) -> Result<PooledConnection, PoolError> {
1531 {
1532 let guard = self.quota_enforcer.lock();
1533 if let Some(ref enforcer) = *guard {
1534 let current = enforcer.current_usage(tenant_id, QuotaResource::Connection);
1535 enforcer
1536 .check_and_record(tenant_id, QuotaResource::Connection, 1)
1537 .map_err(|e| PoolError::Internal(e.to_string()))?;
1538 let _ = current;
1539 }
1540 }
1541 self.acquire().await
1542 }
1543
1544 #[cfg(feature = "tenant-quota-rls-enhanced")]
1549 pub async fn release_with_tenant(&self, tenant_id: &str, pooled: PooledConnection) {
1550 {
1551 let guard = self.quota_enforcer.lock();
1552 if let Some(ref enforcer) = *guard {
1553 enforcer.release_usage(tenant_id, QuotaResource::Connection, 1);
1556 }
1557 }
1558 self.release(pooled).await;
1559 }
1560
1561 fn emit_event(&self, event: PoolEvent) {
1563 if matches!(event, PoolEvent::ConnectionCreated) {
1566 self.connection_created_count
1567 .fetch_add(1, Ordering::Relaxed);
1568 }
1569 if let Some(ref callback) = self.config.on_event {
1570 callback(event);
1571 }
1572 }
1573
1574 async fn close_connection(&self, pooled: PooledConnection) {
1579 let mut pooled = pooled;
1580 let _ = pooled.conn.close().await;
1581 self.connection_closed_count.fetch_add(1, Ordering::Relaxed);
1582 }
1583
1584 #[tracing::instrument(skip(self), fields(max_size = self.config.max_size, acquire_timeout = ?self.config.acquire_timeout))]
1604 pub async fn acquire(&self) -> Result<PooledConnection, PoolError> {
1605 if self.closed.load(Ordering::Acquire) {
1607 self.acquire_failed_count.fetch_add(1, Ordering::Relaxed);
1608 return Err(PoolError::Closed);
1609 }
1610
1611 #[cfg(feature = "circuit-breaker")]
1615 {
1616 let mut guard = self.circuit_breaker.lock();
1617 if !guard.can_execute() {
1618 self.acquire_failed_count.fetch_add(1, Ordering::Relaxed);
1619 return Err(PoolError::CircuitOpen);
1620 }
1621 }
1622
1623 #[cfg(feature = "rate-limit")]
1627 {
1628 let guard = self.rate_limiter.read();
1629 if let Some(ref limiter) = *guard {
1630 match limiter.try_acquire(&self.rate_limit_key) {
1631 Ok(result) if !result.allowed => {
1632 self.acquire_failed_count.fetch_add(1, Ordering::Relaxed);
1633 return Err(PoolError::RateLimited {
1634 remaining: result.remaining,
1635 reset_at: result.reset_at,
1636 });
1637 }
1638 Ok(_) => {} Err(_) => {
1640 }
1642 }
1643 }
1644 }
1645
1646 let mut deadline: Option<Instant> = None;
1647 let mut backoff = Duration::from_millis(1);
1649 const MAX_BACKOFF: Duration = Duration::from_millis(100);
1651 let mut to_close: Vec<PooledConnection> = Vec::with_capacity(4);
1653
1654 loop {
1655 let acquired: Option<PooledConnection> = {
1662 let mut found: Option<PooledConnection> = None;
1663 while let Some(pooled) = self.idle.pop() {
1664 if pooled.is_expired(self.config.max_lifetime) {
1666 to_close.push(pooled);
1667 continue;
1668 }
1669 if pooled.is_idle_too_long(self.config.idle_timeout) {
1671 to_close.push(pooled);
1672 continue;
1673 }
1674 if !pooled.conn.is_connected() {
1677 to_close.push(pooled);
1678 continue;
1679 }
1680 found = Some(pooled);
1681 break;
1682 }
1683 found
1684 };
1685
1686 for pooled in to_close.drain(..) {
1688 self.close_connection(pooled).await;
1689 self.total_count.fetch_sub(1, Ordering::SeqCst);
1691 }
1692
1693 if let Some(mut pooled) = acquired {
1694 if self.config.test_before_acquire {
1696 let ping_timeout = self.config.connection_timeout / 2;
1697 let alive = match tokio::time::timeout(ping_timeout, pooled.conn.ping()).await {
1698 Ok(true) => true,
1699 Ok(false) => false,
1700 Err(_) => false, };
1702 if !alive {
1703 self.close_connection(pooled).await;
1705 self.total_count.fetch_sub(1, Ordering::SeqCst);
1706 continue;
1707 }
1708 }
1709 pooled.pool = Some(self.clone());
1712 self.acquire_count.fetch_add(1, Ordering::Relaxed);
1713 return Ok(pooled);
1714 }
1715
1716 let current_max = self.dynamic_max_size.load(Ordering::Acquire);
1721 let created = loop {
1722 let current = self.total_count.load(Ordering::Acquire);
1723 if current >= current_max {
1724 break None; }
1726 match self.total_count.compare_exchange(
1727 current,
1728 current + 1,
1729 Ordering::SeqCst,
1730 Ordering::Acquire,
1731 ) {
1732 Ok(_) => break Some(()), Err(_) => continue, }
1735 };
1736
1737 if created.is_some() {
1738 match tokio::time::timeout(self.config.connection_timeout, self.factory.create())
1739 .await
1740 {
1741 Ok(Ok(conn)) => {
1742 #[cfg(feature = "circuit-breaker")]
1745 {
1746 self.circuit_breaker.lock().record_success();
1747 }
1748 self.emit_event(PoolEvent::ConnectionCreated);
1749 self.emit_event(PoolEvent::ConnectionAcquired);
1750 self.acquire_count.fetch_add(1, Ordering::Relaxed);
1751 return Ok(PooledConnection::new(conn, self.clone()));
1752 }
1753 Ok(Err(e)) => {
1754 self.total_count.fetch_sub(1, Ordering::SeqCst);
1756 #[cfg(feature = "circuit-breaker")]
1759 {
1760 self.circuit_breaker.lock().record_failure();
1761 }
1762 self.acquire_failed_count.fetch_add(1, Ordering::Relaxed);
1763 return Err(PoolError::ConnectionFailed(e.to_string()));
1764 }
1765 Err(_) => {
1766 self.total_count.fetch_sub(1, Ordering::SeqCst);
1768 #[cfg(feature = "circuit-breaker")]
1771 {
1772 self.circuit_breaker.lock().record_failure();
1773 }
1774 self.acquire_failed_count.fetch_add(1, Ordering::Relaxed);
1775 return Err(PoolError::Timeout);
1776 }
1777 }
1778 }
1779
1780 let now = Instant::now();
1782 let dl = deadline.get_or_insert_with(|| now + self.config.acquire_timeout);
1783 if now >= *dl {
1784 self.emit_event(PoolEvent::AcquireTimeout);
1785 self.acquire_failed_count.fetch_add(1, Ordering::Relaxed);
1786 return Err(PoolError::Timeout);
1787 }
1788 self.waiters_count.fetch_add(1, Ordering::SeqCst);
1790 let wait = std::cmp::min(backoff, *dl - now);
1791 match tokio::time::timeout(wait, self.notify.notified()).await {
1792 Ok(()) => {
1793 backoff = Duration::from_millis(1);
1795 }
1796 Err(_) => {
1797 backoff = std::cmp::min(backoff * 2, MAX_BACKOFF);
1799 }
1800 }
1801 self.waiters_count.fetch_sub(1, Ordering::SeqCst);
1803 self.acquire_wait_time_ns
1805 .fetch_add(wait.as_nanos() as u64, Ordering::Relaxed);
1806 }
1807 }
1808
1809 pub async fn acquire_batch(&self, n: usize) -> Result<Vec<PooledConnection>, PoolError> {
1815 if n == 0 {
1816 return Ok(Vec::new());
1817 }
1818 let max_size = self.dynamic_max_size.load(Ordering::Relaxed) as usize;
1819 if n > max_size {
1820 return Err(PoolError::Exhausted);
1821 }
1822 let mut result = Vec::with_capacity(n);
1823 for _ in 0..n {
1824 match self.acquire().await {
1825 Ok(conn) => result.push(conn),
1826 Err(e) => {
1827 return Err(e);
1828 }
1829 }
1830 }
1831 Ok(result)
1832 }
1833
1834 #[tracing::instrument(skip(self, pooled))]
1842 pub async fn release(&self, mut pooled: PooledConnection) {
1843 pooled.pool = None;
1845 self.release_count.fetch_add(1, Ordering::Relaxed);
1847
1848 if self.closed.load(Ordering::Acquire) {
1850 self.close_connection(pooled).await;
1851 self.total_count.fetch_sub(1, Ordering::SeqCst);
1853 self.emit_event(PoolEvent::ConnectionClosed);
1854 return;
1855 }
1856
1857 if !pooled.conn.is_connected() {
1859 self.close_connection(pooled).await;
1860 self.total_count.fetch_sub(1, Ordering::SeqCst);
1861 self.emit_event(PoolEvent::ConnectionClosed);
1862 return;
1863 }
1864
1865 pooled.last_used_at = Instant::now();
1867
1868 if let Err(rejected) = self.idle.push(pooled) {
1874 self.close_connection(rejected).await;
1876 self.total_count.fetch_sub(1, Ordering::SeqCst);
1877 self.emit_event(PoolEvent::ConnectionClosed);
1878 } else {
1879 self.emit_event(PoolEvent::ConnectionReleased);
1880 }
1881 self.notify.notify_one();
1882 }
1883
1884 pub async fn status(&self) -> PoolStatus {
1889 let idle_count = self.idle.len() as u32;
1890 let active = self.total_count.load(Ordering::Acquire);
1892 let waiters = self.waiters_count.load(Ordering::Acquire);
1893 PoolStatus {
1894 idle: idle_count,
1895 active,
1896 max: self.dynamic_max_size.load(Ordering::Acquire),
1897 min: self.config.min_idle,
1898 waiters,
1899 }
1900 }
1901
1902 pub fn pool_metrics(&self) -> PoolMetrics {
1913 PoolMetrics {
1914 acquire_count: self.acquire_count.load(Ordering::Acquire),
1915 acquire_failed_count: self.acquire_failed_count.load(Ordering::Acquire),
1916 acquire_wait_time: Duration::from_nanos(
1917 self.acquire_wait_time_ns.load(Ordering::Acquire),
1918 ),
1919 release_count: self.release_count.load(Ordering::Acquire),
1920 connection_created_count: self.connection_created_count.load(Ordering::Acquire),
1921 connection_closed_count: self.connection_closed_count.load(Ordering::Acquire),
1922 }
1923 }
1924
1925 #[must_use]
1934 pub fn suggest_tuning(&self) -> PoolTuningAdvice {
1935 let metrics = self.pool_metrics();
1936 let current_max = self.dynamic_max_size.load(Ordering::Acquire);
1937
1938 if metrics.acquire_count == 0 {
1939 return PoolTuningAdvice {
1940 suggested_max_size: None,
1941 suggested_min_idle: None,
1942 suggested_idle_timeout: None,
1943 reason: "数据不足".to_string(),
1944 };
1945 }
1946
1947 let reuse_rate = metrics.connection_reuse_rate();
1948 let mut advice = PoolTuningAdvice {
1949 suggested_max_size: None,
1950 suggested_min_idle: None,
1951 suggested_idle_timeout: None,
1952 reason: String::new(),
1953 };
1954
1955 if reuse_rate < 0.5 {
1956 advice.suggested_max_size = Some(current_max.saturating_mul(2));
1957 advice.reason = "复用率过低,池过小或回收过激".to_string();
1958 } else if reuse_rate < 0.9 {
1959 advice.suggested_min_idle = Some(current_max / 4);
1960 advice.reason = "复用率偏低,预热不足".to_string();
1961 }
1962
1963 let avg_wait = metrics.average_acquire_wait_time();
1964 if avg_wait > Duration::from_millis(100) {
1965 advice.suggested_max_size = Some(current_max.saturating_mul(2));
1966 if !advice.reason.is_empty() {
1967 advice.reason.push(';');
1968 }
1969 advice.reason.push_str("等待时长过高,池容量不足");
1970 }
1971
1972 if metrics.connection_created_count > 0
1973 && metrics.connection_closed_count as f64
1974 > metrics.connection_created_count as f64 * 0.5
1975 {
1976 advice.suggested_idle_timeout = Some(self.config.idle_timeout * 2);
1977 if !advice.reason.is_empty() {
1978 advice.reason.push(';');
1979 }
1980 advice.reason.push_str("连接关闭过快,空闲回收过激");
1981 }
1982
1983 if advice.reason.is_empty() {
1984 advice.reason = "池配置合理".to_string();
1985 }
1986
1987 advice
1988 }
1989
1990 pub fn metrics_snapshot_json(&self) -> String {
1995 serde_json::to_string(&self.pool_metrics()).unwrap_or_else(|_| "{}".to_string())
1996 }
1997
1998 #[tracing::instrument(skip(self))]
2000 pub async fn reap_idle(&self) {
2001 let mut all: Vec<PooledConnection> = Vec::new();
2005 while let Some(pooled) = self.idle.pop() {
2006 all.push(pooled);
2007 }
2008
2009 let mut to_close = Vec::new();
2011 for pooled in all {
2012 if pooled.is_idle_too_long(self.config.idle_timeout)
2013 || pooled.is_expired(self.config.max_lifetime)
2014 {
2015 to_close.push(pooled);
2016 } else {
2017 if let Err(rejected) = self.idle.push(pooled) {
2019 self.close_connection(rejected).await;
2020 self.total_count.fetch_sub(1, Ordering::SeqCst);
2021 }
2022 }
2023 }
2024
2025 for pooled in to_close {
2027 self.close_connection(pooled).await;
2028 self.total_count.fetch_sub(1, Ordering::SeqCst);
2030 }
2031 }
2032
2033 pub async fn close_all(&self) {
2037 self.closed.store(true, Ordering::Release);
2039 let mut to_close: Vec<PooledConnection> = Vec::new();
2042 while let Some(pooled) = self.idle.pop() {
2043 to_close.push(pooled);
2044 }
2045 let closed_count: u32 = to_close.len() as u32;
2047 for pooled in to_close {
2048 self.close_connection(pooled).await;
2049 }
2050 self.total_count.fetch_sub(closed_count, Ordering::SeqCst);
2053 }
2054
2055 pub async fn health_check(&self) -> u32 {
2071 let mut to_check: Vec<PooledConnection> = Vec::new();
2073 while let Some(pooled) = self.idle.pop() {
2074 to_check.push(pooled);
2075 }
2076
2077 let mut removed: u32 = 0;
2078 let mut alive: Vec<PooledConnection> = Vec::with_capacity(to_check.len());
2079 for mut pooled in to_check.drain(..) {
2080 if !pooled.conn.is_connected() {
2082 self.close_connection(pooled).await;
2083 removed += 1;
2084 continue;
2085 }
2086 let ping_timeout = self.config.connection_timeout / 2;
2088 match tokio::time::timeout(ping_timeout, pooled.conn.ping()).await {
2089 Ok(true) => alive.push(pooled),
2090 Ok(false) => {
2091 self.close_connection(pooled).await;
2093 removed += 1;
2094 }
2095 Err(_) => {
2096 self.close_connection(pooled).await;
2098 removed += 1;
2099 }
2100 }
2101 }
2102
2103 let alive_count: u32 = alive.len() as u32;
2105 for pooled in alive {
2106 if let Err(rejected) = self.idle.push(pooled) {
2108 self.close_connection(rejected).await;
2109 removed += 1;
2110 }
2111 }
2112
2113 if removed > 0 {
2115 self.total_count.fetch_sub(removed, Ordering::SeqCst);
2116 }
2117
2118 if alive_count > 0 {
2120 self.notify.notify_one();
2121 }
2122
2123 removed
2124 }
2125
2126 pub async fn shutdown(&self) {
2133 self.shutdown_with_timeout(Duration::from_secs(30)).await;
2134 }
2135
2136 pub async fn shutdown_with_timeout(&self, timeout: Duration) {
2144 if self.closed.swap(true, Ordering::SeqCst) {
2146 return;
2147 }
2148 self.notify.notify_waiters();
2150 self.close_all().await;
2152 let deadline = Instant::now() + timeout;
2154 while self.total_count.load(Ordering::SeqCst) > 0 {
2155 if Instant::now() >= deadline {
2156 let remaining = self.total_count.load(Ordering::SeqCst);
2157 if remaining > 0 {
2158 eprintln!(
2159 "graceful shutdown timeout, {} connections force closed",
2160 remaining
2161 );
2162 }
2163 break;
2164 }
2165 tokio::time::sleep(Duration::from_millis(100)).await;
2166 }
2167 }
2168
2169 pub fn resize(&self, new_max: usize) {
2177 self.set_max_size(new_max as u32);
2178 }
2179
2180 pub fn set_max_size(&self, new_max: u32) {
2182 self.dynamic_max_size.store(new_max, Ordering::SeqCst);
2183 }
2184
2185 pub fn max_size(&self) -> u32 {
2187 self.dynamic_max_size.load(Ordering::Acquire)
2188 }
2189
2190 pub async fn warmup(&self, min_idle: usize) -> Result<(), PoolError> {
2194 for _ in 0..min_idle {
2195 let current_max = self.dynamic_max_size.load(Ordering::Acquire);
2196 let current = self.total_count.load(Ordering::Acquire);
2197 if current >= current_max {
2198 break;
2199 }
2200 match self.total_count.compare_exchange(
2202 current,
2203 current + 1,
2204 Ordering::SeqCst,
2205 Ordering::Acquire,
2206 ) {
2207 Ok(_) => {}
2208 Err(_) => continue, }
2210 match self.factory.create().await {
2211 Ok(conn) => {
2212 let now = Instant::now();
2213 let pooled = PooledConnection {
2214 conn,
2215 created_at: now,
2216 last_used_at: now,
2217 pool: None,
2218 };
2219 if let Err(rejected) = self.idle.push(pooled) {
2220 self.close_connection(rejected).await;
2222 self.total_count.fetch_sub(1, Ordering::SeqCst);
2223 }
2224 self.emit_event(PoolEvent::ConnectionCreated);
2225 }
2226 Err(_) => {
2227 self.total_count.fetch_sub(1, Ordering::SeqCst);
2229 break;
2230 }
2231 }
2232 }
2233 Ok(())
2234 }
2235
2236 pub async fn query_with_timeout(&self, sql: &str) -> Result<QueryRows, crate::DbError> {
2241 let timeout = self.config.query_timeout.unwrap_or(Duration::from_secs(30));
2242 let mut conn = self.acquire().await.map_err(crate::DbError::PoolError)?;
2243 tokio::time::timeout(timeout, conn.query(sql))
2244 .await
2245 .map_err(|_| crate::DbError::QueryError(format!("Query timeout after {:?}", timeout)))?
2246 }
2247}
2248
2249#[cfg(feature = "prod-pool-tuning")]
2254mod pool_prod {
2255 use super::PoolConfig;
2256 use serde::{Deserialize, Serialize};
2257 use std::time::Duration;
2258
2259 #[derive(Debug, Clone, PartialEq, Eq, thiserror::Error)]
2261 pub enum PoolProdError {
2262 #[error("pool max_size must be positive")]
2264 MaxSizeNotPositive,
2265 #[error("pool acquire_timeout must be positive")]
2267 AcquireTimeoutNotPositive,
2268 #[error("pool min_idle cannot exceed max_size")]
2270 MinIdleExceedsMaxSize,
2271 }
2272
2273 #[derive(Debug, Clone, Serialize, Deserialize)]
2275 pub struct PoolProdConfig {
2276 pub max_size: u32,
2278 pub acquire_timeout: Duration,
2280 pub idle_timeout: Duration,
2282 pub connection_timeout: Duration,
2284 pub query_timeout: Duration,
2286 pub min_idle: u32,
2288 pub prewarm: bool,
2290 }
2291
2292 impl Default for PoolProdConfig {
2293 fn default() -> Self {
2294 Self {
2295 max_size: 100,
2296 acquire_timeout: Duration::from_secs(30),
2297 idle_timeout: Duration::from_secs(600),
2298 connection_timeout: Duration::from_secs(10),
2299 query_timeout: Duration::from_secs(30),
2300 min_idle: 0,
2301 prewarm: false,
2302 }
2303 }
2304 }
2305
2306 impl PoolProdConfig {
2307 pub fn new(
2309 max_size: u32,
2310 acquire_timeout: Duration,
2311 idle_timeout: Duration,
2312 connection_timeout: Duration,
2313 query_timeout: Duration,
2314 min_idle: u32,
2315 prewarm: bool,
2316 ) -> Self {
2317 Self {
2318 max_size,
2319 acquire_timeout,
2320 idle_timeout,
2321 connection_timeout,
2322 query_timeout,
2323 min_idle,
2324 prewarm,
2325 }
2326 }
2327
2328 pub fn validate(&self) -> Result<(), PoolProdError> {
2330 if self.max_size == 0 {
2331 return Err(PoolProdError::MaxSizeNotPositive);
2332 }
2333 if self.acquire_timeout.is_zero() {
2334 return Err(PoolProdError::AcquireTimeoutNotPositive);
2335 }
2336 if self.min_idle > self.max_size {
2337 return Err(PoolProdError::MinIdleExceedsMaxSize);
2338 }
2339 Ok(())
2340 }
2341
2342 pub fn to_pool_config(&self) -> PoolConfig {
2344 PoolConfig {
2345 max_size: self.max_size,
2346 min_idle: self.min_idle,
2347 acquire_timeout: self.acquire_timeout,
2348 idle_timeout: self.idle_timeout,
2349 max_lifetime: Duration::from_secs(1800),
2350 connection_timeout: self.connection_timeout,
2351 tls: None,
2352 query_timeout: Some(self.query_timeout),
2353 max_rows: None,
2354 memory_limit: None,
2355 on_event: None,
2356 test_before_acquire: false,
2357 prewarm: self.prewarm,
2358 }
2359 }
2360 }
2361}
2362
2363#[cfg(feature = "prod-pool-tuning")]
2364pub use pool_prod::{PoolProdConfig, PoolProdError};
2365
2366#[cfg(feature = "prod-leak-detection")]
2371mod leak_detection {
2372 use serde::{Deserialize, Serialize};
2373 use std::time::Duration;
2374
2375 #[derive(Debug, Clone, Serialize, Deserialize)]
2377 pub struct LeakDetectionConfig {
2378 pub enabled: bool,
2380 pub interval: Duration,
2382 pub threshold: u32,
2384 pub borrow_timeout: Duration,
2386 }
2387
2388 impl Default for LeakDetectionConfig {
2389 fn default() -> Self {
2390 Self {
2391 enabled: false,
2392 interval: Duration::from_secs(60),
2393 threshold: 5,
2394 borrow_timeout: Duration::from_secs(60),
2395 }
2396 }
2397 }
2398
2399 impl LeakDetectionConfig {
2400 pub fn new(
2402 enabled: bool,
2403 interval: Duration,
2404 threshold: u32,
2405 borrow_timeout: Duration,
2406 ) -> Self {
2407 Self {
2408 enabled,
2409 interval,
2410 threshold,
2411 borrow_timeout,
2412 }
2413 }
2414
2415 pub fn validate(&self) -> Result<(), LeakDetectionError> {
2417 if self.interval.is_zero() {
2418 return Err(LeakDetectionError::IntervalNotPositive);
2419 }
2420 if self.borrow_timeout.is_zero() {
2421 return Err(LeakDetectionError::BorrowTimeoutNotPositive);
2422 }
2423 Ok(())
2424 }
2425 }
2426
2427 #[derive(Debug, Clone, PartialEq, Eq, thiserror::Error)]
2429 pub enum LeakDetectionError {
2430 #[error("leak detection interval must be positive")]
2432 IntervalNotPositive,
2433 #[error("leak detection borrow_timeout must be positive")]
2435 BorrowTimeoutNotPositive,
2436 }
2437
2438 #[derive(Debug, Clone, Serialize, Deserialize)]
2440 pub struct LeakEntry {
2441 pub conn_id: u64,
2443 pub borrowed_at: String,
2445 pub borrow_duration: Duration,
2447 }
2448
2449 #[derive(Debug, Clone, Serialize, Deserialize)]
2451 pub struct LeakReport {
2452 pub borrowed_count: u32,
2454 pub max_borrow_duration: Duration,
2456 pub suspected_leaks: Vec<LeakEntry>,
2458 }
2459
2460 impl LeakReport {
2461 pub fn empty() -> Self {
2463 Self {
2464 borrowed_count: 0,
2465 max_borrow_duration: Duration::ZERO,
2466 suspected_leaks: vec![],
2467 }
2468 }
2469 }
2470}
2471
2472#[cfg(feature = "prod-leak-detection")]
2473pub use leak_detection::{LeakDetectionConfig, LeakDetectionError, LeakEntry, LeakReport};
2474
2475#[cfg(test)]
2476mod tests {
2477 use super::*;
2478
2479 struct MockConnection {
2481 connected: bool,
2482 }
2483
2484 impl MockConnection {
2485 fn new() -> Self {
2486 Self { connected: true }
2487 }
2488 }
2489
2490 impl Connection for MockConnection {
2491 fn execute<'a>(
2492 &'a mut self,
2493 _sql: &'a str,
2494 ) -> Pin<Box<dyn Future<Output = Result<u64, crate::DbError>> + Send + 'a>> {
2495 Box::pin(async move { Ok(1) })
2496 }
2497
2498 fn query<'a>(
2499 &'a mut self,
2500 _sql: &'a str,
2501 ) -> Pin<
2502 Box<
2503 dyn Future<
2504 Output = Result<
2505 Vec<std::collections::HashMap<String, crate::value::Value>>,
2506 crate::DbError,
2507 >,
2508 > + Send
2509 + 'a,
2510 >,
2511 > {
2512 Box::pin(async move { Ok(vec![]) })
2513 }
2514
2515 fn begin_transaction<'a>(
2516 &'a mut self,
2517 ) -> Pin<Box<dyn Future<Output = Result<(), crate::DbError>> + Send + 'a>> {
2518 Box::pin(async move { Ok(()) })
2519 }
2520
2521 fn commit<'a>(
2522 &'a mut self,
2523 ) -> Pin<Box<dyn Future<Output = Result<(), crate::DbError>> + Send + 'a>> {
2524 Box::pin(async move { Ok(()) })
2525 }
2526
2527 fn rollback<'a>(
2528 &'a mut self,
2529 ) -> Pin<Box<dyn Future<Output = Result<(), crate::DbError>> + Send + 'a>> {
2530 Box::pin(async move { Ok(()) })
2531 }
2532
2533 fn is_connected(&self) -> bool {
2534 self.connected
2535 }
2536
2537 fn ping<'a>(&'a mut self) -> Pin<Box<dyn Future<Output = bool> + Send + 'a>> {
2538 Box::pin(async move { true })
2539 }
2540
2541 fn close<'a>(
2542 &'a mut self,
2543 ) -> Pin<Box<dyn Future<Output = Result<(), crate::DbError>> + Send + 'a>> {
2544 Box::pin(async move {
2545 self.connected = false;
2546 Ok(())
2547 })
2548 }
2549 }
2550
2551 struct MockConnectionFactory;
2552
2553 #[async_trait]
2554 impl ConnectionFactory for MockConnectionFactory {
2555 async fn create(&self) -> Result<Box<dyn Connection>, crate::DbError> {
2556 Ok(Box::new(MockConnection::new()))
2557 }
2558 }
2559
2560 #[tokio::test]
2561 async fn test_pool_config_builder() -> Result<(), Box<dyn std::error::Error>> {
2562 let config = PoolConfigBuilder::new().max_size(50).min_idle(10).build()?;
2563
2564 assert_eq!(config.max_size, 50);
2565 assert_eq!(config.min_idle, 10);
2566 Ok(())
2567 }
2568
2569 #[test]
2570 fn test_pool_status_display() {
2571 let status = PoolStatus {
2572 idle: 5,
2573 active: 10,
2574 max: 100,
2575 min: 5,
2576 waiters: 0,
2577 };
2578
2579 let display = format!("{:?}", status);
2580 assert!(display.contains("idle"));
2581 assert!(display.contains("active"));
2582 }
2583
2584 #[test]
2585 fn test_default_pool_config() {
2586 let config = PoolConfig::default();
2587 assert_eq!(config.max_size, 100);
2588 assert_eq!(config.min_idle, 0);
2589 assert_eq!(config.acquire_timeout.as_secs(), 30);
2590 assert_eq!(config.idle_timeout.as_secs(), 600);
2591 assert_eq!(config.max_lifetime.as_secs(), 1800);
2592 }
2593
2594 #[tokio::test]
2595 async fn test_pool_config_clone() {
2596 let config = PoolConfig::default();
2597 let cloned = config.clone();
2598 assert_eq!(cloned.max_size, config.max_size);
2599 assert_eq!(cloned.min_idle, config.min_idle);
2600 }
2601
2602 #[test]
2603 fn test_pool_config_builder_default() -> Result<(), Box<dyn std::error::Error>> {
2604 let builder = PoolConfigBuilder::new();
2605 let config = builder.build()?;
2606 assert_eq!(config.max_size, 100);
2607 Ok(())
2608 }
2609
2610 #[test]
2611 fn test_pool_config_validate() {
2612 let result = PoolConfigBuilder::new().max_size(0).build();
2613 assert!(result.is_err());
2614
2615 let result = PoolConfigBuilder::new().max_size(10).min_idle(20).build();
2616 assert!(result.is_err());
2617 }
2618
2619 #[test]
2620 fn test_pool_config_validate_duration_upper_bound() {
2621 use std::time::Duration;
2622
2623 let config = PoolConfig {
2625 max_size: 10,
2626 min_idle: 1,
2627 acquire_timeout: Duration::from_secs(u64::MAX),
2628 idle_timeout: Duration::from_secs(1),
2629 max_lifetime: Duration::from_secs(1),
2630 connection_timeout: Duration::from_secs(5),
2631 tls: None,
2632 query_timeout: None,
2633 max_rows: None,
2634 memory_limit: None,
2635 on_event: None,
2636 test_before_acquire: false,
2637 prewarm: false,
2638 };
2639 assert!(config.validate().is_err());
2640
2641 let config = PoolConfig {
2643 max_size: 10,
2644 min_idle: 1,
2645 acquire_timeout: Duration::from_secs(u32::MAX as u64),
2646 idle_timeout: Duration::from_secs(1),
2647 max_lifetime: Duration::from_secs(1),
2648 connection_timeout: Duration::from_secs(5),
2649 tls: None,
2650 query_timeout: None,
2651 max_rows: None,
2652 memory_limit: None,
2653 on_event: None,
2654 test_before_acquire: false,
2655 prewarm: false,
2656 };
2657 assert!(config.validate().is_ok());
2658
2659 let config = PoolConfig {
2661 max_size: 10,
2662 min_idle: 1,
2663 acquire_timeout: Duration::from_secs(u32::MAX as u64 + 1),
2664 idle_timeout: Duration::from_secs(1),
2665 max_lifetime: Duration::from_secs(1),
2666 connection_timeout: Duration::from_secs(5),
2667 tls: None,
2668 query_timeout: None,
2669 max_rows: None,
2670 memory_limit: None,
2671 on_event: None,
2672 test_before_acquire: false,
2673 prewarm: false,
2674 };
2675 assert!(config.validate().is_err());
2676 }
2677
2678 #[test]
2679 fn test_pool_config_test_before_acquire_default() {
2680 let config = PoolConfig::default();
2682 assert!(!config.test_before_acquire);
2683 }
2684
2685 #[test]
2686 fn test_pool_config_builder_test_before_acquire() {
2687 let config = PoolConfigBuilder::new()
2689 .test_before_acquire(true)
2690 .build()
2691 .unwrap();
2692 assert!(config.test_before_acquire);
2693 }
2694
2695 #[tokio::test]
2696 async fn test_pool_acquire_and_release() -> Result<(), Box<dyn std::error::Error>> {
2697 let config = PoolConfigBuilder::new().max_size(5).min_idle(1).build()?;
2698 let factory = Arc::new(MockConnectionFactory);
2699 let pool = Pool::new(config, factory)?;
2700
2701 let conn = pool.acquire().await?;
2702 let status = pool.status().await;
2703 assert_eq!(status.active, 1);
2704 assert_eq!(status.idle, 0);
2705
2706 pool.release(conn).await;
2707 let status = pool.status().await;
2708 assert_eq!(status.idle, 1);
2709
2710 let _conn2 = pool.acquire().await?;
2712 let status = pool.status().await;
2713 assert_eq!(status.idle, 0);
2714 Ok(())
2715 }
2716
2717 #[tokio::test]
2718 async fn test_pool_status() -> Result<(), Box<dyn std::error::Error>> {
2719 let config = PoolConfigBuilder::new().max_size(10).min_idle(2).build()?;
2720 let factory = Arc::new(MockConnectionFactory);
2721 let pool = Pool::new(config, factory)?;
2722
2723 let status = pool.status().await;
2724 assert_eq!(status.max, 10);
2725 assert_eq!(status.min, 2);
2726 assert_eq!(status.active, 0);
2727 Ok(())
2728 }
2729
2730 #[tokio::test]
2731 async fn test_pool_close_all() -> Result<(), Box<dyn std::error::Error>> {
2732 let config = PoolConfigBuilder::new().max_size(5).build()?;
2733 let factory = Arc::new(MockConnectionFactory);
2734 let pool = Pool::new(config, factory)?;
2735
2736 let conn1 = pool.acquire().await?;
2738 let conn2 = pool.acquire().await?;
2739 pool.release(conn1).await;
2740 pool.release(conn2).await;
2741
2742 pool.close_all().await;
2743 let status = pool.status().await;
2744 assert_eq!(status.idle, 0);
2745 assert_eq!(status.active, 0);
2746 Ok(())
2747 }
2748
2749 #[tokio::test]
2750 async fn test_pool_reap_idle() -> Result<(), Box<dyn std::error::Error>> {
2751 let config = PoolConfigBuilder::new()
2752 .max_size(5)
2753 .idle_timeout(0) .build()?;
2755 let factory = Arc::new(MockConnectionFactory);
2756 let pool = Pool::new(config, factory)?;
2757
2758 let conn = pool.acquire().await?;
2759 pool.release(conn).await;
2760
2761 tokio::time::sleep(tokio::time::Duration::from_millis(10)).await;
2763
2764 pool.reap_idle().await;
2765 let status = pool.status().await;
2766 assert_eq!(status.idle, 0);
2767 Ok(())
2768 }
2769
2770 #[tokio::test]
2776 async fn test_h7_acquire_timeout_default_30s() {
2777 let config = PoolConfig::default();
2778 assert_eq!(
2779 config.acquire_timeout,
2780 Duration::from_secs(30),
2781 "H-7: acquire_timeout 默认应为 30s"
2782 );
2783 }
2784
2785 #[tokio::test]
2787 async fn test_h7_acquire_timeout_configurable() -> Result<(), Box<dyn std::error::Error>> {
2788 let config = PoolConfigBuilder::new()
2789 .max_size(1)
2790 .acquire_timeout(5) .build()?;
2792 assert_eq!(config.acquire_timeout, Duration::from_secs(5));
2793
2794 let factory = Arc::new(MockConnectionFactory);
2796 let pool = Pool::new(config, factory)?;
2797 let _conn1 = pool.acquire().await?;
2798
2799 let fast_config = PoolConfigBuilder::new()
2801 .max_size(1)
2802 .acquire_timeout(0) .build()?;
2804 let fast_pool = Pool::new(fast_config, Arc::new(MockConnectionFactory))?;
2807 let _fast_conn = fast_pool.acquire().await?; let result = fast_pool.acquire().await;
2809 assert!(
2810 matches!(result, Err(PoolError::Timeout)),
2811 "H-7: 应返回 Timeout"
2812 );
2813 Ok(())
2814 }
2815
2816 #[tokio::test]
2819 async fn test_m7_health_check_removes_nothing_when_all_healthy(
2820 ) -> Result<(), Box<dyn std::error::Error>> {
2821 let config = PoolConfigBuilder::new().max_size(5).build()?;
2823 let factory = Arc::new(MockConnectionFactory);
2824 let pool = Pool::new(config, factory)?;
2825
2826 let conn1 = pool.acquire().await?;
2828 let conn2 = pool.acquire().await?;
2829 let conn3 = pool.acquire().await?;
2830 pool.release(conn1).await;
2831 pool.release(conn2).await;
2832 pool.release(conn3).await;
2833
2834 let removed = pool.health_check().await;
2835 assert_eq!(removed, 0, "Healthy connections should not be removed");
2836
2837 let status = pool.status().await;
2838 assert_eq!(status.idle, 3);
2839 assert_eq!(status.active, 3);
2840 Ok(())
2841 }
2842
2843 #[tokio::test]
2844 async fn test_m7_health_check_returns_zero_for_empty_pool(
2845 ) -> Result<(), Box<dyn std::error::Error>> {
2846 let config = PoolConfigBuilder::new().max_size(5).build()?;
2847 let factory = Arc::new(MockConnectionFactory);
2848 let pool = Pool::new(config, factory)?;
2849
2850 let removed = pool.health_check().await;
2851 assert_eq!(removed, 0);
2852 Ok(())
2853 }
2854
2855 struct CountingFactory {
2859 count: AtomicU32,
2860 }
2861
2862 impl CountingFactory {
2863 fn new() -> Self {
2864 Self {
2865 count: AtomicU32::new(0),
2866 }
2867 }
2868 fn created_count(&self) -> u32 {
2869 self.count.load(Ordering::SeqCst)
2870 }
2871 }
2872
2873 #[async_trait]
2874 impl ConnectionFactory for CountingFactory {
2875 async fn create(&self) -> Result<Box<dyn Connection>, crate::DbError> {
2876 self.count.fetch_add(1, Ordering::SeqCst);
2877 Ok(Box::new(MockConnection::new()))
2878 }
2879 }
2880
2881 #[tokio::test]
2887 async fn test_production_bug_max_lifetime_never_expires(
2888 ) -> Result<(), Box<dyn std::error::Error>> {
2889 let config = PoolConfig {
2892 max_size: 5,
2893 min_idle: 0,
2894 acquire_timeout: Duration::from_secs(30),
2895 idle_timeout: Duration::from_secs(600),
2896 max_lifetime: Duration::from_millis(100), connection_timeout: Duration::from_secs(10),
2898 tls: None,
2899 query_timeout: None,
2900 max_rows: None,
2901 memory_limit: None,
2902 on_event: None,
2903 test_before_acquire: false,
2904 prewarm: false,
2905 };
2906 let factory = Arc::new(CountingFactory::new());
2907 let pool = Pool::new(config, factory.clone())?;
2908
2909 let conn = pool.acquire().await?;
2911 assert_eq!(factory.created_count(), 1, "应创建 1 个连接");
2912
2913 pool.release(conn).await;
2915
2916 tokio::time::sleep(Duration::from_millis(150)).await;
2918
2919 let conn2 = pool.acquire().await?;
2921
2922 assert_eq!(
2925 factory.created_count(),
2926 2,
2927 "超过 max_lifetime 后应创建新连接(旧连接应被回收)"
2928 );
2929
2930 pool.release(conn2).await;
2931 Ok(())
2932 }
2933
2934 #[tokio::test]
2941 async fn test_drop_auto_release_connection() -> Result<(), Box<dyn std::error::Error>> {
2942 let config = PoolConfigBuilder::new().max_size(2).build()?;
2943 let factory = Arc::new(CountingFactory::new());
2944 let pool = Pool::new(config, factory.clone())?;
2945
2946 {
2948 let _conn = pool.acquire().await?;
2949 assert_eq!(factory.created_count(), 1, "应创建 1 个连接");
2950 let status = pool.status().await;
2951 assert_eq!(status.active, 1, "active 应为 1");
2952 assert_eq!(status.idle, 0, "idle 应为 0");
2953 }
2955
2956 tokio::time::sleep(Duration::from_millis(50)).await;
2958
2959 let status = pool.status().await;
2961 assert_eq!(status.idle, 1, "Drop 后连接应自动归还,idle 应为 1");
2962 assert_eq!(status.active, 1, "total_count 应为 1");
2963 assert_eq!(factory.created_count(), 1, "应复用归还的连接,不创建新连接");
2964 Ok(())
2965 }
2966
2967 #[tokio::test]
2969 async fn test_drop_auto_release_then_reuse() -> Result<(), Box<dyn std::error::Error>> {
2970 let config = PoolConfigBuilder::new().max_size(1).build()?;
2971 let factory = Arc::new(CountingFactory::new());
2972 let pool = Pool::new(config, factory.clone())?;
2973
2974 {
2976 let _conn = pool.acquire().await?;
2977 }
2978
2979 tokio::time::sleep(Duration::from_millis(50)).await;
2981
2982 let conn = pool.acquire().await?;
2984 assert_eq!(factory.created_count(), 1, "应复用归还的连接,不创建新连接");
2985
2986 pool.release(conn).await;
2987 Ok(())
2988 }
2989
2990 #[tokio::test]
2992 async fn test_into_inner_does_not_return_to_pool() -> Result<(), Box<dyn std::error::Error>> {
2993 let config = PoolConfigBuilder::new().max_size(2).build()?;
2994 let factory = Arc::new(CountingFactory::new());
2995 let pool = Pool::new(config, factory.clone())?;
2996
2997 let conn = pool.acquire().await?;
2998 assert_eq!(factory.created_count(), 1);
2999
3000 let _raw_conn = conn.into_inner();
3002
3003 tokio::time::sleep(Duration::from_millis(50)).await;
3005
3006 let status = pool.status().await;
3007 assert_eq!(status.idle, 0, "into_inner 后连接不应归还");
3008 assert_eq!(status.active, 1, "total_count 仍为 1(连接被外部持有)");
3009 Ok(())
3010 }
3011
3012 #[tokio::test]
3014 async fn test_explicit_release_no_double_return() -> Result<(), Box<dyn std::error::Error>> {
3015 let config = PoolConfigBuilder::new().max_size(2).build()?;
3016 let factory = Arc::new(CountingFactory::new());
3017 let pool = Pool::new(config, factory.clone())?;
3018
3019 let conn = pool.acquire().await?;
3020 pool.release(conn).await;
3021
3022 let status = pool.status().await;
3023 assert_eq!(status.idle, 1, "release 后 idle 应为 1");
3024
3025 let conn = pool.acquire().await?;
3027 pool.release(conn).await;
3028
3029 let status = pool.status().await;
3030 assert_eq!(status.idle, 1, "再次 release 后 idle 仍应为 1(不重复)");
3031 assert_eq!(status.active, 1, "total_count 应为 1");
3032 Ok(())
3033 }
3034
3035 struct CursorMockConn {
3041 rows: QueryRows,
3042 call_count: usize,
3043 }
3044
3045 impl CursorMockConn {
3046 fn new(rows: QueryRows) -> Self {
3047 Self {
3048 rows,
3049 call_count: 0,
3050 }
3051 }
3052 }
3053
3054 impl Connection for CursorMockConn {
3055 fn execute<'a>(
3056 &'a mut self,
3057 _sql: &'a str,
3058 ) -> Pin<Box<dyn Future<Output = Result<u64, crate::DbError>> + Send + 'a>> {
3059 Box::pin(async move { Ok(1) })
3060 }
3061
3062 fn query<'a>(
3063 &'a mut self,
3064 _sql: &'a str,
3065 ) -> Pin<Box<dyn Future<Output = Result<QueryRows, crate::DbError>> + Send + 'a>> {
3066 Box::pin(async move {
3067 self.call_count += 1;
3068 Ok(self.rows.clone())
3069 })
3070 }
3071
3072 fn begin_transaction<'a>(
3073 &'a mut self,
3074 ) -> Pin<Box<dyn Future<Output = Result<(), crate::DbError>> + Send + 'a>> {
3075 Box::pin(async move { Ok(()) })
3076 }
3077
3078 fn commit<'a>(
3079 &'a mut self,
3080 ) -> Pin<Box<dyn Future<Output = Result<(), crate::DbError>> + Send + 'a>> {
3081 Box::pin(async move { Ok(()) })
3082 }
3083
3084 fn rollback<'a>(
3085 &'a mut self,
3086 ) -> Pin<Box<dyn Future<Output = Result<(), crate::DbError>> + Send + 'a>> {
3087 Box::pin(async move { Ok(()) })
3088 }
3089
3090 fn is_connected(&self) -> bool {
3091 true
3092 }
3093
3094 fn ping<'a>(&'a mut self) -> Pin<Box<dyn Future<Output = bool> + Send + 'a>> {
3095 Box::pin(async move { true })
3096 }
3097
3098 fn close<'a>(
3099 &'a mut self,
3100 ) -> Pin<Box<dyn Future<Output = Result<(), crate::DbError>> + Send + 'a>> {
3101 Box::pin(async move { Ok(()) })
3102 }
3103 }
3104
3105 struct CursorOverrideMockConn {
3107 rows: Vec<crate::value::Value>,
3108 yielded: usize,
3109 }
3110
3111 impl CursorOverrideMockConn {
3112 fn new(rows: Vec<crate::value::Value>) -> Self {
3113 Self { rows, yielded: 0 }
3114 }
3115 }
3116
3117 impl Connection for CursorOverrideMockConn {
3118 fn execute<'a>(
3119 &'a mut self,
3120 _sql: &'a str,
3121 ) -> Pin<Box<dyn Future<Output = Result<u64, crate::DbError>> + Send + 'a>> {
3122 Box::pin(async move { Ok(1) })
3123 }
3124
3125 fn query<'a>(
3126 &'a mut self,
3127 _sql: &'a str,
3128 ) -> Pin<Box<dyn Future<Output = Result<QueryRows, crate::DbError>> + Send + 'a>> {
3129 Box::pin(async move {
3131 Ok(self
3132 .rows
3133 .iter()
3134 .map(|v| {
3135 let mut m = std::collections::HashMap::new();
3136 m.insert("v".to_string(), v.clone());
3137 m
3138 })
3139 .collect())
3140 })
3141 }
3142
3143 fn query_stream<'a>(
3145 &'a mut self,
3146 _sql: &'a str,
3147 ) -> Pin<Box<dyn futures::Stream<Item = QueryStreamItem> + Send + 'a>> {
3148 Box::pin(futures::stream::iter(
3149 self.rows
3150 .iter()
3151 .enumerate()
3152 .map(|(i, v)| {
3153 self.yielded = i + 1;
3154 let mut m = std::collections::HashMap::new();
3155 m.insert("v".to_string(), v.clone());
3156 Ok(m)
3157 })
3158 .collect::<Vec<_>>(),
3159 ))
3160 }
3161
3162 fn begin_transaction<'a>(
3163 &'a mut self,
3164 ) -> Pin<Box<dyn Future<Output = Result<(), crate::DbError>> + Send + 'a>> {
3165 Box::pin(async move { Ok(()) })
3166 }
3167
3168 fn commit<'a>(
3169 &'a mut self,
3170 ) -> Pin<Box<dyn Future<Output = Result<(), crate::DbError>> + Send + 'a>> {
3171 Box::pin(async move { Ok(()) })
3172 }
3173
3174 fn rollback<'a>(
3175 &'a mut self,
3176 ) -> Pin<Box<dyn Future<Output = Result<(), crate::DbError>> + Send + 'a>> {
3177 Box::pin(async move { Ok(()) })
3178 }
3179
3180 fn is_connected(&self) -> bool {
3181 true
3182 }
3183
3184 fn ping<'a>(&'a mut self) -> Pin<Box<dyn Future<Output = bool> + Send + 'a>> {
3185 Box::pin(async move { true })
3186 }
3187
3188 fn close<'a>(
3189 &'a mut self,
3190 ) -> Pin<Box<dyn Future<Output = Result<(), crate::DbError>> + Send + 'a>> {
3191 Box::pin(async move { Ok(()) })
3192 }
3193 }
3194
3195 #[tokio::test]
3197 async fn test_query_stream_default_impl_yields_all_rows() {
3198 use futures::StreamExt;
3199 let rows: QueryRows = vec![
3200 std::collections::HashMap::from([
3201 ("id".to_string(), crate::value::Value::I64(1)),
3202 (
3203 "name".to_string(),
3204 crate::value::Value::String("alice".to_string()),
3205 ),
3206 ]),
3207 std::collections::HashMap::from([
3208 ("id".to_string(), crate::value::Value::I64(2)),
3209 (
3210 "name".to_string(),
3211 crate::value::Value::String("bob".to_string()),
3212 ),
3213 ]),
3214 std::collections::HashMap::from([
3215 ("id".to_string(), crate::value::Value::I64(3)),
3216 (
3217 "name".to_string(),
3218 crate::value::Value::String("carol".to_string()),
3219 ),
3220 ]),
3221 ];
3222 let mut conn = CursorMockConn::new(rows);
3223 let mut stream = conn.query_stream("SELECT id, name FROM users");
3224 let mut received: Vec<QueryStreamItem> = Vec::new();
3225 while let Some(item) = stream.next().await {
3226 received.push(item);
3227 }
3228 assert_eq!(received.len(), 3, "应收到 3 行");
3229 assert!(received.iter().all(|r| r.is_ok()), "所有项应为 Ok");
3230 drop(stream);
3231 assert_eq!(conn.call_count, 1, "默认实现应调用 query() 一次");
3232 }
3233
3234 #[tokio::test]
3236 async fn test_query_stream_default_empty_result() {
3237 use futures::StreamExt;
3238 let mut conn = CursorMockConn::new(Vec::new());
3239 let mut stream = conn.query_stream("SELECT * FROM empty_table");
3240 let mut count = 0;
3241 while let Some(_item) = stream.next().await {
3242 count += 1;
3243 }
3244 assert_eq!(count, 0, "空结果集应产生 0 项");
3245 }
3246
3247 #[tokio::test]
3249 async fn test_query_stream_default_error_propagation() {
3250 use futures::StreamExt;
3251 struct ErrorMockConn;
3253 impl Connection for ErrorMockConn {
3254 fn execute<'a>(
3255 &'a mut self,
3256 _sql: &'a str,
3257 ) -> Pin<Box<dyn Future<Output = Result<u64, crate::DbError>> + Send + 'a>>
3258 {
3259 Box::pin(async move { Ok(1) })
3260 }
3261 fn query<'a>(
3262 &'a mut self,
3263 _sql: &'a str,
3264 ) -> Pin<Box<dyn Future<Output = Result<QueryRows, crate::DbError>> + Send + 'a>>
3265 {
3266 Box::pin(async move { Err(crate::DbError::Internal("query failed".to_string())) })
3267 }
3268 fn begin_transaction<'a>(
3269 &'a mut self,
3270 ) -> Pin<Box<dyn Future<Output = Result<(), crate::DbError>> + Send + 'a>> {
3271 Box::pin(async move { Ok(()) })
3272 }
3273 fn commit<'a>(
3274 &'a mut self,
3275 ) -> Pin<Box<dyn Future<Output = Result<(), crate::DbError>> + Send + 'a>> {
3276 Box::pin(async move { Ok(()) })
3277 }
3278 fn rollback<'a>(
3279 &'a mut self,
3280 ) -> Pin<Box<dyn Future<Output = Result<(), crate::DbError>> + Send + 'a>> {
3281 Box::pin(async move { Ok(()) })
3282 }
3283 fn is_connected(&self) -> bool {
3284 true
3285 }
3286 fn ping<'a>(&'a mut self) -> Pin<Box<dyn Future<Output = bool> + Send + 'a>> {
3287 Box::pin(async move { true })
3288 }
3289 fn close<'a>(
3290 &'a mut self,
3291 ) -> Pin<Box<dyn Future<Output = Result<(), crate::DbError>> + Send + 'a>> {
3292 Box::pin(async move { Ok(()) })
3293 }
3294 }
3295 let mut conn = ErrorMockConn;
3296 let mut stream = conn.query_stream("SELECT * FROM bad_table");
3297 let item = stream.next().await;
3298 assert!(item.is_some(), "应产生一项");
3299 assert!(item.unwrap().is_err(), "该项应为 Err");
3300 }
3301
3302 #[tokio::test]
3304 async fn test_query_stream_override_yields_rows_one_by_one() {
3305 use futures::StreamExt;
3306 let rows = vec![
3307 crate::value::Value::I64(10),
3308 crate::value::Value::I64(20),
3309 crate::value::Value::I64(30),
3310 crate::value::Value::I64(40),
3311 crate::value::Value::I64(50),
3312 ];
3313 let mut conn = CursorOverrideMockConn::new(rows);
3314 let values: Vec<i64> = {
3315 let mut stream = conn.query_stream("SELECT v FROM seq");
3316 let mut vals: Vec<i64> = Vec::new();
3317 while let Some(Ok(row)) = stream.next().await {
3318 if let crate::value::Value::I64(v) = row.get("v").unwrap() {
3319 vals.push(*v);
3320 }
3321 }
3322 vals
3323 };
3324 assert_eq!(values, vec![10, 20, 30, 40, 50], "应按顺序收到全部 5 行");
3325 assert_eq!(conn.yielded, 5, "应逐行 yield 5 次(真游标覆盖)");
3326 }
3327
3328 #[tokio::test]
3330 async fn test_query_stream_override_early_drop() {
3331 use futures::StreamExt;
3332 let rows = vec![
3333 crate::value::Value::I64(1),
3334 crate::value::Value::I64(2),
3335 crate::value::Value::I64(3),
3336 ];
3337 let mut conn = CursorOverrideMockConn::new(rows);
3338 {
3339 let mut stream = conn.query_stream("SELECT v FROM seq");
3340 let first = stream.next().await;
3341 assert!(first.is_some(), "第一项应存在");
3342 drop(stream);
3344 }
3345 assert!(conn.is_connected(), "提前 drop 流后连接仍应可用");
3347 }
3348
3349 #[tokio::test]
3351 async fn test_pool_prewarm() -> Result<(), Box<dyn std::error::Error>> {
3352 use std::sync::atomic::AtomicU32;
3353
3354 let create_count = Arc::new(AtomicU32::new(0));
3356 let create_count_clone = create_count.clone();
3357
3358 struct CountingFactory {
3359 count: Arc<AtomicU32>,
3360 }
3361
3362 #[async_trait]
3363 impl ConnectionFactory for CountingFactory {
3364 async fn create(&self) -> Result<Box<dyn Connection>, crate::DbError> {
3365 self.count.fetch_add(1, Ordering::SeqCst);
3366 Ok(Box::new(MockConnection::new()))
3367 }
3368 }
3369
3370 let config = PoolConfigBuilder::new()
3372 .max_size(10)
3373 .min_idle(5)
3374 .prewarm(true)
3375 .build()?;
3376
3377 let factory = Arc::new(CountingFactory {
3378 count: create_count_clone,
3379 });
3380
3381 let pool = Pool::new(config, factory)?;
3382
3383 let status_before = pool.status().await;
3385 assert_eq!(status_before.idle, 0, "预热前 idle 应为 0");
3386
3387 pool.prewarm().await;
3389
3390 let status_after = pool.status().await;
3392 assert!(
3393 status_after.idle >= 5,
3394 "预热后 idle 应 >= 5,实际: {}",
3395 status_after.idle
3396 );
3397
3398 assert_eq!(
3400 create_count.load(Ordering::SeqCst),
3401 5,
3402 "工厂应被调用 5 次(min_idle)"
3403 );
3404
3405 Ok(())
3406 }
3407
3408 #[tokio::test]
3410 async fn test_pool_prewarm_failure_non_blocking() -> Result<(), Box<dyn std::error::Error>> {
3411 use std::sync::atomic::AtomicBool;
3412
3413 struct FailingFactory {
3414 failed: Arc<AtomicBool>,
3415 }
3416
3417 #[async_trait]
3418 impl ConnectionFactory for FailingFactory {
3419 async fn create(&self) -> Result<Box<dyn Connection>, crate::DbError> {
3420 self.failed.store(true, Ordering::SeqCst);
3421 Err(crate::DbError::Internal(
3423 "simulated connection failure".to_string(),
3424 ))
3425 }
3426 }
3427
3428 let failed = Arc::new(AtomicBool::new(false));
3429 let mut config = PoolConfigBuilder::new()
3430 .max_size(10)
3431 .min_idle(3)
3432 .prewarm(true)
3433 .build()?;
3434 config.connection_timeout = std::time::Duration::from_secs(1); let factory = Arc::new(FailingFactory {
3437 failed: failed.clone(),
3438 });
3439
3440 let pool = Pool::new(config, factory)?;
3442 pool.prewarm().await; assert!(failed.load(Ordering::SeqCst), "工厂应被调用且失败");
3446
3447 let status = pool.status().await;
3449 assert_eq!(status.max, 10, "池配置应正常");
3450
3451 Ok(())
3452 }
3453
3454 #[tokio::test]
3456 async fn test_pool_prewarm_disabled() -> Result<(), Box<dyn std::error::Error>> {
3457 use std::sync::atomic::AtomicU32;
3458
3459 let create_count = Arc::new(AtomicU32::new(0));
3460 let create_count_clone = create_count.clone();
3461
3462 struct CountingFactory {
3463 count: Arc<AtomicU32>,
3464 }
3465
3466 #[async_trait]
3467 impl ConnectionFactory for CountingFactory {
3468 async fn create(&self) -> Result<Box<dyn Connection>, crate::DbError> {
3469 self.count.fetch_add(1, Ordering::SeqCst);
3470 Ok(Box::new(MockConnection::new()))
3471 }
3472 }
3473
3474 let config = PoolConfigBuilder::new()
3476 .max_size(10)
3477 .min_idle(5)
3478 .prewarm(false) .build()?;
3480
3481 let factory = Arc::new(CountingFactory {
3482 count: create_count_clone,
3483 });
3484
3485 let pool = Pool::new(config, factory)?;
3486 pool.prewarm().await; assert_eq!(
3490 create_count.load(Ordering::SeqCst),
3491 0,
3492 "prewarm=false 时工厂不应被调用"
3493 );
3494
3495 let status = pool.status().await;
3496 assert_eq!(status.idle, 0, "idle 应为 0");
3497
3498 Ok(())
3499 }
3500
3501 #[tokio::test]
3503 async fn test_pool_new_async_with_prewarm() -> Result<(), Box<dyn std::error::Error>> {
3504 use std::sync::atomic::AtomicU32;
3505
3506 let create_count = Arc::new(AtomicU32::new(0));
3507 let create_count_clone = create_count.clone();
3508
3509 struct CountingFactory {
3510 count: Arc<AtomicU32>,
3511 }
3512
3513 #[async_trait]
3514 impl ConnectionFactory for CountingFactory {
3515 async fn create(&self) -> Result<Box<dyn Connection>, crate::DbError> {
3516 self.count.fetch_add(1, Ordering::SeqCst);
3517 Ok(Box::new(MockConnection::new()))
3518 }
3519 }
3520
3521 let config = PoolConfigBuilder::new()
3522 .max_size(10)
3523 .min_idle(5)
3524 .prewarm(true)
3525 .build()?;
3526
3527 let factory = Arc::new(CountingFactory {
3528 count: create_count_clone,
3529 });
3530
3531 let pool = Pool::new_async(config, factory).await?;
3532
3533 let status = pool.status().await;
3534 assert!(
3535 status.idle >= 5,
3536 "new_async prewarm=true 后 idle 应 >= 5,实际: {}",
3537 status.idle
3538 );
3539 assert_eq!(create_count.load(Ordering::SeqCst), 5, "工厂应被调用 5 次");
3540
3541 Ok(())
3542 }
3543
3544 #[tokio::test]
3546 async fn test_pool_new_async_without_prewarm() -> Result<(), Box<dyn std::error::Error>> {
3547 use std::sync::atomic::AtomicU32;
3548
3549 let create_count = Arc::new(AtomicU32::new(0));
3550 let create_count_clone = create_count.clone();
3551
3552 struct CountingFactory {
3553 count: Arc<AtomicU32>,
3554 }
3555
3556 #[async_trait]
3557 impl ConnectionFactory for CountingFactory {
3558 async fn create(&self) -> Result<Box<dyn Connection>, crate::DbError> {
3559 self.count.fetch_add(1, Ordering::SeqCst);
3560 Ok(Box::new(MockConnection::new()))
3561 }
3562 }
3563
3564 let config = PoolConfigBuilder::new()
3565 .max_size(10)
3566 .min_idle(5)
3567 .prewarm(false)
3568 .build()?;
3569
3570 let factory = Arc::new(CountingFactory {
3571 count: create_count_clone,
3572 });
3573
3574 let pool = Pool::new_async(config, factory).await?;
3575
3576 let status = pool.status().await;
3577 assert_eq!(status.idle, 0, "prewarm=false 时 idle 应为 0");
3578 assert_eq!(create_count.load(Ordering::SeqCst), 0, "工厂不应被调用");
3579
3580 Ok(())
3581 }
3582
3583 #[tokio::test]
3585 async fn test_pool_new_async_failure_non_blocking() -> Result<(), Box<dyn std::error::Error>> {
3586 struct FailingFactory;
3587
3588 #[async_trait]
3589 impl ConnectionFactory for FailingFactory {
3590 async fn create(&self) -> Result<Box<dyn Connection>, crate::DbError> {
3591 Err(crate::DbError::Internal("simulated failure".to_string()))
3592 }
3593 }
3594
3595 let mut config = PoolConfigBuilder::new()
3596 .max_size(10)
3597 .min_idle(3)
3598 .prewarm(true)
3599 .build()?;
3600 config.connection_timeout = std::time::Duration::from_secs(1);
3601
3602 let pool = Pool::new_async(config, Arc::new(FailingFactory)).await?;
3603
3604 let status = pool.status().await;
3605 assert_eq!(status.max, 10, "池配置应正常");
3606
3607 Ok(())
3608 }
3609
3610 #[cfg(feature = "auto-prewarm")]
3612 #[tokio::test]
3613 async fn test_pool_progressive_prewarm() -> Result<(), Box<dyn std::error::Error>> {
3614 use std::sync::atomic::AtomicU32;
3615
3616 let create_count = Arc::new(AtomicU32::new(0));
3617 let create_count_clone = create_count.clone();
3618
3619 struct CountingFactory {
3620 count: Arc<AtomicU32>,
3621 }
3622
3623 #[async_trait]
3624 impl ConnectionFactory for CountingFactory {
3625 async fn create(&self) -> Result<Box<dyn Connection>, crate::DbError> {
3626 self.count.fetch_add(1, Ordering::SeqCst);
3627 Ok(Box::new(MockConnection::new()))
3628 }
3629 }
3630
3631 let config = PoolConfigBuilder::new()
3632 .max_size(20)
3633 .min_idle(6)
3634 .prewarm(true)
3635 .build()?;
3636
3637 let factory = Arc::new(CountingFactory {
3638 count: create_count_clone,
3639 });
3640
3641 let pool = Pool::new(config, factory)?;
3642
3643 let progress = crate::prewarm::PrewarmProgress::new(6);
3644 pool.progressive_prewarm(
3645 2,
3646 std::time::Duration::from_millis(5),
3647 std::time::Duration::from_secs(10),
3648 &progress,
3649 )
3650 .await;
3651
3652 let snap = progress.snapshot();
3653 assert!(
3654 snap.warmed >= 6,
3655 "progressive_prewarm 后 warmed 应 >= 6,实际: {}",
3656 snap.warmed
3657 );
3658 assert!(snap.is_completed, "应标记完成");
3659 assert_eq!(create_count.load(Ordering::SeqCst), 6, "工厂应被调用 6 次");
3660
3661 let status = pool.status().await;
3662 assert!(status.idle >= 6, "池中 idle 应 >= 6");
3663
3664 Ok(())
3665 }
3666
3667 #[cfg(feature = "auto-prewarm")]
3669 #[tokio::test]
3670 async fn test_pool_progressive_prewarm_timeout_zero() -> Result<(), Box<dyn std::error::Error>>
3671 {
3672 use std::sync::atomic::AtomicU32;
3673
3674 let create_count = Arc::new(AtomicU32::new(0));
3675 let create_count_clone = create_count.clone();
3676
3677 struct CountingFactory {
3678 count: Arc<AtomicU32>,
3679 }
3680
3681 #[async_trait]
3682 impl ConnectionFactory for CountingFactory {
3683 async fn create(&self) -> Result<Box<dyn Connection>, crate::DbError> {
3684 self.count.fetch_add(1, Ordering::SeqCst);
3685 Ok(Box::new(MockConnection::new()))
3686 }
3687 }
3688
3689 let config = PoolConfigBuilder::new()
3690 .max_size(20)
3691 .min_idle(10)
3692 .prewarm(true)
3693 .build()?;
3694
3695 let factory = Arc::new(CountingFactory {
3696 count: create_count_clone,
3697 });
3698
3699 let pool = Pool::new(config, factory)?;
3700
3701 let progress = crate::prewarm::PrewarmProgress::new(10);
3702 pool.progressive_prewarm(
3703 2,
3704 std::time::Duration::from_millis(5),
3705 std::time::Duration::ZERO,
3706 &progress,
3707 )
3708 .await;
3709
3710 let snap = progress.snapshot();
3711 assert!(snap.is_completed, "应标记完成");
3712 assert!(
3713 snap.warmed <= 2,
3714 "total_timeout=0 时最多建一批(batch_size=2),实际: {}",
3715 snap.warmed
3716 );
3717
3718 Ok(())
3719 }
3720
3721 #[cfg(feature = "auto-prewarm")]
3723 #[tokio::test]
3724 async fn test_pool_progressive_prewarm_disabled() -> Result<(), Box<dyn std::error::Error>> {
3725 use std::sync::atomic::AtomicU32;
3726
3727 let create_count = Arc::new(AtomicU32::new(0));
3728 let create_count_clone = create_count.clone();
3729
3730 struct CountingFactory {
3731 count: Arc<AtomicU32>,
3732 }
3733
3734 #[async_trait]
3735 impl ConnectionFactory for CountingFactory {
3736 async fn create(&self) -> Result<Box<dyn Connection>, crate::DbError> {
3737 self.count.fetch_add(1, Ordering::SeqCst);
3738 Ok(Box::new(MockConnection::new()))
3739 }
3740 }
3741
3742 let config = PoolConfigBuilder::new()
3743 .max_size(20)
3744 .min_idle(10)
3745 .prewarm(false)
3746 .build()?;
3747
3748 let factory = Arc::new(CountingFactory {
3749 count: create_count_clone,
3750 });
3751
3752 let pool = Pool::new(config, factory)?;
3753
3754 let progress = crate::prewarm::PrewarmProgress::new(10);
3755 pool.progressive_prewarm(
3756 2,
3757 std::time::Duration::from_millis(5),
3758 std::time::Duration::from_secs(10),
3759 &progress,
3760 )
3761 .await;
3762
3763 let snap = progress.snapshot();
3764 assert!(snap.is_completed, "应标记完成");
3765 assert_eq!(snap.warmed, 0, "prewarm=false 时不应建连");
3766 assert_eq!(create_count.load(Ordering::SeqCst), 0, "工厂不应被调用");
3767
3768 Ok(())
3769 }
3770
3771 #[cfg(feature = "auto-prewarm")]
3773 #[tokio::test]
3774 async fn test_pool_progressive_prewarm_failure_non_blocking(
3775 ) -> Result<(), Box<dyn std::error::Error>> {
3776 struct FailingFactory;
3777
3778 #[async_trait]
3779 impl ConnectionFactory for FailingFactory {
3780 async fn create(&self) -> Result<Box<dyn Connection>, crate::DbError> {
3781 Err(crate::DbError::Internal("simulated failure".to_string()))
3782 }
3783 }
3784
3785 let mut config = PoolConfigBuilder::new()
3786 .max_size(20)
3787 .min_idle(5)
3788 .prewarm(true)
3789 .build()?;
3790 config.connection_timeout = std::time::Duration::from_secs(1);
3791
3792 let pool = Pool::new(config, Arc::new(FailingFactory))?;
3793
3794 let progress = crate::prewarm::PrewarmProgress::new(5);
3795 pool.progressive_prewarm(
3796 2,
3797 std::time::Duration::from_millis(5),
3798 std::time::Duration::from_secs(5),
3799 &progress,
3800 )
3801 .await;
3802
3803 let snap = progress.snapshot();
3804 assert!(snap.is_completed, "应标记完成");
3805 assert_eq!(snap.warmed, 0, "全部失败时 warmed=0");
3806 assert!(snap.failed > 0, "应有失败记录");
3807
3808 Ok(())
3809 }
3810
3811 #[tokio::test]
3813 async fn test_pool_metrics_acquire_release() -> Result<(), Box<dyn std::error::Error>> {
3814 let config = PoolConfigBuilder::new().max_size(10).build()?;
3815 let pool = Pool::new(config, Arc::new(MockConnectionFactory))?;
3816
3817 let metrics = pool.pool_metrics();
3818 assert_eq!(metrics.acquire_count, 0);
3819 assert_eq!(metrics.release_count, 0);
3820 assert_eq!(metrics.connection_created_count, 0);
3821
3822 let conn = pool.acquire().await?;
3823 let metrics = pool.pool_metrics();
3824 assert_eq!(metrics.acquire_count, 1);
3825 assert_eq!(metrics.connection_created_count, 1);
3826 assert_eq!(metrics.acquire_failed_count, 0);
3827
3828 pool.release(conn).await;
3829 let metrics = pool.pool_metrics();
3830 assert_eq!(metrics.release_count, 1);
3831 assert_eq!(metrics.connection_closed_count, 0);
3833
3834 Ok(())
3835 }
3836
3837 #[tokio::test]
3839 async fn test_pool_metrics_acquire_failed() -> Result<(), Box<dyn std::error::Error>> {
3840 struct FailingFactory;
3841
3842 #[async_trait]
3843 impl ConnectionFactory for FailingFactory {
3844 async fn create(&self) -> Result<Box<dyn Connection>, crate::DbError> {
3845 Err(crate::DbError::Internal("simulated failure".to_string()))
3846 }
3847 }
3848
3849 let config = PoolConfigBuilder::new().max_size(10).build()?;
3850 let pool = Pool::new(config, Arc::new(FailingFactory))?;
3851
3852 let result = pool.acquire().await;
3853 assert!(result.is_err());
3854
3855 let metrics = pool.pool_metrics();
3856 assert_eq!(metrics.acquire_failed_count, 1);
3857 assert_eq!(metrics.acquire_count, 0);
3858
3859 Ok(())
3860 }
3861
3862 #[tokio::test]
3864 async fn test_pool_metrics_connection_closed() -> Result<(), Box<dyn std::error::Error>> {
3865 let config = PoolConfigBuilder::new().max_size(10).build()?;
3866 let pool = Pool::new(config, Arc::new(MockConnectionFactory))?;
3867
3868 let conn = pool.acquire().await?;
3869 pool.release(conn).await;
3870
3871 let status = pool.status().await;
3872 assert_eq!(status.idle, 1);
3873
3874 pool.close_all().await;
3875
3876 let metrics = pool.pool_metrics();
3877 assert_eq!(metrics.connection_closed_count, 1);
3878 assert_eq!(metrics.connection_created_count, 1);
3879
3880 Ok(())
3881 }
3882
3883 #[test]
3885 fn test_pool_metrics_average_wait_time() {
3886 let metrics = PoolMetrics {
3887 acquire_count: 4,
3888 acquire_failed_count: 1,
3889 acquire_wait_time: Duration::from_millis(200),
3890 release_count: 4,
3891 connection_created_count: 2,
3892 connection_closed_count: 0,
3893 };
3894 assert_eq!(
3895 metrics.average_acquire_wait_time(),
3896 Duration::from_millis(50)
3897 );
3898
3899 let empty = PoolMetrics::default();
3901 assert_eq!(empty.average_acquire_wait_time(), Duration::ZERO);
3902 }
3903
3904 #[tokio::test]
3905 async fn test_shutdown_with_timeout_fast_return_when_empty() {
3906 let factory = Arc::new(MockConnectionFactory) as Arc<dyn ConnectionFactory>;
3907 let config = PoolConfigBuilder::new().max_size(4).build().unwrap();
3908 let pool = Pool::new(config, factory).unwrap();
3909 pool.shutdown_with_timeout(Duration::from_secs(1)).await;
3910 assert!(pool.closed.load(Ordering::SeqCst));
3911 assert_eq!(pool.total_count.load(Ordering::SeqCst), 0);
3912 }
3913
3914 #[tokio::test]
3915 async fn test_shutdown_delegates_to_shutdown_with_timeout() {
3916 let factory = Arc::new(MockConnectionFactory) as Arc<dyn ConnectionFactory>;
3917 let config = PoolConfigBuilder::new().max_size(4).build().unwrap();
3918 let pool = Pool::new(config, factory).unwrap();
3919 pool.shutdown().await;
3920 assert!(pool.closed.load(Ordering::SeqCst));
3921 }
3922
3923 #[tokio::test]
3924 async fn test_shutdown_with_timeout_idempotent() {
3925 let factory = Arc::new(MockConnectionFactory) as Arc<dyn ConnectionFactory>;
3926 let config = PoolConfigBuilder::new().max_size(4).build().unwrap();
3927 let pool = Pool::new(config, factory).unwrap();
3928 pool.shutdown_with_timeout(Duration::from_secs(1)).await;
3929 let count_after_first = pool.total_count.load(Ordering::SeqCst);
3930 pool.shutdown_with_timeout(Duration::from_secs(1)).await;
3931 let count_after_second = pool.total_count.load(Ordering::SeqCst);
3932 assert_eq!(count_after_first, count_after_second);
3933 }
3934
3935 #[tokio::test]
3936 async fn test_shutdown_with_timeout_rejects_new_acquire() {
3937 let factory = Arc::new(MockConnectionFactory) as Arc<dyn ConnectionFactory>;
3938 let config = PoolConfigBuilder::new().max_size(4).build().unwrap();
3939 let pool = Pool::new(config, factory).unwrap();
3940 pool.shutdown_with_timeout(Duration::from_secs(1)).await;
3941 let result = pool.acquire().await;
3942 assert!(result.is_err());
3943 }
3944}
3945
3946#[cfg(all(test, feature = "prod-pool-tuning"))]
3947mod pool_prod_tests {
3948 use super::*;
3949
3950 struct MockFactory;
3951
3952 #[async_trait]
3953 impl ConnectionFactory for MockFactory {
3954 async fn create(&self) -> Result<Box<dyn Connection>, crate::DbError> {
3955 Ok(Box::new(MockConn))
3956 }
3957 }
3958
3959 struct MockConn;
3960
3961 impl Connection for MockConn {
3962 fn execute<'a>(
3963 &'a mut self,
3964 _sql: &'a str,
3965 ) -> Pin<Box<dyn Future<Output = Result<u64, crate::DbError>> + Send + 'a>> {
3966 Box::pin(async move { Ok(1) })
3967 }
3968 fn query<'a>(
3969 &'a mut self,
3970 _sql: &'a str,
3971 ) -> Pin<
3972 Box<
3973 dyn Future<
3974 Output = Result<
3975 Vec<std::collections::HashMap<String, crate::value::Value>>,
3976 crate::DbError,
3977 >,
3978 > + Send
3979 + 'a,
3980 >,
3981 > {
3982 Box::pin(async move { Ok(vec![]) })
3983 }
3984 fn begin_transaction<'a>(
3985 &'a mut self,
3986 ) -> Pin<Box<dyn Future<Output = Result<(), crate::DbError>> + Send + 'a>> {
3987 Box::pin(async move { Ok(()) })
3988 }
3989 fn commit<'a>(
3990 &'a mut self,
3991 ) -> Pin<Box<dyn Future<Output = Result<(), crate::DbError>> + Send + 'a>> {
3992 Box::pin(async move { Ok(()) })
3993 }
3994 fn rollback<'a>(
3995 &'a mut self,
3996 ) -> Pin<Box<dyn Future<Output = Result<(), crate::DbError>> + Send + 'a>> {
3997 Box::pin(async move { Ok(()) })
3998 }
3999 fn is_connected(&self) -> bool {
4000 true
4001 }
4002 fn ping<'a>(&'a mut self) -> Pin<Box<dyn Future<Output = bool> + Send + 'a>> {
4003 Box::pin(async move { true })
4004 }
4005 fn close<'a>(
4006 &'a mut self,
4007 ) -> Pin<Box<dyn Future<Output = Result<(), crate::DbError>> + Send + 'a>> {
4008 Box::pin(async move { Ok(()) })
4009 }
4010 }
4011
4012 #[test]
4013 fn test_pool_prod_config_validate_ok() {
4014 let config = PoolProdConfig::new(
4015 50,
4016 Duration::from_secs(10),
4017 Duration::from_secs(600),
4018 Duration::from_secs(5),
4019 Duration::from_secs(30),
4020 5,
4021 true,
4022 );
4023 assert!(config.validate().is_ok());
4024 }
4025
4026 #[test]
4027 fn test_pool_prod_config_max_size_zero_rejected() {
4028 let config = PoolProdConfig::default();
4029 let mut c = config;
4030 c.max_size = 0;
4031 let err = c.validate().unwrap_err();
4032 assert!(err.to_string().contains("max_size must be positive"));
4033 }
4034
4035 #[test]
4036 fn test_pool_prod_config_min_idle_exceeds_max_size() {
4037 let config = PoolProdConfig::new(
4038 10,
4039 Duration::from_secs(10),
4040 Duration::from_secs(600),
4041 Duration::from_secs(5),
4042 Duration::from_secs(30),
4043 20,
4044 false,
4045 );
4046 let err = config.validate().unwrap_err();
4047 assert!(err.to_string().contains("min_idle cannot exceed max_size"));
4048 }
4049
4050 #[test]
4051 fn test_pool_prod_config_to_pool_config() {
4052 let config = PoolProdConfig::new(
4053 50,
4054 Duration::from_secs(10),
4055 Duration::from_secs(600),
4056 Duration::from_secs(5),
4057 Duration::from_secs(30),
4058 5,
4059 true,
4060 );
4061 let pool_config = config.to_pool_config();
4062 assert_eq!(pool_config.max_size, 50);
4063 assert_eq!(pool_config.min_idle, 5);
4064 assert_eq!(pool_config.acquire_timeout, Duration::from_secs(10));
4065 assert!(pool_config.prewarm);
4066 }
4067
4068 #[tokio::test]
4069 async fn test_pool_prod_config_runtime_resize() {
4070 let factory = Arc::new(MockFactory) as Arc<dyn ConnectionFactory>;
4071 let config = PoolProdConfig::default();
4072 let pool = Pool::new(config.to_pool_config(), factory).unwrap();
4073 assert_eq!(pool.max_size(), 100);
4074 pool.resize(50);
4075 assert_eq!(pool.max_size(), 50);
4076 }
4077}
4078
4079#[cfg(all(test, feature = "prod-leak-detection"))]
4080mod leak_prod_tests {
4081 use super::*;
4082
4083 #[test]
4084 fn test_leak_config_default() {
4085 let config = LeakDetectionConfig::default();
4086 assert!(!config.enabled);
4087 assert_eq!(config.interval, Duration::from_secs(60));
4088 assert_eq!(config.threshold, 5);
4089 }
4090
4091 #[test]
4092 fn test_leak_config_validate_ok() {
4093 let config =
4094 LeakDetectionConfig::new(true, Duration::from_secs(30), 10, Duration::from_secs(60));
4095 assert!(config.validate().is_ok());
4096 }
4097
4098 #[test]
4099 fn test_leak_config_interval_zero_rejected() {
4100 let config = LeakDetectionConfig::new(true, Duration::ZERO, 10, Duration::from_secs(60));
4101 assert!(config.validate().is_err());
4102 }
4103
4104 #[test]
4105 fn test_leak_report_empty() {
4106 let report = LeakReport::empty();
4107 assert_eq!(report.borrowed_count, 0);
4108 assert!(report.suspected_leaks.is_empty());
4109 }
4110
4111 #[test]
4112 fn connection_reuse_rate_zero() {
4113 let metrics = PoolMetrics::default();
4114 assert_eq!(metrics.connection_reuse_rate(), 0.0);
4115 }
4116
4117 #[test]
4118 fn connection_reuse_rate_full() {
4119 let metrics = PoolMetrics {
4120 acquire_count: 100,
4121 connection_created_count: 1,
4122 ..Default::default()
4123 };
4124 let rate = metrics.connection_reuse_rate();
4125 assert!(
4126 (rate - 0.99).abs() < 0.001,
4127 "复用率应接近 0.99,实际 {rate}"
4128 );
4129 }
4130
4131 #[test]
4132 fn connection_reuse_rate_partial() {
4133 let metrics = PoolMetrics {
4134 acquire_count: 10,
4135 connection_created_count: 2,
4136 ..Default::default()
4137 };
4138 assert!((metrics.connection_reuse_rate() - 0.8).abs() < 0.001);
4139 }
4140
4141 #[test]
4142 fn pool_tuning_advice_is_optimal() {
4143 let advice = PoolTuningAdvice {
4144 suggested_max_size: None,
4145 suggested_min_idle: None,
4146 suggested_idle_timeout: None,
4147 reason: "池配置合理".to_string(),
4148 };
4149 assert!(advice.is_optimal());
4150
4151 let not_optimal = PoolTuningAdvice {
4152 suggested_max_size: Some(20),
4153 suggested_min_idle: None,
4154 suggested_idle_timeout: None,
4155 reason: "test".to_string(),
4156 };
4157 assert!(!not_optimal.is_optimal());
4158 }
4159
4160 #[test]
4161 fn suggest_tuning_low_reuse() {
4162 let metrics = PoolMetrics {
4163 acquire_count: 100,
4164 connection_created_count: 60,
4165 ..Default::default()
4166 };
4167 let reuse = metrics.connection_reuse_rate();
4168 assert!(reuse < 0.5, "复用率 {reuse} 应 < 0.5");
4169 }
4170
4171 #[test]
4172 fn suggest_tuning_optimal() {
4173 let metrics = PoolMetrics {
4174 acquire_count: 1000,
4175 connection_created_count: 10,
4176 acquire_wait_time: Duration::from_millis(10),
4177 ..Default::default()
4178 };
4179 let reuse = metrics.connection_reuse_rate();
4180 assert!(reuse >= 0.9, "复用率 {reuse} 应 >= 0.9");
4181 let avg_wait = metrics.average_acquire_wait_time();
4182 assert!(avg_wait <= Duration::from_millis(100));
4183 }
4184
4185 #[test]
4186 fn suggest_tuning_high_wait() {
4187 let metrics = PoolMetrics {
4188 acquire_count: 100,
4189 acquire_wait_time: Duration::from_millis(200 * 100),
4190 ..Default::default()
4191 };
4192 let avg_wait = metrics.average_acquire_wait_time();
4193 assert!(
4194 avg_wait > Duration::from_millis(100),
4195 "平均等待 {avg_wait:?} 应 > 100ms"
4196 );
4197 }
4198}