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.closed.store(true, Ordering::SeqCst);
1698 self.notify.notify_waiters();
1700 self.close_all().await;
1702 let deadline = Instant::now() + Duration::from_secs(30);
1704 while self.total_count.load(Ordering::SeqCst) > 0 {
1705 if Instant::now() >= deadline {
1706 break;
1707 }
1708 tokio::time::sleep(Duration::from_millis(100)).await;
1709 }
1710 }
1711
1712 pub fn resize(&self, new_max: usize) {
1720 self.set_max_size(new_max as u32);
1721 }
1722
1723 pub fn set_max_size(&self, new_max: u32) {
1725 self.dynamic_max_size.store(new_max, Ordering::SeqCst);
1726 }
1727
1728 pub fn max_size(&self) -> u32 {
1730 self.dynamic_max_size.load(Ordering::Acquire)
1731 }
1732
1733 pub async fn warmup(&self, min_idle: usize) -> Result<(), PoolError> {
1737 for _ in 0..min_idle {
1738 let current_max = self.dynamic_max_size.load(Ordering::Acquire);
1739 let current = self.total_count.load(Ordering::Acquire);
1740 if current >= current_max {
1741 break;
1742 }
1743 match self.total_count.compare_exchange(
1745 current,
1746 current + 1,
1747 Ordering::SeqCst,
1748 Ordering::Acquire,
1749 ) {
1750 Ok(_) => {}
1751 Err(_) => continue, }
1753 match self.factory.create().await {
1754 Ok(conn) => {
1755 let now = Instant::now();
1756 let pooled = PooledConnection {
1757 conn,
1758 created_at: now,
1759 last_used_at: now,
1760 pool: None,
1761 };
1762 if let Err(rejected) = self.idle.push(pooled) {
1763 self.close_connection(rejected).await;
1765 self.total_count.fetch_sub(1, Ordering::SeqCst);
1766 }
1767 self.emit_event(PoolEvent::ConnectionCreated);
1768 }
1769 Err(_) => {
1770 self.total_count.fetch_sub(1, Ordering::SeqCst);
1772 break;
1773 }
1774 }
1775 }
1776 Ok(())
1777 }
1778
1779 pub async fn query_with_timeout(&self, sql: &str) -> Result<QueryRows, crate::DbError> {
1784 let timeout = self.config.query_timeout.unwrap_or(Duration::from_secs(30));
1785 let mut conn = self.acquire().await.map_err(crate::DbError::PoolError)?;
1786 tokio::time::timeout(timeout, conn.query(sql))
1787 .await
1788 .map_err(|_| crate::DbError::QueryError(format!("Query timeout after {:?}", timeout)))?
1789 }
1790}
1791
1792#[cfg(test)]
1793mod tests {
1794 use super::*;
1795
1796 struct MockConnection {
1798 connected: bool,
1799 }
1800
1801 impl MockConnection {
1802 fn new() -> Self {
1803 Self { connected: true }
1804 }
1805 }
1806
1807 impl Connection for MockConnection {
1808 fn execute<'a>(
1809 &'a mut self,
1810 _sql: &'a str,
1811 ) -> Pin<Box<dyn Future<Output = Result<u64, crate::DbError>> + Send + 'a>> {
1812 Box::pin(async move { Ok(1) })
1813 }
1814
1815 fn query<'a>(
1816 &'a mut self,
1817 _sql: &'a str,
1818 ) -> Pin<
1819 Box<
1820 dyn Future<
1821 Output = Result<
1822 Vec<std::collections::HashMap<String, crate::value::Value>>,
1823 crate::DbError,
1824 >,
1825 > + Send
1826 + 'a,
1827 >,
1828 > {
1829 Box::pin(async move { Ok(vec![]) })
1830 }
1831
1832 fn begin_transaction<'a>(
1833 &'a mut self,
1834 ) -> Pin<Box<dyn Future<Output = Result<(), crate::DbError>> + Send + 'a>> {
1835 Box::pin(async move { Ok(()) })
1836 }
1837
1838 fn commit<'a>(
1839 &'a mut self,
1840 ) -> Pin<Box<dyn Future<Output = Result<(), crate::DbError>> + Send + 'a>> {
1841 Box::pin(async move { Ok(()) })
1842 }
1843
1844 fn rollback<'a>(
1845 &'a mut self,
1846 ) -> Pin<Box<dyn Future<Output = Result<(), crate::DbError>> + Send + 'a>> {
1847 Box::pin(async move { Ok(()) })
1848 }
1849
1850 fn is_connected(&self) -> bool {
1851 self.connected
1852 }
1853
1854 fn ping<'a>(&'a mut self) -> Pin<Box<dyn Future<Output = bool> + Send + 'a>> {
1855 Box::pin(async move { true })
1856 }
1857
1858 fn close<'a>(
1859 &'a mut self,
1860 ) -> Pin<Box<dyn Future<Output = Result<(), crate::DbError>> + Send + 'a>> {
1861 Box::pin(async move {
1862 self.connected = false;
1863 Ok(())
1864 })
1865 }
1866 }
1867
1868 struct MockConnectionFactory;
1869
1870 #[async_trait]
1871 impl ConnectionFactory for MockConnectionFactory {
1872 async fn create(&self) -> Result<Box<dyn Connection>, crate::DbError> {
1873 Ok(Box::new(MockConnection::new()))
1874 }
1875 }
1876
1877 #[tokio::test]
1878 async fn test_pool_config_builder() -> Result<(), Box<dyn std::error::Error>> {
1879 let config = PoolConfigBuilder::new().max_size(50).min_idle(10).build()?;
1880
1881 assert_eq!(config.max_size, 50);
1882 assert_eq!(config.min_idle, 10);
1883 Ok(())
1884 }
1885
1886 #[test]
1887 fn test_pool_status_display() {
1888 let status = PoolStatus {
1889 idle: 5,
1890 active: 10,
1891 max: 100,
1892 min: 5,
1893 waiters: 0,
1894 };
1895
1896 let display = format!("{:?}", status);
1897 assert!(display.contains("idle"));
1898 assert!(display.contains("active"));
1899 }
1900
1901 #[test]
1902 fn test_default_pool_config() {
1903 let config = PoolConfig::default();
1904 assert_eq!(config.max_size, 100);
1905 assert_eq!(config.min_idle, 0);
1906 assert_eq!(config.acquire_timeout.as_secs(), 30);
1907 assert_eq!(config.idle_timeout.as_secs(), 600);
1908 assert_eq!(config.max_lifetime.as_secs(), 1800);
1909 }
1910
1911 #[tokio::test]
1912 async fn test_pool_config_clone() {
1913 let config = PoolConfig::default();
1914 let cloned = config.clone();
1915 assert_eq!(cloned.max_size, config.max_size);
1916 assert_eq!(cloned.min_idle, config.min_idle);
1917 }
1918
1919 #[test]
1920 fn test_pool_config_builder_default() -> Result<(), Box<dyn std::error::Error>> {
1921 let builder = PoolConfigBuilder::new();
1922 let config = builder.build()?;
1923 assert_eq!(config.max_size, 100);
1924 Ok(())
1925 }
1926
1927 #[test]
1928 fn test_pool_config_validate() {
1929 let result = PoolConfigBuilder::new().max_size(0).build();
1930 assert!(result.is_err());
1931
1932 let result = PoolConfigBuilder::new().max_size(10).min_idle(20).build();
1933 assert!(result.is_err());
1934 }
1935
1936 #[test]
1937 fn test_pool_config_validate_duration_upper_bound() {
1938 use std::time::Duration;
1939
1940 let config = PoolConfig {
1942 max_size: 10,
1943 min_idle: 1,
1944 acquire_timeout: Duration::from_secs(u64::MAX),
1945 idle_timeout: Duration::from_secs(1),
1946 max_lifetime: Duration::from_secs(1),
1947 connection_timeout: Duration::from_secs(5),
1948 tls: None,
1949 query_timeout: None,
1950 max_rows: None,
1951 memory_limit: None,
1952 on_event: None,
1953 test_before_acquire: false,
1954 prewarm: false,
1955 };
1956 assert!(config.validate().is_err());
1957
1958 let config = PoolConfig {
1960 max_size: 10,
1961 min_idle: 1,
1962 acquire_timeout: Duration::from_secs(u32::MAX as u64),
1963 idle_timeout: Duration::from_secs(1),
1964 max_lifetime: Duration::from_secs(1),
1965 connection_timeout: Duration::from_secs(5),
1966 tls: None,
1967 query_timeout: None,
1968 max_rows: None,
1969 memory_limit: None,
1970 on_event: None,
1971 test_before_acquire: false,
1972 prewarm: false,
1973 };
1974 assert!(config.validate().is_ok());
1975
1976 let config = PoolConfig {
1978 max_size: 10,
1979 min_idle: 1,
1980 acquire_timeout: Duration::from_secs(u32::MAX as u64 + 1),
1981 idle_timeout: Duration::from_secs(1),
1982 max_lifetime: Duration::from_secs(1),
1983 connection_timeout: Duration::from_secs(5),
1984 tls: None,
1985 query_timeout: None,
1986 max_rows: None,
1987 memory_limit: None,
1988 on_event: None,
1989 test_before_acquire: false,
1990 prewarm: false,
1991 };
1992 assert!(config.validate().is_err());
1993 }
1994
1995 #[test]
1996 fn test_pool_config_test_before_acquire_default() {
1997 let config = PoolConfig::default();
1999 assert!(!config.test_before_acquire);
2000 }
2001
2002 #[test]
2003 fn test_pool_config_builder_test_before_acquire() {
2004 let config = PoolConfigBuilder::new()
2006 .test_before_acquire(true)
2007 .build()
2008 .unwrap();
2009 assert!(config.test_before_acquire);
2010 }
2011
2012 #[tokio::test]
2013 async fn test_pool_acquire_and_release() -> Result<(), Box<dyn std::error::Error>> {
2014 let config = PoolConfigBuilder::new().max_size(5).min_idle(1).build()?;
2015 let factory = Arc::new(MockConnectionFactory);
2016 let pool = Pool::new(config, factory)?;
2017
2018 let conn = pool.acquire().await?;
2019 let status = pool.status().await;
2020 assert_eq!(status.active, 1);
2021 assert_eq!(status.idle, 0);
2022
2023 pool.release(conn).await;
2024 let status = pool.status().await;
2025 assert_eq!(status.idle, 1);
2026
2027 let _conn2 = pool.acquire().await?;
2029 let status = pool.status().await;
2030 assert_eq!(status.idle, 0);
2031 Ok(())
2032 }
2033
2034 #[tokio::test]
2035 async fn test_pool_status() -> Result<(), Box<dyn std::error::Error>> {
2036 let config = PoolConfigBuilder::new().max_size(10).min_idle(2).build()?;
2037 let factory = Arc::new(MockConnectionFactory);
2038 let pool = Pool::new(config, factory)?;
2039
2040 let status = pool.status().await;
2041 assert_eq!(status.max, 10);
2042 assert_eq!(status.min, 2);
2043 assert_eq!(status.active, 0);
2044 Ok(())
2045 }
2046
2047 #[tokio::test]
2048 async fn test_pool_close_all() -> Result<(), Box<dyn std::error::Error>> {
2049 let config = PoolConfigBuilder::new().max_size(5).build()?;
2050 let factory = Arc::new(MockConnectionFactory);
2051 let pool = Pool::new(config, factory)?;
2052
2053 let conn1 = pool.acquire().await?;
2055 let conn2 = pool.acquire().await?;
2056 pool.release(conn1).await;
2057 pool.release(conn2).await;
2058
2059 pool.close_all().await;
2060 let status = pool.status().await;
2061 assert_eq!(status.idle, 0);
2062 assert_eq!(status.active, 0);
2063 Ok(())
2064 }
2065
2066 #[tokio::test]
2067 async fn test_pool_reap_idle() -> Result<(), Box<dyn std::error::Error>> {
2068 let config = PoolConfigBuilder::new()
2069 .max_size(5)
2070 .idle_timeout(0) .build()?;
2072 let factory = Arc::new(MockConnectionFactory);
2073 let pool = Pool::new(config, factory)?;
2074
2075 let conn = pool.acquire().await?;
2076 pool.release(conn).await;
2077
2078 tokio::time::sleep(tokio::time::Duration::from_millis(10)).await;
2080
2081 pool.reap_idle().await;
2082 let status = pool.status().await;
2083 assert_eq!(status.idle, 0);
2084 Ok(())
2085 }
2086
2087 #[tokio::test]
2093 async fn test_h7_acquire_timeout_default_30s() {
2094 let config = PoolConfig::default();
2095 assert_eq!(
2096 config.acquire_timeout,
2097 Duration::from_secs(30),
2098 "H-7: acquire_timeout 默认应为 30s"
2099 );
2100 }
2101
2102 #[tokio::test]
2104 async fn test_h7_acquire_timeout_configurable() -> Result<(), Box<dyn std::error::Error>> {
2105 let config = PoolConfigBuilder::new()
2106 .max_size(1)
2107 .acquire_timeout(5) .build()?;
2109 assert_eq!(config.acquire_timeout, Duration::from_secs(5));
2110
2111 let factory = Arc::new(MockConnectionFactory);
2113 let pool = Pool::new(config, factory)?;
2114 let _conn1 = pool.acquire().await?;
2115
2116 let fast_config = PoolConfigBuilder::new()
2118 .max_size(1)
2119 .acquire_timeout(0) .build()?;
2121 let fast_pool = Pool::new(fast_config, Arc::new(MockConnectionFactory))?;
2124 let _fast_conn = fast_pool.acquire().await?; let result = fast_pool.acquire().await;
2126 assert!(
2127 matches!(result, Err(PoolError::Timeout)),
2128 "H-7: 应返回 Timeout"
2129 );
2130 Ok(())
2131 }
2132
2133 #[tokio::test]
2136 async fn test_m7_health_check_removes_nothing_when_all_healthy(
2137 ) -> Result<(), Box<dyn std::error::Error>> {
2138 let config = PoolConfigBuilder::new().max_size(5).build()?;
2140 let factory = Arc::new(MockConnectionFactory);
2141 let pool = Pool::new(config, factory)?;
2142
2143 let conn1 = pool.acquire().await?;
2145 let conn2 = pool.acquire().await?;
2146 let conn3 = pool.acquire().await?;
2147 pool.release(conn1).await;
2148 pool.release(conn2).await;
2149 pool.release(conn3).await;
2150
2151 let removed = pool.health_check().await;
2152 assert_eq!(removed, 0, "Healthy connections should not be removed");
2153
2154 let status = pool.status().await;
2155 assert_eq!(status.idle, 3);
2156 assert_eq!(status.active, 3);
2157 Ok(())
2158 }
2159
2160 #[tokio::test]
2161 async fn test_m7_health_check_returns_zero_for_empty_pool(
2162 ) -> Result<(), Box<dyn std::error::Error>> {
2163 let config = PoolConfigBuilder::new().max_size(5).build()?;
2164 let factory = Arc::new(MockConnectionFactory);
2165 let pool = Pool::new(config, factory)?;
2166
2167 let removed = pool.health_check().await;
2168 assert_eq!(removed, 0);
2169 Ok(())
2170 }
2171
2172 struct CountingFactory {
2176 count: AtomicU32,
2177 }
2178
2179 impl CountingFactory {
2180 fn new() -> Self {
2181 Self {
2182 count: AtomicU32::new(0),
2183 }
2184 }
2185 fn created_count(&self) -> u32 {
2186 self.count.load(Ordering::SeqCst)
2187 }
2188 }
2189
2190 #[async_trait]
2191 impl ConnectionFactory for CountingFactory {
2192 async fn create(&self) -> Result<Box<dyn Connection>, crate::DbError> {
2193 self.count.fetch_add(1, Ordering::SeqCst);
2194 Ok(Box::new(MockConnection::new()))
2195 }
2196 }
2197
2198 #[tokio::test]
2204 async fn test_production_bug_max_lifetime_never_expires(
2205 ) -> Result<(), Box<dyn std::error::Error>> {
2206 let config = PoolConfig {
2209 max_size: 5,
2210 min_idle: 0,
2211 acquire_timeout: Duration::from_secs(30),
2212 idle_timeout: Duration::from_secs(600),
2213 max_lifetime: Duration::from_millis(100), connection_timeout: Duration::from_secs(10),
2215 tls: None,
2216 query_timeout: None,
2217 max_rows: None,
2218 memory_limit: None,
2219 on_event: None,
2220 test_before_acquire: false,
2221 prewarm: false,
2222 };
2223 let factory = Arc::new(CountingFactory::new());
2224 let pool = Pool::new(config, factory.clone())?;
2225
2226 let conn = pool.acquire().await?;
2228 assert_eq!(factory.created_count(), 1, "应创建 1 个连接");
2229
2230 pool.release(conn).await;
2232
2233 tokio::time::sleep(Duration::from_millis(150)).await;
2235
2236 let conn2 = pool.acquire().await?;
2238
2239 assert_eq!(
2242 factory.created_count(),
2243 2,
2244 "超过 max_lifetime 后应创建新连接(旧连接应被回收)"
2245 );
2246
2247 pool.release(conn2).await;
2248 Ok(())
2249 }
2250
2251 #[tokio::test]
2258 async fn test_drop_auto_release_connection() -> Result<(), Box<dyn std::error::Error>> {
2259 let config = PoolConfigBuilder::new().max_size(2).build()?;
2260 let factory = Arc::new(CountingFactory::new());
2261 let pool = Pool::new(config, factory.clone())?;
2262
2263 {
2265 let _conn = pool.acquire().await?;
2266 assert_eq!(factory.created_count(), 1, "应创建 1 个连接");
2267 let status = pool.status().await;
2268 assert_eq!(status.active, 1, "active 应为 1");
2269 assert_eq!(status.idle, 0, "idle 应为 0");
2270 }
2272
2273 tokio::time::sleep(Duration::from_millis(50)).await;
2275
2276 let status = pool.status().await;
2278 assert_eq!(status.idle, 1, "Drop 后连接应自动归还,idle 应为 1");
2279 assert_eq!(status.active, 1, "total_count 应为 1");
2280 assert_eq!(factory.created_count(), 1, "应复用归还的连接,不创建新连接");
2281 Ok(())
2282 }
2283
2284 #[tokio::test]
2286 async fn test_drop_auto_release_then_reuse() -> Result<(), Box<dyn std::error::Error>> {
2287 let config = PoolConfigBuilder::new().max_size(1).build()?;
2288 let factory = Arc::new(CountingFactory::new());
2289 let pool = Pool::new(config, factory.clone())?;
2290
2291 {
2293 let _conn = pool.acquire().await?;
2294 }
2295
2296 tokio::time::sleep(Duration::from_millis(50)).await;
2298
2299 let conn = pool.acquire().await?;
2301 assert_eq!(factory.created_count(), 1, "应复用归还的连接,不创建新连接");
2302
2303 pool.release(conn).await;
2304 Ok(())
2305 }
2306
2307 #[tokio::test]
2309 async fn test_into_inner_does_not_return_to_pool() -> Result<(), Box<dyn std::error::Error>> {
2310 let config = PoolConfigBuilder::new().max_size(2).build()?;
2311 let factory = Arc::new(CountingFactory::new());
2312 let pool = Pool::new(config, factory.clone())?;
2313
2314 let conn = pool.acquire().await?;
2315 assert_eq!(factory.created_count(), 1);
2316
2317 let _raw_conn = conn.into_inner();
2319
2320 tokio::time::sleep(Duration::from_millis(50)).await;
2322
2323 let status = pool.status().await;
2324 assert_eq!(status.idle, 0, "into_inner 后连接不应归还");
2325 assert_eq!(status.active, 1, "total_count 仍为 1(连接被外部持有)");
2326 Ok(())
2327 }
2328
2329 #[tokio::test]
2331 async fn test_explicit_release_no_double_return() -> Result<(), Box<dyn std::error::Error>> {
2332 let config = PoolConfigBuilder::new().max_size(2).build()?;
2333 let factory = Arc::new(CountingFactory::new());
2334 let pool = Pool::new(config, factory.clone())?;
2335
2336 let conn = pool.acquire().await?;
2337 pool.release(conn).await;
2338
2339 let status = pool.status().await;
2340 assert_eq!(status.idle, 1, "release 后 idle 应为 1");
2341
2342 let conn = pool.acquire().await?;
2344 pool.release(conn).await;
2345
2346 let status = pool.status().await;
2347 assert_eq!(status.idle, 1, "再次 release 后 idle 仍应为 1(不重复)");
2348 assert_eq!(status.active, 1, "total_count 应为 1");
2349 Ok(())
2350 }
2351
2352 struct CursorMockConn {
2358 rows: QueryRows,
2359 call_count: usize,
2360 }
2361
2362 impl CursorMockConn {
2363 fn new(rows: QueryRows) -> Self {
2364 Self {
2365 rows,
2366 call_count: 0,
2367 }
2368 }
2369 }
2370
2371 impl Connection for CursorMockConn {
2372 fn execute<'a>(
2373 &'a mut self,
2374 _sql: &'a str,
2375 ) -> Pin<Box<dyn Future<Output = Result<u64, crate::DbError>> + Send + 'a>> {
2376 Box::pin(async move { Ok(1) })
2377 }
2378
2379 fn query<'a>(
2380 &'a mut self,
2381 _sql: &'a str,
2382 ) -> Pin<Box<dyn Future<Output = Result<QueryRows, crate::DbError>> + Send + 'a>> {
2383 Box::pin(async move {
2384 self.call_count += 1;
2385 Ok(self.rows.clone())
2386 })
2387 }
2388
2389 fn begin_transaction<'a>(
2390 &'a mut self,
2391 ) -> Pin<Box<dyn Future<Output = Result<(), crate::DbError>> + Send + 'a>> {
2392 Box::pin(async move { Ok(()) })
2393 }
2394
2395 fn commit<'a>(
2396 &'a mut self,
2397 ) -> Pin<Box<dyn Future<Output = Result<(), crate::DbError>> + Send + 'a>> {
2398 Box::pin(async move { Ok(()) })
2399 }
2400
2401 fn rollback<'a>(
2402 &'a mut self,
2403 ) -> Pin<Box<dyn Future<Output = Result<(), crate::DbError>> + Send + 'a>> {
2404 Box::pin(async move { Ok(()) })
2405 }
2406
2407 fn is_connected(&self) -> bool {
2408 true
2409 }
2410
2411 fn ping<'a>(&'a mut self) -> Pin<Box<dyn Future<Output = bool> + Send + 'a>> {
2412 Box::pin(async move { true })
2413 }
2414
2415 fn close<'a>(
2416 &'a mut self,
2417 ) -> Pin<Box<dyn Future<Output = Result<(), crate::DbError>> + Send + 'a>> {
2418 Box::pin(async move { Ok(()) })
2419 }
2420 }
2421
2422 struct CursorOverrideMockConn {
2424 rows: Vec<crate::value::Value>,
2425 yielded: usize,
2426 }
2427
2428 impl CursorOverrideMockConn {
2429 fn new(rows: Vec<crate::value::Value>) -> Self {
2430 Self { rows, yielded: 0 }
2431 }
2432 }
2433
2434 impl Connection for CursorOverrideMockConn {
2435 fn execute<'a>(
2436 &'a mut self,
2437 _sql: &'a str,
2438 ) -> Pin<Box<dyn Future<Output = Result<u64, crate::DbError>> + Send + 'a>> {
2439 Box::pin(async move { Ok(1) })
2440 }
2441
2442 fn query<'a>(
2443 &'a mut self,
2444 _sql: &'a str,
2445 ) -> Pin<Box<dyn Future<Output = Result<QueryRows, crate::DbError>> + Send + 'a>> {
2446 Box::pin(async move {
2448 Ok(self
2449 .rows
2450 .iter()
2451 .map(|v| {
2452 let mut m = std::collections::HashMap::new();
2453 m.insert("v".to_string(), v.clone());
2454 m
2455 })
2456 .collect())
2457 })
2458 }
2459
2460 fn query_stream<'a>(
2462 &'a mut self,
2463 _sql: &'a str,
2464 ) -> Pin<Box<dyn futures::Stream<Item = QueryStreamItem> + Send + 'a>> {
2465 Box::pin(futures::stream::iter(
2466 self.rows
2467 .iter()
2468 .enumerate()
2469 .map(|(i, v)| {
2470 self.yielded = i + 1;
2471 let mut m = std::collections::HashMap::new();
2472 m.insert("v".to_string(), v.clone());
2473 Ok(m)
2474 })
2475 .collect::<Vec<_>>(),
2476 ))
2477 }
2478
2479 fn begin_transaction<'a>(
2480 &'a mut self,
2481 ) -> Pin<Box<dyn Future<Output = Result<(), crate::DbError>> + Send + 'a>> {
2482 Box::pin(async move { Ok(()) })
2483 }
2484
2485 fn commit<'a>(
2486 &'a mut self,
2487 ) -> Pin<Box<dyn Future<Output = Result<(), crate::DbError>> + Send + 'a>> {
2488 Box::pin(async move { Ok(()) })
2489 }
2490
2491 fn rollback<'a>(
2492 &'a mut self,
2493 ) -> Pin<Box<dyn Future<Output = Result<(), crate::DbError>> + Send + 'a>> {
2494 Box::pin(async move { Ok(()) })
2495 }
2496
2497 fn is_connected(&self) -> bool {
2498 true
2499 }
2500
2501 fn ping<'a>(&'a mut self) -> Pin<Box<dyn Future<Output = bool> + Send + 'a>> {
2502 Box::pin(async move { true })
2503 }
2504
2505 fn close<'a>(
2506 &'a mut self,
2507 ) -> Pin<Box<dyn Future<Output = Result<(), crate::DbError>> + Send + 'a>> {
2508 Box::pin(async move { Ok(()) })
2509 }
2510 }
2511
2512 #[tokio::test]
2514 async fn test_query_stream_default_impl_yields_all_rows() {
2515 use futures::StreamExt;
2516 let rows: QueryRows = vec![
2517 std::collections::HashMap::from([
2518 ("id".to_string(), crate::value::Value::I64(1)),
2519 (
2520 "name".to_string(),
2521 crate::value::Value::String("alice".to_string()),
2522 ),
2523 ]),
2524 std::collections::HashMap::from([
2525 ("id".to_string(), crate::value::Value::I64(2)),
2526 (
2527 "name".to_string(),
2528 crate::value::Value::String("bob".to_string()),
2529 ),
2530 ]),
2531 std::collections::HashMap::from([
2532 ("id".to_string(), crate::value::Value::I64(3)),
2533 (
2534 "name".to_string(),
2535 crate::value::Value::String("carol".to_string()),
2536 ),
2537 ]),
2538 ];
2539 let mut conn = CursorMockConn::new(rows);
2540 let mut stream = conn.query_stream("SELECT id, name FROM users");
2541 let mut received: Vec<QueryStreamItem> = Vec::new();
2542 while let Some(item) = stream.next().await {
2543 received.push(item);
2544 }
2545 assert_eq!(received.len(), 3, "应收到 3 行");
2546 assert!(received.iter().all(|r| r.is_ok()), "所有项应为 Ok");
2547 drop(stream);
2548 assert_eq!(conn.call_count, 1, "默认实现应调用 query() 一次");
2549 }
2550
2551 #[tokio::test]
2553 async fn test_query_stream_default_empty_result() {
2554 use futures::StreamExt;
2555 let mut conn = CursorMockConn::new(Vec::new());
2556 let mut stream = conn.query_stream("SELECT * FROM empty_table");
2557 let mut count = 0;
2558 while let Some(_item) = stream.next().await {
2559 count += 1;
2560 }
2561 assert_eq!(count, 0, "空结果集应产生 0 项");
2562 }
2563
2564 #[tokio::test]
2566 async fn test_query_stream_default_error_propagation() {
2567 use futures::StreamExt;
2568 struct ErrorMockConn;
2570 impl Connection for ErrorMockConn {
2571 fn execute<'a>(
2572 &'a mut self,
2573 _sql: &'a str,
2574 ) -> Pin<Box<dyn Future<Output = Result<u64, crate::DbError>> + Send + 'a>>
2575 {
2576 Box::pin(async move { Ok(1) })
2577 }
2578 fn query<'a>(
2579 &'a mut self,
2580 _sql: &'a str,
2581 ) -> Pin<Box<dyn Future<Output = Result<QueryRows, crate::DbError>> + Send + 'a>>
2582 {
2583 Box::pin(async move { Err(crate::DbError::Internal("query failed".to_string())) })
2584 }
2585 fn begin_transaction<'a>(
2586 &'a mut self,
2587 ) -> Pin<Box<dyn Future<Output = Result<(), crate::DbError>> + Send + 'a>> {
2588 Box::pin(async move { Ok(()) })
2589 }
2590 fn commit<'a>(
2591 &'a mut self,
2592 ) -> Pin<Box<dyn Future<Output = Result<(), crate::DbError>> + Send + 'a>> {
2593 Box::pin(async move { Ok(()) })
2594 }
2595 fn rollback<'a>(
2596 &'a mut self,
2597 ) -> Pin<Box<dyn Future<Output = Result<(), crate::DbError>> + Send + 'a>> {
2598 Box::pin(async move { Ok(()) })
2599 }
2600 fn is_connected(&self) -> bool {
2601 true
2602 }
2603 fn ping<'a>(&'a mut self) -> Pin<Box<dyn Future<Output = bool> + Send + 'a>> {
2604 Box::pin(async move { true })
2605 }
2606 fn close<'a>(
2607 &'a mut self,
2608 ) -> Pin<Box<dyn Future<Output = Result<(), crate::DbError>> + Send + 'a>> {
2609 Box::pin(async move { Ok(()) })
2610 }
2611 }
2612 let mut conn = ErrorMockConn;
2613 let mut stream = conn.query_stream("SELECT * FROM bad_table");
2614 let item = stream.next().await;
2615 assert!(item.is_some(), "应产生一项");
2616 assert!(item.unwrap().is_err(), "该项应为 Err");
2617 }
2618
2619 #[tokio::test]
2621 async fn test_query_stream_override_yields_rows_one_by_one() {
2622 use futures::StreamExt;
2623 let rows = vec![
2624 crate::value::Value::I64(10),
2625 crate::value::Value::I64(20),
2626 crate::value::Value::I64(30),
2627 crate::value::Value::I64(40),
2628 crate::value::Value::I64(50),
2629 ];
2630 let mut conn = CursorOverrideMockConn::new(rows);
2631 let values: Vec<i64> = {
2632 let mut stream = conn.query_stream("SELECT v FROM seq");
2633 let mut vals: Vec<i64> = Vec::new();
2634 while let Some(Ok(row)) = stream.next().await {
2635 if let crate::value::Value::I64(v) = row.get("v").unwrap() {
2636 vals.push(*v);
2637 }
2638 }
2639 vals
2640 };
2641 assert_eq!(values, vec![10, 20, 30, 40, 50], "应按顺序收到全部 5 行");
2642 assert_eq!(conn.yielded, 5, "应逐行 yield 5 次(真游标覆盖)");
2643 }
2644
2645 #[tokio::test]
2647 async fn test_query_stream_override_early_drop() {
2648 use futures::StreamExt;
2649 let rows = vec![
2650 crate::value::Value::I64(1),
2651 crate::value::Value::I64(2),
2652 crate::value::Value::I64(3),
2653 ];
2654 let mut conn = CursorOverrideMockConn::new(rows);
2655 {
2656 let mut stream = conn.query_stream("SELECT v FROM seq");
2657 let first = stream.next().await;
2658 assert!(first.is_some(), "第一项应存在");
2659 drop(stream);
2661 }
2662 assert!(conn.is_connected(), "提前 drop 流后连接仍应可用");
2664 }
2665
2666 #[tokio::test]
2668 async fn test_pool_prewarm() -> Result<(), Box<dyn std::error::Error>> {
2669 use std::sync::atomic::AtomicU32;
2670
2671 let create_count = Arc::new(AtomicU32::new(0));
2673 let create_count_clone = create_count.clone();
2674
2675 struct CountingFactory {
2676 count: Arc<AtomicU32>,
2677 }
2678
2679 #[async_trait]
2680 impl ConnectionFactory for CountingFactory {
2681 async fn create(&self) -> Result<Box<dyn Connection>, crate::DbError> {
2682 self.count.fetch_add(1, Ordering::SeqCst);
2683 Ok(Box::new(MockConnection::new()))
2684 }
2685 }
2686
2687 let config = PoolConfigBuilder::new()
2689 .max_size(10)
2690 .min_idle(5)
2691 .prewarm(true)
2692 .build()?;
2693
2694 let factory = Arc::new(CountingFactory {
2695 count: create_count_clone,
2696 });
2697
2698 let pool = Pool::new(config, factory)?;
2699
2700 let status_before = pool.status().await;
2702 assert_eq!(status_before.idle, 0, "预热前 idle 应为 0");
2703
2704 pool.prewarm().await;
2706
2707 let status_after = pool.status().await;
2709 assert!(
2710 status_after.idle >= 5,
2711 "预热后 idle 应 >= 5,实际: {}",
2712 status_after.idle
2713 );
2714
2715 assert_eq!(
2717 create_count.load(Ordering::SeqCst),
2718 5,
2719 "工厂应被调用 5 次(min_idle)"
2720 );
2721
2722 Ok(())
2723 }
2724
2725 #[tokio::test]
2727 async fn test_pool_prewarm_failure_non_blocking() -> Result<(), Box<dyn std::error::Error>> {
2728 use std::sync::atomic::AtomicBool;
2729
2730 struct FailingFactory {
2731 failed: Arc<AtomicBool>,
2732 }
2733
2734 #[async_trait]
2735 impl ConnectionFactory for FailingFactory {
2736 async fn create(&self) -> Result<Box<dyn Connection>, crate::DbError> {
2737 self.failed.store(true, Ordering::SeqCst);
2738 Err(crate::DbError::Internal(
2740 "simulated connection failure".to_string(),
2741 ))
2742 }
2743 }
2744
2745 let failed = Arc::new(AtomicBool::new(false));
2746 let mut config = PoolConfigBuilder::new()
2747 .max_size(10)
2748 .min_idle(3)
2749 .prewarm(true)
2750 .build()?;
2751 config.connection_timeout = std::time::Duration::from_secs(1); let factory = Arc::new(FailingFactory {
2754 failed: failed.clone(),
2755 });
2756
2757 let pool = Pool::new(config, factory)?;
2759 pool.prewarm().await; assert!(failed.load(Ordering::SeqCst), "工厂应被调用且失败");
2763
2764 let status = pool.status().await;
2766 assert_eq!(status.max, 10, "池配置应正常");
2767
2768 Ok(())
2769 }
2770
2771 #[tokio::test]
2773 async fn test_pool_prewarm_disabled() -> Result<(), Box<dyn std::error::Error>> {
2774 use std::sync::atomic::AtomicU32;
2775
2776 let create_count = Arc::new(AtomicU32::new(0));
2777 let create_count_clone = create_count.clone();
2778
2779 struct CountingFactory {
2780 count: Arc<AtomicU32>,
2781 }
2782
2783 #[async_trait]
2784 impl ConnectionFactory for CountingFactory {
2785 async fn create(&self) -> Result<Box<dyn Connection>, crate::DbError> {
2786 self.count.fetch_add(1, Ordering::SeqCst);
2787 Ok(Box::new(MockConnection::new()))
2788 }
2789 }
2790
2791 let config = PoolConfigBuilder::new()
2793 .max_size(10)
2794 .min_idle(5)
2795 .prewarm(false) .build()?;
2797
2798 let factory = Arc::new(CountingFactory {
2799 count: create_count_clone,
2800 });
2801
2802 let pool = Pool::new(config, factory)?;
2803 pool.prewarm().await; assert_eq!(
2807 create_count.load(Ordering::SeqCst),
2808 0,
2809 "prewarm=false 时工厂不应被调用"
2810 );
2811
2812 let status = pool.status().await;
2813 assert_eq!(status.idle, 0, "idle 应为 0");
2814
2815 Ok(())
2816 }
2817
2818 #[tokio::test]
2820 async fn test_pool_new_async_with_prewarm() -> Result<(), Box<dyn std::error::Error>> {
2821 use std::sync::atomic::AtomicU32;
2822
2823 let create_count = Arc::new(AtomicU32::new(0));
2824 let create_count_clone = create_count.clone();
2825
2826 struct CountingFactory {
2827 count: Arc<AtomicU32>,
2828 }
2829
2830 #[async_trait]
2831 impl ConnectionFactory for CountingFactory {
2832 async fn create(&self) -> Result<Box<dyn Connection>, crate::DbError> {
2833 self.count.fetch_add(1, Ordering::SeqCst);
2834 Ok(Box::new(MockConnection::new()))
2835 }
2836 }
2837
2838 let config = PoolConfigBuilder::new()
2839 .max_size(10)
2840 .min_idle(5)
2841 .prewarm(true)
2842 .build()?;
2843
2844 let factory = Arc::new(CountingFactory {
2845 count: create_count_clone,
2846 });
2847
2848 let pool = Pool::new_async(config, factory).await?;
2849
2850 let status = pool.status().await;
2851 assert!(
2852 status.idle >= 5,
2853 "new_async prewarm=true 后 idle 应 >= 5,实际: {}",
2854 status.idle
2855 );
2856 assert_eq!(create_count.load(Ordering::SeqCst), 5, "工厂应被调用 5 次");
2857
2858 Ok(())
2859 }
2860
2861 #[tokio::test]
2863 async fn test_pool_new_async_without_prewarm() -> Result<(), Box<dyn std::error::Error>> {
2864 use std::sync::atomic::AtomicU32;
2865
2866 let create_count = Arc::new(AtomicU32::new(0));
2867 let create_count_clone = create_count.clone();
2868
2869 struct CountingFactory {
2870 count: Arc<AtomicU32>,
2871 }
2872
2873 #[async_trait]
2874 impl ConnectionFactory for CountingFactory {
2875 async fn create(&self) -> Result<Box<dyn Connection>, crate::DbError> {
2876 self.count.fetch_add(1, Ordering::SeqCst);
2877 Ok(Box::new(MockConnection::new()))
2878 }
2879 }
2880
2881 let config = PoolConfigBuilder::new()
2882 .max_size(10)
2883 .min_idle(5)
2884 .prewarm(false)
2885 .build()?;
2886
2887 let factory = Arc::new(CountingFactory {
2888 count: create_count_clone,
2889 });
2890
2891 let pool = Pool::new_async(config, factory).await?;
2892
2893 let status = pool.status().await;
2894 assert_eq!(status.idle, 0, "prewarm=false 时 idle 应为 0");
2895 assert_eq!(create_count.load(Ordering::SeqCst), 0, "工厂不应被调用");
2896
2897 Ok(())
2898 }
2899
2900 #[tokio::test]
2902 async fn test_pool_new_async_failure_non_blocking() -> Result<(), Box<dyn std::error::Error>> {
2903 struct FailingFactory;
2904
2905 #[async_trait]
2906 impl ConnectionFactory for FailingFactory {
2907 async fn create(&self) -> Result<Box<dyn Connection>, crate::DbError> {
2908 Err(crate::DbError::Internal("simulated failure".to_string()))
2909 }
2910 }
2911
2912 let mut config = PoolConfigBuilder::new()
2913 .max_size(10)
2914 .min_idle(3)
2915 .prewarm(true)
2916 .build()?;
2917 config.connection_timeout = std::time::Duration::from_secs(1);
2918
2919 let pool = Pool::new_async(config, Arc::new(FailingFactory)).await?;
2920
2921 let status = pool.status().await;
2922 assert_eq!(status.max, 10, "池配置应正常");
2923
2924 Ok(())
2925 }
2926
2927 #[cfg(feature = "auto-prewarm")]
2929 #[tokio::test]
2930 async fn test_pool_progressive_prewarm() -> Result<(), Box<dyn std::error::Error>> {
2931 use std::sync::atomic::AtomicU32;
2932
2933 let create_count = Arc::new(AtomicU32::new(0));
2934 let create_count_clone = create_count.clone();
2935
2936 struct CountingFactory {
2937 count: Arc<AtomicU32>,
2938 }
2939
2940 #[async_trait]
2941 impl ConnectionFactory for CountingFactory {
2942 async fn create(&self) -> Result<Box<dyn Connection>, crate::DbError> {
2943 self.count.fetch_add(1, Ordering::SeqCst);
2944 Ok(Box::new(MockConnection::new()))
2945 }
2946 }
2947
2948 let config = PoolConfigBuilder::new()
2949 .max_size(20)
2950 .min_idle(6)
2951 .prewarm(true)
2952 .build()?;
2953
2954 let factory = Arc::new(CountingFactory {
2955 count: create_count_clone,
2956 });
2957
2958 let pool = Pool::new(config, factory)?;
2959
2960 let progress = crate::prewarm::PrewarmProgress::new(6);
2961 pool.progressive_prewarm(
2962 2,
2963 std::time::Duration::from_millis(5),
2964 std::time::Duration::from_secs(10),
2965 &progress,
2966 )
2967 .await;
2968
2969 let snap = progress.snapshot();
2970 assert!(
2971 snap.warmed >= 6,
2972 "progressive_prewarm 后 warmed 应 >= 6,实际: {}",
2973 snap.warmed
2974 );
2975 assert!(snap.is_completed, "应标记完成");
2976 assert_eq!(create_count.load(Ordering::SeqCst), 6, "工厂应被调用 6 次");
2977
2978 let status = pool.status().await;
2979 assert!(status.idle >= 6, "池中 idle 应 >= 6");
2980
2981 Ok(())
2982 }
2983
2984 #[cfg(feature = "auto-prewarm")]
2986 #[tokio::test]
2987 async fn test_pool_progressive_prewarm_timeout_zero() -> Result<(), Box<dyn std::error::Error>>
2988 {
2989 use std::sync::atomic::AtomicU32;
2990
2991 let create_count = Arc::new(AtomicU32::new(0));
2992 let create_count_clone = create_count.clone();
2993
2994 struct CountingFactory {
2995 count: Arc<AtomicU32>,
2996 }
2997
2998 #[async_trait]
2999 impl ConnectionFactory for CountingFactory {
3000 async fn create(&self) -> Result<Box<dyn Connection>, crate::DbError> {
3001 self.count.fetch_add(1, Ordering::SeqCst);
3002 Ok(Box::new(MockConnection::new()))
3003 }
3004 }
3005
3006 let config = PoolConfigBuilder::new()
3007 .max_size(20)
3008 .min_idle(10)
3009 .prewarm(true)
3010 .build()?;
3011
3012 let factory = Arc::new(CountingFactory {
3013 count: create_count_clone,
3014 });
3015
3016 let pool = Pool::new(config, factory)?;
3017
3018 let progress = crate::prewarm::PrewarmProgress::new(10);
3019 pool.progressive_prewarm(
3020 2,
3021 std::time::Duration::from_millis(5),
3022 std::time::Duration::ZERO,
3023 &progress,
3024 )
3025 .await;
3026
3027 let snap = progress.snapshot();
3028 assert!(snap.is_completed, "应标记完成");
3029 assert!(
3030 snap.warmed <= 2,
3031 "total_timeout=0 时最多建一批(batch_size=2),实际: {}",
3032 snap.warmed
3033 );
3034
3035 Ok(())
3036 }
3037
3038 #[cfg(feature = "auto-prewarm")]
3040 #[tokio::test]
3041 async fn test_pool_progressive_prewarm_disabled() -> Result<(), Box<dyn std::error::Error>> {
3042 use std::sync::atomic::AtomicU32;
3043
3044 let create_count = Arc::new(AtomicU32::new(0));
3045 let create_count_clone = create_count.clone();
3046
3047 struct CountingFactory {
3048 count: Arc<AtomicU32>,
3049 }
3050
3051 #[async_trait]
3052 impl ConnectionFactory for CountingFactory {
3053 async fn create(&self) -> Result<Box<dyn Connection>, crate::DbError> {
3054 self.count.fetch_add(1, Ordering::SeqCst);
3055 Ok(Box::new(MockConnection::new()))
3056 }
3057 }
3058
3059 let config = PoolConfigBuilder::new()
3060 .max_size(20)
3061 .min_idle(10)
3062 .prewarm(false)
3063 .build()?;
3064
3065 let factory = Arc::new(CountingFactory {
3066 count: create_count_clone,
3067 });
3068
3069 let pool = Pool::new(config, factory)?;
3070
3071 let progress = crate::prewarm::PrewarmProgress::new(10);
3072 pool.progressive_prewarm(
3073 2,
3074 std::time::Duration::from_millis(5),
3075 std::time::Duration::from_secs(10),
3076 &progress,
3077 )
3078 .await;
3079
3080 let snap = progress.snapshot();
3081 assert!(snap.is_completed, "应标记完成");
3082 assert_eq!(snap.warmed, 0, "prewarm=false 时不应建连");
3083 assert_eq!(create_count.load(Ordering::SeqCst), 0, "工厂不应被调用");
3084
3085 Ok(())
3086 }
3087
3088 #[cfg(feature = "auto-prewarm")]
3090 #[tokio::test]
3091 async fn test_pool_progressive_prewarm_failure_non_blocking(
3092 ) -> Result<(), Box<dyn std::error::Error>> {
3093 struct FailingFactory;
3094
3095 #[async_trait]
3096 impl ConnectionFactory for FailingFactory {
3097 async fn create(&self) -> Result<Box<dyn Connection>, crate::DbError> {
3098 Err(crate::DbError::Internal("simulated failure".to_string()))
3099 }
3100 }
3101
3102 let mut config = PoolConfigBuilder::new()
3103 .max_size(20)
3104 .min_idle(5)
3105 .prewarm(true)
3106 .build()?;
3107 config.connection_timeout = std::time::Duration::from_secs(1);
3108
3109 let pool = Pool::new(config, Arc::new(FailingFactory))?;
3110
3111 let progress = crate::prewarm::PrewarmProgress::new(5);
3112 pool.progressive_prewarm(
3113 2,
3114 std::time::Duration::from_millis(5),
3115 std::time::Duration::from_secs(5),
3116 &progress,
3117 )
3118 .await;
3119
3120 let snap = progress.snapshot();
3121 assert!(snap.is_completed, "应标记完成");
3122 assert_eq!(snap.warmed, 0, "全部失败时 warmed=0");
3123 assert!(snap.failed > 0, "应有失败记录");
3124
3125 Ok(())
3126 }
3127
3128 #[tokio::test]
3130 async fn test_pool_metrics_acquire_release() -> Result<(), Box<dyn std::error::Error>> {
3131 let config = PoolConfigBuilder::new().max_size(10).build()?;
3132 let pool = Pool::new(config, Arc::new(MockConnectionFactory))?;
3133
3134 let metrics = pool.pool_metrics();
3135 assert_eq!(metrics.acquire_count, 0);
3136 assert_eq!(metrics.release_count, 0);
3137 assert_eq!(metrics.connection_created_count, 0);
3138
3139 let conn = pool.acquire().await?;
3140 let metrics = pool.pool_metrics();
3141 assert_eq!(metrics.acquire_count, 1);
3142 assert_eq!(metrics.connection_created_count, 1);
3143 assert_eq!(metrics.acquire_failed_count, 0);
3144
3145 pool.release(conn).await;
3146 let metrics = pool.pool_metrics();
3147 assert_eq!(metrics.release_count, 1);
3148 assert_eq!(metrics.connection_closed_count, 0);
3150
3151 Ok(())
3152 }
3153
3154 #[tokio::test]
3156 async fn test_pool_metrics_acquire_failed() -> Result<(), Box<dyn std::error::Error>> {
3157 struct FailingFactory;
3158
3159 #[async_trait]
3160 impl ConnectionFactory for FailingFactory {
3161 async fn create(&self) -> Result<Box<dyn Connection>, crate::DbError> {
3162 Err(crate::DbError::Internal("simulated failure".to_string()))
3163 }
3164 }
3165
3166 let config = PoolConfigBuilder::new().max_size(10).build()?;
3167 let pool = Pool::new(config, Arc::new(FailingFactory))?;
3168
3169 let result = pool.acquire().await;
3170 assert!(result.is_err());
3171
3172 let metrics = pool.pool_metrics();
3173 assert_eq!(metrics.acquire_failed_count, 1);
3174 assert_eq!(metrics.acquire_count, 0);
3175
3176 Ok(())
3177 }
3178
3179 #[tokio::test]
3181 async fn test_pool_metrics_connection_closed() -> Result<(), Box<dyn std::error::Error>> {
3182 let config = PoolConfigBuilder::new().max_size(10).build()?;
3183 let pool = Pool::new(config, Arc::new(MockConnectionFactory))?;
3184
3185 let conn = pool.acquire().await?;
3186 pool.release(conn).await;
3187
3188 let status = pool.status().await;
3189 assert_eq!(status.idle, 1);
3190
3191 pool.close_all().await;
3192
3193 let metrics = pool.pool_metrics();
3194 assert_eq!(metrics.connection_closed_count, 1);
3195 assert_eq!(metrics.connection_created_count, 1);
3196
3197 Ok(())
3198 }
3199
3200 #[test]
3202 fn test_pool_metrics_average_wait_time() {
3203 let metrics = PoolMetrics {
3204 acquire_count: 4,
3205 acquire_failed_count: 1,
3206 acquire_wait_time: Duration::from_millis(200),
3207 release_count: 4,
3208 connection_created_count: 2,
3209 connection_closed_count: 0,
3210 };
3211 assert_eq!(
3212 metrics.average_acquire_wait_time(),
3213 Duration::from_millis(50)
3214 );
3215
3216 let empty = PoolMetrics::default();
3218 assert_eq!(empty.average_acquire_wait_time(), Duration::ZERO);
3219 }
3220}