1use async_trait::async_trait;
6use crossbeam_queue::ArrayQueue;
7use futures::StreamExt;
8#[cfg(feature = "circuit-breaker")]
12use parking_lot::Mutex as PlMutex;
13#[cfg(feature = "rate-limit")]
14use parking_lot::RwLock as PlRwLock;
15use std::future::Future;
16use std::ops::{Deref, DerefMut};
17use std::pin::Pin;
18use std::sync::atomic::{AtomicBool, AtomicU32, AtomicU64, Ordering};
19use std::sync::Arc;
20use std::time::{Duration, Instant};
21use tokio::sync::Notify;
22
23#[cfg(feature = "circuit-breaker")]
27use crate::circuit_breaker::{CircuitBreaker, CircuitState, DefaultCircuitBreaker};
28use crate::error::PoolError;
29#[cfg(feature = "rate-limit")]
30use crate::rate_limiter::RateLimiter;
31
32pub type QueryRows = Vec<std::collections::HashMap<String, crate::value::Value>>;
34
35pub type QueryStreamItem =
37 Result<std::collections::HashMap<String, crate::value::Value>, crate::DbError>;
38
39pub trait Connection: Send + Sync {
46 fn execute<'a>(
48 &'a mut self,
49 sql: &'a str,
50 ) -> Pin<Box<dyn Future<Output = Result<u64, crate::DbError>> + Send + 'a>>;
51 fn query<'a>(
53 &'a mut self,
54 sql: &'a str,
55 ) -> Pin<Box<dyn Future<Output = Result<QueryRows, crate::DbError>> + Send + 'a>>;
56 fn begin_transaction<'a>(
58 &'a mut self,
59 ) -> Pin<Box<dyn Future<Output = Result<(), crate::DbError>> + Send + 'a>>;
60 fn commit<'a>(
62 &'a mut self,
63 ) -> Pin<Box<dyn Future<Output = Result<(), crate::DbError>> + Send + 'a>>;
64 fn rollback<'a>(
66 &'a mut self,
67 ) -> Pin<Box<dyn Future<Output = Result<(), crate::DbError>> + Send + 'a>>;
68 fn is_connected(&self) -> bool;
70 fn ping<'a>(&'a mut self) -> Pin<Box<dyn Future<Output = bool> + Send + 'a>>;
72 fn close<'a>(
74 &'a mut self,
75 ) -> Pin<Box<dyn Future<Output = Result<(), crate::DbError>> + Send + 'a>>;
76
77 fn execute_with_params<'a>(
83 &'a mut self,
84 sql: &'a str,
85 params: &'a [crate::value::Value],
86 ) -> Pin<Box<dyn Future<Output = Result<u64, crate::DbError>> + Send + 'a>> {
87 let _ = (sql, params);
88 Box::pin(async move {
89 Err(crate::DbError::Internal(
90 "execute_with_params not implemented for this adapter".to_string(),
91 ))
92 })
93 }
94
95 fn query_with_params<'a>(
101 &'a mut self,
102 sql: &'a str,
103 params: &'a [crate::value::Value],
104 ) -> Pin<Box<dyn Future<Output = Result<QueryRows, crate::DbError>> + Send + 'a>> {
105 let _ = (sql, params);
106 Box::pin(async move {
107 Err(crate::DbError::Internal(
108 "query_with_params not implemented for this adapter".to_string(),
109 ))
110 })
111 }
112
113 fn query_values<'a>(
118 &'a mut self,
119 sql: &'a str,
120 ) -> Pin<Box<dyn Future<Output = Result<crate::value::QueryValues, crate::DbError>> + Send + 'a>>
121 {
122 let _ = sql;
123 Box::pin(async move {
124 Err(crate::DbError::Internal(
125 "query_values not implemented for this adapter".to_string(),
126 ))
127 })
128 }
129
130 fn query_values_with_params<'a>(
134 &'a mut self,
135 sql: &'a str,
136 params: &'a [crate::value::Value],
137 ) -> Pin<Box<dyn Future<Output = Result<crate::value::QueryValues, crate::DbError>> + Send + 'a>>
138 {
139 let _ = (sql, params);
140 Box::pin(async move {
141 Err(crate::DbError::Internal(
142 "query_values_with_params not implemented for this adapter".to_string(),
143 ))
144 })
145 }
146
147 fn query_stream<'a>(
159 &'a mut self,
160 sql: &'a str,
161 ) -> Pin<Box<dyn futures::Stream<Item = QueryStreamItem> + Send + 'a>> {
162 let sql_owned = sql.to_string();
164 let stream = futures::stream::once(async move { self.query(&sql_owned).await })
166 .map(|result| {
168 let items: Vec<QueryStreamItem> = match result {
169 Ok(rows) => rows.into_iter().map(Ok).collect(),
170 Err(e) => vec![Err(e)],
171 };
172 futures::stream::iter(items)
173 })
174 .flatten();
175 Box::pin(stream)
176 }
177
178 fn query_stream_cursor<'a>(
189 &'a mut self,
190 sql: &'a str,
191 _batch_size: usize,
192 ) -> Pin<Box<dyn futures::Stream<Item = QueryStreamItem> + Send + 'a>> {
193 self.query_stream(sql)
194 }
195
196 fn execute_batch<'a>(
201 &'a mut self,
202 sqls: &'a [String],
203 ) -> Pin<Box<dyn Future<Output = Result<u64, crate::DbError>> + Send + 'a>> {
204 Box::pin(async move {
205 let mut total = 0u64;
206 for sql in sqls {
207 total += self.execute(sql).await?;
208 }
209 Ok(total)
210 })
211 }
212
213 fn execute_batch_params<'a>(
218 &'a mut self,
219 sql: &'a str,
220 params_batch: &'a [Vec<crate::value::Value>],
221 ) -> Pin<Box<dyn Future<Output = Result<u64, crate::DbError>> + Send + 'a>> {
222 Box::pin(async move {
223 let mut total = 0u64;
224 for params in params_batch {
225 total += self.execute_with_params(sql, params).await?;
226 }
227 Ok(total)
228 })
229 }
230}
231
232pub struct PooledConnection {
240 conn: Box<dyn Connection>,
241 created_at: Instant,
242 last_used_at: Instant,
243 pool: Option<Pool>,
244}
245
246impl PooledConnection {
247 fn new(conn: Box<dyn Connection>, pool: Pool) -> Self {
248 let now = Instant::now();
249 Self {
250 conn,
251 created_at: now,
252 last_used_at: now,
253 pool: Some(pool),
254 }
255 }
256
257 fn is_expired(&self, max_lifetime: Duration) -> bool {
258 self.created_at.elapsed() >= max_lifetime
259 }
260
261 fn is_idle_too_long(&self, idle_timeout: Duration) -> bool {
262 self.last_used_at.elapsed() >= idle_timeout
263 }
264
265 pub fn created_at(&self) -> Instant {
267 self.created_at
268 }
269
270 pub fn into_inner(mut self) -> Box<dyn Connection> {
275 self.pool = None; std::mem::replace(&mut self.conn, Box::new(ClosedConnection))
279 }
280}
281
282impl Drop for PooledConnection {
294 fn drop(&mut self) {
295 if let Some(pool) = self.pool.take() {
296 let conn = std::mem::replace(&mut self.conn, Box::new(ClosedConnection));
298 let pooled = PooledConnection {
299 conn,
300 created_at: self.created_at,
301 last_used_at: self.last_used_at,
302 pool: None,
303 };
304 if let Ok(handle) = tokio::runtime::Handle::try_current() {
306 handle.spawn(async move {
307 pool.release(pooled).await;
308 });
309 } else {
310 drop(pooled);
314 pool.total_count.fetch_sub(1, Ordering::SeqCst);
315 }
316 }
317 }
318}
319
320struct ClosedConnection;
324
325impl Connection for ClosedConnection {
326 fn execute<'a>(
327 &'a mut self,
328 _sql: &'a str,
329 ) -> Pin<Box<dyn Future<Output = Result<u64, crate::DbError>> + Send + 'a>> {
330 Box::pin(async {
331 Err(crate::DbError::ConnectionError(
332 "connection already returned to pool".to_string(),
333 ))
334 })
335 }
336
337 fn query<'a>(
338 &'a mut self,
339 _sql: &'a str,
340 ) -> Pin<Box<dyn Future<Output = Result<QueryRows, crate::DbError>> + Send + 'a>> {
341 Box::pin(async {
342 Err(crate::DbError::ConnectionError(
343 "connection already returned to pool".to_string(),
344 ))
345 })
346 }
347
348 fn begin_transaction<'a>(
349 &'a mut self,
350 ) -> Pin<Box<dyn Future<Output = Result<(), crate::DbError>> + Send + 'a>> {
351 Box::pin(async {
352 Err(crate::DbError::ConnectionError(
353 "connection already returned to pool".to_string(),
354 ))
355 })
356 }
357
358 fn commit<'a>(
359 &'a mut self,
360 ) -> Pin<Box<dyn Future<Output = Result<(), crate::DbError>> + Send + 'a>> {
361 Box::pin(async { Ok(()) })
362 }
363
364 fn rollback<'a>(
365 &'a mut self,
366 ) -> Pin<Box<dyn Future<Output = Result<(), crate::DbError>> + Send + 'a>> {
367 Box::pin(async { Ok(()) })
368 }
369
370 fn is_connected(&self) -> bool {
371 false
372 }
373
374 fn ping<'a>(&'a mut self) -> Pin<Box<dyn Future<Output = bool> + Send + 'a>> {
375 Box::pin(async { false })
376 }
377
378 fn close<'a>(
379 &'a mut self,
380 ) -> Pin<Box<dyn Future<Output = Result<(), crate::DbError>> + Send + 'a>> {
381 Box::pin(async { Ok(()) })
382 }
383}
384
385impl Deref for PooledConnection {
386 type Target = dyn Connection;
387
388 fn deref(&self) -> &Self::Target {
389 self.conn.as_ref()
390 }
391}
392
393impl DerefMut for PooledConnection {
394 fn deref_mut(&mut self) -> &mut Self::Target {
395 self.conn.as_mut()
396 }
397}
398
399#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
401pub enum TlsVersion {
402 #[default]
404 Tls12,
405 Tls13,
407}
408
409#[derive(Debug, Clone, Default)]
411pub struct TlsConfig {
412 pub enabled: bool,
414 pub ca_cert_path: Option<String>,
416 pub client_cert_path: Option<String>,
418 pub client_key_path: Option<String>,
420 pub min_version: TlsVersion,
422}
423
424#[derive(Debug, Clone)]
426pub enum PoolEvent {
427 ConnectionCreated,
429 ConnectionClosed,
431 ConnectionAcquired,
433 ConnectionReleased,
435 AcquireTimeout,
437}
438
439pub type PoolEventCallback = Arc<dyn Fn(PoolEvent) + Send + Sync>;
441
442pub struct PoolConfig {
444 pub max_size: u32,
446 pub min_idle: u32,
448 pub acquire_timeout: Duration,
450 pub idle_timeout: Duration,
452 pub max_lifetime: Duration,
454 pub connection_timeout: Duration,
456 pub tls: Option<TlsConfig>,
458 pub query_timeout: Option<Duration>,
460 pub max_rows: Option<usize>,
462 pub memory_limit: Option<usize>,
464 pub on_event: Option<PoolEventCallback>,
466 pub test_before_acquire: bool,
474 pub prewarm: bool,
486}
487
488impl Default for PoolConfig {
489 fn default() -> Self {
490 Self {
491 max_size: 100,
492 min_idle: 0,
493 acquire_timeout: Duration::from_secs(30),
494 idle_timeout: Duration::from_secs(600),
495 max_lifetime: Duration::from_secs(1800),
496 connection_timeout: Duration::from_secs(10),
497 tls: None,
498 query_timeout: Some(Duration::from_secs(30)),
499 max_rows: None,
500 memory_limit: None,
501 on_event: None,
502 test_before_acquire: false,
503 prewarm: false,
504 }
505 }
506}
507
508impl Clone for PoolConfig {
509 fn clone(&self) -> Self {
510 Self {
511 max_size: self.max_size,
512 min_idle: self.min_idle,
513 acquire_timeout: self.acquire_timeout,
514 idle_timeout: self.idle_timeout,
515 max_lifetime: self.max_lifetime,
516 connection_timeout: self.connection_timeout,
517 tls: self.tls.clone(),
518 query_timeout: self.query_timeout,
519 max_rows: self.max_rows,
520 memory_limit: self.memory_limit,
521 on_event: self.on_event.clone(),
522 test_before_acquire: self.test_before_acquire,
523 prewarm: self.prewarm,
524 }
525 }
526}
527
528impl PoolConfig {
529 pub fn validate(&self) -> Result<(), PoolError> {
531 if self.max_size == 0 {
532 return Err(PoolError::InvalidConfig("max_size cannot be 0".to_string()));
533 }
534 if self.min_idle > self.max_size {
535 return Err(PoolError::InvalidConfig(
536 "min_idle cannot exceed max_size".to_string(),
537 ));
538 }
539 const MAX_DURATION_SECS: u64 = u32::MAX as u64; for (name, dur) in [
545 ("acquire_timeout", self.acquire_timeout),
546 ("idle_timeout", self.idle_timeout),
547 ("max_lifetime", self.max_lifetime),
548 ("connection_timeout", self.connection_timeout),
549 ] {
550 if dur.as_secs() > MAX_DURATION_SECS {
551 return Err(PoolError::InvalidConfig(format!(
552 "{name} ({:?}) exceeds maximum allowed duration ({} seconds)",
553 dur, MAX_DURATION_SECS
554 )));
555 }
556 }
557 Ok(())
558 }
559
560 #[must_use]
562 pub fn with_prewarm(mut self, prewarm: bool) -> Self {
563 self.prewarm = prewarm;
564 self
565 }
566}
567
568pub struct PoolStatus {
570 pub idle: u32,
572 pub active: u32,
574 pub max: u32,
576 pub min: u32,
578 pub waiters: u32,
580}
581
582impl std::fmt::Debug for PoolStatus {
583 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
584 f.debug_struct("PoolStatus")
585 .field("idle", &self.idle)
586 .field("active", &self.active)
587 .field("max", &self.max)
588 .field("min", &self.min)
589 .field("waiters", &self.waiters)
590 .finish()
591 }
592}
593
594#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
600pub struct PoolMetrics {
601 pub acquire_count: u64,
603 pub acquire_failed_count: u64,
605 pub acquire_wait_time: Duration,
607 pub release_count: u64,
609 pub connection_created_count: u64,
611 pub connection_closed_count: u64,
613}
614
615impl PoolMetrics {
616 #[must_use]
618 pub fn average_acquire_wait_time(&self) -> Duration {
619 if self.acquire_count == 0 {
620 Duration::ZERO
621 } else {
622 self.acquire_wait_time / self.acquire_count as u32
623 }
624 }
625}
626
627pub struct PoolConfigBuilder {
629 config: PoolConfig,
630}
631
632impl PoolConfigBuilder {
633 pub fn new() -> Self {
635 Self {
636 config: PoolConfig::default(),
637 }
638 }
639
640 pub fn max_size(mut self, size: u32) -> Self {
642 self.config.max_size = size;
643 self
644 }
645
646 pub fn min_idle(mut self, count: u32) -> Self {
648 self.config.min_idle = count;
649 self
650 }
651
652 pub fn acquire_timeout(mut self, timeout_secs: u64) -> Self {
654 self.config.acquire_timeout = Duration::from_secs(timeout_secs);
655 self
656 }
657
658 pub fn idle_timeout(mut self, timeout_secs: u64) -> Self {
660 self.config.idle_timeout = Duration::from_secs(timeout_secs);
661 self
662 }
663
664 pub fn max_lifetime(mut self, lifetime_secs: u64) -> Self {
666 self.config.max_lifetime = Duration::from_secs(lifetime_secs);
667 self
668 }
669
670 pub fn tls(mut self, tls: TlsConfig) -> Self {
672 self.config.tls = Some(tls);
673 self
674 }
675
676 pub fn query_timeout(mut self, timeout: Duration) -> Self {
678 self.config.query_timeout = Some(timeout);
679 self
680 }
681
682 pub fn max_rows(mut self, max_rows: usize) -> Self {
684 self.config.max_rows = Some(max_rows);
685 self
686 }
687
688 pub fn memory_limit(mut self, memory_limit: usize) -> Self {
690 self.config.memory_limit = Some(memory_limit);
691 self
692 }
693
694 pub fn on_event(mut self, callback: PoolEventCallback) -> Self {
696 self.config.on_event = Some(callback);
697 self
698 }
699
700 pub fn test_before_acquire(mut self, enabled: bool) -> Self {
705 self.config.test_before_acquire = enabled;
706 self
707 }
708
709 pub fn prewarm(mut self, enabled: bool) -> Self {
714 self.config.prewarm = enabled;
715 self
716 }
717
718 pub fn build(self) -> Result<PoolConfig, PoolError> {
720 self.config.validate()?;
721 Ok(self.config)
722 }
723}
724
725impl Default for PoolConfigBuilder {
726 fn default() -> Self {
727 Self::new()
728 }
729}
730
731#[async_trait]
733pub trait ConnectionFactory: Send + Sync {
734 async fn create(&self) -> Result<Box<dyn Connection>, crate::DbError>;
736}
737
738pub struct Pool {
744 config: PoolConfig,
745 factory: Arc<dyn ConnectionFactory>,
746 idle: Arc<ArrayQueue<PooledConnection>>,
752 total_count: Arc<AtomicU32>,
762 closed: Arc<AtomicBool>,
764 notify: Arc<Notify>,
765 waiters_count: Arc<AtomicU32>,
767 dynamic_max_size: Arc<AtomicU32>,
769 #[cfg(feature = "circuit-breaker")]
775 circuit_breaker: Arc<PlMutex<DefaultCircuitBreaker>>,
776 #[cfg(feature = "rate-limit")]
785 rate_limiter: Arc<PlRwLock<Option<Arc<dyn RateLimiter>>>>,
786 #[cfg(feature = "rate-limit")]
788 rate_limit_key: String,
789 acquire_count: Arc<AtomicU64>,
791 acquire_failed_count: Arc<AtomicU64>,
793 acquire_wait_time_ns: Arc<AtomicU64>,
795 release_count: Arc<AtomicU64>,
797 connection_created_count: Arc<AtomicU64>,
799 connection_closed_count: Arc<AtomicU64>,
801}
802
803impl Clone for Pool {
807 fn clone(&self) -> Self {
808 Self {
809 config: self.config.clone(),
810 factory: self.factory.clone(),
811 idle: self.idle.clone(),
812 total_count: self.total_count.clone(),
813 closed: self.closed.clone(),
814 notify: Arc::clone(&self.notify),
815 waiters_count: self.waiters_count.clone(),
816 dynamic_max_size: self.dynamic_max_size.clone(),
817 #[cfg(feature = "circuit-breaker")]
818 circuit_breaker: Arc::clone(&self.circuit_breaker),
819 #[cfg(feature = "rate-limit")]
820 rate_limiter: Arc::clone(&self.rate_limiter),
821 #[cfg(feature = "rate-limit")]
822 rate_limit_key: self.rate_limit_key.clone(),
823 acquire_count: self.acquire_count.clone(),
824 acquire_failed_count: self.acquire_failed_count.clone(),
825 acquire_wait_time_ns: self.acquire_wait_time_ns.clone(),
826 release_count: self.release_count.clone(),
827 connection_created_count: self.connection_created_count.clone(),
828 connection_closed_count: self.connection_closed_count.clone(),
829 }
830 }
831}
832
833impl Pool {
834 pub fn new(config: PoolConfig, factory: Arc<dyn ConnectionFactory>) -> Result<Self, PoolError> {
858 config.validate()?;
859 let max_size = config.max_size as usize;
862 let dynamic_max = config.max_size;
863 Ok(Self {
864 config,
865 factory,
866 idle: Arc::new(ArrayQueue::new(max_size)),
867 total_count: Arc::new(AtomicU32::new(0)),
868 closed: Arc::new(AtomicBool::new(false)),
869 notify: Arc::new(Notify::new()),
870 waiters_count: Arc::new(AtomicU32::new(0)),
871 dynamic_max_size: Arc::new(AtomicU32::new(dynamic_max)),
872 #[cfg(feature = "circuit-breaker")]
875 circuit_breaker: Arc::new(PlMutex::new(DefaultCircuitBreaker::new(
876 5,
877 std::time::Duration::from_secs(30),
878 ))),
879 #[cfg(feature = "rate-limit")]
882 rate_limiter: Arc::new(PlRwLock::new(None)),
883 #[cfg(feature = "rate-limit")]
884 rate_limit_key: "pool".to_string(),
885 acquire_count: Arc::new(AtomicU64::new(0)),
886 acquire_failed_count: Arc::new(AtomicU64::new(0)),
887 acquire_wait_time_ns: Arc::new(AtomicU64::new(0)),
888 release_count: Arc::new(AtomicU64::new(0)),
889 connection_created_count: Arc::new(AtomicU64::new(0)),
890 connection_closed_count: Arc::new(AtomicU64::new(0)),
891 })
892 }
893
894 pub async fn new_async(
901 config: PoolConfig,
902 factory: Arc<dyn ConnectionFactory>,
903 ) -> Result<Self, PoolError> {
904 let pool = Self::new(config, factory)?;
905 if pool.config.prewarm {
906 pool.prewarm().await;
907 }
908 Ok(pool)
909 }
910
911 pub async fn prewarm(&self) {
928 if !self.config.prewarm {
929 return;
930 }
931
932 let min_idle = self.config.min_idle as usize;
933 let mut warmed = 0;
934
935 for i in 0..min_idle {
936 if self.closed.load(Ordering::Acquire) {
938 break;
939 }
940
941 let current_max = self.dynamic_max_size.load(Ordering::Acquire);
943 let current = self.total_count.load(Ordering::Acquire);
944 if current >= current_max {
945 break;
946 }
947
948 let created = loop {
950 let current = self.total_count.load(Ordering::Acquire);
951 if current >= current_max {
952 break None;
953 }
954 match self.total_count.compare_exchange(
955 current,
956 current + 1,
957 Ordering::SeqCst,
958 Ordering::Acquire,
959 ) {
960 Ok(_) => break Some(()),
961 Err(_) => continue,
962 }
963 };
964
965 if created.is_some() {
966 match tokio::time::timeout(self.config.connection_timeout, self.factory.create())
967 .await
968 {
969 Ok(Ok(conn)) => {
970 #[cfg(feature = "circuit-breaker")]
971 {
972 self.circuit_breaker.lock().record_success();
973 }
974 self.emit_event(PoolEvent::ConnectionCreated);
975 let pooled = PooledConnection::new(conn, self.clone());
976 if self.idle.push(pooled).is_err() {
978 let _ = self.total_count.fetch_sub(1, Ordering::SeqCst);
980 tracing::warn!(
981 target: "sz_orm::pool::prewarm",
982 "prewarm connection {} failed: idle queue full",
983 i
984 );
985 } else {
986 warmed += 1;
987 self.notify.notify_one();
988 }
989 }
990 Ok(Err(e)) => {
991 let _ = self.total_count.fetch_sub(1, Ordering::SeqCst);
992 #[cfg(feature = "circuit-breaker")]
993 {
994 self.circuit_breaker.lock().record_failure();
995 }
996 tracing::warn!(
997 target: "sz_orm::pool::prewarm",
998 "prewarm connection {} failed: {}",
999 i,
1000 e
1001 );
1002 }
1003 Err(_) => {
1004 let _ = self.total_count.fetch_sub(1, Ordering::SeqCst);
1005 #[cfg(feature = "circuit-breaker")]
1006 {
1007 self.circuit_breaker.lock().record_failure();
1008 }
1009 tracing::warn!(
1010 target: "sz_orm::pool::prewarm",
1011 "prewarm connection {} timeout",
1012 i
1013 );
1014 }
1015 }
1016 }
1017 }
1018
1019 if warmed > 0 {
1020 tracing::info!(
1021 target: "sz_orm::pool::prewarm",
1022 "pool prewarm completed: {}/{} connections established",
1023 warmed,
1024 min_idle
1025 );
1026 }
1027 }
1028
1029 #[cfg(feature = "auto-prewarm")]
1034 pub async fn progressive_prewarm(
1035 &self,
1036 batch_size: u32,
1037 interval: std::time::Duration,
1038 total_timeout: std::time::Duration,
1039 progress: &crate::prewarm::PrewarmProgress,
1040 ) {
1041 use std::time::Instant;
1042
1043 let min_idle = self.config.min_idle;
1044 if min_idle == 0 || !self.config.prewarm {
1045 progress.mark_completed();
1046 return;
1047 }
1048
1049 let start = Instant::now();
1050 let batch = batch_size.max(1);
1051 let mut warmed_total: u32 = 0;
1052
1053 while warmed_total < min_idle {
1054 if start.elapsed() >= total_timeout {
1055 tracing::warn!(
1056 target: "sz_orm::pool::prewarm",
1057 "progressive prewarm timeout: {}/{} connections established",
1058 warmed_total,
1059 min_idle
1060 );
1061 break;
1062 }
1063
1064 if self.closed.load(Ordering::Acquire) {
1065 break;
1066 }
1067
1068 let remaining = min_idle - warmed_total;
1069 let this_batch = batch.min(remaining);
1070
1071 for _ in 0..this_batch {
1072 let current_max = self.dynamic_max_size.load(Ordering::Acquire);
1073 let current = self.total_count.load(Ordering::Acquire);
1074 if current >= current_max {
1075 break;
1076 }
1077
1078 let created = loop {
1079 let current = self.total_count.load(Ordering::Acquire);
1080 if current >= current_max {
1081 break None;
1082 }
1083 match self.total_count.compare_exchange(
1084 current,
1085 current + 1,
1086 Ordering::SeqCst,
1087 Ordering::Acquire,
1088 ) {
1089 Ok(_) => break Some(()),
1090 Err(_) => continue,
1091 }
1092 };
1093
1094 if created.is_some() {
1095 match tokio::time::timeout(
1096 self.config.connection_timeout,
1097 self.factory.create(),
1098 )
1099 .await
1100 {
1101 Ok(Ok(conn)) => {
1102 #[cfg(feature = "circuit-breaker")]
1103 {
1104 self.circuit_breaker.lock().record_success();
1105 }
1106 self.emit_event(PoolEvent::ConnectionCreated);
1107 let pooled = PooledConnection::new(conn, self.clone());
1108 if self.idle.push(pooled).is_err() {
1109 let _ = self.total_count.fetch_sub(1, Ordering::SeqCst);
1110 progress.record_failure();
1111 } else {
1112 progress.record_success();
1113 warmed_total += 1;
1114 self.notify.notify_one();
1115 }
1116 }
1117 Ok(Err(_)) => {
1118 let _ = self.total_count.fetch_sub(1, Ordering::SeqCst);
1119 progress.record_failure();
1120 #[cfg(feature = "circuit-breaker")]
1121 {
1122 self.circuit_breaker.lock().record_failure();
1123 }
1124 }
1125 Err(_) => {
1126 let _ = self.total_count.fetch_sub(1, Ordering::SeqCst);
1127 progress.record_failure();
1128 #[cfg(feature = "circuit-breaker")]
1129 {
1130 self.circuit_breaker.lock().record_failure();
1131 }
1132 }
1133 }
1134 }
1135 }
1136
1137 if warmed_total < min_idle && interval > std::time::Duration::ZERO {
1138 tokio::time::sleep(interval).await;
1139 }
1140 }
1141
1142 progress.set_elapsed(start.elapsed());
1143 progress.mark_completed();
1144
1145 tracing::info!(
1146 target: "sz_orm::pool::prewarm",
1147 "progressive prewarm completed: {} warmed, {} failed, elapsed {:?}",
1148 progress.snapshot().warmed,
1149 progress.snapshot().failed,
1150 start.elapsed()
1151 );
1152 }
1153
1154 pub fn config(&self) -> &PoolConfig {
1156 &self.config
1157 }
1158
1159 #[cfg(feature = "circuit-breaker")]
1173 pub fn configure_circuit_breaker(
1174 &self,
1175 failure_threshold: usize,
1176 reset_timeout: std::time::Duration,
1177 ) {
1178 let new_cb = DefaultCircuitBreaker::new(failure_threshold, reset_timeout);
1179 let mut guard = self.circuit_breaker.lock();
1181 *guard = new_cb;
1182 }
1183
1184 #[cfg(feature = "circuit-breaker")]
1189 pub fn reset_circuit_breaker(&self) -> bool {
1190 let mut guard = self.circuit_breaker.lock();
1192 guard.reset()
1193 }
1194
1195 #[cfg(feature = "circuit-breaker")]
1197 pub fn circuit_state(&self) -> CircuitState {
1198 let guard = self.circuit_breaker.lock();
1200 guard.state()
1201 }
1202
1203 #[cfg(feature = "rate-limit")]
1212 pub fn set_rate_limiter(&self, limiter: Option<Arc<dyn RateLimiter>>) {
1213 let mut guard = self.rate_limiter.write();
1215 *guard = limiter;
1216 }
1217
1218 #[cfg(feature = "rate-limit")]
1220 pub fn with_rate_limit_key(mut self, key: impl Into<String>) -> Self {
1221 self.rate_limit_key = key.into();
1222 self
1223 }
1224
1225 fn emit_event(&self, event: PoolEvent) {
1227 if matches!(event, PoolEvent::ConnectionCreated) {
1230 self.connection_created_count
1231 .fetch_add(1, Ordering::Relaxed);
1232 }
1233 if let Some(ref callback) = self.config.on_event {
1234 callback(event);
1235 }
1236 }
1237
1238 async fn close_connection(&self, pooled: PooledConnection) {
1243 let mut pooled = pooled;
1244 let _ = pooled.conn.close().await;
1245 self.connection_closed_count.fetch_add(1, Ordering::Relaxed);
1246 }
1247
1248 #[tracing::instrument(skip(self), fields(max_size = self.config.max_size, acquire_timeout = ?self.config.acquire_timeout))]
1268 pub async fn acquire(&self) -> Result<PooledConnection, PoolError> {
1269 if self.closed.load(Ordering::Acquire) {
1271 self.acquire_failed_count.fetch_add(1, Ordering::Relaxed);
1272 return Err(PoolError::Closed);
1273 }
1274
1275 #[cfg(feature = "circuit-breaker")]
1279 {
1280 let mut guard = self.circuit_breaker.lock();
1281 if !guard.can_execute() {
1282 self.acquire_failed_count.fetch_add(1, Ordering::Relaxed);
1283 return Err(PoolError::CircuitOpen);
1284 }
1285 }
1286
1287 #[cfg(feature = "rate-limit")]
1291 {
1292 let guard = self.rate_limiter.read();
1293 if let Some(ref limiter) = *guard {
1294 match limiter.try_acquire(&self.rate_limit_key) {
1295 Ok(result) if !result.allowed => {
1296 self.acquire_failed_count.fetch_add(1, Ordering::Relaxed);
1297 return Err(PoolError::RateLimited {
1298 remaining: result.remaining,
1299 reset_at: result.reset_at,
1300 });
1301 }
1302 Ok(_) => {} Err(_) => {
1304 }
1306 }
1307 }
1308 }
1309
1310 let deadline = Instant::now() + self.config.acquire_timeout;
1311 let mut backoff = Duration::from_millis(1);
1313 const MAX_BACKOFF: Duration = Duration::from_millis(100);
1315
1316 loop {
1317 let mut to_close: Vec<PooledConnection> = Vec::new();
1323 let acquired: Option<PooledConnection> = {
1324 let mut found: Option<PooledConnection> = None;
1325 while let Some(pooled) = self.idle.pop() {
1326 if pooled.is_expired(self.config.max_lifetime) {
1328 to_close.push(pooled);
1329 continue;
1330 }
1331 if pooled.is_idle_too_long(self.config.idle_timeout) {
1333 to_close.push(pooled);
1334 continue;
1335 }
1336 if !pooled.conn.is_connected() {
1339 to_close.push(pooled);
1340 continue;
1341 }
1342 found = Some(pooled);
1343 break;
1344 }
1345 found
1346 };
1347
1348 for pooled in to_close {
1350 self.close_connection(pooled).await;
1351 self.total_count.fetch_sub(1, Ordering::SeqCst);
1353 }
1354
1355 if let Some(mut pooled) = acquired {
1356 if self.config.test_before_acquire {
1358 let ping_timeout = self.config.connection_timeout / 2;
1359 let alive = match tokio::time::timeout(ping_timeout, pooled.conn.ping()).await {
1360 Ok(true) => true,
1361 Ok(false) => false,
1362 Err(_) => false, };
1364 if !alive {
1365 self.close_connection(pooled).await;
1367 self.total_count.fetch_sub(1, Ordering::SeqCst);
1368 continue;
1369 }
1370 }
1371 pooled.pool = Some(self.clone());
1374 self.acquire_count.fetch_add(1, Ordering::Relaxed);
1375 return Ok(pooled);
1376 }
1377
1378 let current_max = self.dynamic_max_size.load(Ordering::Acquire);
1383 let created = loop {
1384 let current = self.total_count.load(Ordering::Acquire);
1385 if current >= current_max {
1386 break None; }
1388 match self.total_count.compare_exchange(
1389 current,
1390 current + 1,
1391 Ordering::SeqCst,
1392 Ordering::Acquire,
1393 ) {
1394 Ok(_) => break Some(()), Err(_) => continue, }
1397 };
1398
1399 if created.is_some() {
1400 match tokio::time::timeout(self.config.connection_timeout, self.factory.create())
1401 .await
1402 {
1403 Ok(Ok(conn)) => {
1404 #[cfg(feature = "circuit-breaker")]
1407 {
1408 self.circuit_breaker.lock().record_success();
1409 }
1410 self.emit_event(PoolEvent::ConnectionCreated);
1411 self.emit_event(PoolEvent::ConnectionAcquired);
1412 self.acquire_count.fetch_add(1, Ordering::Relaxed);
1413 return Ok(PooledConnection::new(conn, self.clone()));
1414 }
1415 Ok(Err(e)) => {
1416 self.total_count.fetch_sub(1, Ordering::SeqCst);
1418 #[cfg(feature = "circuit-breaker")]
1421 {
1422 self.circuit_breaker.lock().record_failure();
1423 }
1424 self.acquire_failed_count.fetch_add(1, Ordering::Relaxed);
1425 return Err(PoolError::ConnectionFailed(e.to_string()));
1426 }
1427 Err(_) => {
1428 self.total_count.fetch_sub(1, Ordering::SeqCst);
1430 #[cfg(feature = "circuit-breaker")]
1433 {
1434 self.circuit_breaker.lock().record_failure();
1435 }
1436 self.acquire_failed_count.fetch_add(1, Ordering::Relaxed);
1437 return Err(PoolError::Timeout);
1438 }
1439 }
1440 }
1441
1442 let now = Instant::now();
1444 if now >= deadline {
1445 self.emit_event(PoolEvent::AcquireTimeout);
1446 self.acquire_failed_count.fetch_add(1, Ordering::Relaxed);
1447 return Err(PoolError::Timeout);
1448 }
1449 self.waiters_count.fetch_add(1, Ordering::SeqCst);
1451 let wait = std::cmp::min(backoff, deadline - now);
1452 match tokio::time::timeout(wait, self.notify.notified()).await {
1453 Ok(()) => {
1454 backoff = Duration::from_millis(1);
1456 }
1457 Err(_) => {
1458 backoff = std::cmp::min(backoff * 2, MAX_BACKOFF);
1460 }
1461 }
1462 self.waiters_count.fetch_sub(1, Ordering::SeqCst);
1464 self.acquire_wait_time_ns
1466 .fetch_add(wait.as_nanos() as u64, Ordering::Relaxed);
1467 }
1468 }
1469
1470 #[tracing::instrument(skip(self, pooled))]
1478 pub async fn release(&self, mut pooled: PooledConnection) {
1479 pooled.pool = None;
1481 self.release_count.fetch_add(1, Ordering::Relaxed);
1483
1484 if self.closed.load(Ordering::Acquire) {
1486 self.close_connection(pooled).await;
1487 self.total_count.fetch_sub(1, Ordering::SeqCst);
1489 self.emit_event(PoolEvent::ConnectionClosed);
1490 return;
1491 }
1492
1493 if !pooled.conn.is_connected() {
1495 self.close_connection(pooled).await;
1496 self.total_count.fetch_sub(1, Ordering::SeqCst);
1497 self.emit_event(PoolEvent::ConnectionClosed);
1498 return;
1499 }
1500
1501 pooled.last_used_at = Instant::now();
1503
1504 if let Err(rejected) = self.idle.push(pooled) {
1510 self.close_connection(rejected).await;
1512 self.total_count.fetch_sub(1, Ordering::SeqCst);
1513 self.emit_event(PoolEvent::ConnectionClosed);
1514 } else {
1515 self.emit_event(PoolEvent::ConnectionReleased);
1516 }
1517 self.notify.notify_one();
1518 }
1519
1520 pub async fn status(&self) -> PoolStatus {
1525 let idle_count = self.idle.len() as u32;
1526 let active = self.total_count.load(Ordering::Acquire);
1528 let waiters = self.waiters_count.load(Ordering::Acquire);
1529 PoolStatus {
1530 idle: idle_count,
1531 active,
1532 max: self.dynamic_max_size.load(Ordering::Acquire),
1533 min: self.config.min_idle,
1534 waiters,
1535 }
1536 }
1537
1538 pub fn pool_metrics(&self) -> PoolMetrics {
1549 PoolMetrics {
1550 acquire_count: self.acquire_count.load(Ordering::Acquire),
1551 acquire_failed_count: self.acquire_failed_count.load(Ordering::Acquire),
1552 acquire_wait_time: Duration::from_nanos(
1553 self.acquire_wait_time_ns.load(Ordering::Acquire),
1554 ),
1555 release_count: self.release_count.load(Ordering::Acquire),
1556 connection_created_count: self.connection_created_count.load(Ordering::Acquire),
1557 connection_closed_count: self.connection_closed_count.load(Ordering::Acquire),
1558 }
1559 }
1560
1561 #[tracing::instrument(skip(self))]
1563 pub async fn reap_idle(&self) {
1564 let mut all: Vec<PooledConnection> = Vec::new();
1568 while let Some(pooled) = self.idle.pop() {
1569 all.push(pooled);
1570 }
1571
1572 let mut to_close = Vec::new();
1574 for pooled in all {
1575 if pooled.is_idle_too_long(self.config.idle_timeout)
1576 || pooled.is_expired(self.config.max_lifetime)
1577 {
1578 to_close.push(pooled);
1579 } else {
1580 if let Err(rejected) = self.idle.push(pooled) {
1582 self.close_connection(rejected).await;
1583 self.total_count.fetch_sub(1, Ordering::SeqCst);
1584 }
1585 }
1586 }
1587
1588 for pooled in to_close {
1590 self.close_connection(pooled).await;
1591 self.total_count.fetch_sub(1, Ordering::SeqCst);
1593 }
1594 }
1595
1596 pub async fn close_all(&self) {
1600 self.closed.store(true, Ordering::Release);
1602 let mut to_close: Vec<PooledConnection> = Vec::new();
1605 while let Some(pooled) = self.idle.pop() {
1606 to_close.push(pooled);
1607 }
1608 let closed_count: u32 = to_close.len() as u32;
1610 for pooled in to_close {
1611 self.close_connection(pooled).await;
1612 }
1613 self.total_count.fetch_sub(closed_count, Ordering::SeqCst);
1616 }
1617
1618 pub async fn health_check(&self) -> u32 {
1634 let mut to_check: Vec<PooledConnection> = Vec::new();
1636 while let Some(pooled) = self.idle.pop() {
1637 to_check.push(pooled);
1638 }
1639
1640 let mut removed: u32 = 0;
1641 let mut alive: Vec<PooledConnection> = Vec::with_capacity(to_check.len());
1642 for mut pooled in to_check.drain(..) {
1643 if !pooled.conn.is_connected() {
1645 self.close_connection(pooled).await;
1646 removed += 1;
1647 continue;
1648 }
1649 let ping_timeout = self.config.connection_timeout / 2;
1651 match tokio::time::timeout(ping_timeout, pooled.conn.ping()).await {
1652 Ok(true) => alive.push(pooled),
1653 Ok(false) => {
1654 self.close_connection(pooled).await;
1656 removed += 1;
1657 }
1658 Err(_) => {
1659 self.close_connection(pooled).await;
1661 removed += 1;
1662 }
1663 }
1664 }
1665
1666 let alive_count: u32 = alive.len() as u32;
1668 for pooled in alive {
1669 if let Err(rejected) = self.idle.push(pooled) {
1671 self.close_connection(rejected).await;
1672 removed += 1;
1673 }
1674 }
1675
1676 if removed > 0 {
1678 self.total_count.fetch_sub(removed, Ordering::SeqCst);
1679 }
1680
1681 if alive_count > 0 {
1683 self.notify.notify_one();
1684 }
1685
1686 removed
1687 }
1688
1689 pub async fn shutdown(&self) {
1696 self.shutdown_with_timeout(Duration::from_secs(30)).await;
1697 }
1698
1699 pub async fn shutdown_with_timeout(&self, timeout: Duration) {
1707 if self.closed.swap(true, Ordering::SeqCst) {
1709 return;
1710 }
1711 self.notify.notify_waiters();
1713 self.close_all().await;
1715 let deadline = Instant::now() + timeout;
1717 while self.total_count.load(Ordering::SeqCst) > 0 {
1718 if Instant::now() >= deadline {
1719 let remaining = self.total_count.load(Ordering::SeqCst);
1720 if remaining > 0 {
1721 eprintln!(
1722 "graceful shutdown timeout, {} connections force closed",
1723 remaining
1724 );
1725 }
1726 break;
1727 }
1728 tokio::time::sleep(Duration::from_millis(100)).await;
1729 }
1730 }
1731
1732 pub fn resize(&self, new_max: usize) {
1740 self.set_max_size(new_max as u32);
1741 }
1742
1743 pub fn set_max_size(&self, new_max: u32) {
1745 self.dynamic_max_size.store(new_max, Ordering::SeqCst);
1746 }
1747
1748 pub fn max_size(&self) -> u32 {
1750 self.dynamic_max_size.load(Ordering::Acquire)
1751 }
1752
1753 pub async fn warmup(&self, min_idle: usize) -> Result<(), PoolError> {
1757 for _ in 0..min_idle {
1758 let current_max = self.dynamic_max_size.load(Ordering::Acquire);
1759 let current = self.total_count.load(Ordering::Acquire);
1760 if current >= current_max {
1761 break;
1762 }
1763 match self.total_count.compare_exchange(
1765 current,
1766 current + 1,
1767 Ordering::SeqCst,
1768 Ordering::Acquire,
1769 ) {
1770 Ok(_) => {}
1771 Err(_) => continue, }
1773 match self.factory.create().await {
1774 Ok(conn) => {
1775 let now = Instant::now();
1776 let pooled = PooledConnection {
1777 conn,
1778 created_at: now,
1779 last_used_at: now,
1780 pool: None,
1781 };
1782 if let Err(rejected) = self.idle.push(pooled) {
1783 self.close_connection(rejected).await;
1785 self.total_count.fetch_sub(1, Ordering::SeqCst);
1786 }
1787 self.emit_event(PoolEvent::ConnectionCreated);
1788 }
1789 Err(_) => {
1790 self.total_count.fetch_sub(1, Ordering::SeqCst);
1792 break;
1793 }
1794 }
1795 }
1796 Ok(())
1797 }
1798
1799 pub async fn query_with_timeout(&self, sql: &str) -> Result<QueryRows, crate::DbError> {
1804 let timeout = self.config.query_timeout.unwrap_or(Duration::from_secs(30));
1805 let mut conn = self.acquire().await.map_err(crate::DbError::PoolError)?;
1806 tokio::time::timeout(timeout, conn.query(sql))
1807 .await
1808 .map_err(|_| crate::DbError::QueryError(format!("Query timeout after {:?}", timeout)))?
1809 }
1810}
1811
1812#[cfg(feature = "prod-pool-tuning")]
1817mod pool_prod {
1818 use super::PoolConfig;
1819 use serde::{Deserialize, Serialize};
1820 use std::time::Duration;
1821
1822 #[derive(Debug, Clone, PartialEq, Eq, thiserror::Error)]
1824 pub enum PoolProdError {
1825 #[error("pool max_size must be positive")]
1827 MaxSizeNotPositive,
1828 #[error("pool acquire_timeout must be positive")]
1830 AcquireTimeoutNotPositive,
1831 #[error("pool min_idle cannot exceed max_size")]
1833 MinIdleExceedsMaxSize,
1834 }
1835
1836 #[derive(Debug, Clone, Serialize, Deserialize)]
1838 pub struct PoolProdConfig {
1839 pub max_size: u32,
1841 pub acquire_timeout: Duration,
1843 pub idle_timeout: Duration,
1845 pub connection_timeout: Duration,
1847 pub query_timeout: Duration,
1849 pub min_idle: u32,
1851 pub prewarm: bool,
1853 }
1854
1855 impl Default for PoolProdConfig {
1856 fn default() -> Self {
1857 Self {
1858 max_size: 100,
1859 acquire_timeout: Duration::from_secs(30),
1860 idle_timeout: Duration::from_secs(600),
1861 connection_timeout: Duration::from_secs(10),
1862 query_timeout: Duration::from_secs(30),
1863 min_idle: 0,
1864 prewarm: false,
1865 }
1866 }
1867 }
1868
1869 impl PoolProdConfig {
1870 pub fn new(
1872 max_size: u32,
1873 acquire_timeout: Duration,
1874 idle_timeout: Duration,
1875 connection_timeout: Duration,
1876 query_timeout: Duration,
1877 min_idle: u32,
1878 prewarm: bool,
1879 ) -> Self {
1880 Self {
1881 max_size,
1882 acquire_timeout,
1883 idle_timeout,
1884 connection_timeout,
1885 query_timeout,
1886 min_idle,
1887 prewarm,
1888 }
1889 }
1890
1891 pub fn validate(&self) -> Result<(), PoolProdError> {
1893 if self.max_size == 0 {
1894 return Err(PoolProdError::MaxSizeNotPositive);
1895 }
1896 if self.acquire_timeout.is_zero() {
1897 return Err(PoolProdError::AcquireTimeoutNotPositive);
1898 }
1899 if self.min_idle > self.max_size {
1900 return Err(PoolProdError::MinIdleExceedsMaxSize);
1901 }
1902 Ok(())
1903 }
1904
1905 pub fn to_pool_config(&self) -> PoolConfig {
1907 PoolConfig {
1908 max_size: self.max_size,
1909 min_idle: self.min_idle,
1910 acquire_timeout: self.acquire_timeout,
1911 idle_timeout: self.idle_timeout,
1912 max_lifetime: Duration::from_secs(1800),
1913 connection_timeout: self.connection_timeout,
1914 tls: None,
1915 query_timeout: Some(self.query_timeout),
1916 max_rows: None,
1917 memory_limit: None,
1918 on_event: None,
1919 test_before_acquire: false,
1920 prewarm: self.prewarm,
1921 }
1922 }
1923 }
1924}
1925
1926#[cfg(feature = "prod-pool-tuning")]
1927pub use pool_prod::{PoolProdConfig, PoolProdError};
1928
1929#[cfg(feature = "prod-leak-detection")]
1934mod leak_detection {
1935 use serde::{Deserialize, Serialize};
1936 use std::time::Duration;
1937
1938 #[derive(Debug, Clone, Serialize, Deserialize)]
1940 pub struct LeakDetectionConfig {
1941 pub enabled: bool,
1943 pub interval: Duration,
1945 pub threshold: u32,
1947 pub borrow_timeout: Duration,
1949 }
1950
1951 impl Default for LeakDetectionConfig {
1952 fn default() -> Self {
1953 Self {
1954 enabled: false,
1955 interval: Duration::from_secs(60),
1956 threshold: 5,
1957 borrow_timeout: Duration::from_secs(60),
1958 }
1959 }
1960 }
1961
1962 impl LeakDetectionConfig {
1963 pub fn new(
1965 enabled: bool,
1966 interval: Duration,
1967 threshold: u32,
1968 borrow_timeout: Duration,
1969 ) -> Self {
1970 Self {
1971 enabled,
1972 interval,
1973 threshold,
1974 borrow_timeout,
1975 }
1976 }
1977
1978 pub fn validate(&self) -> Result<(), LeakDetectionError> {
1980 if self.interval.is_zero() {
1981 return Err(LeakDetectionError::IntervalNotPositive);
1982 }
1983 if self.borrow_timeout.is_zero() {
1984 return Err(LeakDetectionError::BorrowTimeoutNotPositive);
1985 }
1986 Ok(())
1987 }
1988 }
1989
1990 #[derive(Debug, Clone, PartialEq, Eq, thiserror::Error)]
1992 pub enum LeakDetectionError {
1993 #[error("leak detection interval must be positive")]
1995 IntervalNotPositive,
1996 #[error("leak detection borrow_timeout must be positive")]
1998 BorrowTimeoutNotPositive,
1999 }
2000
2001 #[derive(Debug, Clone, Serialize, Deserialize)]
2003 pub struct LeakEntry {
2004 pub conn_id: u64,
2006 pub borrowed_at: String,
2008 pub borrow_duration: Duration,
2010 }
2011
2012 #[derive(Debug, Clone, Serialize, Deserialize)]
2014 pub struct LeakReport {
2015 pub borrowed_count: u32,
2017 pub max_borrow_duration: Duration,
2019 pub suspected_leaks: Vec<LeakEntry>,
2021 }
2022
2023 impl LeakReport {
2024 pub fn empty() -> Self {
2026 Self {
2027 borrowed_count: 0,
2028 max_borrow_duration: Duration::ZERO,
2029 suspected_leaks: vec![],
2030 }
2031 }
2032 }
2033}
2034
2035#[cfg(feature = "prod-leak-detection")]
2036pub use leak_detection::{LeakDetectionConfig, LeakDetectionError, LeakEntry, LeakReport};
2037
2038#[cfg(test)]
2039mod tests {
2040 use super::*;
2041
2042 struct MockConnection {
2044 connected: bool,
2045 }
2046
2047 impl MockConnection {
2048 fn new() -> Self {
2049 Self { connected: true }
2050 }
2051 }
2052
2053 impl Connection for MockConnection {
2054 fn execute<'a>(
2055 &'a mut self,
2056 _sql: &'a str,
2057 ) -> Pin<Box<dyn Future<Output = Result<u64, crate::DbError>> + Send + 'a>> {
2058 Box::pin(async move { Ok(1) })
2059 }
2060
2061 fn query<'a>(
2062 &'a mut self,
2063 _sql: &'a str,
2064 ) -> Pin<
2065 Box<
2066 dyn Future<
2067 Output = Result<
2068 Vec<std::collections::HashMap<String, crate::value::Value>>,
2069 crate::DbError,
2070 >,
2071 > + Send
2072 + 'a,
2073 >,
2074 > {
2075 Box::pin(async move { Ok(vec![]) })
2076 }
2077
2078 fn begin_transaction<'a>(
2079 &'a mut self,
2080 ) -> Pin<Box<dyn Future<Output = Result<(), crate::DbError>> + Send + 'a>> {
2081 Box::pin(async move { Ok(()) })
2082 }
2083
2084 fn commit<'a>(
2085 &'a mut self,
2086 ) -> Pin<Box<dyn Future<Output = Result<(), crate::DbError>> + Send + 'a>> {
2087 Box::pin(async move { Ok(()) })
2088 }
2089
2090 fn rollback<'a>(
2091 &'a mut self,
2092 ) -> Pin<Box<dyn Future<Output = Result<(), crate::DbError>> + Send + 'a>> {
2093 Box::pin(async move { Ok(()) })
2094 }
2095
2096 fn is_connected(&self) -> bool {
2097 self.connected
2098 }
2099
2100 fn ping<'a>(&'a mut self) -> Pin<Box<dyn Future<Output = bool> + Send + 'a>> {
2101 Box::pin(async move { true })
2102 }
2103
2104 fn close<'a>(
2105 &'a mut self,
2106 ) -> Pin<Box<dyn Future<Output = Result<(), crate::DbError>> + Send + 'a>> {
2107 Box::pin(async move {
2108 self.connected = false;
2109 Ok(())
2110 })
2111 }
2112 }
2113
2114 struct MockConnectionFactory;
2115
2116 #[async_trait]
2117 impl ConnectionFactory for MockConnectionFactory {
2118 async fn create(&self) -> Result<Box<dyn Connection>, crate::DbError> {
2119 Ok(Box::new(MockConnection::new()))
2120 }
2121 }
2122
2123 #[tokio::test]
2124 async fn test_pool_config_builder() -> Result<(), Box<dyn std::error::Error>> {
2125 let config = PoolConfigBuilder::new().max_size(50).min_idle(10).build()?;
2126
2127 assert_eq!(config.max_size, 50);
2128 assert_eq!(config.min_idle, 10);
2129 Ok(())
2130 }
2131
2132 #[test]
2133 fn test_pool_status_display() {
2134 let status = PoolStatus {
2135 idle: 5,
2136 active: 10,
2137 max: 100,
2138 min: 5,
2139 waiters: 0,
2140 };
2141
2142 let display = format!("{:?}", status);
2143 assert!(display.contains("idle"));
2144 assert!(display.contains("active"));
2145 }
2146
2147 #[test]
2148 fn test_default_pool_config() {
2149 let config = PoolConfig::default();
2150 assert_eq!(config.max_size, 100);
2151 assert_eq!(config.min_idle, 0);
2152 assert_eq!(config.acquire_timeout.as_secs(), 30);
2153 assert_eq!(config.idle_timeout.as_secs(), 600);
2154 assert_eq!(config.max_lifetime.as_secs(), 1800);
2155 }
2156
2157 #[tokio::test]
2158 async fn test_pool_config_clone() {
2159 let config = PoolConfig::default();
2160 let cloned = config.clone();
2161 assert_eq!(cloned.max_size, config.max_size);
2162 assert_eq!(cloned.min_idle, config.min_idle);
2163 }
2164
2165 #[test]
2166 fn test_pool_config_builder_default() -> Result<(), Box<dyn std::error::Error>> {
2167 let builder = PoolConfigBuilder::new();
2168 let config = builder.build()?;
2169 assert_eq!(config.max_size, 100);
2170 Ok(())
2171 }
2172
2173 #[test]
2174 fn test_pool_config_validate() {
2175 let result = PoolConfigBuilder::new().max_size(0).build();
2176 assert!(result.is_err());
2177
2178 let result = PoolConfigBuilder::new().max_size(10).min_idle(20).build();
2179 assert!(result.is_err());
2180 }
2181
2182 #[test]
2183 fn test_pool_config_validate_duration_upper_bound() {
2184 use std::time::Duration;
2185
2186 let config = PoolConfig {
2188 max_size: 10,
2189 min_idle: 1,
2190 acquire_timeout: Duration::from_secs(u64::MAX),
2191 idle_timeout: Duration::from_secs(1),
2192 max_lifetime: Duration::from_secs(1),
2193 connection_timeout: Duration::from_secs(5),
2194 tls: None,
2195 query_timeout: None,
2196 max_rows: None,
2197 memory_limit: None,
2198 on_event: None,
2199 test_before_acquire: false,
2200 prewarm: false,
2201 };
2202 assert!(config.validate().is_err());
2203
2204 let config = PoolConfig {
2206 max_size: 10,
2207 min_idle: 1,
2208 acquire_timeout: Duration::from_secs(u32::MAX as u64),
2209 idle_timeout: Duration::from_secs(1),
2210 max_lifetime: Duration::from_secs(1),
2211 connection_timeout: Duration::from_secs(5),
2212 tls: None,
2213 query_timeout: None,
2214 max_rows: None,
2215 memory_limit: None,
2216 on_event: None,
2217 test_before_acquire: false,
2218 prewarm: false,
2219 };
2220 assert!(config.validate().is_ok());
2221
2222 let config = PoolConfig {
2224 max_size: 10,
2225 min_idle: 1,
2226 acquire_timeout: Duration::from_secs(u32::MAX as u64 + 1),
2227 idle_timeout: Duration::from_secs(1),
2228 max_lifetime: Duration::from_secs(1),
2229 connection_timeout: Duration::from_secs(5),
2230 tls: None,
2231 query_timeout: None,
2232 max_rows: None,
2233 memory_limit: None,
2234 on_event: None,
2235 test_before_acquire: false,
2236 prewarm: false,
2237 };
2238 assert!(config.validate().is_err());
2239 }
2240
2241 #[test]
2242 fn test_pool_config_test_before_acquire_default() {
2243 let config = PoolConfig::default();
2245 assert!(!config.test_before_acquire);
2246 }
2247
2248 #[test]
2249 fn test_pool_config_builder_test_before_acquire() {
2250 let config = PoolConfigBuilder::new()
2252 .test_before_acquire(true)
2253 .build()
2254 .unwrap();
2255 assert!(config.test_before_acquire);
2256 }
2257
2258 #[tokio::test]
2259 async fn test_pool_acquire_and_release() -> Result<(), Box<dyn std::error::Error>> {
2260 let config = PoolConfigBuilder::new().max_size(5).min_idle(1).build()?;
2261 let factory = Arc::new(MockConnectionFactory);
2262 let pool = Pool::new(config, factory)?;
2263
2264 let conn = pool.acquire().await?;
2265 let status = pool.status().await;
2266 assert_eq!(status.active, 1);
2267 assert_eq!(status.idle, 0);
2268
2269 pool.release(conn).await;
2270 let status = pool.status().await;
2271 assert_eq!(status.idle, 1);
2272
2273 let _conn2 = pool.acquire().await?;
2275 let status = pool.status().await;
2276 assert_eq!(status.idle, 0);
2277 Ok(())
2278 }
2279
2280 #[tokio::test]
2281 async fn test_pool_status() -> Result<(), Box<dyn std::error::Error>> {
2282 let config = PoolConfigBuilder::new().max_size(10).min_idle(2).build()?;
2283 let factory = Arc::new(MockConnectionFactory);
2284 let pool = Pool::new(config, factory)?;
2285
2286 let status = pool.status().await;
2287 assert_eq!(status.max, 10);
2288 assert_eq!(status.min, 2);
2289 assert_eq!(status.active, 0);
2290 Ok(())
2291 }
2292
2293 #[tokio::test]
2294 async fn test_pool_close_all() -> Result<(), Box<dyn std::error::Error>> {
2295 let config = PoolConfigBuilder::new().max_size(5).build()?;
2296 let factory = Arc::new(MockConnectionFactory);
2297 let pool = Pool::new(config, factory)?;
2298
2299 let conn1 = pool.acquire().await?;
2301 let conn2 = pool.acquire().await?;
2302 pool.release(conn1).await;
2303 pool.release(conn2).await;
2304
2305 pool.close_all().await;
2306 let status = pool.status().await;
2307 assert_eq!(status.idle, 0);
2308 assert_eq!(status.active, 0);
2309 Ok(())
2310 }
2311
2312 #[tokio::test]
2313 async fn test_pool_reap_idle() -> Result<(), Box<dyn std::error::Error>> {
2314 let config = PoolConfigBuilder::new()
2315 .max_size(5)
2316 .idle_timeout(0) .build()?;
2318 let factory = Arc::new(MockConnectionFactory);
2319 let pool = Pool::new(config, factory)?;
2320
2321 let conn = pool.acquire().await?;
2322 pool.release(conn).await;
2323
2324 tokio::time::sleep(tokio::time::Duration::from_millis(10)).await;
2326
2327 pool.reap_idle().await;
2328 let status = pool.status().await;
2329 assert_eq!(status.idle, 0);
2330 Ok(())
2331 }
2332
2333 #[tokio::test]
2339 async fn test_h7_acquire_timeout_default_30s() {
2340 let config = PoolConfig::default();
2341 assert_eq!(
2342 config.acquire_timeout,
2343 Duration::from_secs(30),
2344 "H-7: acquire_timeout 默认应为 30s"
2345 );
2346 }
2347
2348 #[tokio::test]
2350 async fn test_h7_acquire_timeout_configurable() -> Result<(), Box<dyn std::error::Error>> {
2351 let config = PoolConfigBuilder::new()
2352 .max_size(1)
2353 .acquire_timeout(5) .build()?;
2355 assert_eq!(config.acquire_timeout, Duration::from_secs(5));
2356
2357 let factory = Arc::new(MockConnectionFactory);
2359 let pool = Pool::new(config, factory)?;
2360 let _conn1 = pool.acquire().await?;
2361
2362 let fast_config = PoolConfigBuilder::new()
2364 .max_size(1)
2365 .acquire_timeout(0) .build()?;
2367 let fast_pool = Pool::new(fast_config, Arc::new(MockConnectionFactory))?;
2370 let _fast_conn = fast_pool.acquire().await?; let result = fast_pool.acquire().await;
2372 assert!(
2373 matches!(result, Err(PoolError::Timeout)),
2374 "H-7: 应返回 Timeout"
2375 );
2376 Ok(())
2377 }
2378
2379 #[tokio::test]
2382 async fn test_m7_health_check_removes_nothing_when_all_healthy(
2383 ) -> Result<(), Box<dyn std::error::Error>> {
2384 let config = PoolConfigBuilder::new().max_size(5).build()?;
2386 let factory = Arc::new(MockConnectionFactory);
2387 let pool = Pool::new(config, factory)?;
2388
2389 let conn1 = pool.acquire().await?;
2391 let conn2 = pool.acquire().await?;
2392 let conn3 = pool.acquire().await?;
2393 pool.release(conn1).await;
2394 pool.release(conn2).await;
2395 pool.release(conn3).await;
2396
2397 let removed = pool.health_check().await;
2398 assert_eq!(removed, 0, "Healthy connections should not be removed");
2399
2400 let status = pool.status().await;
2401 assert_eq!(status.idle, 3);
2402 assert_eq!(status.active, 3);
2403 Ok(())
2404 }
2405
2406 #[tokio::test]
2407 async fn test_m7_health_check_returns_zero_for_empty_pool(
2408 ) -> Result<(), Box<dyn std::error::Error>> {
2409 let config = PoolConfigBuilder::new().max_size(5).build()?;
2410 let factory = Arc::new(MockConnectionFactory);
2411 let pool = Pool::new(config, factory)?;
2412
2413 let removed = pool.health_check().await;
2414 assert_eq!(removed, 0);
2415 Ok(())
2416 }
2417
2418 struct CountingFactory {
2422 count: AtomicU32,
2423 }
2424
2425 impl CountingFactory {
2426 fn new() -> Self {
2427 Self {
2428 count: AtomicU32::new(0),
2429 }
2430 }
2431 fn created_count(&self) -> u32 {
2432 self.count.load(Ordering::SeqCst)
2433 }
2434 }
2435
2436 #[async_trait]
2437 impl ConnectionFactory for CountingFactory {
2438 async fn create(&self) -> Result<Box<dyn Connection>, crate::DbError> {
2439 self.count.fetch_add(1, Ordering::SeqCst);
2440 Ok(Box::new(MockConnection::new()))
2441 }
2442 }
2443
2444 #[tokio::test]
2450 async fn test_production_bug_max_lifetime_never_expires(
2451 ) -> Result<(), Box<dyn std::error::Error>> {
2452 let config = PoolConfig {
2455 max_size: 5,
2456 min_idle: 0,
2457 acquire_timeout: Duration::from_secs(30),
2458 idle_timeout: Duration::from_secs(600),
2459 max_lifetime: Duration::from_millis(100), connection_timeout: Duration::from_secs(10),
2461 tls: None,
2462 query_timeout: None,
2463 max_rows: None,
2464 memory_limit: None,
2465 on_event: None,
2466 test_before_acquire: false,
2467 prewarm: false,
2468 };
2469 let factory = Arc::new(CountingFactory::new());
2470 let pool = Pool::new(config, factory.clone())?;
2471
2472 let conn = pool.acquire().await?;
2474 assert_eq!(factory.created_count(), 1, "应创建 1 个连接");
2475
2476 pool.release(conn).await;
2478
2479 tokio::time::sleep(Duration::from_millis(150)).await;
2481
2482 let conn2 = pool.acquire().await?;
2484
2485 assert_eq!(
2488 factory.created_count(),
2489 2,
2490 "超过 max_lifetime 后应创建新连接(旧连接应被回收)"
2491 );
2492
2493 pool.release(conn2).await;
2494 Ok(())
2495 }
2496
2497 #[tokio::test]
2504 async fn test_drop_auto_release_connection() -> Result<(), Box<dyn std::error::Error>> {
2505 let config = PoolConfigBuilder::new().max_size(2).build()?;
2506 let factory = Arc::new(CountingFactory::new());
2507 let pool = Pool::new(config, factory.clone())?;
2508
2509 {
2511 let _conn = pool.acquire().await?;
2512 assert_eq!(factory.created_count(), 1, "应创建 1 个连接");
2513 let status = pool.status().await;
2514 assert_eq!(status.active, 1, "active 应为 1");
2515 assert_eq!(status.idle, 0, "idle 应为 0");
2516 }
2518
2519 tokio::time::sleep(Duration::from_millis(50)).await;
2521
2522 let status = pool.status().await;
2524 assert_eq!(status.idle, 1, "Drop 后连接应自动归还,idle 应为 1");
2525 assert_eq!(status.active, 1, "total_count 应为 1");
2526 assert_eq!(factory.created_count(), 1, "应复用归还的连接,不创建新连接");
2527 Ok(())
2528 }
2529
2530 #[tokio::test]
2532 async fn test_drop_auto_release_then_reuse() -> Result<(), Box<dyn std::error::Error>> {
2533 let config = PoolConfigBuilder::new().max_size(1).build()?;
2534 let factory = Arc::new(CountingFactory::new());
2535 let pool = Pool::new(config, factory.clone())?;
2536
2537 {
2539 let _conn = pool.acquire().await?;
2540 }
2541
2542 tokio::time::sleep(Duration::from_millis(50)).await;
2544
2545 let conn = pool.acquire().await?;
2547 assert_eq!(factory.created_count(), 1, "应复用归还的连接,不创建新连接");
2548
2549 pool.release(conn).await;
2550 Ok(())
2551 }
2552
2553 #[tokio::test]
2555 async fn test_into_inner_does_not_return_to_pool() -> Result<(), Box<dyn std::error::Error>> {
2556 let config = PoolConfigBuilder::new().max_size(2).build()?;
2557 let factory = Arc::new(CountingFactory::new());
2558 let pool = Pool::new(config, factory.clone())?;
2559
2560 let conn = pool.acquire().await?;
2561 assert_eq!(factory.created_count(), 1);
2562
2563 let _raw_conn = conn.into_inner();
2565
2566 tokio::time::sleep(Duration::from_millis(50)).await;
2568
2569 let status = pool.status().await;
2570 assert_eq!(status.idle, 0, "into_inner 后连接不应归还");
2571 assert_eq!(status.active, 1, "total_count 仍为 1(连接被外部持有)");
2572 Ok(())
2573 }
2574
2575 #[tokio::test]
2577 async fn test_explicit_release_no_double_return() -> Result<(), Box<dyn std::error::Error>> {
2578 let config = PoolConfigBuilder::new().max_size(2).build()?;
2579 let factory = Arc::new(CountingFactory::new());
2580 let pool = Pool::new(config, factory.clone())?;
2581
2582 let conn = pool.acquire().await?;
2583 pool.release(conn).await;
2584
2585 let status = pool.status().await;
2586 assert_eq!(status.idle, 1, "release 后 idle 应为 1");
2587
2588 let conn = pool.acquire().await?;
2590 pool.release(conn).await;
2591
2592 let status = pool.status().await;
2593 assert_eq!(status.idle, 1, "再次 release 后 idle 仍应为 1(不重复)");
2594 assert_eq!(status.active, 1, "total_count 应为 1");
2595 Ok(())
2596 }
2597
2598 struct CursorMockConn {
2604 rows: QueryRows,
2605 call_count: usize,
2606 }
2607
2608 impl CursorMockConn {
2609 fn new(rows: QueryRows) -> Self {
2610 Self {
2611 rows,
2612 call_count: 0,
2613 }
2614 }
2615 }
2616
2617 impl Connection for CursorMockConn {
2618 fn execute<'a>(
2619 &'a mut self,
2620 _sql: &'a str,
2621 ) -> Pin<Box<dyn Future<Output = Result<u64, crate::DbError>> + Send + 'a>> {
2622 Box::pin(async move { Ok(1) })
2623 }
2624
2625 fn query<'a>(
2626 &'a mut self,
2627 _sql: &'a str,
2628 ) -> Pin<Box<dyn Future<Output = Result<QueryRows, crate::DbError>> + Send + 'a>> {
2629 Box::pin(async move {
2630 self.call_count += 1;
2631 Ok(self.rows.clone())
2632 })
2633 }
2634
2635 fn begin_transaction<'a>(
2636 &'a mut self,
2637 ) -> Pin<Box<dyn Future<Output = Result<(), crate::DbError>> + Send + 'a>> {
2638 Box::pin(async move { Ok(()) })
2639 }
2640
2641 fn commit<'a>(
2642 &'a mut self,
2643 ) -> Pin<Box<dyn Future<Output = Result<(), crate::DbError>> + Send + 'a>> {
2644 Box::pin(async move { Ok(()) })
2645 }
2646
2647 fn rollback<'a>(
2648 &'a mut self,
2649 ) -> Pin<Box<dyn Future<Output = Result<(), crate::DbError>> + Send + 'a>> {
2650 Box::pin(async move { Ok(()) })
2651 }
2652
2653 fn is_connected(&self) -> bool {
2654 true
2655 }
2656
2657 fn ping<'a>(&'a mut self) -> Pin<Box<dyn Future<Output = bool> + Send + 'a>> {
2658 Box::pin(async move { true })
2659 }
2660
2661 fn close<'a>(
2662 &'a mut self,
2663 ) -> Pin<Box<dyn Future<Output = Result<(), crate::DbError>> + Send + 'a>> {
2664 Box::pin(async move { Ok(()) })
2665 }
2666 }
2667
2668 struct CursorOverrideMockConn {
2670 rows: Vec<crate::value::Value>,
2671 yielded: usize,
2672 }
2673
2674 impl CursorOverrideMockConn {
2675 fn new(rows: Vec<crate::value::Value>) -> Self {
2676 Self { rows, yielded: 0 }
2677 }
2678 }
2679
2680 impl Connection for CursorOverrideMockConn {
2681 fn execute<'a>(
2682 &'a mut self,
2683 _sql: &'a str,
2684 ) -> Pin<Box<dyn Future<Output = Result<u64, crate::DbError>> + Send + 'a>> {
2685 Box::pin(async move { Ok(1) })
2686 }
2687
2688 fn query<'a>(
2689 &'a mut self,
2690 _sql: &'a str,
2691 ) -> Pin<Box<dyn Future<Output = Result<QueryRows, crate::DbError>> + Send + 'a>> {
2692 Box::pin(async move {
2694 Ok(self
2695 .rows
2696 .iter()
2697 .map(|v| {
2698 let mut m = std::collections::HashMap::new();
2699 m.insert("v".to_string(), v.clone());
2700 m
2701 })
2702 .collect())
2703 })
2704 }
2705
2706 fn query_stream<'a>(
2708 &'a mut self,
2709 _sql: &'a str,
2710 ) -> Pin<Box<dyn futures::Stream<Item = QueryStreamItem> + Send + 'a>> {
2711 Box::pin(futures::stream::iter(
2712 self.rows
2713 .iter()
2714 .enumerate()
2715 .map(|(i, v)| {
2716 self.yielded = i + 1;
2717 let mut m = std::collections::HashMap::new();
2718 m.insert("v".to_string(), v.clone());
2719 Ok(m)
2720 })
2721 .collect::<Vec<_>>(),
2722 ))
2723 }
2724
2725 fn begin_transaction<'a>(
2726 &'a mut self,
2727 ) -> Pin<Box<dyn Future<Output = Result<(), crate::DbError>> + Send + 'a>> {
2728 Box::pin(async move { Ok(()) })
2729 }
2730
2731 fn commit<'a>(
2732 &'a mut self,
2733 ) -> Pin<Box<dyn Future<Output = Result<(), crate::DbError>> + Send + 'a>> {
2734 Box::pin(async move { Ok(()) })
2735 }
2736
2737 fn rollback<'a>(
2738 &'a mut self,
2739 ) -> Pin<Box<dyn Future<Output = Result<(), crate::DbError>> + Send + 'a>> {
2740 Box::pin(async move { Ok(()) })
2741 }
2742
2743 fn is_connected(&self) -> bool {
2744 true
2745 }
2746
2747 fn ping<'a>(&'a mut self) -> Pin<Box<dyn Future<Output = bool> + Send + 'a>> {
2748 Box::pin(async move { true })
2749 }
2750
2751 fn close<'a>(
2752 &'a mut self,
2753 ) -> Pin<Box<dyn Future<Output = Result<(), crate::DbError>> + Send + 'a>> {
2754 Box::pin(async move { Ok(()) })
2755 }
2756 }
2757
2758 #[tokio::test]
2760 async fn test_query_stream_default_impl_yields_all_rows() {
2761 use futures::StreamExt;
2762 let rows: QueryRows = vec![
2763 std::collections::HashMap::from([
2764 ("id".to_string(), crate::value::Value::I64(1)),
2765 (
2766 "name".to_string(),
2767 crate::value::Value::String("alice".to_string()),
2768 ),
2769 ]),
2770 std::collections::HashMap::from([
2771 ("id".to_string(), crate::value::Value::I64(2)),
2772 (
2773 "name".to_string(),
2774 crate::value::Value::String("bob".to_string()),
2775 ),
2776 ]),
2777 std::collections::HashMap::from([
2778 ("id".to_string(), crate::value::Value::I64(3)),
2779 (
2780 "name".to_string(),
2781 crate::value::Value::String("carol".to_string()),
2782 ),
2783 ]),
2784 ];
2785 let mut conn = CursorMockConn::new(rows);
2786 let mut stream = conn.query_stream("SELECT id, name FROM users");
2787 let mut received: Vec<QueryStreamItem> = Vec::new();
2788 while let Some(item) = stream.next().await {
2789 received.push(item);
2790 }
2791 assert_eq!(received.len(), 3, "应收到 3 行");
2792 assert!(received.iter().all(|r| r.is_ok()), "所有项应为 Ok");
2793 drop(stream);
2794 assert_eq!(conn.call_count, 1, "默认实现应调用 query() 一次");
2795 }
2796
2797 #[tokio::test]
2799 async fn test_query_stream_default_empty_result() {
2800 use futures::StreamExt;
2801 let mut conn = CursorMockConn::new(Vec::new());
2802 let mut stream = conn.query_stream("SELECT * FROM empty_table");
2803 let mut count = 0;
2804 while let Some(_item) = stream.next().await {
2805 count += 1;
2806 }
2807 assert_eq!(count, 0, "空结果集应产生 0 项");
2808 }
2809
2810 #[tokio::test]
2812 async fn test_query_stream_default_error_propagation() {
2813 use futures::StreamExt;
2814 struct ErrorMockConn;
2816 impl Connection for ErrorMockConn {
2817 fn execute<'a>(
2818 &'a mut self,
2819 _sql: &'a str,
2820 ) -> Pin<Box<dyn Future<Output = Result<u64, crate::DbError>> + Send + 'a>>
2821 {
2822 Box::pin(async move { Ok(1) })
2823 }
2824 fn query<'a>(
2825 &'a mut self,
2826 _sql: &'a str,
2827 ) -> Pin<Box<dyn Future<Output = Result<QueryRows, crate::DbError>> + Send + 'a>>
2828 {
2829 Box::pin(async move { Err(crate::DbError::Internal("query failed".to_string())) })
2830 }
2831 fn begin_transaction<'a>(
2832 &'a mut self,
2833 ) -> Pin<Box<dyn Future<Output = Result<(), crate::DbError>> + Send + 'a>> {
2834 Box::pin(async move { Ok(()) })
2835 }
2836 fn commit<'a>(
2837 &'a mut self,
2838 ) -> Pin<Box<dyn Future<Output = Result<(), crate::DbError>> + Send + 'a>> {
2839 Box::pin(async move { Ok(()) })
2840 }
2841 fn rollback<'a>(
2842 &'a mut self,
2843 ) -> Pin<Box<dyn Future<Output = Result<(), crate::DbError>> + Send + 'a>> {
2844 Box::pin(async move { Ok(()) })
2845 }
2846 fn is_connected(&self) -> bool {
2847 true
2848 }
2849 fn ping<'a>(&'a mut self) -> Pin<Box<dyn Future<Output = bool> + Send + 'a>> {
2850 Box::pin(async move { true })
2851 }
2852 fn close<'a>(
2853 &'a mut self,
2854 ) -> Pin<Box<dyn Future<Output = Result<(), crate::DbError>> + Send + 'a>> {
2855 Box::pin(async move { Ok(()) })
2856 }
2857 }
2858 let mut conn = ErrorMockConn;
2859 let mut stream = conn.query_stream("SELECT * FROM bad_table");
2860 let item = stream.next().await;
2861 assert!(item.is_some(), "应产生一项");
2862 assert!(item.unwrap().is_err(), "该项应为 Err");
2863 }
2864
2865 #[tokio::test]
2867 async fn test_query_stream_override_yields_rows_one_by_one() {
2868 use futures::StreamExt;
2869 let rows = vec![
2870 crate::value::Value::I64(10),
2871 crate::value::Value::I64(20),
2872 crate::value::Value::I64(30),
2873 crate::value::Value::I64(40),
2874 crate::value::Value::I64(50),
2875 ];
2876 let mut conn = CursorOverrideMockConn::new(rows);
2877 let values: Vec<i64> = {
2878 let mut stream = conn.query_stream("SELECT v FROM seq");
2879 let mut vals: Vec<i64> = Vec::new();
2880 while let Some(Ok(row)) = stream.next().await {
2881 if let crate::value::Value::I64(v) = row.get("v").unwrap() {
2882 vals.push(*v);
2883 }
2884 }
2885 vals
2886 };
2887 assert_eq!(values, vec![10, 20, 30, 40, 50], "应按顺序收到全部 5 行");
2888 assert_eq!(conn.yielded, 5, "应逐行 yield 5 次(真游标覆盖)");
2889 }
2890
2891 #[tokio::test]
2893 async fn test_query_stream_override_early_drop() {
2894 use futures::StreamExt;
2895 let rows = vec![
2896 crate::value::Value::I64(1),
2897 crate::value::Value::I64(2),
2898 crate::value::Value::I64(3),
2899 ];
2900 let mut conn = CursorOverrideMockConn::new(rows);
2901 {
2902 let mut stream = conn.query_stream("SELECT v FROM seq");
2903 let first = stream.next().await;
2904 assert!(first.is_some(), "第一项应存在");
2905 drop(stream);
2907 }
2908 assert!(conn.is_connected(), "提前 drop 流后连接仍应可用");
2910 }
2911
2912 #[tokio::test]
2914 async fn test_pool_prewarm() -> Result<(), Box<dyn std::error::Error>> {
2915 use std::sync::atomic::AtomicU32;
2916
2917 let create_count = Arc::new(AtomicU32::new(0));
2919 let create_count_clone = create_count.clone();
2920
2921 struct CountingFactory {
2922 count: Arc<AtomicU32>,
2923 }
2924
2925 #[async_trait]
2926 impl ConnectionFactory for CountingFactory {
2927 async fn create(&self) -> Result<Box<dyn Connection>, crate::DbError> {
2928 self.count.fetch_add(1, Ordering::SeqCst);
2929 Ok(Box::new(MockConnection::new()))
2930 }
2931 }
2932
2933 let config = PoolConfigBuilder::new()
2935 .max_size(10)
2936 .min_idle(5)
2937 .prewarm(true)
2938 .build()?;
2939
2940 let factory = Arc::new(CountingFactory {
2941 count: create_count_clone,
2942 });
2943
2944 let pool = Pool::new(config, factory)?;
2945
2946 let status_before = pool.status().await;
2948 assert_eq!(status_before.idle, 0, "预热前 idle 应为 0");
2949
2950 pool.prewarm().await;
2952
2953 let status_after = pool.status().await;
2955 assert!(
2956 status_after.idle >= 5,
2957 "预热后 idle 应 >= 5,实际: {}",
2958 status_after.idle
2959 );
2960
2961 assert_eq!(
2963 create_count.load(Ordering::SeqCst),
2964 5,
2965 "工厂应被调用 5 次(min_idle)"
2966 );
2967
2968 Ok(())
2969 }
2970
2971 #[tokio::test]
2973 async fn test_pool_prewarm_failure_non_blocking() -> Result<(), Box<dyn std::error::Error>> {
2974 use std::sync::atomic::AtomicBool;
2975
2976 struct FailingFactory {
2977 failed: Arc<AtomicBool>,
2978 }
2979
2980 #[async_trait]
2981 impl ConnectionFactory for FailingFactory {
2982 async fn create(&self) -> Result<Box<dyn Connection>, crate::DbError> {
2983 self.failed.store(true, Ordering::SeqCst);
2984 Err(crate::DbError::Internal(
2986 "simulated connection failure".to_string(),
2987 ))
2988 }
2989 }
2990
2991 let failed = Arc::new(AtomicBool::new(false));
2992 let mut config = PoolConfigBuilder::new()
2993 .max_size(10)
2994 .min_idle(3)
2995 .prewarm(true)
2996 .build()?;
2997 config.connection_timeout = std::time::Duration::from_secs(1); let factory = Arc::new(FailingFactory {
3000 failed: failed.clone(),
3001 });
3002
3003 let pool = Pool::new(config, factory)?;
3005 pool.prewarm().await; assert!(failed.load(Ordering::SeqCst), "工厂应被调用且失败");
3009
3010 let status = pool.status().await;
3012 assert_eq!(status.max, 10, "池配置应正常");
3013
3014 Ok(())
3015 }
3016
3017 #[tokio::test]
3019 async fn test_pool_prewarm_disabled() -> Result<(), Box<dyn std::error::Error>> {
3020 use std::sync::atomic::AtomicU32;
3021
3022 let create_count = Arc::new(AtomicU32::new(0));
3023 let create_count_clone = create_count.clone();
3024
3025 struct CountingFactory {
3026 count: Arc<AtomicU32>,
3027 }
3028
3029 #[async_trait]
3030 impl ConnectionFactory for CountingFactory {
3031 async fn create(&self) -> Result<Box<dyn Connection>, crate::DbError> {
3032 self.count.fetch_add(1, Ordering::SeqCst);
3033 Ok(Box::new(MockConnection::new()))
3034 }
3035 }
3036
3037 let config = PoolConfigBuilder::new()
3039 .max_size(10)
3040 .min_idle(5)
3041 .prewarm(false) .build()?;
3043
3044 let factory = Arc::new(CountingFactory {
3045 count: create_count_clone,
3046 });
3047
3048 let pool = Pool::new(config, factory)?;
3049 pool.prewarm().await; assert_eq!(
3053 create_count.load(Ordering::SeqCst),
3054 0,
3055 "prewarm=false 时工厂不应被调用"
3056 );
3057
3058 let status = pool.status().await;
3059 assert_eq!(status.idle, 0, "idle 应为 0");
3060
3061 Ok(())
3062 }
3063
3064 #[tokio::test]
3066 async fn test_pool_new_async_with_prewarm() -> Result<(), Box<dyn std::error::Error>> {
3067 use std::sync::atomic::AtomicU32;
3068
3069 let create_count = Arc::new(AtomicU32::new(0));
3070 let create_count_clone = create_count.clone();
3071
3072 struct CountingFactory {
3073 count: Arc<AtomicU32>,
3074 }
3075
3076 #[async_trait]
3077 impl ConnectionFactory for CountingFactory {
3078 async fn create(&self) -> Result<Box<dyn Connection>, crate::DbError> {
3079 self.count.fetch_add(1, Ordering::SeqCst);
3080 Ok(Box::new(MockConnection::new()))
3081 }
3082 }
3083
3084 let config = PoolConfigBuilder::new()
3085 .max_size(10)
3086 .min_idle(5)
3087 .prewarm(true)
3088 .build()?;
3089
3090 let factory = Arc::new(CountingFactory {
3091 count: create_count_clone,
3092 });
3093
3094 let pool = Pool::new_async(config, factory).await?;
3095
3096 let status = pool.status().await;
3097 assert!(
3098 status.idle >= 5,
3099 "new_async prewarm=true 后 idle 应 >= 5,实际: {}",
3100 status.idle
3101 );
3102 assert_eq!(create_count.load(Ordering::SeqCst), 5, "工厂应被调用 5 次");
3103
3104 Ok(())
3105 }
3106
3107 #[tokio::test]
3109 async fn test_pool_new_async_without_prewarm() -> Result<(), Box<dyn std::error::Error>> {
3110 use std::sync::atomic::AtomicU32;
3111
3112 let create_count = Arc::new(AtomicU32::new(0));
3113 let create_count_clone = create_count.clone();
3114
3115 struct CountingFactory {
3116 count: Arc<AtomicU32>,
3117 }
3118
3119 #[async_trait]
3120 impl ConnectionFactory for CountingFactory {
3121 async fn create(&self) -> Result<Box<dyn Connection>, crate::DbError> {
3122 self.count.fetch_add(1, Ordering::SeqCst);
3123 Ok(Box::new(MockConnection::new()))
3124 }
3125 }
3126
3127 let config = PoolConfigBuilder::new()
3128 .max_size(10)
3129 .min_idle(5)
3130 .prewarm(false)
3131 .build()?;
3132
3133 let factory = Arc::new(CountingFactory {
3134 count: create_count_clone,
3135 });
3136
3137 let pool = Pool::new_async(config, factory).await?;
3138
3139 let status = pool.status().await;
3140 assert_eq!(status.idle, 0, "prewarm=false 时 idle 应为 0");
3141 assert_eq!(create_count.load(Ordering::SeqCst), 0, "工厂不应被调用");
3142
3143 Ok(())
3144 }
3145
3146 #[tokio::test]
3148 async fn test_pool_new_async_failure_non_blocking() -> Result<(), Box<dyn std::error::Error>> {
3149 struct FailingFactory;
3150
3151 #[async_trait]
3152 impl ConnectionFactory for FailingFactory {
3153 async fn create(&self) -> Result<Box<dyn Connection>, crate::DbError> {
3154 Err(crate::DbError::Internal("simulated failure".to_string()))
3155 }
3156 }
3157
3158 let mut config = PoolConfigBuilder::new()
3159 .max_size(10)
3160 .min_idle(3)
3161 .prewarm(true)
3162 .build()?;
3163 config.connection_timeout = std::time::Duration::from_secs(1);
3164
3165 let pool = Pool::new_async(config, Arc::new(FailingFactory)).await?;
3166
3167 let status = pool.status().await;
3168 assert_eq!(status.max, 10, "池配置应正常");
3169
3170 Ok(())
3171 }
3172
3173 #[cfg(feature = "auto-prewarm")]
3175 #[tokio::test]
3176 async fn test_pool_progressive_prewarm() -> Result<(), Box<dyn std::error::Error>> {
3177 use std::sync::atomic::AtomicU32;
3178
3179 let create_count = Arc::new(AtomicU32::new(0));
3180 let create_count_clone = create_count.clone();
3181
3182 struct CountingFactory {
3183 count: Arc<AtomicU32>,
3184 }
3185
3186 #[async_trait]
3187 impl ConnectionFactory for CountingFactory {
3188 async fn create(&self) -> Result<Box<dyn Connection>, crate::DbError> {
3189 self.count.fetch_add(1, Ordering::SeqCst);
3190 Ok(Box::new(MockConnection::new()))
3191 }
3192 }
3193
3194 let config = PoolConfigBuilder::new()
3195 .max_size(20)
3196 .min_idle(6)
3197 .prewarm(true)
3198 .build()?;
3199
3200 let factory = Arc::new(CountingFactory {
3201 count: create_count_clone,
3202 });
3203
3204 let pool = Pool::new(config, factory)?;
3205
3206 let progress = crate::prewarm::PrewarmProgress::new(6);
3207 pool.progressive_prewarm(
3208 2,
3209 std::time::Duration::from_millis(5),
3210 std::time::Duration::from_secs(10),
3211 &progress,
3212 )
3213 .await;
3214
3215 let snap = progress.snapshot();
3216 assert!(
3217 snap.warmed >= 6,
3218 "progressive_prewarm 后 warmed 应 >= 6,实际: {}",
3219 snap.warmed
3220 );
3221 assert!(snap.is_completed, "应标记完成");
3222 assert_eq!(create_count.load(Ordering::SeqCst), 6, "工厂应被调用 6 次");
3223
3224 let status = pool.status().await;
3225 assert!(status.idle >= 6, "池中 idle 应 >= 6");
3226
3227 Ok(())
3228 }
3229
3230 #[cfg(feature = "auto-prewarm")]
3232 #[tokio::test]
3233 async fn test_pool_progressive_prewarm_timeout_zero() -> Result<(), Box<dyn std::error::Error>>
3234 {
3235 use std::sync::atomic::AtomicU32;
3236
3237 let create_count = Arc::new(AtomicU32::new(0));
3238 let create_count_clone = create_count.clone();
3239
3240 struct CountingFactory {
3241 count: Arc<AtomicU32>,
3242 }
3243
3244 #[async_trait]
3245 impl ConnectionFactory for CountingFactory {
3246 async fn create(&self) -> Result<Box<dyn Connection>, crate::DbError> {
3247 self.count.fetch_add(1, Ordering::SeqCst);
3248 Ok(Box::new(MockConnection::new()))
3249 }
3250 }
3251
3252 let config = PoolConfigBuilder::new()
3253 .max_size(20)
3254 .min_idle(10)
3255 .prewarm(true)
3256 .build()?;
3257
3258 let factory = Arc::new(CountingFactory {
3259 count: create_count_clone,
3260 });
3261
3262 let pool = Pool::new(config, factory)?;
3263
3264 let progress = crate::prewarm::PrewarmProgress::new(10);
3265 pool.progressive_prewarm(
3266 2,
3267 std::time::Duration::from_millis(5),
3268 std::time::Duration::ZERO,
3269 &progress,
3270 )
3271 .await;
3272
3273 let snap = progress.snapshot();
3274 assert!(snap.is_completed, "应标记完成");
3275 assert!(
3276 snap.warmed <= 2,
3277 "total_timeout=0 时最多建一批(batch_size=2),实际: {}",
3278 snap.warmed
3279 );
3280
3281 Ok(())
3282 }
3283
3284 #[cfg(feature = "auto-prewarm")]
3286 #[tokio::test]
3287 async fn test_pool_progressive_prewarm_disabled() -> Result<(), Box<dyn std::error::Error>> {
3288 use std::sync::atomic::AtomicU32;
3289
3290 let create_count = Arc::new(AtomicU32::new(0));
3291 let create_count_clone = create_count.clone();
3292
3293 struct CountingFactory {
3294 count: Arc<AtomicU32>,
3295 }
3296
3297 #[async_trait]
3298 impl ConnectionFactory for CountingFactory {
3299 async fn create(&self) -> Result<Box<dyn Connection>, crate::DbError> {
3300 self.count.fetch_add(1, Ordering::SeqCst);
3301 Ok(Box::new(MockConnection::new()))
3302 }
3303 }
3304
3305 let config = PoolConfigBuilder::new()
3306 .max_size(20)
3307 .min_idle(10)
3308 .prewarm(false)
3309 .build()?;
3310
3311 let factory = Arc::new(CountingFactory {
3312 count: create_count_clone,
3313 });
3314
3315 let pool = Pool::new(config, factory)?;
3316
3317 let progress = crate::prewarm::PrewarmProgress::new(10);
3318 pool.progressive_prewarm(
3319 2,
3320 std::time::Duration::from_millis(5),
3321 std::time::Duration::from_secs(10),
3322 &progress,
3323 )
3324 .await;
3325
3326 let snap = progress.snapshot();
3327 assert!(snap.is_completed, "应标记完成");
3328 assert_eq!(snap.warmed, 0, "prewarm=false 时不应建连");
3329 assert_eq!(create_count.load(Ordering::SeqCst), 0, "工厂不应被调用");
3330
3331 Ok(())
3332 }
3333
3334 #[cfg(feature = "auto-prewarm")]
3336 #[tokio::test]
3337 async fn test_pool_progressive_prewarm_failure_non_blocking(
3338 ) -> Result<(), Box<dyn std::error::Error>> {
3339 struct FailingFactory;
3340
3341 #[async_trait]
3342 impl ConnectionFactory for FailingFactory {
3343 async fn create(&self) -> Result<Box<dyn Connection>, crate::DbError> {
3344 Err(crate::DbError::Internal("simulated failure".to_string()))
3345 }
3346 }
3347
3348 let mut config = PoolConfigBuilder::new()
3349 .max_size(20)
3350 .min_idle(5)
3351 .prewarm(true)
3352 .build()?;
3353 config.connection_timeout = std::time::Duration::from_secs(1);
3354
3355 let pool = Pool::new(config, Arc::new(FailingFactory))?;
3356
3357 let progress = crate::prewarm::PrewarmProgress::new(5);
3358 pool.progressive_prewarm(
3359 2,
3360 std::time::Duration::from_millis(5),
3361 std::time::Duration::from_secs(5),
3362 &progress,
3363 )
3364 .await;
3365
3366 let snap = progress.snapshot();
3367 assert!(snap.is_completed, "应标记完成");
3368 assert_eq!(snap.warmed, 0, "全部失败时 warmed=0");
3369 assert!(snap.failed > 0, "应有失败记录");
3370
3371 Ok(())
3372 }
3373
3374 #[tokio::test]
3376 async fn test_pool_metrics_acquire_release() -> Result<(), Box<dyn std::error::Error>> {
3377 let config = PoolConfigBuilder::new().max_size(10).build()?;
3378 let pool = Pool::new(config, Arc::new(MockConnectionFactory))?;
3379
3380 let metrics = pool.pool_metrics();
3381 assert_eq!(metrics.acquire_count, 0);
3382 assert_eq!(metrics.release_count, 0);
3383 assert_eq!(metrics.connection_created_count, 0);
3384
3385 let conn = pool.acquire().await?;
3386 let metrics = pool.pool_metrics();
3387 assert_eq!(metrics.acquire_count, 1);
3388 assert_eq!(metrics.connection_created_count, 1);
3389 assert_eq!(metrics.acquire_failed_count, 0);
3390
3391 pool.release(conn).await;
3392 let metrics = pool.pool_metrics();
3393 assert_eq!(metrics.release_count, 1);
3394 assert_eq!(metrics.connection_closed_count, 0);
3396
3397 Ok(())
3398 }
3399
3400 #[tokio::test]
3402 async fn test_pool_metrics_acquire_failed() -> Result<(), Box<dyn std::error::Error>> {
3403 struct FailingFactory;
3404
3405 #[async_trait]
3406 impl ConnectionFactory for FailingFactory {
3407 async fn create(&self) -> Result<Box<dyn Connection>, crate::DbError> {
3408 Err(crate::DbError::Internal("simulated failure".to_string()))
3409 }
3410 }
3411
3412 let config = PoolConfigBuilder::new().max_size(10).build()?;
3413 let pool = Pool::new(config, Arc::new(FailingFactory))?;
3414
3415 let result = pool.acquire().await;
3416 assert!(result.is_err());
3417
3418 let metrics = pool.pool_metrics();
3419 assert_eq!(metrics.acquire_failed_count, 1);
3420 assert_eq!(metrics.acquire_count, 0);
3421
3422 Ok(())
3423 }
3424
3425 #[tokio::test]
3427 async fn test_pool_metrics_connection_closed() -> Result<(), Box<dyn std::error::Error>> {
3428 let config = PoolConfigBuilder::new().max_size(10).build()?;
3429 let pool = Pool::new(config, Arc::new(MockConnectionFactory))?;
3430
3431 let conn = pool.acquire().await?;
3432 pool.release(conn).await;
3433
3434 let status = pool.status().await;
3435 assert_eq!(status.idle, 1);
3436
3437 pool.close_all().await;
3438
3439 let metrics = pool.pool_metrics();
3440 assert_eq!(metrics.connection_closed_count, 1);
3441 assert_eq!(metrics.connection_created_count, 1);
3442
3443 Ok(())
3444 }
3445
3446 #[test]
3448 fn test_pool_metrics_average_wait_time() {
3449 let metrics = PoolMetrics {
3450 acquire_count: 4,
3451 acquire_failed_count: 1,
3452 acquire_wait_time: Duration::from_millis(200),
3453 release_count: 4,
3454 connection_created_count: 2,
3455 connection_closed_count: 0,
3456 };
3457 assert_eq!(
3458 metrics.average_acquire_wait_time(),
3459 Duration::from_millis(50)
3460 );
3461
3462 let empty = PoolMetrics::default();
3464 assert_eq!(empty.average_acquire_wait_time(), Duration::ZERO);
3465 }
3466
3467 #[tokio::test]
3468 async fn test_shutdown_with_timeout_fast_return_when_empty() {
3469 let factory = Arc::new(MockConnectionFactory) as Arc<dyn ConnectionFactory>;
3470 let config = PoolConfigBuilder::new().max_size(4).build().unwrap();
3471 let pool = Pool::new(config, factory).unwrap();
3472 pool.shutdown_with_timeout(Duration::from_secs(1)).await;
3473 assert!(pool.closed.load(Ordering::SeqCst));
3474 assert_eq!(pool.total_count.load(Ordering::SeqCst), 0);
3475 }
3476
3477 #[tokio::test]
3478 async fn test_shutdown_delegates_to_shutdown_with_timeout() {
3479 let factory = Arc::new(MockConnectionFactory) as Arc<dyn ConnectionFactory>;
3480 let config = PoolConfigBuilder::new().max_size(4).build().unwrap();
3481 let pool = Pool::new(config, factory).unwrap();
3482 pool.shutdown().await;
3483 assert!(pool.closed.load(Ordering::SeqCst));
3484 }
3485
3486 #[tokio::test]
3487 async fn test_shutdown_with_timeout_idempotent() {
3488 let factory = Arc::new(MockConnectionFactory) as Arc<dyn ConnectionFactory>;
3489 let config = PoolConfigBuilder::new().max_size(4).build().unwrap();
3490 let pool = Pool::new(config, factory).unwrap();
3491 pool.shutdown_with_timeout(Duration::from_secs(1)).await;
3492 let count_after_first = pool.total_count.load(Ordering::SeqCst);
3493 pool.shutdown_with_timeout(Duration::from_secs(1)).await;
3494 let count_after_second = pool.total_count.load(Ordering::SeqCst);
3495 assert_eq!(count_after_first, count_after_second);
3496 }
3497
3498 #[tokio::test]
3499 async fn test_shutdown_with_timeout_rejects_new_acquire() {
3500 let factory = Arc::new(MockConnectionFactory) as Arc<dyn ConnectionFactory>;
3501 let config = PoolConfigBuilder::new().max_size(4).build().unwrap();
3502 let pool = Pool::new(config, factory).unwrap();
3503 pool.shutdown_with_timeout(Duration::from_secs(1)).await;
3504 let result = pool.acquire().await;
3505 assert!(result.is_err());
3506 }
3507}
3508
3509#[cfg(all(test, feature = "prod-pool-tuning"))]
3510mod pool_prod_tests {
3511 use super::*;
3512
3513 struct MockFactory;
3514
3515 #[async_trait]
3516 impl ConnectionFactory for MockFactory {
3517 async fn create(&self) -> Result<Box<dyn Connection>, crate::DbError> {
3518 Ok(Box::new(MockConn))
3519 }
3520 }
3521
3522 struct MockConn;
3523
3524 impl Connection for MockConn {
3525 fn execute<'a>(
3526 &'a mut self,
3527 _sql: &'a str,
3528 ) -> Pin<Box<dyn Future<Output = Result<u64, crate::DbError>> + Send + 'a>> {
3529 Box::pin(async move { Ok(1) })
3530 }
3531 fn query<'a>(
3532 &'a mut self,
3533 _sql: &'a str,
3534 ) -> Pin<
3535 Box<
3536 dyn Future<
3537 Output = Result<
3538 Vec<std::collections::HashMap<String, crate::value::Value>>,
3539 crate::DbError,
3540 >,
3541 > + Send
3542 + 'a,
3543 >,
3544 > {
3545 Box::pin(async move { Ok(vec![]) })
3546 }
3547 fn begin_transaction<'a>(
3548 &'a mut self,
3549 ) -> Pin<Box<dyn Future<Output = Result<(), crate::DbError>> + Send + 'a>> {
3550 Box::pin(async move { Ok(()) })
3551 }
3552 fn commit<'a>(
3553 &'a mut self,
3554 ) -> Pin<Box<dyn Future<Output = Result<(), crate::DbError>> + Send + 'a>> {
3555 Box::pin(async move { Ok(()) })
3556 }
3557 fn rollback<'a>(
3558 &'a mut self,
3559 ) -> Pin<Box<dyn Future<Output = Result<(), crate::DbError>> + Send + 'a>> {
3560 Box::pin(async move { Ok(()) })
3561 }
3562 fn is_connected(&self) -> bool {
3563 true
3564 }
3565 fn ping<'a>(&'a mut self) -> Pin<Box<dyn Future<Output = bool> + Send + 'a>> {
3566 Box::pin(async move { true })
3567 }
3568 fn close<'a>(
3569 &'a mut self,
3570 ) -> Pin<Box<dyn Future<Output = Result<(), crate::DbError>> + Send + 'a>> {
3571 Box::pin(async move { Ok(()) })
3572 }
3573 }
3574
3575 #[test]
3576 fn test_pool_prod_config_validate_ok() {
3577 let config = PoolProdConfig::new(
3578 50,
3579 Duration::from_secs(10),
3580 Duration::from_secs(600),
3581 Duration::from_secs(5),
3582 Duration::from_secs(30),
3583 5,
3584 true,
3585 );
3586 assert!(config.validate().is_ok());
3587 }
3588
3589 #[test]
3590 fn test_pool_prod_config_max_size_zero_rejected() {
3591 let config = PoolProdConfig::default();
3592 let mut c = config;
3593 c.max_size = 0;
3594 let err = c.validate().unwrap_err();
3595 assert!(err.to_string().contains("max_size must be positive"));
3596 }
3597
3598 #[test]
3599 fn test_pool_prod_config_min_idle_exceeds_max_size() {
3600 let config = PoolProdConfig::new(
3601 10,
3602 Duration::from_secs(10),
3603 Duration::from_secs(600),
3604 Duration::from_secs(5),
3605 Duration::from_secs(30),
3606 20,
3607 false,
3608 );
3609 let err = config.validate().unwrap_err();
3610 assert!(err.to_string().contains("min_idle cannot exceed max_size"));
3611 }
3612
3613 #[test]
3614 fn test_pool_prod_config_to_pool_config() {
3615 let config = PoolProdConfig::new(
3616 50,
3617 Duration::from_secs(10),
3618 Duration::from_secs(600),
3619 Duration::from_secs(5),
3620 Duration::from_secs(30),
3621 5,
3622 true,
3623 );
3624 let pool_config = config.to_pool_config();
3625 assert_eq!(pool_config.max_size, 50);
3626 assert_eq!(pool_config.min_idle, 5);
3627 assert_eq!(pool_config.acquire_timeout, Duration::from_secs(10));
3628 assert!(pool_config.prewarm);
3629 }
3630
3631 #[tokio::test]
3632 async fn test_pool_prod_config_runtime_resize() {
3633 let factory = Arc::new(MockFactory) as Arc<dyn ConnectionFactory>;
3634 let config = PoolProdConfig::default();
3635 let pool = Pool::new(config.to_pool_config(), factory).unwrap();
3636 assert_eq!(pool.max_size(), 100);
3637 pool.resize(50);
3638 assert_eq!(pool.max_size(), 50);
3639 }
3640}
3641
3642#[cfg(all(test, feature = "prod-leak-detection"))]
3643mod leak_prod_tests {
3644 use super::*;
3645
3646 #[test]
3647 fn test_leak_config_default() {
3648 let config = LeakDetectionConfig::default();
3649 assert!(!config.enabled);
3650 assert_eq!(config.interval, Duration::from_secs(60));
3651 assert_eq!(config.threshold, 5);
3652 }
3653
3654 #[test]
3655 fn test_leak_config_validate_ok() {
3656 let config =
3657 LeakDetectionConfig::new(true, Duration::from_secs(30), 10, Duration::from_secs(60));
3658 assert!(config.validate().is_ok());
3659 }
3660
3661 #[test]
3662 fn test_leak_config_interval_zero_rejected() {
3663 let config = LeakDetectionConfig::new(true, Duration::ZERO, 10, Duration::from_secs(60));
3664 assert!(config.validate().is_err());
3665 }
3666
3667 #[test]
3668 fn test_leak_report_empty() {
3669 let report = LeakReport::empty();
3670 assert_eq!(report.borrowed_count, 0);
3671 assert!(report.suspected_leaks.is_empty());
3672 }
3673}