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
633pub struct PoolConfigBuilder {
635 config: PoolConfig,
636}
637
638impl PoolConfigBuilder {
639 pub fn new() -> Self {
641 Self {
642 config: PoolConfig::default(),
643 }
644 }
645
646 pub fn max_size(mut self, size: u32) -> Self {
648 self.config.max_size = size;
649 self
650 }
651
652 pub fn min_idle(mut self, count: u32) -> Self {
654 self.config.min_idle = count;
655 self
656 }
657
658 pub fn acquire_timeout(mut self, timeout_secs: u64) -> Self {
660 self.config.acquire_timeout = Duration::from_secs(timeout_secs);
661 self
662 }
663
664 pub fn idle_timeout(mut self, timeout_secs: u64) -> Self {
666 self.config.idle_timeout = Duration::from_secs(timeout_secs);
667 self
668 }
669
670 pub fn max_lifetime(mut self, lifetime_secs: u64) -> Self {
672 self.config.max_lifetime = Duration::from_secs(lifetime_secs);
673 self
674 }
675
676 pub fn tls(mut self, tls: TlsConfig) -> Self {
678 self.config.tls = Some(tls);
679 self
680 }
681
682 pub fn query_timeout(mut self, timeout: Duration) -> Self {
684 self.config.query_timeout = Some(timeout);
685 self
686 }
687
688 pub fn max_rows(mut self, max_rows: usize) -> Self {
690 self.config.max_rows = Some(max_rows);
691 self
692 }
693
694 pub fn memory_limit(mut self, memory_limit: usize) -> Self {
696 self.config.memory_limit = Some(memory_limit);
697 self
698 }
699
700 pub fn on_event(mut self, callback: PoolEventCallback) -> Self {
702 self.config.on_event = Some(callback);
703 self
704 }
705
706 pub fn test_before_acquire(mut self, enabled: bool) -> Self {
711 self.config.test_before_acquire = enabled;
712 self
713 }
714
715 pub fn prewarm(mut self, enabled: bool) -> Self {
720 self.config.prewarm = enabled;
721 self
722 }
723
724 pub fn build(self) -> Result<PoolConfig, PoolError> {
726 self.config.validate()?;
727 Ok(self.config)
728 }
729}
730
731impl Default for PoolConfigBuilder {
732 fn default() -> Self {
733 Self::new()
734 }
735}
736
737#[async_trait]
739pub trait ConnectionFactory: Send + Sync {
740 async fn create(&self) -> Result<Box<dyn Connection>, crate::DbError>;
742}
743
744pub struct Pool {
750 config: PoolConfig,
751 factory: Arc<dyn ConnectionFactory>,
752 idle: Arc<ArrayQueue<PooledConnection>>,
758 total_count: Arc<AtomicU32>,
768 closed: Arc<AtomicBool>,
770 notify: Arc<Notify>,
771 waiters_count: Arc<AtomicU32>,
773 dynamic_max_size: Arc<AtomicU32>,
775 #[cfg(feature = "circuit-breaker")]
781 circuit_breaker: Arc<PlMutex<DefaultCircuitBreaker>>,
782 #[cfg(feature = "rate-limit")]
791 rate_limiter: Arc<PlRwLock<Option<Arc<dyn RateLimiter>>>>,
792 #[cfg(feature = "rate-limit")]
794 rate_limit_key: String,
795 #[cfg(feature = "tenant-quota-rls-enhanced")]
800 quota_enforcer: Arc<PlMutex<Option<Arc<QuotaEnforcer>>>>,
801 acquire_count: Arc<AtomicU64>,
803 acquire_failed_count: Arc<AtomicU64>,
805 acquire_wait_time_ns: Arc<AtomicU64>,
807 release_count: Arc<AtomicU64>,
809 connection_created_count: Arc<AtomicU64>,
811 connection_closed_count: Arc<AtomicU64>,
813}
814
815impl Clone for Pool {
819 fn clone(&self) -> Self {
820 Self {
821 config: self.config.clone(),
822 factory: self.factory.clone(),
823 idle: self.idle.clone(),
824 total_count: self.total_count.clone(),
825 closed: self.closed.clone(),
826 notify: Arc::clone(&self.notify),
827 waiters_count: self.waiters_count.clone(),
828 dynamic_max_size: self.dynamic_max_size.clone(),
829 #[cfg(feature = "circuit-breaker")]
830 circuit_breaker: Arc::clone(&self.circuit_breaker),
831 #[cfg(feature = "rate-limit")]
832 rate_limiter: Arc::clone(&self.rate_limiter),
833 #[cfg(feature = "rate-limit")]
834 rate_limit_key: self.rate_limit_key.clone(),
835 #[cfg(feature = "tenant-quota-rls-enhanced")]
836 quota_enforcer: Arc::clone(&self.quota_enforcer),
837 acquire_count: self.acquire_count.clone(),
838 acquire_failed_count: self.acquire_failed_count.clone(),
839 acquire_wait_time_ns: self.acquire_wait_time_ns.clone(),
840 release_count: self.release_count.clone(),
841 connection_created_count: self.connection_created_count.clone(),
842 connection_closed_count: self.connection_closed_count.clone(),
843 }
844 }
845}
846
847impl Pool {
848 pub fn new(config: PoolConfig, factory: Arc<dyn ConnectionFactory>) -> Result<Self, PoolError> {
872 config.validate()?;
873 let max_size = config.max_size as usize;
876 let dynamic_max = config.max_size;
877 Ok(Self {
878 config,
879 factory,
880 idle: Arc::new(ArrayQueue::new(max_size)),
881 total_count: Arc::new(AtomicU32::new(0)),
882 closed: Arc::new(AtomicBool::new(false)),
883 notify: Arc::new(Notify::new()),
884 waiters_count: Arc::new(AtomicU32::new(0)),
885 dynamic_max_size: Arc::new(AtomicU32::new(dynamic_max)),
886 #[cfg(feature = "circuit-breaker")]
889 circuit_breaker: Arc::new(PlMutex::new(DefaultCircuitBreaker::new(
890 5,
891 std::time::Duration::from_secs(30),
892 ))),
893 #[cfg(feature = "rate-limit")]
896 rate_limiter: Arc::new(PlRwLock::new(None)),
897 #[cfg(feature = "rate-limit")]
898 rate_limit_key: "pool".to_string(),
899 #[cfg(feature = "tenant-quota-rls-enhanced")]
900 quota_enforcer: Arc::new(PlMutex::new(None)),
901 acquire_count: Arc::new(AtomicU64::new(0)),
902 acquire_failed_count: Arc::new(AtomicU64::new(0)),
903 acquire_wait_time_ns: Arc::new(AtomicU64::new(0)),
904 release_count: Arc::new(AtomicU64::new(0)),
905 connection_created_count: Arc::new(AtomicU64::new(0)),
906 connection_closed_count: Arc::new(AtomicU64::new(0)),
907 })
908 }
909
910 pub async fn new_async(
917 config: PoolConfig,
918 factory: Arc<dyn ConnectionFactory>,
919 ) -> Result<Self, PoolError> {
920 let pool = Self::new(config, factory)?;
921 if pool.config.prewarm {
922 pool.prewarm().await;
923 }
924 Ok(pool)
925 }
926
927 pub async fn prewarm(&self) {
944 if !self.config.prewarm {
945 return;
946 }
947
948 let min_idle = self.config.min_idle as usize;
949 let mut warmed = 0;
950
951 for i in 0..min_idle {
952 if self.closed.load(Ordering::Acquire) {
954 break;
955 }
956
957 let current_max = self.dynamic_max_size.load(Ordering::Acquire);
959 let current = self.total_count.load(Ordering::Acquire);
960 if current >= current_max {
961 break;
962 }
963
964 let created = loop {
966 let current = self.total_count.load(Ordering::Acquire);
967 if current >= current_max {
968 break None;
969 }
970 match self.total_count.compare_exchange(
971 current,
972 current + 1,
973 Ordering::SeqCst,
974 Ordering::Acquire,
975 ) {
976 Ok(_) => break Some(()),
977 Err(_) => continue,
978 }
979 };
980
981 if created.is_some() {
982 match tokio::time::timeout(self.config.connection_timeout, self.factory.create())
983 .await
984 {
985 Ok(Ok(conn)) => {
986 #[cfg(feature = "circuit-breaker")]
987 {
988 self.circuit_breaker.lock().record_success();
989 }
990 self.emit_event(PoolEvent::ConnectionCreated);
991 let pooled = PooledConnection::new(conn, self.clone());
992 if self.idle.push(pooled).is_err() {
994 let _ = self.total_count.fetch_sub(1, Ordering::SeqCst);
996 tracing::warn!(
997 target: "sz_orm::pool::prewarm",
998 "prewarm connection {} failed: idle queue full",
999 i
1000 );
1001 } else {
1002 warmed += 1;
1003 self.notify.notify_one();
1004 }
1005 }
1006 Ok(Err(e)) => {
1007 let _ = self.total_count.fetch_sub(1, Ordering::SeqCst);
1008 #[cfg(feature = "circuit-breaker")]
1009 {
1010 self.circuit_breaker.lock().record_failure();
1011 }
1012 tracing::warn!(
1013 target: "sz_orm::pool::prewarm",
1014 "prewarm connection {} failed: {}",
1015 i,
1016 e
1017 );
1018 }
1019 Err(_) => {
1020 let _ = self.total_count.fetch_sub(1, Ordering::SeqCst);
1021 #[cfg(feature = "circuit-breaker")]
1022 {
1023 self.circuit_breaker.lock().record_failure();
1024 }
1025 tracing::warn!(
1026 target: "sz_orm::pool::prewarm",
1027 "prewarm connection {} timeout",
1028 i
1029 );
1030 }
1031 }
1032 }
1033 }
1034
1035 if warmed > 0 {
1036 tracing::info!(
1037 target: "sz_orm::pool::prewarm",
1038 "pool prewarm completed: {}/{} connections established",
1039 warmed,
1040 min_idle
1041 );
1042 }
1043 }
1044
1045 #[cfg(feature = "auto-prewarm")]
1050 pub async fn progressive_prewarm(
1051 &self,
1052 batch_size: u32,
1053 interval: std::time::Duration,
1054 total_timeout: std::time::Duration,
1055 progress: &crate::prewarm::PrewarmProgress,
1056 ) {
1057 use std::time::Instant;
1058
1059 let min_idle = self.config.min_idle;
1060 if min_idle == 0 || !self.config.prewarm {
1061 progress.mark_completed();
1062 return;
1063 }
1064
1065 let start = Instant::now();
1066 let batch = batch_size.max(1);
1067 let mut warmed_total: u32 = 0;
1068
1069 while warmed_total < min_idle {
1070 if start.elapsed() >= total_timeout {
1071 tracing::warn!(
1072 target: "sz_orm::pool::prewarm",
1073 "progressive prewarm timeout: {}/{} connections established",
1074 warmed_total,
1075 min_idle
1076 );
1077 break;
1078 }
1079
1080 if self.closed.load(Ordering::Acquire) {
1081 break;
1082 }
1083
1084 let remaining = min_idle - warmed_total;
1085 let this_batch = batch.min(remaining);
1086
1087 for _ in 0..this_batch {
1088 let current_max = self.dynamic_max_size.load(Ordering::Acquire);
1089 let current = self.total_count.load(Ordering::Acquire);
1090 if current >= current_max {
1091 break;
1092 }
1093
1094 let created = loop {
1095 let current = self.total_count.load(Ordering::Acquire);
1096 if current >= current_max {
1097 break None;
1098 }
1099 match self.total_count.compare_exchange(
1100 current,
1101 current + 1,
1102 Ordering::SeqCst,
1103 Ordering::Acquire,
1104 ) {
1105 Ok(_) => break Some(()),
1106 Err(_) => continue,
1107 }
1108 };
1109
1110 if created.is_some() {
1111 match tokio::time::timeout(
1112 self.config.connection_timeout,
1113 self.factory.create(),
1114 )
1115 .await
1116 {
1117 Ok(Ok(conn)) => {
1118 #[cfg(feature = "circuit-breaker")]
1119 {
1120 self.circuit_breaker.lock().record_success();
1121 }
1122 self.emit_event(PoolEvent::ConnectionCreated);
1123 let pooled = PooledConnection::new(conn, self.clone());
1124 if self.idle.push(pooled).is_err() {
1125 let _ = self.total_count.fetch_sub(1, Ordering::SeqCst);
1126 progress.record_failure();
1127 } else {
1128 progress.record_success();
1129 warmed_total += 1;
1130 self.notify.notify_one();
1131 }
1132 }
1133 Ok(Err(_)) => {
1134 let _ = self.total_count.fetch_sub(1, Ordering::SeqCst);
1135 progress.record_failure();
1136 #[cfg(feature = "circuit-breaker")]
1137 {
1138 self.circuit_breaker.lock().record_failure();
1139 }
1140 }
1141 Err(_) => {
1142 let _ = self.total_count.fetch_sub(1, Ordering::SeqCst);
1143 progress.record_failure();
1144 #[cfg(feature = "circuit-breaker")]
1145 {
1146 self.circuit_breaker.lock().record_failure();
1147 }
1148 }
1149 }
1150 }
1151 }
1152
1153 if warmed_total < min_idle && interval > std::time::Duration::ZERO {
1154 tokio::time::sleep(interval).await;
1155 }
1156 }
1157
1158 progress.set_elapsed(start.elapsed());
1159 progress.mark_completed();
1160
1161 tracing::info!(
1162 target: "sz_orm::pool::prewarm",
1163 "progressive prewarm completed: {} warmed, {} failed, elapsed {:?}",
1164 progress.snapshot().warmed,
1165 progress.snapshot().failed,
1166 start.elapsed()
1167 );
1168 }
1169
1170 pub fn config(&self) -> &PoolConfig {
1172 &self.config
1173 }
1174
1175 #[cfg(feature = "circuit-breaker")]
1189 pub fn configure_circuit_breaker(
1190 &self,
1191 failure_threshold: usize,
1192 reset_timeout: std::time::Duration,
1193 ) {
1194 let new_cb = DefaultCircuitBreaker::new(failure_threshold, reset_timeout);
1195 let mut guard = self.circuit_breaker.lock();
1197 *guard = new_cb;
1198 }
1199
1200 #[cfg(feature = "circuit-breaker")]
1205 pub fn reset_circuit_breaker(&self) -> bool {
1206 let mut guard = self.circuit_breaker.lock();
1208 guard.reset()
1209 }
1210
1211 #[cfg(feature = "circuit-breaker")]
1213 pub fn circuit_state(&self) -> CircuitState {
1214 let guard = self.circuit_breaker.lock();
1216 guard.state()
1217 }
1218
1219 #[cfg(feature = "rate-limit")]
1228 pub fn set_rate_limiter(&self, limiter: Option<Arc<dyn RateLimiter>>) {
1229 let mut guard = self.rate_limiter.write();
1231 *guard = limiter;
1232 }
1233
1234 #[cfg(feature = "rate-limit")]
1236 pub fn with_rate_limit_key(mut self, key: impl Into<String>) -> Self {
1237 self.rate_limit_key = key.into();
1238 self
1239 }
1240
1241 #[cfg(feature = "tenant-quota-rls-enhanced")]
1246 pub fn set_quota_enforcer(&self, enforcer: Option<Arc<QuotaEnforcer>>) {
1247 let mut guard = self.quota_enforcer.lock();
1248 *guard = enforcer;
1249 }
1250
1251 #[cfg(feature = "tenant-quota-rls-enhanced")]
1260 pub async fn acquire_with_tenant(
1261 &self,
1262 tenant_id: &str,
1263 ) -> Result<PooledConnection, PoolError> {
1264 {
1265 let guard = self.quota_enforcer.lock();
1266 if let Some(ref enforcer) = *guard {
1267 let current = enforcer.current_usage(tenant_id, QuotaResource::Connection);
1268 enforcer
1269 .check_and_record(tenant_id, QuotaResource::Connection, 1)
1270 .map_err(|e| PoolError::Internal(e.to_string()))?;
1271 let _ = current;
1272 }
1273 }
1274 self.acquire().await
1275 }
1276
1277 #[cfg(feature = "tenant-quota-rls-enhanced")]
1282 pub async fn release_with_tenant(&self, tenant_id: &str, pooled: PooledConnection) {
1283 {
1284 let guard = self.quota_enforcer.lock();
1285 if let Some(ref enforcer) = *guard {
1286 enforcer.release_usage(tenant_id, QuotaResource::Connection, 1);
1289 }
1290 }
1291 self.release(pooled).await;
1292 }
1293
1294 fn emit_event(&self, event: PoolEvent) {
1296 if matches!(event, PoolEvent::ConnectionCreated) {
1299 self.connection_created_count
1300 .fetch_add(1, Ordering::Relaxed);
1301 }
1302 if let Some(ref callback) = self.config.on_event {
1303 callback(event);
1304 }
1305 }
1306
1307 async fn close_connection(&self, pooled: PooledConnection) {
1312 let mut pooled = pooled;
1313 let _ = pooled.conn.close().await;
1314 self.connection_closed_count.fetch_add(1, Ordering::Relaxed);
1315 }
1316
1317 #[tracing::instrument(skip(self), fields(max_size = self.config.max_size, acquire_timeout = ?self.config.acquire_timeout))]
1337 pub async fn acquire(&self) -> Result<PooledConnection, PoolError> {
1338 if self.closed.load(Ordering::Acquire) {
1340 self.acquire_failed_count.fetch_add(1, Ordering::Relaxed);
1341 return Err(PoolError::Closed);
1342 }
1343
1344 #[cfg(feature = "circuit-breaker")]
1348 {
1349 let mut guard = self.circuit_breaker.lock();
1350 if !guard.can_execute() {
1351 self.acquire_failed_count.fetch_add(1, Ordering::Relaxed);
1352 return Err(PoolError::CircuitOpen);
1353 }
1354 }
1355
1356 #[cfg(feature = "rate-limit")]
1360 {
1361 let guard = self.rate_limiter.read();
1362 if let Some(ref limiter) = *guard {
1363 match limiter.try_acquire(&self.rate_limit_key) {
1364 Ok(result) if !result.allowed => {
1365 self.acquire_failed_count.fetch_add(1, Ordering::Relaxed);
1366 return Err(PoolError::RateLimited {
1367 remaining: result.remaining,
1368 reset_at: result.reset_at,
1369 });
1370 }
1371 Ok(_) => {} Err(_) => {
1373 }
1375 }
1376 }
1377 }
1378
1379 let deadline = Instant::now() + self.config.acquire_timeout;
1380 let mut backoff = Duration::from_millis(1);
1382 const MAX_BACKOFF: Duration = Duration::from_millis(100);
1384
1385 loop {
1386 let mut to_close: Vec<PooledConnection> = Vec::new();
1392 let acquired: Option<PooledConnection> = {
1393 let mut found: Option<PooledConnection> = None;
1394 while let Some(pooled) = self.idle.pop() {
1395 if pooled.is_expired(self.config.max_lifetime) {
1397 to_close.push(pooled);
1398 continue;
1399 }
1400 if pooled.is_idle_too_long(self.config.idle_timeout) {
1402 to_close.push(pooled);
1403 continue;
1404 }
1405 if !pooled.conn.is_connected() {
1408 to_close.push(pooled);
1409 continue;
1410 }
1411 found = Some(pooled);
1412 break;
1413 }
1414 found
1415 };
1416
1417 for pooled in to_close {
1419 self.close_connection(pooled).await;
1420 self.total_count.fetch_sub(1, Ordering::SeqCst);
1422 }
1423
1424 if let Some(mut pooled) = acquired {
1425 if self.config.test_before_acquire {
1427 let ping_timeout = self.config.connection_timeout / 2;
1428 let alive = match tokio::time::timeout(ping_timeout, pooled.conn.ping()).await {
1429 Ok(true) => true,
1430 Ok(false) => false,
1431 Err(_) => false, };
1433 if !alive {
1434 self.close_connection(pooled).await;
1436 self.total_count.fetch_sub(1, Ordering::SeqCst);
1437 continue;
1438 }
1439 }
1440 pooled.pool = Some(self.clone());
1443 self.acquire_count.fetch_add(1, Ordering::Relaxed);
1444 return Ok(pooled);
1445 }
1446
1447 let current_max = self.dynamic_max_size.load(Ordering::Acquire);
1452 let created = loop {
1453 let current = self.total_count.load(Ordering::Acquire);
1454 if current >= current_max {
1455 break None; }
1457 match self.total_count.compare_exchange(
1458 current,
1459 current + 1,
1460 Ordering::SeqCst,
1461 Ordering::Acquire,
1462 ) {
1463 Ok(_) => break Some(()), Err(_) => continue, }
1466 };
1467
1468 if created.is_some() {
1469 match tokio::time::timeout(self.config.connection_timeout, self.factory.create())
1470 .await
1471 {
1472 Ok(Ok(conn)) => {
1473 #[cfg(feature = "circuit-breaker")]
1476 {
1477 self.circuit_breaker.lock().record_success();
1478 }
1479 self.emit_event(PoolEvent::ConnectionCreated);
1480 self.emit_event(PoolEvent::ConnectionAcquired);
1481 self.acquire_count.fetch_add(1, Ordering::Relaxed);
1482 return Ok(PooledConnection::new(conn, self.clone()));
1483 }
1484 Ok(Err(e)) => {
1485 self.total_count.fetch_sub(1, Ordering::SeqCst);
1487 #[cfg(feature = "circuit-breaker")]
1490 {
1491 self.circuit_breaker.lock().record_failure();
1492 }
1493 self.acquire_failed_count.fetch_add(1, Ordering::Relaxed);
1494 return Err(PoolError::ConnectionFailed(e.to_string()));
1495 }
1496 Err(_) => {
1497 self.total_count.fetch_sub(1, Ordering::SeqCst);
1499 #[cfg(feature = "circuit-breaker")]
1502 {
1503 self.circuit_breaker.lock().record_failure();
1504 }
1505 self.acquire_failed_count.fetch_add(1, Ordering::Relaxed);
1506 return Err(PoolError::Timeout);
1507 }
1508 }
1509 }
1510
1511 let now = Instant::now();
1513 if now >= deadline {
1514 self.emit_event(PoolEvent::AcquireTimeout);
1515 self.acquire_failed_count.fetch_add(1, Ordering::Relaxed);
1516 return Err(PoolError::Timeout);
1517 }
1518 self.waiters_count.fetch_add(1, Ordering::SeqCst);
1520 let wait = std::cmp::min(backoff, deadline - now);
1521 match tokio::time::timeout(wait, self.notify.notified()).await {
1522 Ok(()) => {
1523 backoff = Duration::from_millis(1);
1525 }
1526 Err(_) => {
1527 backoff = std::cmp::min(backoff * 2, MAX_BACKOFF);
1529 }
1530 }
1531 self.waiters_count.fetch_sub(1, Ordering::SeqCst);
1533 self.acquire_wait_time_ns
1535 .fetch_add(wait.as_nanos() as u64, Ordering::Relaxed);
1536 }
1537 }
1538
1539 #[tracing::instrument(skip(self, pooled))]
1547 pub async fn release(&self, mut pooled: PooledConnection) {
1548 pooled.pool = None;
1550 self.release_count.fetch_add(1, Ordering::Relaxed);
1552
1553 if self.closed.load(Ordering::Acquire) {
1555 self.close_connection(pooled).await;
1556 self.total_count.fetch_sub(1, Ordering::SeqCst);
1558 self.emit_event(PoolEvent::ConnectionClosed);
1559 return;
1560 }
1561
1562 if !pooled.conn.is_connected() {
1564 self.close_connection(pooled).await;
1565 self.total_count.fetch_sub(1, Ordering::SeqCst);
1566 self.emit_event(PoolEvent::ConnectionClosed);
1567 return;
1568 }
1569
1570 pooled.last_used_at = Instant::now();
1572
1573 if let Err(rejected) = self.idle.push(pooled) {
1579 self.close_connection(rejected).await;
1581 self.total_count.fetch_sub(1, Ordering::SeqCst);
1582 self.emit_event(PoolEvent::ConnectionClosed);
1583 } else {
1584 self.emit_event(PoolEvent::ConnectionReleased);
1585 }
1586 self.notify.notify_one();
1587 }
1588
1589 pub async fn status(&self) -> PoolStatus {
1594 let idle_count = self.idle.len() as u32;
1595 let active = self.total_count.load(Ordering::Acquire);
1597 let waiters = self.waiters_count.load(Ordering::Acquire);
1598 PoolStatus {
1599 idle: idle_count,
1600 active,
1601 max: self.dynamic_max_size.load(Ordering::Acquire),
1602 min: self.config.min_idle,
1603 waiters,
1604 }
1605 }
1606
1607 pub fn pool_metrics(&self) -> PoolMetrics {
1618 PoolMetrics {
1619 acquire_count: self.acquire_count.load(Ordering::Acquire),
1620 acquire_failed_count: self.acquire_failed_count.load(Ordering::Acquire),
1621 acquire_wait_time: Duration::from_nanos(
1622 self.acquire_wait_time_ns.load(Ordering::Acquire),
1623 ),
1624 release_count: self.release_count.load(Ordering::Acquire),
1625 connection_created_count: self.connection_created_count.load(Ordering::Acquire),
1626 connection_closed_count: self.connection_closed_count.load(Ordering::Acquire),
1627 }
1628 }
1629
1630 pub fn metrics_snapshot_json(&self) -> String {
1635 serde_json::to_string(&self.pool_metrics()).unwrap_or_else(|_| "{}".to_string())
1636 }
1637
1638 #[tracing::instrument(skip(self))]
1640 pub async fn reap_idle(&self) {
1641 let mut all: Vec<PooledConnection> = Vec::new();
1645 while let Some(pooled) = self.idle.pop() {
1646 all.push(pooled);
1647 }
1648
1649 let mut to_close = Vec::new();
1651 for pooled in all {
1652 if pooled.is_idle_too_long(self.config.idle_timeout)
1653 || pooled.is_expired(self.config.max_lifetime)
1654 {
1655 to_close.push(pooled);
1656 } else {
1657 if let Err(rejected) = self.idle.push(pooled) {
1659 self.close_connection(rejected).await;
1660 self.total_count.fetch_sub(1, Ordering::SeqCst);
1661 }
1662 }
1663 }
1664
1665 for pooled in to_close {
1667 self.close_connection(pooled).await;
1668 self.total_count.fetch_sub(1, Ordering::SeqCst);
1670 }
1671 }
1672
1673 pub async fn close_all(&self) {
1677 self.closed.store(true, Ordering::Release);
1679 let mut to_close: Vec<PooledConnection> = Vec::new();
1682 while let Some(pooled) = self.idle.pop() {
1683 to_close.push(pooled);
1684 }
1685 let closed_count: u32 = to_close.len() as u32;
1687 for pooled in to_close {
1688 self.close_connection(pooled).await;
1689 }
1690 self.total_count.fetch_sub(closed_count, Ordering::SeqCst);
1693 }
1694
1695 pub async fn health_check(&self) -> u32 {
1711 let mut to_check: Vec<PooledConnection> = Vec::new();
1713 while let Some(pooled) = self.idle.pop() {
1714 to_check.push(pooled);
1715 }
1716
1717 let mut removed: u32 = 0;
1718 let mut alive: Vec<PooledConnection> = Vec::with_capacity(to_check.len());
1719 for mut pooled in to_check.drain(..) {
1720 if !pooled.conn.is_connected() {
1722 self.close_connection(pooled).await;
1723 removed += 1;
1724 continue;
1725 }
1726 let ping_timeout = self.config.connection_timeout / 2;
1728 match tokio::time::timeout(ping_timeout, pooled.conn.ping()).await {
1729 Ok(true) => alive.push(pooled),
1730 Ok(false) => {
1731 self.close_connection(pooled).await;
1733 removed += 1;
1734 }
1735 Err(_) => {
1736 self.close_connection(pooled).await;
1738 removed += 1;
1739 }
1740 }
1741 }
1742
1743 let alive_count: u32 = alive.len() as u32;
1745 for pooled in alive {
1746 if let Err(rejected) = self.idle.push(pooled) {
1748 self.close_connection(rejected).await;
1749 removed += 1;
1750 }
1751 }
1752
1753 if removed > 0 {
1755 self.total_count.fetch_sub(removed, Ordering::SeqCst);
1756 }
1757
1758 if alive_count > 0 {
1760 self.notify.notify_one();
1761 }
1762
1763 removed
1764 }
1765
1766 pub async fn shutdown(&self) {
1773 self.shutdown_with_timeout(Duration::from_secs(30)).await;
1774 }
1775
1776 pub async fn shutdown_with_timeout(&self, timeout: Duration) {
1784 if self.closed.swap(true, Ordering::SeqCst) {
1786 return;
1787 }
1788 self.notify.notify_waiters();
1790 self.close_all().await;
1792 let deadline = Instant::now() + timeout;
1794 while self.total_count.load(Ordering::SeqCst) > 0 {
1795 if Instant::now() >= deadline {
1796 let remaining = self.total_count.load(Ordering::SeqCst);
1797 if remaining > 0 {
1798 eprintln!(
1799 "graceful shutdown timeout, {} connections force closed",
1800 remaining
1801 );
1802 }
1803 break;
1804 }
1805 tokio::time::sleep(Duration::from_millis(100)).await;
1806 }
1807 }
1808
1809 pub fn resize(&self, new_max: usize) {
1817 self.set_max_size(new_max as u32);
1818 }
1819
1820 pub fn set_max_size(&self, new_max: u32) {
1822 self.dynamic_max_size.store(new_max, Ordering::SeqCst);
1823 }
1824
1825 pub fn max_size(&self) -> u32 {
1827 self.dynamic_max_size.load(Ordering::Acquire)
1828 }
1829
1830 pub async fn warmup(&self, min_idle: usize) -> Result<(), PoolError> {
1834 for _ in 0..min_idle {
1835 let current_max = self.dynamic_max_size.load(Ordering::Acquire);
1836 let current = self.total_count.load(Ordering::Acquire);
1837 if current >= current_max {
1838 break;
1839 }
1840 match self.total_count.compare_exchange(
1842 current,
1843 current + 1,
1844 Ordering::SeqCst,
1845 Ordering::Acquire,
1846 ) {
1847 Ok(_) => {}
1848 Err(_) => continue, }
1850 match self.factory.create().await {
1851 Ok(conn) => {
1852 let now = Instant::now();
1853 let pooled = PooledConnection {
1854 conn,
1855 created_at: now,
1856 last_used_at: now,
1857 pool: None,
1858 };
1859 if let Err(rejected) = self.idle.push(pooled) {
1860 self.close_connection(rejected).await;
1862 self.total_count.fetch_sub(1, Ordering::SeqCst);
1863 }
1864 self.emit_event(PoolEvent::ConnectionCreated);
1865 }
1866 Err(_) => {
1867 self.total_count.fetch_sub(1, Ordering::SeqCst);
1869 break;
1870 }
1871 }
1872 }
1873 Ok(())
1874 }
1875
1876 pub async fn query_with_timeout(&self, sql: &str) -> Result<QueryRows, crate::DbError> {
1881 let timeout = self.config.query_timeout.unwrap_or(Duration::from_secs(30));
1882 let mut conn = self.acquire().await.map_err(crate::DbError::PoolError)?;
1883 tokio::time::timeout(timeout, conn.query(sql))
1884 .await
1885 .map_err(|_| crate::DbError::QueryError(format!("Query timeout after {:?}", timeout)))?
1886 }
1887}
1888
1889#[cfg(feature = "prod-pool-tuning")]
1894mod pool_prod {
1895 use super::PoolConfig;
1896 use serde::{Deserialize, Serialize};
1897 use std::time::Duration;
1898
1899 #[derive(Debug, Clone, PartialEq, Eq, thiserror::Error)]
1901 pub enum PoolProdError {
1902 #[error("pool max_size must be positive")]
1904 MaxSizeNotPositive,
1905 #[error("pool acquire_timeout must be positive")]
1907 AcquireTimeoutNotPositive,
1908 #[error("pool min_idle cannot exceed max_size")]
1910 MinIdleExceedsMaxSize,
1911 }
1912
1913 #[derive(Debug, Clone, Serialize, Deserialize)]
1915 pub struct PoolProdConfig {
1916 pub max_size: u32,
1918 pub acquire_timeout: Duration,
1920 pub idle_timeout: Duration,
1922 pub connection_timeout: Duration,
1924 pub query_timeout: Duration,
1926 pub min_idle: u32,
1928 pub prewarm: bool,
1930 }
1931
1932 impl Default for PoolProdConfig {
1933 fn default() -> Self {
1934 Self {
1935 max_size: 100,
1936 acquire_timeout: Duration::from_secs(30),
1937 idle_timeout: Duration::from_secs(600),
1938 connection_timeout: Duration::from_secs(10),
1939 query_timeout: Duration::from_secs(30),
1940 min_idle: 0,
1941 prewarm: false,
1942 }
1943 }
1944 }
1945
1946 impl PoolProdConfig {
1947 pub fn new(
1949 max_size: u32,
1950 acquire_timeout: Duration,
1951 idle_timeout: Duration,
1952 connection_timeout: Duration,
1953 query_timeout: Duration,
1954 min_idle: u32,
1955 prewarm: bool,
1956 ) -> Self {
1957 Self {
1958 max_size,
1959 acquire_timeout,
1960 idle_timeout,
1961 connection_timeout,
1962 query_timeout,
1963 min_idle,
1964 prewarm,
1965 }
1966 }
1967
1968 pub fn validate(&self) -> Result<(), PoolProdError> {
1970 if self.max_size == 0 {
1971 return Err(PoolProdError::MaxSizeNotPositive);
1972 }
1973 if self.acquire_timeout.is_zero() {
1974 return Err(PoolProdError::AcquireTimeoutNotPositive);
1975 }
1976 if self.min_idle > self.max_size {
1977 return Err(PoolProdError::MinIdleExceedsMaxSize);
1978 }
1979 Ok(())
1980 }
1981
1982 pub fn to_pool_config(&self) -> PoolConfig {
1984 PoolConfig {
1985 max_size: self.max_size,
1986 min_idle: self.min_idle,
1987 acquire_timeout: self.acquire_timeout,
1988 idle_timeout: self.idle_timeout,
1989 max_lifetime: Duration::from_secs(1800),
1990 connection_timeout: self.connection_timeout,
1991 tls: None,
1992 query_timeout: Some(self.query_timeout),
1993 max_rows: None,
1994 memory_limit: None,
1995 on_event: None,
1996 test_before_acquire: false,
1997 prewarm: self.prewarm,
1998 }
1999 }
2000 }
2001}
2002
2003#[cfg(feature = "prod-pool-tuning")]
2004pub use pool_prod::{PoolProdConfig, PoolProdError};
2005
2006#[cfg(feature = "prod-leak-detection")]
2011mod leak_detection {
2012 use serde::{Deserialize, Serialize};
2013 use std::time::Duration;
2014
2015 #[derive(Debug, Clone, Serialize, Deserialize)]
2017 pub struct LeakDetectionConfig {
2018 pub enabled: bool,
2020 pub interval: Duration,
2022 pub threshold: u32,
2024 pub borrow_timeout: Duration,
2026 }
2027
2028 impl Default for LeakDetectionConfig {
2029 fn default() -> Self {
2030 Self {
2031 enabled: false,
2032 interval: Duration::from_secs(60),
2033 threshold: 5,
2034 borrow_timeout: Duration::from_secs(60),
2035 }
2036 }
2037 }
2038
2039 impl LeakDetectionConfig {
2040 pub fn new(
2042 enabled: bool,
2043 interval: Duration,
2044 threshold: u32,
2045 borrow_timeout: Duration,
2046 ) -> Self {
2047 Self {
2048 enabled,
2049 interval,
2050 threshold,
2051 borrow_timeout,
2052 }
2053 }
2054
2055 pub fn validate(&self) -> Result<(), LeakDetectionError> {
2057 if self.interval.is_zero() {
2058 return Err(LeakDetectionError::IntervalNotPositive);
2059 }
2060 if self.borrow_timeout.is_zero() {
2061 return Err(LeakDetectionError::BorrowTimeoutNotPositive);
2062 }
2063 Ok(())
2064 }
2065 }
2066
2067 #[derive(Debug, Clone, PartialEq, Eq, thiserror::Error)]
2069 pub enum LeakDetectionError {
2070 #[error("leak detection interval must be positive")]
2072 IntervalNotPositive,
2073 #[error("leak detection borrow_timeout must be positive")]
2075 BorrowTimeoutNotPositive,
2076 }
2077
2078 #[derive(Debug, Clone, Serialize, Deserialize)]
2080 pub struct LeakEntry {
2081 pub conn_id: u64,
2083 pub borrowed_at: String,
2085 pub borrow_duration: Duration,
2087 }
2088
2089 #[derive(Debug, Clone, Serialize, Deserialize)]
2091 pub struct LeakReport {
2092 pub borrowed_count: u32,
2094 pub max_borrow_duration: Duration,
2096 pub suspected_leaks: Vec<LeakEntry>,
2098 }
2099
2100 impl LeakReport {
2101 pub fn empty() -> Self {
2103 Self {
2104 borrowed_count: 0,
2105 max_borrow_duration: Duration::ZERO,
2106 suspected_leaks: vec![],
2107 }
2108 }
2109 }
2110}
2111
2112#[cfg(feature = "prod-leak-detection")]
2113pub use leak_detection::{LeakDetectionConfig, LeakDetectionError, LeakEntry, LeakReport};
2114
2115#[cfg(test)]
2116mod tests {
2117 use super::*;
2118
2119 struct MockConnection {
2121 connected: bool,
2122 }
2123
2124 impl MockConnection {
2125 fn new() -> Self {
2126 Self { connected: true }
2127 }
2128 }
2129
2130 impl Connection for MockConnection {
2131 fn execute<'a>(
2132 &'a mut self,
2133 _sql: &'a str,
2134 ) -> Pin<Box<dyn Future<Output = Result<u64, crate::DbError>> + Send + 'a>> {
2135 Box::pin(async move { Ok(1) })
2136 }
2137
2138 fn query<'a>(
2139 &'a mut self,
2140 _sql: &'a str,
2141 ) -> Pin<
2142 Box<
2143 dyn Future<
2144 Output = Result<
2145 Vec<std::collections::HashMap<String, crate::value::Value>>,
2146 crate::DbError,
2147 >,
2148 > + Send
2149 + 'a,
2150 >,
2151 > {
2152 Box::pin(async move { Ok(vec![]) })
2153 }
2154
2155 fn begin_transaction<'a>(
2156 &'a mut self,
2157 ) -> Pin<Box<dyn Future<Output = Result<(), crate::DbError>> + Send + 'a>> {
2158 Box::pin(async move { Ok(()) })
2159 }
2160
2161 fn commit<'a>(
2162 &'a mut self,
2163 ) -> Pin<Box<dyn Future<Output = Result<(), crate::DbError>> + Send + 'a>> {
2164 Box::pin(async move { Ok(()) })
2165 }
2166
2167 fn rollback<'a>(
2168 &'a mut self,
2169 ) -> Pin<Box<dyn Future<Output = Result<(), crate::DbError>> + Send + 'a>> {
2170 Box::pin(async move { Ok(()) })
2171 }
2172
2173 fn is_connected(&self) -> bool {
2174 self.connected
2175 }
2176
2177 fn ping<'a>(&'a mut self) -> Pin<Box<dyn Future<Output = bool> + Send + 'a>> {
2178 Box::pin(async move { true })
2179 }
2180
2181 fn close<'a>(
2182 &'a mut self,
2183 ) -> Pin<Box<dyn Future<Output = Result<(), crate::DbError>> + Send + 'a>> {
2184 Box::pin(async move {
2185 self.connected = false;
2186 Ok(())
2187 })
2188 }
2189 }
2190
2191 struct MockConnectionFactory;
2192
2193 #[async_trait]
2194 impl ConnectionFactory for MockConnectionFactory {
2195 async fn create(&self) -> Result<Box<dyn Connection>, crate::DbError> {
2196 Ok(Box::new(MockConnection::new()))
2197 }
2198 }
2199
2200 #[tokio::test]
2201 async fn test_pool_config_builder() -> Result<(), Box<dyn std::error::Error>> {
2202 let config = PoolConfigBuilder::new().max_size(50).min_idle(10).build()?;
2203
2204 assert_eq!(config.max_size, 50);
2205 assert_eq!(config.min_idle, 10);
2206 Ok(())
2207 }
2208
2209 #[test]
2210 fn test_pool_status_display() {
2211 let status = PoolStatus {
2212 idle: 5,
2213 active: 10,
2214 max: 100,
2215 min: 5,
2216 waiters: 0,
2217 };
2218
2219 let display = format!("{:?}", status);
2220 assert!(display.contains("idle"));
2221 assert!(display.contains("active"));
2222 }
2223
2224 #[test]
2225 fn test_default_pool_config() {
2226 let config = PoolConfig::default();
2227 assert_eq!(config.max_size, 100);
2228 assert_eq!(config.min_idle, 0);
2229 assert_eq!(config.acquire_timeout.as_secs(), 30);
2230 assert_eq!(config.idle_timeout.as_secs(), 600);
2231 assert_eq!(config.max_lifetime.as_secs(), 1800);
2232 }
2233
2234 #[tokio::test]
2235 async fn test_pool_config_clone() {
2236 let config = PoolConfig::default();
2237 let cloned = config.clone();
2238 assert_eq!(cloned.max_size, config.max_size);
2239 assert_eq!(cloned.min_idle, config.min_idle);
2240 }
2241
2242 #[test]
2243 fn test_pool_config_builder_default() -> Result<(), Box<dyn std::error::Error>> {
2244 let builder = PoolConfigBuilder::new();
2245 let config = builder.build()?;
2246 assert_eq!(config.max_size, 100);
2247 Ok(())
2248 }
2249
2250 #[test]
2251 fn test_pool_config_validate() {
2252 let result = PoolConfigBuilder::new().max_size(0).build();
2253 assert!(result.is_err());
2254
2255 let result = PoolConfigBuilder::new().max_size(10).min_idle(20).build();
2256 assert!(result.is_err());
2257 }
2258
2259 #[test]
2260 fn test_pool_config_validate_duration_upper_bound() {
2261 use std::time::Duration;
2262
2263 let config = PoolConfig {
2265 max_size: 10,
2266 min_idle: 1,
2267 acquire_timeout: Duration::from_secs(u64::MAX),
2268 idle_timeout: Duration::from_secs(1),
2269 max_lifetime: Duration::from_secs(1),
2270 connection_timeout: Duration::from_secs(5),
2271 tls: None,
2272 query_timeout: None,
2273 max_rows: None,
2274 memory_limit: None,
2275 on_event: None,
2276 test_before_acquire: false,
2277 prewarm: false,
2278 };
2279 assert!(config.validate().is_err());
2280
2281 let config = PoolConfig {
2283 max_size: 10,
2284 min_idle: 1,
2285 acquire_timeout: Duration::from_secs(u32::MAX as u64),
2286 idle_timeout: Duration::from_secs(1),
2287 max_lifetime: Duration::from_secs(1),
2288 connection_timeout: Duration::from_secs(5),
2289 tls: None,
2290 query_timeout: None,
2291 max_rows: None,
2292 memory_limit: None,
2293 on_event: None,
2294 test_before_acquire: false,
2295 prewarm: false,
2296 };
2297 assert!(config.validate().is_ok());
2298
2299 let config = PoolConfig {
2301 max_size: 10,
2302 min_idle: 1,
2303 acquire_timeout: Duration::from_secs(u32::MAX as u64 + 1),
2304 idle_timeout: Duration::from_secs(1),
2305 max_lifetime: Duration::from_secs(1),
2306 connection_timeout: Duration::from_secs(5),
2307 tls: None,
2308 query_timeout: None,
2309 max_rows: None,
2310 memory_limit: None,
2311 on_event: None,
2312 test_before_acquire: false,
2313 prewarm: false,
2314 };
2315 assert!(config.validate().is_err());
2316 }
2317
2318 #[test]
2319 fn test_pool_config_test_before_acquire_default() {
2320 let config = PoolConfig::default();
2322 assert!(!config.test_before_acquire);
2323 }
2324
2325 #[test]
2326 fn test_pool_config_builder_test_before_acquire() {
2327 let config = PoolConfigBuilder::new()
2329 .test_before_acquire(true)
2330 .build()
2331 .unwrap();
2332 assert!(config.test_before_acquire);
2333 }
2334
2335 #[tokio::test]
2336 async fn test_pool_acquire_and_release() -> Result<(), Box<dyn std::error::Error>> {
2337 let config = PoolConfigBuilder::new().max_size(5).min_idle(1).build()?;
2338 let factory = Arc::new(MockConnectionFactory);
2339 let pool = Pool::new(config, factory)?;
2340
2341 let conn = pool.acquire().await?;
2342 let status = pool.status().await;
2343 assert_eq!(status.active, 1);
2344 assert_eq!(status.idle, 0);
2345
2346 pool.release(conn).await;
2347 let status = pool.status().await;
2348 assert_eq!(status.idle, 1);
2349
2350 let _conn2 = pool.acquire().await?;
2352 let status = pool.status().await;
2353 assert_eq!(status.idle, 0);
2354 Ok(())
2355 }
2356
2357 #[tokio::test]
2358 async fn test_pool_status() -> Result<(), Box<dyn std::error::Error>> {
2359 let config = PoolConfigBuilder::new().max_size(10).min_idle(2).build()?;
2360 let factory = Arc::new(MockConnectionFactory);
2361 let pool = Pool::new(config, factory)?;
2362
2363 let status = pool.status().await;
2364 assert_eq!(status.max, 10);
2365 assert_eq!(status.min, 2);
2366 assert_eq!(status.active, 0);
2367 Ok(())
2368 }
2369
2370 #[tokio::test]
2371 async fn test_pool_close_all() -> Result<(), Box<dyn std::error::Error>> {
2372 let config = PoolConfigBuilder::new().max_size(5).build()?;
2373 let factory = Arc::new(MockConnectionFactory);
2374 let pool = Pool::new(config, factory)?;
2375
2376 let conn1 = pool.acquire().await?;
2378 let conn2 = pool.acquire().await?;
2379 pool.release(conn1).await;
2380 pool.release(conn2).await;
2381
2382 pool.close_all().await;
2383 let status = pool.status().await;
2384 assert_eq!(status.idle, 0);
2385 assert_eq!(status.active, 0);
2386 Ok(())
2387 }
2388
2389 #[tokio::test]
2390 async fn test_pool_reap_idle() -> Result<(), Box<dyn std::error::Error>> {
2391 let config = PoolConfigBuilder::new()
2392 .max_size(5)
2393 .idle_timeout(0) .build()?;
2395 let factory = Arc::new(MockConnectionFactory);
2396 let pool = Pool::new(config, factory)?;
2397
2398 let conn = pool.acquire().await?;
2399 pool.release(conn).await;
2400
2401 tokio::time::sleep(tokio::time::Duration::from_millis(10)).await;
2403
2404 pool.reap_idle().await;
2405 let status = pool.status().await;
2406 assert_eq!(status.idle, 0);
2407 Ok(())
2408 }
2409
2410 #[tokio::test]
2416 async fn test_h7_acquire_timeout_default_30s() {
2417 let config = PoolConfig::default();
2418 assert_eq!(
2419 config.acquire_timeout,
2420 Duration::from_secs(30),
2421 "H-7: acquire_timeout 默认应为 30s"
2422 );
2423 }
2424
2425 #[tokio::test]
2427 async fn test_h7_acquire_timeout_configurable() -> Result<(), Box<dyn std::error::Error>> {
2428 let config = PoolConfigBuilder::new()
2429 .max_size(1)
2430 .acquire_timeout(5) .build()?;
2432 assert_eq!(config.acquire_timeout, Duration::from_secs(5));
2433
2434 let factory = Arc::new(MockConnectionFactory);
2436 let pool = Pool::new(config, factory)?;
2437 let _conn1 = pool.acquire().await?;
2438
2439 let fast_config = PoolConfigBuilder::new()
2441 .max_size(1)
2442 .acquire_timeout(0) .build()?;
2444 let fast_pool = Pool::new(fast_config, Arc::new(MockConnectionFactory))?;
2447 let _fast_conn = fast_pool.acquire().await?; let result = fast_pool.acquire().await;
2449 assert!(
2450 matches!(result, Err(PoolError::Timeout)),
2451 "H-7: 应返回 Timeout"
2452 );
2453 Ok(())
2454 }
2455
2456 #[tokio::test]
2459 async fn test_m7_health_check_removes_nothing_when_all_healthy(
2460 ) -> Result<(), Box<dyn std::error::Error>> {
2461 let config = PoolConfigBuilder::new().max_size(5).build()?;
2463 let factory = Arc::new(MockConnectionFactory);
2464 let pool = Pool::new(config, factory)?;
2465
2466 let conn1 = pool.acquire().await?;
2468 let conn2 = pool.acquire().await?;
2469 let conn3 = pool.acquire().await?;
2470 pool.release(conn1).await;
2471 pool.release(conn2).await;
2472 pool.release(conn3).await;
2473
2474 let removed = pool.health_check().await;
2475 assert_eq!(removed, 0, "Healthy connections should not be removed");
2476
2477 let status = pool.status().await;
2478 assert_eq!(status.idle, 3);
2479 assert_eq!(status.active, 3);
2480 Ok(())
2481 }
2482
2483 #[tokio::test]
2484 async fn test_m7_health_check_returns_zero_for_empty_pool(
2485 ) -> Result<(), Box<dyn std::error::Error>> {
2486 let config = PoolConfigBuilder::new().max_size(5).build()?;
2487 let factory = Arc::new(MockConnectionFactory);
2488 let pool = Pool::new(config, factory)?;
2489
2490 let removed = pool.health_check().await;
2491 assert_eq!(removed, 0);
2492 Ok(())
2493 }
2494
2495 struct CountingFactory {
2499 count: AtomicU32,
2500 }
2501
2502 impl CountingFactory {
2503 fn new() -> Self {
2504 Self {
2505 count: AtomicU32::new(0),
2506 }
2507 }
2508 fn created_count(&self) -> u32 {
2509 self.count.load(Ordering::SeqCst)
2510 }
2511 }
2512
2513 #[async_trait]
2514 impl ConnectionFactory for CountingFactory {
2515 async fn create(&self) -> Result<Box<dyn Connection>, crate::DbError> {
2516 self.count.fetch_add(1, Ordering::SeqCst);
2517 Ok(Box::new(MockConnection::new()))
2518 }
2519 }
2520
2521 #[tokio::test]
2527 async fn test_production_bug_max_lifetime_never_expires(
2528 ) -> Result<(), Box<dyn std::error::Error>> {
2529 let config = PoolConfig {
2532 max_size: 5,
2533 min_idle: 0,
2534 acquire_timeout: Duration::from_secs(30),
2535 idle_timeout: Duration::from_secs(600),
2536 max_lifetime: Duration::from_millis(100), connection_timeout: Duration::from_secs(10),
2538 tls: None,
2539 query_timeout: None,
2540 max_rows: None,
2541 memory_limit: None,
2542 on_event: None,
2543 test_before_acquire: false,
2544 prewarm: false,
2545 };
2546 let factory = Arc::new(CountingFactory::new());
2547 let pool = Pool::new(config, factory.clone())?;
2548
2549 let conn = pool.acquire().await?;
2551 assert_eq!(factory.created_count(), 1, "应创建 1 个连接");
2552
2553 pool.release(conn).await;
2555
2556 tokio::time::sleep(Duration::from_millis(150)).await;
2558
2559 let conn2 = pool.acquire().await?;
2561
2562 assert_eq!(
2565 factory.created_count(),
2566 2,
2567 "超过 max_lifetime 后应创建新连接(旧连接应被回收)"
2568 );
2569
2570 pool.release(conn2).await;
2571 Ok(())
2572 }
2573
2574 #[tokio::test]
2581 async fn test_drop_auto_release_connection() -> Result<(), Box<dyn std::error::Error>> {
2582 let config = PoolConfigBuilder::new().max_size(2).build()?;
2583 let factory = Arc::new(CountingFactory::new());
2584 let pool = Pool::new(config, factory.clone())?;
2585
2586 {
2588 let _conn = pool.acquire().await?;
2589 assert_eq!(factory.created_count(), 1, "应创建 1 个连接");
2590 let status = pool.status().await;
2591 assert_eq!(status.active, 1, "active 应为 1");
2592 assert_eq!(status.idle, 0, "idle 应为 0");
2593 }
2595
2596 tokio::time::sleep(Duration::from_millis(50)).await;
2598
2599 let status = pool.status().await;
2601 assert_eq!(status.idle, 1, "Drop 后连接应自动归还,idle 应为 1");
2602 assert_eq!(status.active, 1, "total_count 应为 1");
2603 assert_eq!(factory.created_count(), 1, "应复用归还的连接,不创建新连接");
2604 Ok(())
2605 }
2606
2607 #[tokio::test]
2609 async fn test_drop_auto_release_then_reuse() -> Result<(), Box<dyn std::error::Error>> {
2610 let config = PoolConfigBuilder::new().max_size(1).build()?;
2611 let factory = Arc::new(CountingFactory::new());
2612 let pool = Pool::new(config, factory.clone())?;
2613
2614 {
2616 let _conn = pool.acquire().await?;
2617 }
2618
2619 tokio::time::sleep(Duration::from_millis(50)).await;
2621
2622 let conn = pool.acquire().await?;
2624 assert_eq!(factory.created_count(), 1, "应复用归还的连接,不创建新连接");
2625
2626 pool.release(conn).await;
2627 Ok(())
2628 }
2629
2630 #[tokio::test]
2632 async fn test_into_inner_does_not_return_to_pool() -> Result<(), Box<dyn std::error::Error>> {
2633 let config = PoolConfigBuilder::new().max_size(2).build()?;
2634 let factory = Arc::new(CountingFactory::new());
2635 let pool = Pool::new(config, factory.clone())?;
2636
2637 let conn = pool.acquire().await?;
2638 assert_eq!(factory.created_count(), 1);
2639
2640 let _raw_conn = conn.into_inner();
2642
2643 tokio::time::sleep(Duration::from_millis(50)).await;
2645
2646 let status = pool.status().await;
2647 assert_eq!(status.idle, 0, "into_inner 后连接不应归还");
2648 assert_eq!(status.active, 1, "total_count 仍为 1(连接被外部持有)");
2649 Ok(())
2650 }
2651
2652 #[tokio::test]
2654 async fn test_explicit_release_no_double_return() -> Result<(), Box<dyn std::error::Error>> {
2655 let config = PoolConfigBuilder::new().max_size(2).build()?;
2656 let factory = Arc::new(CountingFactory::new());
2657 let pool = Pool::new(config, factory.clone())?;
2658
2659 let conn = pool.acquire().await?;
2660 pool.release(conn).await;
2661
2662 let status = pool.status().await;
2663 assert_eq!(status.idle, 1, "release 后 idle 应为 1");
2664
2665 let conn = pool.acquire().await?;
2667 pool.release(conn).await;
2668
2669 let status = pool.status().await;
2670 assert_eq!(status.idle, 1, "再次 release 后 idle 仍应为 1(不重复)");
2671 assert_eq!(status.active, 1, "total_count 应为 1");
2672 Ok(())
2673 }
2674
2675 struct CursorMockConn {
2681 rows: QueryRows,
2682 call_count: usize,
2683 }
2684
2685 impl CursorMockConn {
2686 fn new(rows: QueryRows) -> Self {
2687 Self {
2688 rows,
2689 call_count: 0,
2690 }
2691 }
2692 }
2693
2694 impl Connection for CursorMockConn {
2695 fn execute<'a>(
2696 &'a mut self,
2697 _sql: &'a str,
2698 ) -> Pin<Box<dyn Future<Output = Result<u64, crate::DbError>> + Send + 'a>> {
2699 Box::pin(async move { Ok(1) })
2700 }
2701
2702 fn query<'a>(
2703 &'a mut self,
2704 _sql: &'a str,
2705 ) -> Pin<Box<dyn Future<Output = Result<QueryRows, crate::DbError>> + Send + 'a>> {
2706 Box::pin(async move {
2707 self.call_count += 1;
2708 Ok(self.rows.clone())
2709 })
2710 }
2711
2712 fn begin_transaction<'a>(
2713 &'a mut self,
2714 ) -> Pin<Box<dyn Future<Output = Result<(), crate::DbError>> + Send + 'a>> {
2715 Box::pin(async move { Ok(()) })
2716 }
2717
2718 fn commit<'a>(
2719 &'a mut self,
2720 ) -> Pin<Box<dyn Future<Output = Result<(), crate::DbError>> + Send + 'a>> {
2721 Box::pin(async move { Ok(()) })
2722 }
2723
2724 fn rollback<'a>(
2725 &'a mut self,
2726 ) -> Pin<Box<dyn Future<Output = Result<(), crate::DbError>> + Send + 'a>> {
2727 Box::pin(async move { Ok(()) })
2728 }
2729
2730 fn is_connected(&self) -> bool {
2731 true
2732 }
2733
2734 fn ping<'a>(&'a mut self) -> Pin<Box<dyn Future<Output = bool> + Send + 'a>> {
2735 Box::pin(async move { true })
2736 }
2737
2738 fn close<'a>(
2739 &'a mut self,
2740 ) -> Pin<Box<dyn Future<Output = Result<(), crate::DbError>> + Send + 'a>> {
2741 Box::pin(async move { Ok(()) })
2742 }
2743 }
2744
2745 struct CursorOverrideMockConn {
2747 rows: Vec<crate::value::Value>,
2748 yielded: usize,
2749 }
2750
2751 impl CursorOverrideMockConn {
2752 fn new(rows: Vec<crate::value::Value>) -> Self {
2753 Self { rows, yielded: 0 }
2754 }
2755 }
2756
2757 impl Connection for CursorOverrideMockConn {
2758 fn execute<'a>(
2759 &'a mut self,
2760 _sql: &'a str,
2761 ) -> Pin<Box<dyn Future<Output = Result<u64, crate::DbError>> + Send + 'a>> {
2762 Box::pin(async move { Ok(1) })
2763 }
2764
2765 fn query<'a>(
2766 &'a mut self,
2767 _sql: &'a str,
2768 ) -> Pin<Box<dyn Future<Output = Result<QueryRows, crate::DbError>> + Send + 'a>> {
2769 Box::pin(async move {
2771 Ok(self
2772 .rows
2773 .iter()
2774 .map(|v| {
2775 let mut m = std::collections::HashMap::new();
2776 m.insert("v".to_string(), v.clone());
2777 m
2778 })
2779 .collect())
2780 })
2781 }
2782
2783 fn query_stream<'a>(
2785 &'a mut self,
2786 _sql: &'a str,
2787 ) -> Pin<Box<dyn futures::Stream<Item = QueryStreamItem> + Send + 'a>> {
2788 Box::pin(futures::stream::iter(
2789 self.rows
2790 .iter()
2791 .enumerate()
2792 .map(|(i, v)| {
2793 self.yielded = i + 1;
2794 let mut m = std::collections::HashMap::new();
2795 m.insert("v".to_string(), v.clone());
2796 Ok(m)
2797 })
2798 .collect::<Vec<_>>(),
2799 ))
2800 }
2801
2802 fn begin_transaction<'a>(
2803 &'a mut self,
2804 ) -> Pin<Box<dyn Future<Output = Result<(), crate::DbError>> + Send + 'a>> {
2805 Box::pin(async move { Ok(()) })
2806 }
2807
2808 fn commit<'a>(
2809 &'a mut self,
2810 ) -> Pin<Box<dyn Future<Output = Result<(), crate::DbError>> + Send + 'a>> {
2811 Box::pin(async move { Ok(()) })
2812 }
2813
2814 fn rollback<'a>(
2815 &'a mut self,
2816 ) -> Pin<Box<dyn Future<Output = Result<(), crate::DbError>> + Send + 'a>> {
2817 Box::pin(async move { Ok(()) })
2818 }
2819
2820 fn is_connected(&self) -> bool {
2821 true
2822 }
2823
2824 fn ping<'a>(&'a mut self) -> Pin<Box<dyn Future<Output = bool> + Send + 'a>> {
2825 Box::pin(async move { true })
2826 }
2827
2828 fn close<'a>(
2829 &'a mut self,
2830 ) -> Pin<Box<dyn Future<Output = Result<(), crate::DbError>> + Send + 'a>> {
2831 Box::pin(async move { Ok(()) })
2832 }
2833 }
2834
2835 #[tokio::test]
2837 async fn test_query_stream_default_impl_yields_all_rows() {
2838 use futures::StreamExt;
2839 let rows: QueryRows = vec![
2840 std::collections::HashMap::from([
2841 ("id".to_string(), crate::value::Value::I64(1)),
2842 (
2843 "name".to_string(),
2844 crate::value::Value::String("alice".to_string()),
2845 ),
2846 ]),
2847 std::collections::HashMap::from([
2848 ("id".to_string(), crate::value::Value::I64(2)),
2849 (
2850 "name".to_string(),
2851 crate::value::Value::String("bob".to_string()),
2852 ),
2853 ]),
2854 std::collections::HashMap::from([
2855 ("id".to_string(), crate::value::Value::I64(3)),
2856 (
2857 "name".to_string(),
2858 crate::value::Value::String("carol".to_string()),
2859 ),
2860 ]),
2861 ];
2862 let mut conn = CursorMockConn::new(rows);
2863 let mut stream = conn.query_stream("SELECT id, name FROM users");
2864 let mut received: Vec<QueryStreamItem> = Vec::new();
2865 while let Some(item) = stream.next().await {
2866 received.push(item);
2867 }
2868 assert_eq!(received.len(), 3, "应收到 3 行");
2869 assert!(received.iter().all(|r| r.is_ok()), "所有项应为 Ok");
2870 drop(stream);
2871 assert_eq!(conn.call_count, 1, "默认实现应调用 query() 一次");
2872 }
2873
2874 #[tokio::test]
2876 async fn test_query_stream_default_empty_result() {
2877 use futures::StreamExt;
2878 let mut conn = CursorMockConn::new(Vec::new());
2879 let mut stream = conn.query_stream("SELECT * FROM empty_table");
2880 let mut count = 0;
2881 while let Some(_item) = stream.next().await {
2882 count += 1;
2883 }
2884 assert_eq!(count, 0, "空结果集应产生 0 项");
2885 }
2886
2887 #[tokio::test]
2889 async fn test_query_stream_default_error_propagation() {
2890 use futures::StreamExt;
2891 struct ErrorMockConn;
2893 impl Connection for ErrorMockConn {
2894 fn execute<'a>(
2895 &'a mut self,
2896 _sql: &'a str,
2897 ) -> Pin<Box<dyn Future<Output = Result<u64, crate::DbError>> + Send + 'a>>
2898 {
2899 Box::pin(async move { Ok(1) })
2900 }
2901 fn query<'a>(
2902 &'a mut self,
2903 _sql: &'a str,
2904 ) -> Pin<Box<dyn Future<Output = Result<QueryRows, crate::DbError>> + Send + 'a>>
2905 {
2906 Box::pin(async move { Err(crate::DbError::Internal("query failed".to_string())) })
2907 }
2908 fn begin_transaction<'a>(
2909 &'a mut self,
2910 ) -> Pin<Box<dyn Future<Output = Result<(), crate::DbError>> + Send + 'a>> {
2911 Box::pin(async move { Ok(()) })
2912 }
2913 fn commit<'a>(
2914 &'a mut self,
2915 ) -> Pin<Box<dyn Future<Output = Result<(), crate::DbError>> + Send + 'a>> {
2916 Box::pin(async move { Ok(()) })
2917 }
2918 fn rollback<'a>(
2919 &'a mut self,
2920 ) -> Pin<Box<dyn Future<Output = Result<(), crate::DbError>> + Send + 'a>> {
2921 Box::pin(async move { Ok(()) })
2922 }
2923 fn is_connected(&self) -> bool {
2924 true
2925 }
2926 fn ping<'a>(&'a mut self) -> Pin<Box<dyn Future<Output = bool> + Send + 'a>> {
2927 Box::pin(async move { true })
2928 }
2929 fn close<'a>(
2930 &'a mut self,
2931 ) -> Pin<Box<dyn Future<Output = Result<(), crate::DbError>> + Send + 'a>> {
2932 Box::pin(async move { Ok(()) })
2933 }
2934 }
2935 let mut conn = ErrorMockConn;
2936 let mut stream = conn.query_stream("SELECT * FROM bad_table");
2937 let item = stream.next().await;
2938 assert!(item.is_some(), "应产生一项");
2939 assert!(item.unwrap().is_err(), "该项应为 Err");
2940 }
2941
2942 #[tokio::test]
2944 async fn test_query_stream_override_yields_rows_one_by_one() {
2945 use futures::StreamExt;
2946 let rows = vec![
2947 crate::value::Value::I64(10),
2948 crate::value::Value::I64(20),
2949 crate::value::Value::I64(30),
2950 crate::value::Value::I64(40),
2951 crate::value::Value::I64(50),
2952 ];
2953 let mut conn = CursorOverrideMockConn::new(rows);
2954 let values: Vec<i64> = {
2955 let mut stream = conn.query_stream("SELECT v FROM seq");
2956 let mut vals: Vec<i64> = Vec::new();
2957 while let Some(Ok(row)) = stream.next().await {
2958 if let crate::value::Value::I64(v) = row.get("v").unwrap() {
2959 vals.push(*v);
2960 }
2961 }
2962 vals
2963 };
2964 assert_eq!(values, vec![10, 20, 30, 40, 50], "应按顺序收到全部 5 行");
2965 assert_eq!(conn.yielded, 5, "应逐行 yield 5 次(真游标覆盖)");
2966 }
2967
2968 #[tokio::test]
2970 async fn test_query_stream_override_early_drop() {
2971 use futures::StreamExt;
2972 let rows = vec![
2973 crate::value::Value::I64(1),
2974 crate::value::Value::I64(2),
2975 crate::value::Value::I64(3),
2976 ];
2977 let mut conn = CursorOverrideMockConn::new(rows);
2978 {
2979 let mut stream = conn.query_stream("SELECT v FROM seq");
2980 let first = stream.next().await;
2981 assert!(first.is_some(), "第一项应存在");
2982 drop(stream);
2984 }
2985 assert!(conn.is_connected(), "提前 drop 流后连接仍应可用");
2987 }
2988
2989 #[tokio::test]
2991 async fn test_pool_prewarm() -> Result<(), Box<dyn std::error::Error>> {
2992 use std::sync::atomic::AtomicU32;
2993
2994 let create_count = Arc::new(AtomicU32::new(0));
2996 let create_count_clone = create_count.clone();
2997
2998 struct CountingFactory {
2999 count: Arc<AtomicU32>,
3000 }
3001
3002 #[async_trait]
3003 impl ConnectionFactory for CountingFactory {
3004 async fn create(&self) -> Result<Box<dyn Connection>, crate::DbError> {
3005 self.count.fetch_add(1, Ordering::SeqCst);
3006 Ok(Box::new(MockConnection::new()))
3007 }
3008 }
3009
3010 let config = PoolConfigBuilder::new()
3012 .max_size(10)
3013 .min_idle(5)
3014 .prewarm(true)
3015 .build()?;
3016
3017 let factory = Arc::new(CountingFactory {
3018 count: create_count_clone,
3019 });
3020
3021 let pool = Pool::new(config, factory)?;
3022
3023 let status_before = pool.status().await;
3025 assert_eq!(status_before.idle, 0, "预热前 idle 应为 0");
3026
3027 pool.prewarm().await;
3029
3030 let status_after = pool.status().await;
3032 assert!(
3033 status_after.idle >= 5,
3034 "预热后 idle 应 >= 5,实际: {}",
3035 status_after.idle
3036 );
3037
3038 assert_eq!(
3040 create_count.load(Ordering::SeqCst),
3041 5,
3042 "工厂应被调用 5 次(min_idle)"
3043 );
3044
3045 Ok(())
3046 }
3047
3048 #[tokio::test]
3050 async fn test_pool_prewarm_failure_non_blocking() -> Result<(), Box<dyn std::error::Error>> {
3051 use std::sync::atomic::AtomicBool;
3052
3053 struct FailingFactory {
3054 failed: Arc<AtomicBool>,
3055 }
3056
3057 #[async_trait]
3058 impl ConnectionFactory for FailingFactory {
3059 async fn create(&self) -> Result<Box<dyn Connection>, crate::DbError> {
3060 self.failed.store(true, Ordering::SeqCst);
3061 Err(crate::DbError::Internal(
3063 "simulated connection failure".to_string(),
3064 ))
3065 }
3066 }
3067
3068 let failed = Arc::new(AtomicBool::new(false));
3069 let mut config = PoolConfigBuilder::new()
3070 .max_size(10)
3071 .min_idle(3)
3072 .prewarm(true)
3073 .build()?;
3074 config.connection_timeout = std::time::Duration::from_secs(1); let factory = Arc::new(FailingFactory {
3077 failed: failed.clone(),
3078 });
3079
3080 let pool = Pool::new(config, factory)?;
3082 pool.prewarm().await; assert!(failed.load(Ordering::SeqCst), "工厂应被调用且失败");
3086
3087 let status = pool.status().await;
3089 assert_eq!(status.max, 10, "池配置应正常");
3090
3091 Ok(())
3092 }
3093
3094 #[tokio::test]
3096 async fn test_pool_prewarm_disabled() -> Result<(), Box<dyn std::error::Error>> {
3097 use std::sync::atomic::AtomicU32;
3098
3099 let create_count = Arc::new(AtomicU32::new(0));
3100 let create_count_clone = create_count.clone();
3101
3102 struct CountingFactory {
3103 count: Arc<AtomicU32>,
3104 }
3105
3106 #[async_trait]
3107 impl ConnectionFactory for CountingFactory {
3108 async fn create(&self) -> Result<Box<dyn Connection>, crate::DbError> {
3109 self.count.fetch_add(1, Ordering::SeqCst);
3110 Ok(Box::new(MockConnection::new()))
3111 }
3112 }
3113
3114 let config = PoolConfigBuilder::new()
3116 .max_size(10)
3117 .min_idle(5)
3118 .prewarm(false) .build()?;
3120
3121 let factory = Arc::new(CountingFactory {
3122 count: create_count_clone,
3123 });
3124
3125 let pool = Pool::new(config, factory)?;
3126 pool.prewarm().await; assert_eq!(
3130 create_count.load(Ordering::SeqCst),
3131 0,
3132 "prewarm=false 时工厂不应被调用"
3133 );
3134
3135 let status = pool.status().await;
3136 assert_eq!(status.idle, 0, "idle 应为 0");
3137
3138 Ok(())
3139 }
3140
3141 #[tokio::test]
3143 async fn test_pool_new_async_with_prewarm() -> Result<(), Box<dyn std::error::Error>> {
3144 use std::sync::atomic::AtomicU32;
3145
3146 let create_count = Arc::new(AtomicU32::new(0));
3147 let create_count_clone = create_count.clone();
3148
3149 struct CountingFactory {
3150 count: Arc<AtomicU32>,
3151 }
3152
3153 #[async_trait]
3154 impl ConnectionFactory for CountingFactory {
3155 async fn create(&self) -> Result<Box<dyn Connection>, crate::DbError> {
3156 self.count.fetch_add(1, Ordering::SeqCst);
3157 Ok(Box::new(MockConnection::new()))
3158 }
3159 }
3160
3161 let config = PoolConfigBuilder::new()
3162 .max_size(10)
3163 .min_idle(5)
3164 .prewarm(true)
3165 .build()?;
3166
3167 let factory = Arc::new(CountingFactory {
3168 count: create_count_clone,
3169 });
3170
3171 let pool = Pool::new_async(config, factory).await?;
3172
3173 let status = pool.status().await;
3174 assert!(
3175 status.idle >= 5,
3176 "new_async prewarm=true 后 idle 应 >= 5,实际: {}",
3177 status.idle
3178 );
3179 assert_eq!(create_count.load(Ordering::SeqCst), 5, "工厂应被调用 5 次");
3180
3181 Ok(())
3182 }
3183
3184 #[tokio::test]
3186 async fn test_pool_new_async_without_prewarm() -> Result<(), Box<dyn std::error::Error>> {
3187 use std::sync::atomic::AtomicU32;
3188
3189 let create_count = Arc::new(AtomicU32::new(0));
3190 let create_count_clone = create_count.clone();
3191
3192 struct CountingFactory {
3193 count: Arc<AtomicU32>,
3194 }
3195
3196 #[async_trait]
3197 impl ConnectionFactory for CountingFactory {
3198 async fn create(&self) -> Result<Box<dyn Connection>, crate::DbError> {
3199 self.count.fetch_add(1, Ordering::SeqCst);
3200 Ok(Box::new(MockConnection::new()))
3201 }
3202 }
3203
3204 let config = PoolConfigBuilder::new()
3205 .max_size(10)
3206 .min_idle(5)
3207 .prewarm(false)
3208 .build()?;
3209
3210 let factory = Arc::new(CountingFactory {
3211 count: create_count_clone,
3212 });
3213
3214 let pool = Pool::new_async(config, factory).await?;
3215
3216 let status = pool.status().await;
3217 assert_eq!(status.idle, 0, "prewarm=false 时 idle 应为 0");
3218 assert_eq!(create_count.load(Ordering::SeqCst), 0, "工厂不应被调用");
3219
3220 Ok(())
3221 }
3222
3223 #[tokio::test]
3225 async fn test_pool_new_async_failure_non_blocking() -> Result<(), Box<dyn std::error::Error>> {
3226 struct FailingFactory;
3227
3228 #[async_trait]
3229 impl ConnectionFactory for FailingFactory {
3230 async fn create(&self) -> Result<Box<dyn Connection>, crate::DbError> {
3231 Err(crate::DbError::Internal("simulated failure".to_string()))
3232 }
3233 }
3234
3235 let mut config = PoolConfigBuilder::new()
3236 .max_size(10)
3237 .min_idle(3)
3238 .prewarm(true)
3239 .build()?;
3240 config.connection_timeout = std::time::Duration::from_secs(1);
3241
3242 let pool = Pool::new_async(config, Arc::new(FailingFactory)).await?;
3243
3244 let status = pool.status().await;
3245 assert_eq!(status.max, 10, "池配置应正常");
3246
3247 Ok(())
3248 }
3249
3250 #[cfg(feature = "auto-prewarm")]
3252 #[tokio::test]
3253 async fn test_pool_progressive_prewarm() -> Result<(), Box<dyn std::error::Error>> {
3254 use std::sync::atomic::AtomicU32;
3255
3256 let create_count = Arc::new(AtomicU32::new(0));
3257 let create_count_clone = create_count.clone();
3258
3259 struct CountingFactory {
3260 count: Arc<AtomicU32>,
3261 }
3262
3263 #[async_trait]
3264 impl ConnectionFactory for CountingFactory {
3265 async fn create(&self) -> Result<Box<dyn Connection>, crate::DbError> {
3266 self.count.fetch_add(1, Ordering::SeqCst);
3267 Ok(Box::new(MockConnection::new()))
3268 }
3269 }
3270
3271 let config = PoolConfigBuilder::new()
3272 .max_size(20)
3273 .min_idle(6)
3274 .prewarm(true)
3275 .build()?;
3276
3277 let factory = Arc::new(CountingFactory {
3278 count: create_count_clone,
3279 });
3280
3281 let pool = Pool::new(config, factory)?;
3282
3283 let progress = crate::prewarm::PrewarmProgress::new(6);
3284 pool.progressive_prewarm(
3285 2,
3286 std::time::Duration::from_millis(5),
3287 std::time::Duration::from_secs(10),
3288 &progress,
3289 )
3290 .await;
3291
3292 let snap = progress.snapshot();
3293 assert!(
3294 snap.warmed >= 6,
3295 "progressive_prewarm 后 warmed 应 >= 6,实际: {}",
3296 snap.warmed
3297 );
3298 assert!(snap.is_completed, "应标记完成");
3299 assert_eq!(create_count.load(Ordering::SeqCst), 6, "工厂应被调用 6 次");
3300
3301 let status = pool.status().await;
3302 assert!(status.idle >= 6, "池中 idle 应 >= 6");
3303
3304 Ok(())
3305 }
3306
3307 #[cfg(feature = "auto-prewarm")]
3309 #[tokio::test]
3310 async fn test_pool_progressive_prewarm_timeout_zero() -> Result<(), Box<dyn std::error::Error>>
3311 {
3312 use std::sync::atomic::AtomicU32;
3313
3314 let create_count = Arc::new(AtomicU32::new(0));
3315 let create_count_clone = create_count.clone();
3316
3317 struct CountingFactory {
3318 count: Arc<AtomicU32>,
3319 }
3320
3321 #[async_trait]
3322 impl ConnectionFactory for CountingFactory {
3323 async fn create(&self) -> Result<Box<dyn Connection>, crate::DbError> {
3324 self.count.fetch_add(1, Ordering::SeqCst);
3325 Ok(Box::new(MockConnection::new()))
3326 }
3327 }
3328
3329 let config = PoolConfigBuilder::new()
3330 .max_size(20)
3331 .min_idle(10)
3332 .prewarm(true)
3333 .build()?;
3334
3335 let factory = Arc::new(CountingFactory {
3336 count: create_count_clone,
3337 });
3338
3339 let pool = Pool::new(config, factory)?;
3340
3341 let progress = crate::prewarm::PrewarmProgress::new(10);
3342 pool.progressive_prewarm(
3343 2,
3344 std::time::Duration::from_millis(5),
3345 std::time::Duration::ZERO,
3346 &progress,
3347 )
3348 .await;
3349
3350 let snap = progress.snapshot();
3351 assert!(snap.is_completed, "应标记完成");
3352 assert!(
3353 snap.warmed <= 2,
3354 "total_timeout=0 时最多建一批(batch_size=2),实际: {}",
3355 snap.warmed
3356 );
3357
3358 Ok(())
3359 }
3360
3361 #[cfg(feature = "auto-prewarm")]
3363 #[tokio::test]
3364 async fn test_pool_progressive_prewarm_disabled() -> Result<(), Box<dyn std::error::Error>> {
3365 use std::sync::atomic::AtomicU32;
3366
3367 let create_count = Arc::new(AtomicU32::new(0));
3368 let create_count_clone = create_count.clone();
3369
3370 struct CountingFactory {
3371 count: Arc<AtomicU32>,
3372 }
3373
3374 #[async_trait]
3375 impl ConnectionFactory for CountingFactory {
3376 async fn create(&self) -> Result<Box<dyn Connection>, crate::DbError> {
3377 self.count.fetch_add(1, Ordering::SeqCst);
3378 Ok(Box::new(MockConnection::new()))
3379 }
3380 }
3381
3382 let config = PoolConfigBuilder::new()
3383 .max_size(20)
3384 .min_idle(10)
3385 .prewarm(false)
3386 .build()?;
3387
3388 let factory = Arc::new(CountingFactory {
3389 count: create_count_clone,
3390 });
3391
3392 let pool = Pool::new(config, factory)?;
3393
3394 let progress = crate::prewarm::PrewarmProgress::new(10);
3395 pool.progressive_prewarm(
3396 2,
3397 std::time::Duration::from_millis(5),
3398 std::time::Duration::from_secs(10),
3399 &progress,
3400 )
3401 .await;
3402
3403 let snap = progress.snapshot();
3404 assert!(snap.is_completed, "应标记完成");
3405 assert_eq!(snap.warmed, 0, "prewarm=false 时不应建连");
3406 assert_eq!(create_count.load(Ordering::SeqCst), 0, "工厂不应被调用");
3407
3408 Ok(())
3409 }
3410
3411 #[cfg(feature = "auto-prewarm")]
3413 #[tokio::test]
3414 async fn test_pool_progressive_prewarm_failure_non_blocking(
3415 ) -> Result<(), Box<dyn std::error::Error>> {
3416 struct FailingFactory;
3417
3418 #[async_trait]
3419 impl ConnectionFactory for FailingFactory {
3420 async fn create(&self) -> Result<Box<dyn Connection>, crate::DbError> {
3421 Err(crate::DbError::Internal("simulated failure".to_string()))
3422 }
3423 }
3424
3425 let mut config = PoolConfigBuilder::new()
3426 .max_size(20)
3427 .min_idle(5)
3428 .prewarm(true)
3429 .build()?;
3430 config.connection_timeout = std::time::Duration::from_secs(1);
3431
3432 let pool = Pool::new(config, Arc::new(FailingFactory))?;
3433
3434 let progress = crate::prewarm::PrewarmProgress::new(5);
3435 pool.progressive_prewarm(
3436 2,
3437 std::time::Duration::from_millis(5),
3438 std::time::Duration::from_secs(5),
3439 &progress,
3440 )
3441 .await;
3442
3443 let snap = progress.snapshot();
3444 assert!(snap.is_completed, "应标记完成");
3445 assert_eq!(snap.warmed, 0, "全部失败时 warmed=0");
3446 assert!(snap.failed > 0, "应有失败记录");
3447
3448 Ok(())
3449 }
3450
3451 #[tokio::test]
3453 async fn test_pool_metrics_acquire_release() -> Result<(), Box<dyn std::error::Error>> {
3454 let config = PoolConfigBuilder::new().max_size(10).build()?;
3455 let pool = Pool::new(config, Arc::new(MockConnectionFactory))?;
3456
3457 let metrics = pool.pool_metrics();
3458 assert_eq!(metrics.acquire_count, 0);
3459 assert_eq!(metrics.release_count, 0);
3460 assert_eq!(metrics.connection_created_count, 0);
3461
3462 let conn = pool.acquire().await?;
3463 let metrics = pool.pool_metrics();
3464 assert_eq!(metrics.acquire_count, 1);
3465 assert_eq!(metrics.connection_created_count, 1);
3466 assert_eq!(metrics.acquire_failed_count, 0);
3467
3468 pool.release(conn).await;
3469 let metrics = pool.pool_metrics();
3470 assert_eq!(metrics.release_count, 1);
3471 assert_eq!(metrics.connection_closed_count, 0);
3473
3474 Ok(())
3475 }
3476
3477 #[tokio::test]
3479 async fn test_pool_metrics_acquire_failed() -> Result<(), Box<dyn std::error::Error>> {
3480 struct FailingFactory;
3481
3482 #[async_trait]
3483 impl ConnectionFactory for FailingFactory {
3484 async fn create(&self) -> Result<Box<dyn Connection>, crate::DbError> {
3485 Err(crate::DbError::Internal("simulated failure".to_string()))
3486 }
3487 }
3488
3489 let config = PoolConfigBuilder::new().max_size(10).build()?;
3490 let pool = Pool::new(config, Arc::new(FailingFactory))?;
3491
3492 let result = pool.acquire().await;
3493 assert!(result.is_err());
3494
3495 let metrics = pool.pool_metrics();
3496 assert_eq!(metrics.acquire_failed_count, 1);
3497 assert_eq!(metrics.acquire_count, 0);
3498
3499 Ok(())
3500 }
3501
3502 #[tokio::test]
3504 async fn test_pool_metrics_connection_closed() -> Result<(), Box<dyn std::error::Error>> {
3505 let config = PoolConfigBuilder::new().max_size(10).build()?;
3506 let pool = Pool::new(config, Arc::new(MockConnectionFactory))?;
3507
3508 let conn = pool.acquire().await?;
3509 pool.release(conn).await;
3510
3511 let status = pool.status().await;
3512 assert_eq!(status.idle, 1);
3513
3514 pool.close_all().await;
3515
3516 let metrics = pool.pool_metrics();
3517 assert_eq!(metrics.connection_closed_count, 1);
3518 assert_eq!(metrics.connection_created_count, 1);
3519
3520 Ok(())
3521 }
3522
3523 #[test]
3525 fn test_pool_metrics_average_wait_time() {
3526 let metrics = PoolMetrics {
3527 acquire_count: 4,
3528 acquire_failed_count: 1,
3529 acquire_wait_time: Duration::from_millis(200),
3530 release_count: 4,
3531 connection_created_count: 2,
3532 connection_closed_count: 0,
3533 };
3534 assert_eq!(
3535 metrics.average_acquire_wait_time(),
3536 Duration::from_millis(50)
3537 );
3538
3539 let empty = PoolMetrics::default();
3541 assert_eq!(empty.average_acquire_wait_time(), Duration::ZERO);
3542 }
3543
3544 #[tokio::test]
3545 async fn test_shutdown_with_timeout_fast_return_when_empty() {
3546 let factory = Arc::new(MockConnectionFactory) as Arc<dyn ConnectionFactory>;
3547 let config = PoolConfigBuilder::new().max_size(4).build().unwrap();
3548 let pool = Pool::new(config, factory).unwrap();
3549 pool.shutdown_with_timeout(Duration::from_secs(1)).await;
3550 assert!(pool.closed.load(Ordering::SeqCst));
3551 assert_eq!(pool.total_count.load(Ordering::SeqCst), 0);
3552 }
3553
3554 #[tokio::test]
3555 async fn test_shutdown_delegates_to_shutdown_with_timeout() {
3556 let factory = Arc::new(MockConnectionFactory) as Arc<dyn ConnectionFactory>;
3557 let config = PoolConfigBuilder::new().max_size(4).build().unwrap();
3558 let pool = Pool::new(config, factory).unwrap();
3559 pool.shutdown().await;
3560 assert!(pool.closed.load(Ordering::SeqCst));
3561 }
3562
3563 #[tokio::test]
3564 async fn test_shutdown_with_timeout_idempotent() {
3565 let factory = Arc::new(MockConnectionFactory) as Arc<dyn ConnectionFactory>;
3566 let config = PoolConfigBuilder::new().max_size(4).build().unwrap();
3567 let pool = Pool::new(config, factory).unwrap();
3568 pool.shutdown_with_timeout(Duration::from_secs(1)).await;
3569 let count_after_first = pool.total_count.load(Ordering::SeqCst);
3570 pool.shutdown_with_timeout(Duration::from_secs(1)).await;
3571 let count_after_second = pool.total_count.load(Ordering::SeqCst);
3572 assert_eq!(count_after_first, count_after_second);
3573 }
3574
3575 #[tokio::test]
3576 async fn test_shutdown_with_timeout_rejects_new_acquire() {
3577 let factory = Arc::new(MockConnectionFactory) as Arc<dyn ConnectionFactory>;
3578 let config = PoolConfigBuilder::new().max_size(4).build().unwrap();
3579 let pool = Pool::new(config, factory).unwrap();
3580 pool.shutdown_with_timeout(Duration::from_secs(1)).await;
3581 let result = pool.acquire().await;
3582 assert!(result.is_err());
3583 }
3584}
3585
3586#[cfg(all(test, feature = "prod-pool-tuning"))]
3587mod pool_prod_tests {
3588 use super::*;
3589
3590 struct MockFactory;
3591
3592 #[async_trait]
3593 impl ConnectionFactory for MockFactory {
3594 async fn create(&self) -> Result<Box<dyn Connection>, crate::DbError> {
3595 Ok(Box::new(MockConn))
3596 }
3597 }
3598
3599 struct MockConn;
3600
3601 impl Connection for MockConn {
3602 fn execute<'a>(
3603 &'a mut self,
3604 _sql: &'a str,
3605 ) -> Pin<Box<dyn Future<Output = Result<u64, crate::DbError>> + Send + 'a>> {
3606 Box::pin(async move { Ok(1) })
3607 }
3608 fn query<'a>(
3609 &'a mut self,
3610 _sql: &'a str,
3611 ) -> Pin<
3612 Box<
3613 dyn Future<
3614 Output = Result<
3615 Vec<std::collections::HashMap<String, crate::value::Value>>,
3616 crate::DbError,
3617 >,
3618 > + Send
3619 + 'a,
3620 >,
3621 > {
3622 Box::pin(async move { Ok(vec![]) })
3623 }
3624 fn begin_transaction<'a>(
3625 &'a mut self,
3626 ) -> Pin<Box<dyn Future<Output = Result<(), crate::DbError>> + Send + 'a>> {
3627 Box::pin(async move { Ok(()) })
3628 }
3629 fn commit<'a>(
3630 &'a mut self,
3631 ) -> Pin<Box<dyn Future<Output = Result<(), crate::DbError>> + Send + 'a>> {
3632 Box::pin(async move { Ok(()) })
3633 }
3634 fn rollback<'a>(
3635 &'a mut self,
3636 ) -> Pin<Box<dyn Future<Output = Result<(), crate::DbError>> + Send + 'a>> {
3637 Box::pin(async move { Ok(()) })
3638 }
3639 fn is_connected(&self) -> bool {
3640 true
3641 }
3642 fn ping<'a>(&'a mut self) -> Pin<Box<dyn Future<Output = bool> + Send + 'a>> {
3643 Box::pin(async move { true })
3644 }
3645 fn close<'a>(
3646 &'a mut self,
3647 ) -> Pin<Box<dyn Future<Output = Result<(), crate::DbError>> + Send + 'a>> {
3648 Box::pin(async move { Ok(()) })
3649 }
3650 }
3651
3652 #[test]
3653 fn test_pool_prod_config_validate_ok() {
3654 let config = PoolProdConfig::new(
3655 50,
3656 Duration::from_secs(10),
3657 Duration::from_secs(600),
3658 Duration::from_secs(5),
3659 Duration::from_secs(30),
3660 5,
3661 true,
3662 );
3663 assert!(config.validate().is_ok());
3664 }
3665
3666 #[test]
3667 fn test_pool_prod_config_max_size_zero_rejected() {
3668 let config = PoolProdConfig::default();
3669 let mut c = config;
3670 c.max_size = 0;
3671 let err = c.validate().unwrap_err();
3672 assert!(err.to_string().contains("max_size must be positive"));
3673 }
3674
3675 #[test]
3676 fn test_pool_prod_config_min_idle_exceeds_max_size() {
3677 let config = PoolProdConfig::new(
3678 10,
3679 Duration::from_secs(10),
3680 Duration::from_secs(600),
3681 Duration::from_secs(5),
3682 Duration::from_secs(30),
3683 20,
3684 false,
3685 );
3686 let err = config.validate().unwrap_err();
3687 assert!(err.to_string().contains("min_idle cannot exceed max_size"));
3688 }
3689
3690 #[test]
3691 fn test_pool_prod_config_to_pool_config() {
3692 let config = PoolProdConfig::new(
3693 50,
3694 Duration::from_secs(10),
3695 Duration::from_secs(600),
3696 Duration::from_secs(5),
3697 Duration::from_secs(30),
3698 5,
3699 true,
3700 );
3701 let pool_config = config.to_pool_config();
3702 assert_eq!(pool_config.max_size, 50);
3703 assert_eq!(pool_config.min_idle, 5);
3704 assert_eq!(pool_config.acquire_timeout, Duration::from_secs(10));
3705 assert!(pool_config.prewarm);
3706 }
3707
3708 #[tokio::test]
3709 async fn test_pool_prod_config_runtime_resize() {
3710 let factory = Arc::new(MockFactory) as Arc<dyn ConnectionFactory>;
3711 let config = PoolProdConfig::default();
3712 let pool = Pool::new(config.to_pool_config(), factory).unwrap();
3713 assert_eq!(pool.max_size(), 100);
3714 pool.resize(50);
3715 assert_eq!(pool.max_size(), 50);
3716 }
3717}
3718
3719#[cfg(all(test, feature = "prod-leak-detection"))]
3720mod leak_prod_tests {
3721 use super::*;
3722
3723 #[test]
3724 fn test_leak_config_default() {
3725 let config = LeakDetectionConfig::default();
3726 assert!(!config.enabled);
3727 assert_eq!(config.interval, Duration::from_secs(60));
3728 assert_eq!(config.threshold, 5);
3729 }
3730
3731 #[test]
3732 fn test_leak_config_validate_ok() {
3733 let config =
3734 LeakDetectionConfig::new(true, Duration::from_secs(30), 10, Duration::from_secs(60));
3735 assert!(config.validate().is_ok());
3736 }
3737
3738 #[test]
3739 fn test_leak_config_interval_zero_rejected() {
3740 let config = LeakDetectionConfig::new(true, Duration::ZERO, 10, Duration::from_secs(60));
3741 assert!(config.validate().is_err());
3742 }
3743
3744 #[test]
3745 fn test_leak_report_empty() {
3746 let report = LeakReport::empty();
3747 assert_eq!(report.borrowed_count, 0);
3748 assert!(report.suspected_leaks.is_empty());
3749 }
3750}