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>(
47 &'a mut self,
48 sql: &'a str,
49 ) -> Pin<Box<dyn Future<Output = Result<u64, crate::DbError>> + Send + 'a>>;
50 fn query<'a>(
51 &'a mut self,
52 sql: &'a str,
53 ) -> Pin<Box<dyn Future<Output = Result<QueryRows, crate::DbError>> + Send + 'a>>;
54 fn begin_transaction<'a>(
55 &'a mut self,
56 ) -> Pin<Box<dyn Future<Output = Result<(), crate::DbError>> + Send + 'a>>;
57 fn commit<'a>(
58 &'a mut self,
59 ) -> Pin<Box<dyn Future<Output = Result<(), crate::DbError>> + Send + 'a>>;
60 fn rollback<'a>(
61 &'a mut self,
62 ) -> Pin<Box<dyn Future<Output = Result<(), crate::DbError>> + Send + 'a>>;
63 fn is_connected(&self) -> bool;
64 fn ping<'a>(&'a mut self) -> Pin<Box<dyn Future<Output = bool> + Send + 'a>>;
65 fn close<'a>(
66 &'a mut self,
67 ) -> Pin<Box<dyn Future<Output = Result<(), crate::DbError>> + Send + 'a>>;
68
69 fn execute_with_params<'a>(
75 &'a mut self,
76 sql: &'a str,
77 params: &'a [crate::value::Value],
78 ) -> Pin<Box<dyn Future<Output = Result<u64, crate::DbError>> + Send + 'a>> {
79 let _ = (sql, params);
80 Box::pin(async move {
81 Err(crate::DbError::Internal(
82 "execute_with_params not implemented for this adapter".to_string(),
83 ))
84 })
85 }
86
87 fn query_with_params<'a>(
93 &'a mut self,
94 sql: &'a str,
95 params: &'a [crate::value::Value],
96 ) -> Pin<Box<dyn Future<Output = Result<QueryRows, crate::DbError>> + Send + 'a>> {
97 let _ = (sql, params);
98 Box::pin(async move {
99 Err(crate::DbError::Internal(
100 "query_with_params not implemented for this adapter".to_string(),
101 ))
102 })
103 }
104
105 fn query_values<'a>(
110 &'a mut self,
111 sql: &'a str,
112 ) -> Pin<Box<dyn Future<Output = Result<crate::value::QueryValues, crate::DbError>> + Send + 'a>>
113 {
114 let _ = sql;
115 Box::pin(async move {
116 Err(crate::DbError::Internal(
117 "query_values not implemented for this adapter".to_string(),
118 ))
119 })
120 }
121
122 fn query_values_with_params<'a>(
126 &'a mut self,
127 sql: &'a str,
128 params: &'a [crate::value::Value],
129 ) -> Pin<Box<dyn Future<Output = Result<crate::value::QueryValues, crate::DbError>> + Send + 'a>>
130 {
131 let _ = (sql, params);
132 Box::pin(async move {
133 Err(crate::DbError::Internal(
134 "query_values_with_params not implemented for this adapter".to_string(),
135 ))
136 })
137 }
138
139 fn query_stream<'a>(
151 &'a mut self,
152 sql: &'a str,
153 ) -> Pin<Box<dyn futures::Stream<Item = QueryStreamItem> + Send + 'a>> {
154 let sql_owned = sql.to_string();
156 let stream = futures::stream::once(async move { self.query(&sql_owned).await })
158 .map(|result| {
160 let items: Vec<QueryStreamItem> = match result {
161 Ok(rows) => rows.into_iter().map(Ok).collect(),
162 Err(e) => vec![Err(e)],
163 };
164 futures::stream::iter(items)
165 })
166 .flatten();
167 Box::pin(stream)
168 }
169
170 fn query_stream_cursor<'a>(
181 &'a mut self,
182 sql: &'a str,
183 _batch_size: usize,
184 ) -> Pin<Box<dyn futures::Stream<Item = QueryStreamItem> + Send + 'a>> {
185 self.query_stream(sql)
186 }
187
188 fn execute_batch<'a>(
193 &'a mut self,
194 sqls: &'a [String],
195 ) -> Pin<Box<dyn Future<Output = Result<u64, crate::DbError>> + Send + 'a>> {
196 Box::pin(async move {
197 let mut total = 0u64;
198 for sql in sqls {
199 total += self.execute(sql).await?;
200 }
201 Ok(total)
202 })
203 }
204
205 fn execute_batch_params<'a>(
210 &'a mut self,
211 sql: &'a str,
212 params_batch: &'a [Vec<crate::value::Value>],
213 ) -> Pin<Box<dyn Future<Output = Result<u64, crate::DbError>> + Send + 'a>> {
214 Box::pin(async move {
215 let mut total = 0u64;
216 for params in params_batch {
217 total += self.execute_with_params(sql, params).await?;
218 }
219 Ok(total)
220 })
221 }
222}
223
224pub struct PooledConnection {
232 conn: Box<dyn Connection>,
233 created_at: Instant,
234 last_used_at: Instant,
235 pool: Option<Pool>,
236}
237
238impl PooledConnection {
239 fn new(conn: Box<dyn Connection>, pool: Pool) -> Self {
240 let now = Instant::now();
241 Self {
242 conn,
243 created_at: now,
244 last_used_at: now,
245 pool: Some(pool),
246 }
247 }
248
249 fn is_expired(&self, max_lifetime: Duration) -> bool {
250 self.created_at.elapsed() >= max_lifetime
251 }
252
253 fn is_idle_too_long(&self, idle_timeout: Duration) -> bool {
254 self.last_used_at.elapsed() >= idle_timeout
255 }
256
257 pub fn created_at(&self) -> Instant {
259 self.created_at
260 }
261
262 pub fn into_inner(mut self) -> Box<dyn Connection> {
267 self.pool = None; std::mem::replace(&mut self.conn, Box::new(ClosedConnection))
271 }
272}
273
274impl Drop for PooledConnection {
286 fn drop(&mut self) {
287 if let Some(pool) = self.pool.take() {
288 let conn = std::mem::replace(&mut self.conn, Box::new(ClosedConnection));
290 let pooled = PooledConnection {
291 conn,
292 created_at: self.created_at,
293 last_used_at: self.last_used_at,
294 pool: None,
295 };
296 if let Ok(handle) = tokio::runtime::Handle::try_current() {
298 handle.spawn(async move {
299 pool.release(pooled).await;
300 });
301 } else {
302 drop(pooled);
306 pool.total_count.fetch_sub(1, Ordering::SeqCst);
307 }
308 }
309 }
310}
311
312struct ClosedConnection;
316
317impl Connection for ClosedConnection {
318 fn execute<'a>(
319 &'a mut self,
320 _sql: &'a str,
321 ) -> Pin<Box<dyn Future<Output = Result<u64, crate::DbError>> + Send + 'a>> {
322 Box::pin(async {
323 Err(crate::DbError::ConnectionError(
324 "connection already returned to pool".to_string(),
325 ))
326 })
327 }
328
329 fn query<'a>(
330 &'a mut self,
331 _sql: &'a str,
332 ) -> Pin<Box<dyn Future<Output = Result<QueryRows, crate::DbError>> + Send + 'a>> {
333 Box::pin(async {
334 Err(crate::DbError::ConnectionError(
335 "connection already returned to pool".to_string(),
336 ))
337 })
338 }
339
340 fn begin_transaction<'a>(
341 &'a mut self,
342 ) -> Pin<Box<dyn Future<Output = Result<(), crate::DbError>> + Send + 'a>> {
343 Box::pin(async {
344 Err(crate::DbError::ConnectionError(
345 "connection already returned to pool".to_string(),
346 ))
347 })
348 }
349
350 fn commit<'a>(
351 &'a mut self,
352 ) -> Pin<Box<dyn Future<Output = Result<(), crate::DbError>> + Send + 'a>> {
353 Box::pin(async { Ok(()) })
354 }
355
356 fn rollback<'a>(
357 &'a mut self,
358 ) -> Pin<Box<dyn Future<Output = Result<(), crate::DbError>> + Send + 'a>> {
359 Box::pin(async { Ok(()) })
360 }
361
362 fn is_connected(&self) -> bool {
363 false
364 }
365
366 fn ping<'a>(&'a mut self) -> Pin<Box<dyn Future<Output = bool> + Send + 'a>> {
367 Box::pin(async { false })
368 }
369
370 fn close<'a>(
371 &'a mut self,
372 ) -> Pin<Box<dyn Future<Output = Result<(), crate::DbError>> + Send + 'a>> {
373 Box::pin(async { Ok(()) })
374 }
375}
376
377impl Deref for PooledConnection {
378 type Target = dyn Connection;
379
380 fn deref(&self) -> &Self::Target {
381 self.conn.as_ref()
382 }
383}
384
385impl DerefMut for PooledConnection {
386 fn deref_mut(&mut self) -> &mut Self::Target {
387 self.conn.as_mut()
388 }
389}
390
391#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
393pub enum TlsVersion {
394 #[default]
395 Tls12,
396 Tls13,
397}
398
399#[derive(Debug, Clone, Default)]
401pub struct TlsConfig {
402 pub enabled: bool,
404 pub ca_cert_path: Option<String>,
406 pub client_cert_path: Option<String>,
408 pub client_key_path: Option<String>,
410 pub min_version: TlsVersion,
412}
413
414#[derive(Debug, Clone)]
416pub enum PoolEvent {
417 ConnectionCreated,
419 ConnectionClosed,
421 ConnectionAcquired,
423 ConnectionReleased,
425 AcquireTimeout,
427}
428
429pub type PoolEventCallback = Arc<dyn Fn(PoolEvent) + Send + Sync>;
431
432pub struct PoolConfig {
433 pub max_size: u32,
434 pub min_idle: u32,
435 pub acquire_timeout: Duration,
436 pub idle_timeout: Duration,
437 pub max_lifetime: Duration,
438 pub connection_timeout: Duration,
439 pub tls: Option<TlsConfig>,
441 pub query_timeout: Option<Duration>,
443 pub max_rows: Option<usize>,
445 pub memory_limit: Option<usize>,
447 pub on_event: Option<PoolEventCallback>,
449 pub test_before_acquire: bool,
457 pub prewarm: bool,
469}
470
471impl Default for PoolConfig {
472 fn default() -> Self {
473 Self {
474 max_size: 100,
475 min_idle: 0,
476 acquire_timeout: Duration::from_secs(30),
477 idle_timeout: Duration::from_secs(600),
478 max_lifetime: Duration::from_secs(1800),
479 connection_timeout: Duration::from_secs(10),
480 tls: None,
481 query_timeout: Some(Duration::from_secs(30)),
482 max_rows: None,
483 memory_limit: None,
484 on_event: None,
485 test_before_acquire: false,
486 prewarm: false,
487 }
488 }
489}
490
491impl Clone for PoolConfig {
492 fn clone(&self) -> Self {
493 Self {
494 max_size: self.max_size,
495 min_idle: self.min_idle,
496 acquire_timeout: self.acquire_timeout,
497 idle_timeout: self.idle_timeout,
498 max_lifetime: self.max_lifetime,
499 connection_timeout: self.connection_timeout,
500 tls: self.tls.clone(),
501 query_timeout: self.query_timeout,
502 max_rows: self.max_rows,
503 memory_limit: self.memory_limit,
504 on_event: self.on_event.clone(),
505 test_before_acquire: self.test_before_acquire,
506 prewarm: self.prewarm,
507 }
508 }
509}
510
511impl PoolConfig {
512 pub fn validate(&self) -> Result<(), PoolError> {
514 if self.max_size == 0 {
515 return Err(PoolError::InvalidConfig("max_size cannot be 0".to_string()));
516 }
517 if self.min_idle > self.max_size {
518 return Err(PoolError::InvalidConfig(
519 "min_idle cannot exceed max_size".to_string(),
520 ));
521 }
522 const MAX_DURATION_SECS: u64 = u32::MAX as u64; for (name, dur) in [
528 ("acquire_timeout", self.acquire_timeout),
529 ("idle_timeout", self.idle_timeout),
530 ("max_lifetime", self.max_lifetime),
531 ("connection_timeout", self.connection_timeout),
532 ] {
533 if dur.as_secs() > MAX_DURATION_SECS {
534 return Err(PoolError::InvalidConfig(format!(
535 "{name} ({:?}) exceeds maximum allowed duration ({} seconds)",
536 dur, MAX_DURATION_SECS
537 )));
538 }
539 }
540 Ok(())
541 }
542
543 #[must_use]
545 pub fn with_prewarm(mut self, prewarm: bool) -> Self {
546 self.prewarm = prewarm;
547 self
548 }
549}
550
551pub struct PoolStatus {
552 pub idle: u32,
553 pub active: u32,
554 pub max: u32,
555 pub min: u32,
556 pub waiters: u32,
558}
559
560impl std::fmt::Debug for PoolStatus {
561 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
562 f.debug_struct("PoolStatus")
563 .field("idle", &self.idle)
564 .field("active", &self.active)
565 .field("max", &self.max)
566 .field("min", &self.min)
567 .field("waiters", &self.waiters)
568 .finish()
569 }
570}
571
572#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
578pub struct PoolMetrics {
579 pub acquire_count: u64,
581 pub acquire_failed_count: u64,
583 pub acquire_wait_time: Duration,
585 pub release_count: u64,
587 pub connection_created_count: u64,
589 pub connection_closed_count: u64,
591}
592
593impl PoolMetrics {
594 #[must_use]
596 pub fn average_acquire_wait_time(&self) -> Duration {
597 if self.acquire_count == 0 {
598 Duration::ZERO
599 } else {
600 self.acquire_wait_time / self.acquire_count as u32
601 }
602 }
603}
604
605pub struct PoolConfigBuilder {
606 config: PoolConfig,
607}
608
609impl PoolConfigBuilder {
610 pub fn new() -> Self {
611 Self {
612 config: PoolConfig::default(),
613 }
614 }
615
616 pub fn max_size(mut self, size: u32) -> Self {
617 self.config.max_size = size;
618 self
619 }
620
621 pub fn min_idle(mut self, count: u32) -> Self {
622 self.config.min_idle = count;
623 self
624 }
625
626 pub fn acquire_timeout(mut self, timeout_secs: u64) -> Self {
627 self.config.acquire_timeout = Duration::from_secs(timeout_secs);
628 self
629 }
630
631 pub fn idle_timeout(mut self, timeout_secs: u64) -> Self {
632 self.config.idle_timeout = Duration::from_secs(timeout_secs);
633 self
634 }
635
636 pub fn max_lifetime(mut self, lifetime_secs: u64) -> Self {
637 self.config.max_lifetime = Duration::from_secs(lifetime_secs);
638 self
639 }
640
641 pub fn tls(mut self, tls: TlsConfig) -> Self {
643 self.config.tls = Some(tls);
644 self
645 }
646
647 pub fn query_timeout(mut self, timeout: Duration) -> Self {
649 self.config.query_timeout = Some(timeout);
650 self
651 }
652
653 pub fn max_rows(mut self, max_rows: usize) -> Self {
655 self.config.max_rows = Some(max_rows);
656 self
657 }
658
659 pub fn memory_limit(mut self, memory_limit: usize) -> Self {
661 self.config.memory_limit = Some(memory_limit);
662 self
663 }
664
665 pub fn on_event(mut self, callback: PoolEventCallback) -> Self {
667 self.config.on_event = Some(callback);
668 self
669 }
670
671 pub fn test_before_acquire(mut self, enabled: bool) -> Self {
676 self.config.test_before_acquire = enabled;
677 self
678 }
679
680 pub fn prewarm(mut self, enabled: bool) -> Self {
685 self.config.prewarm = enabled;
686 self
687 }
688
689 pub fn build(self) -> Result<PoolConfig, PoolError> {
690 self.config.validate()?;
691 Ok(self.config)
692 }
693}
694
695impl Default for PoolConfigBuilder {
696 fn default() -> Self {
697 Self::new()
698 }
699}
700
701#[async_trait]
703pub trait ConnectionFactory: Send + Sync {
704 async fn create(&self) -> Result<Box<dyn Connection>, crate::DbError>;
705}
706
707pub struct Pool {
713 config: PoolConfig,
714 factory: Arc<dyn ConnectionFactory>,
715 idle: Arc<ArrayQueue<PooledConnection>>,
721 total_count: Arc<AtomicU32>,
731 closed: Arc<AtomicBool>,
733 notify: Arc<Notify>,
734 waiters_count: Arc<AtomicU32>,
736 dynamic_max_size: Arc<AtomicU32>,
738 #[cfg(feature = "circuit-breaker")]
744 circuit_breaker: Arc<PlMutex<DefaultCircuitBreaker>>,
745 #[cfg(feature = "rate-limit")]
754 rate_limiter: Arc<PlRwLock<Option<Arc<dyn RateLimiter>>>>,
755 #[cfg(feature = "rate-limit")]
757 rate_limit_key: String,
758 acquire_count: Arc<AtomicU64>,
760 acquire_failed_count: Arc<AtomicU64>,
762 acquire_wait_time_ns: Arc<AtomicU64>,
764 release_count: Arc<AtomicU64>,
766 connection_created_count: Arc<AtomicU64>,
768 connection_closed_count: Arc<AtomicU64>,
770}
771
772impl Clone for Pool {
776 fn clone(&self) -> Self {
777 Self {
778 config: self.config.clone(),
779 factory: self.factory.clone(),
780 idle: self.idle.clone(),
781 total_count: self.total_count.clone(),
782 closed: self.closed.clone(),
783 notify: Arc::clone(&self.notify),
784 waiters_count: self.waiters_count.clone(),
785 dynamic_max_size: self.dynamic_max_size.clone(),
786 #[cfg(feature = "circuit-breaker")]
787 circuit_breaker: Arc::clone(&self.circuit_breaker),
788 #[cfg(feature = "rate-limit")]
789 rate_limiter: Arc::clone(&self.rate_limiter),
790 #[cfg(feature = "rate-limit")]
791 rate_limit_key: self.rate_limit_key.clone(),
792 acquire_count: self.acquire_count.clone(),
793 acquire_failed_count: self.acquire_failed_count.clone(),
794 acquire_wait_time_ns: self.acquire_wait_time_ns.clone(),
795 release_count: self.release_count.clone(),
796 connection_created_count: self.connection_created_count.clone(),
797 connection_closed_count: self.connection_closed_count.clone(),
798 }
799 }
800}
801
802impl Pool {
803 pub fn new(config: PoolConfig, factory: Arc<dyn ConnectionFactory>) -> Result<Self, PoolError> {
827 config.validate()?;
828 let max_size = config.max_size as usize;
831 let dynamic_max = config.max_size;
832 Ok(Self {
833 config,
834 factory,
835 idle: Arc::new(ArrayQueue::new(max_size)),
836 total_count: Arc::new(AtomicU32::new(0)),
837 closed: Arc::new(AtomicBool::new(false)),
838 notify: Arc::new(Notify::new()),
839 waiters_count: Arc::new(AtomicU32::new(0)),
840 dynamic_max_size: Arc::new(AtomicU32::new(dynamic_max)),
841 #[cfg(feature = "circuit-breaker")]
844 circuit_breaker: Arc::new(PlMutex::new(DefaultCircuitBreaker::new(
845 5,
846 std::time::Duration::from_secs(30),
847 ))),
848 #[cfg(feature = "rate-limit")]
851 rate_limiter: Arc::new(PlRwLock::new(None)),
852 #[cfg(feature = "rate-limit")]
853 rate_limit_key: "pool".to_string(),
854 acquire_count: Arc::new(AtomicU64::new(0)),
855 acquire_failed_count: Arc::new(AtomicU64::new(0)),
856 acquire_wait_time_ns: Arc::new(AtomicU64::new(0)),
857 release_count: Arc::new(AtomicU64::new(0)),
858 connection_created_count: Arc::new(AtomicU64::new(0)),
859 connection_closed_count: Arc::new(AtomicU64::new(0)),
860 })
861 }
862
863 pub async fn prewarm(&self) {
880 if !self.config.prewarm {
881 return;
882 }
883
884 let min_idle = self.config.min_idle as usize;
885 let mut warmed = 0;
886
887 for i in 0..min_idle {
888 if self.closed.load(Ordering::Acquire) {
890 break;
891 }
892
893 let current_max = self.dynamic_max_size.load(Ordering::Acquire);
895 let current = self.total_count.load(Ordering::Acquire);
896 if current >= current_max {
897 break;
898 }
899
900 let created = loop {
902 let current = self.total_count.load(Ordering::Acquire);
903 if current >= current_max {
904 break None;
905 }
906 match self.total_count.compare_exchange(
907 current,
908 current + 1,
909 Ordering::SeqCst,
910 Ordering::Acquire,
911 ) {
912 Ok(_) => break Some(()),
913 Err(_) => continue,
914 }
915 };
916
917 if created.is_some() {
918 match tokio::time::timeout(self.config.connection_timeout, self.factory.create())
919 .await
920 {
921 Ok(Ok(conn)) => {
922 #[cfg(feature = "circuit-breaker")]
923 {
924 self.circuit_breaker.lock().record_success();
925 }
926 self.emit_event(PoolEvent::ConnectionCreated);
927 let pooled = PooledConnection::new(conn, self.clone());
928 if self.idle.push(pooled).is_err() {
930 let _ = self.total_count.fetch_sub(1, Ordering::SeqCst);
932 tracing::warn!(
933 target: "sz_orm::pool::prewarm",
934 "prewarm connection {} failed: idle queue full",
935 i
936 );
937 } else {
938 warmed += 1;
939 self.notify.notify_one();
940 }
941 }
942 Ok(Err(e)) => {
943 let _ = self.total_count.fetch_sub(1, Ordering::SeqCst);
944 #[cfg(feature = "circuit-breaker")]
945 {
946 self.circuit_breaker.lock().record_failure();
947 }
948 tracing::warn!(
949 target: "sz_orm::pool::prewarm",
950 "prewarm connection {} failed: {}",
951 i,
952 e
953 );
954 }
955 Err(_) => {
956 let _ = self.total_count.fetch_sub(1, Ordering::SeqCst);
957 #[cfg(feature = "circuit-breaker")]
958 {
959 self.circuit_breaker.lock().record_failure();
960 }
961 tracing::warn!(
962 target: "sz_orm::pool::prewarm",
963 "prewarm connection {} timeout",
964 i
965 );
966 }
967 }
968 }
969 }
970
971 if warmed > 0 {
972 tracing::info!(
973 target: "sz_orm::pool::prewarm",
974 "pool prewarm completed: {}/{} connections established",
975 warmed,
976 min_idle
977 );
978 }
979 }
980
981 pub fn config(&self) -> &PoolConfig {
983 &self.config
984 }
985
986 #[cfg(feature = "circuit-breaker")]
1000 pub fn configure_circuit_breaker(
1001 &self,
1002 failure_threshold: usize,
1003 reset_timeout: std::time::Duration,
1004 ) {
1005 let new_cb = DefaultCircuitBreaker::new(failure_threshold, reset_timeout);
1006 let mut guard = self.circuit_breaker.lock();
1008 *guard = new_cb;
1009 }
1010
1011 #[cfg(feature = "circuit-breaker")]
1016 pub fn reset_circuit_breaker(&self) -> bool {
1017 let mut guard = self.circuit_breaker.lock();
1019 guard.reset()
1020 }
1021
1022 #[cfg(feature = "circuit-breaker")]
1024 pub fn circuit_state(&self) -> CircuitState {
1025 let guard = self.circuit_breaker.lock();
1027 guard.state()
1028 }
1029
1030 #[cfg(feature = "rate-limit")]
1039 pub fn set_rate_limiter(&self, limiter: Option<Arc<dyn RateLimiter>>) {
1040 let mut guard = self.rate_limiter.write();
1042 *guard = limiter;
1043 }
1044
1045 #[cfg(feature = "rate-limit")]
1047 pub fn with_rate_limit_key(mut self, key: impl Into<String>) -> Self {
1048 self.rate_limit_key = key.into();
1049 self
1050 }
1051
1052 fn emit_event(&self, event: PoolEvent) {
1054 if matches!(event, PoolEvent::ConnectionCreated) {
1057 self.connection_created_count
1058 .fetch_add(1, Ordering::Relaxed);
1059 }
1060 if let Some(ref callback) = self.config.on_event {
1061 callback(event);
1062 }
1063 }
1064
1065 async fn close_connection(&self, pooled: PooledConnection) {
1070 let mut pooled = pooled;
1071 let _ = pooled.conn.close().await;
1072 self.connection_closed_count.fetch_add(1, Ordering::Relaxed);
1073 }
1074
1075 #[tracing::instrument(skip(self), fields(max_size = self.config.max_size, acquire_timeout = ?self.config.acquire_timeout))]
1095 pub async fn acquire(&self) -> Result<PooledConnection, PoolError> {
1096 if self.closed.load(Ordering::Acquire) {
1098 self.acquire_failed_count.fetch_add(1, Ordering::Relaxed);
1099 return Err(PoolError::Closed);
1100 }
1101
1102 #[cfg(feature = "circuit-breaker")]
1106 {
1107 let mut guard = self.circuit_breaker.lock();
1108 if !guard.can_execute() {
1109 self.acquire_failed_count.fetch_add(1, Ordering::Relaxed);
1110 return Err(PoolError::CircuitOpen);
1111 }
1112 }
1113
1114 #[cfg(feature = "rate-limit")]
1118 {
1119 let guard = self.rate_limiter.read();
1120 if let Some(ref limiter) = *guard {
1121 match limiter.try_acquire(&self.rate_limit_key) {
1122 Ok(result) if !result.allowed => {
1123 self.acquire_failed_count.fetch_add(1, Ordering::Relaxed);
1124 return Err(PoolError::RateLimited {
1125 remaining: result.remaining,
1126 reset_at: result.reset_at,
1127 });
1128 }
1129 Ok(_) => {} Err(_) => {
1131 }
1133 }
1134 }
1135 }
1136
1137 let deadline = Instant::now() + self.config.acquire_timeout;
1138 let mut backoff = Duration::from_millis(1);
1140 const MAX_BACKOFF: Duration = Duration::from_millis(100);
1142
1143 loop {
1144 let mut to_close: Vec<PooledConnection> = Vec::new();
1150 let acquired: Option<PooledConnection> = {
1151 let mut found: Option<PooledConnection> = None;
1152 while let Some(pooled) = self.idle.pop() {
1153 if pooled.is_expired(self.config.max_lifetime) {
1155 to_close.push(pooled);
1156 continue;
1157 }
1158 if pooled.is_idle_too_long(self.config.idle_timeout) {
1160 to_close.push(pooled);
1161 continue;
1162 }
1163 if !pooled.conn.is_connected() {
1166 to_close.push(pooled);
1167 continue;
1168 }
1169 found = Some(pooled);
1170 break;
1171 }
1172 found
1173 };
1174
1175 for pooled in to_close {
1177 self.close_connection(pooled).await;
1178 self.total_count.fetch_sub(1, Ordering::SeqCst);
1180 }
1181
1182 if let Some(mut pooled) = acquired {
1183 if self.config.test_before_acquire {
1185 let ping_timeout = self.config.connection_timeout / 2;
1186 let alive = match tokio::time::timeout(ping_timeout, pooled.conn.ping()).await {
1187 Ok(true) => true,
1188 Ok(false) => false,
1189 Err(_) => false, };
1191 if !alive {
1192 self.close_connection(pooled).await;
1194 self.total_count.fetch_sub(1, Ordering::SeqCst);
1195 continue;
1196 }
1197 }
1198 pooled.pool = Some(self.clone());
1201 self.acquire_count.fetch_add(1, Ordering::Relaxed);
1202 return Ok(pooled);
1203 }
1204
1205 let current_max = self.dynamic_max_size.load(Ordering::Acquire);
1210 let created = loop {
1211 let current = self.total_count.load(Ordering::Acquire);
1212 if current >= current_max {
1213 break None; }
1215 match self.total_count.compare_exchange(
1216 current,
1217 current + 1,
1218 Ordering::SeqCst,
1219 Ordering::Acquire,
1220 ) {
1221 Ok(_) => break Some(()), Err(_) => continue, }
1224 };
1225
1226 if created.is_some() {
1227 match tokio::time::timeout(self.config.connection_timeout, self.factory.create())
1228 .await
1229 {
1230 Ok(Ok(conn)) => {
1231 #[cfg(feature = "circuit-breaker")]
1234 {
1235 self.circuit_breaker.lock().record_success();
1236 }
1237 self.emit_event(PoolEvent::ConnectionCreated);
1238 self.emit_event(PoolEvent::ConnectionAcquired);
1239 self.acquire_count.fetch_add(1, Ordering::Relaxed);
1240 return Ok(PooledConnection::new(conn, self.clone()));
1241 }
1242 Ok(Err(e)) => {
1243 self.total_count.fetch_sub(1, Ordering::SeqCst);
1245 #[cfg(feature = "circuit-breaker")]
1248 {
1249 self.circuit_breaker.lock().record_failure();
1250 }
1251 self.acquire_failed_count.fetch_add(1, Ordering::Relaxed);
1252 return Err(PoolError::ConnectionFailed(e.to_string()));
1253 }
1254 Err(_) => {
1255 self.total_count.fetch_sub(1, Ordering::SeqCst);
1257 #[cfg(feature = "circuit-breaker")]
1260 {
1261 self.circuit_breaker.lock().record_failure();
1262 }
1263 self.acquire_failed_count.fetch_add(1, Ordering::Relaxed);
1264 return Err(PoolError::Timeout);
1265 }
1266 }
1267 }
1268
1269 let now = Instant::now();
1271 if now >= deadline {
1272 self.emit_event(PoolEvent::AcquireTimeout);
1273 self.acquire_failed_count.fetch_add(1, Ordering::Relaxed);
1274 return Err(PoolError::Timeout);
1275 }
1276 self.waiters_count.fetch_add(1, Ordering::SeqCst);
1278 let wait = std::cmp::min(backoff, deadline - now);
1279 match tokio::time::timeout(wait, self.notify.notified()).await {
1280 Ok(()) => {
1281 backoff = Duration::from_millis(1);
1283 }
1284 Err(_) => {
1285 backoff = std::cmp::min(backoff * 2, MAX_BACKOFF);
1287 }
1288 }
1289 self.waiters_count.fetch_sub(1, Ordering::SeqCst);
1291 self.acquire_wait_time_ns
1293 .fetch_add(wait.as_nanos() as u64, Ordering::Relaxed);
1294 }
1295 }
1296
1297 #[tracing::instrument(skip(self, pooled))]
1305 pub async fn release(&self, mut pooled: PooledConnection) {
1306 pooled.pool = None;
1308 self.release_count.fetch_add(1, Ordering::Relaxed);
1310
1311 if self.closed.load(Ordering::Acquire) {
1313 self.close_connection(pooled).await;
1314 self.total_count.fetch_sub(1, Ordering::SeqCst);
1316 self.emit_event(PoolEvent::ConnectionClosed);
1317 return;
1318 }
1319
1320 if !pooled.conn.is_connected() {
1322 self.close_connection(pooled).await;
1323 self.total_count.fetch_sub(1, Ordering::SeqCst);
1324 self.emit_event(PoolEvent::ConnectionClosed);
1325 return;
1326 }
1327
1328 pooled.last_used_at = Instant::now();
1330
1331 if let Err(rejected) = self.idle.push(pooled) {
1337 self.close_connection(rejected).await;
1339 self.total_count.fetch_sub(1, Ordering::SeqCst);
1340 self.emit_event(PoolEvent::ConnectionClosed);
1341 } else {
1342 self.emit_event(PoolEvent::ConnectionReleased);
1343 }
1344 self.notify.notify_one();
1345 }
1346
1347 pub async fn status(&self) -> PoolStatus {
1352 let idle_count = self.idle.len() as u32;
1353 let active = self.total_count.load(Ordering::Acquire);
1355 let waiters = self.waiters_count.load(Ordering::Acquire);
1356 PoolStatus {
1357 idle: idle_count,
1358 active,
1359 max: self.dynamic_max_size.load(Ordering::Acquire),
1360 min: self.config.min_idle,
1361 waiters,
1362 }
1363 }
1364
1365 pub fn pool_metrics(&self) -> PoolMetrics {
1376 PoolMetrics {
1377 acquire_count: self.acquire_count.load(Ordering::Acquire),
1378 acquire_failed_count: self.acquire_failed_count.load(Ordering::Acquire),
1379 acquire_wait_time: Duration::from_nanos(
1380 self.acquire_wait_time_ns.load(Ordering::Acquire),
1381 ),
1382 release_count: self.release_count.load(Ordering::Acquire),
1383 connection_created_count: self.connection_created_count.load(Ordering::Acquire),
1384 connection_closed_count: self.connection_closed_count.load(Ordering::Acquire),
1385 }
1386 }
1387
1388 #[tracing::instrument(skip(self))]
1390 pub async fn reap_idle(&self) {
1391 let mut all: Vec<PooledConnection> = Vec::new();
1395 while let Some(pooled) = self.idle.pop() {
1396 all.push(pooled);
1397 }
1398
1399 let mut to_close = Vec::new();
1401 for pooled in all {
1402 if pooled.is_idle_too_long(self.config.idle_timeout)
1403 || pooled.is_expired(self.config.max_lifetime)
1404 {
1405 to_close.push(pooled);
1406 } else {
1407 if let Err(rejected) = self.idle.push(pooled) {
1409 self.close_connection(rejected).await;
1410 self.total_count.fetch_sub(1, Ordering::SeqCst);
1411 }
1412 }
1413 }
1414
1415 for pooled in to_close {
1417 self.close_connection(pooled).await;
1418 self.total_count.fetch_sub(1, Ordering::SeqCst);
1420 }
1421 }
1422
1423 pub async fn close_all(&self) {
1427 self.closed.store(true, Ordering::Release);
1429 let mut to_close: Vec<PooledConnection> = Vec::new();
1432 while let Some(pooled) = self.idle.pop() {
1433 to_close.push(pooled);
1434 }
1435 let closed_count: u32 = to_close.len() as u32;
1437 for pooled in to_close {
1438 self.close_connection(pooled).await;
1439 }
1440 self.total_count.fetch_sub(closed_count, Ordering::SeqCst);
1443 }
1444
1445 pub async fn health_check(&self) -> u32 {
1461 let mut to_check: Vec<PooledConnection> = Vec::new();
1463 while let Some(pooled) = self.idle.pop() {
1464 to_check.push(pooled);
1465 }
1466
1467 let mut removed: u32 = 0;
1468 let mut alive: Vec<PooledConnection> = Vec::with_capacity(to_check.len());
1469 for mut pooled in to_check.drain(..) {
1470 if !pooled.conn.is_connected() {
1472 self.close_connection(pooled).await;
1473 removed += 1;
1474 continue;
1475 }
1476 let ping_timeout = self.config.connection_timeout / 2;
1478 match tokio::time::timeout(ping_timeout, pooled.conn.ping()).await {
1479 Ok(true) => alive.push(pooled),
1480 Ok(false) => {
1481 self.close_connection(pooled).await;
1483 removed += 1;
1484 }
1485 Err(_) => {
1486 self.close_connection(pooled).await;
1488 removed += 1;
1489 }
1490 }
1491 }
1492
1493 let alive_count: u32 = alive.len() as u32;
1495 for pooled in alive {
1496 if let Err(rejected) = self.idle.push(pooled) {
1498 self.close_connection(rejected).await;
1499 removed += 1;
1500 }
1501 }
1502
1503 if removed > 0 {
1505 self.total_count.fetch_sub(removed, Ordering::SeqCst);
1506 }
1507
1508 if alive_count > 0 {
1510 self.notify.notify_one();
1511 }
1512
1513 removed
1514 }
1515
1516 pub async fn shutdown(&self) {
1523 self.closed.store(true, Ordering::SeqCst);
1525 self.notify.notify_waiters();
1527 self.close_all().await;
1529 let deadline = Instant::now() + Duration::from_secs(30);
1531 while self.total_count.load(Ordering::SeqCst) > 0 {
1532 if Instant::now() >= deadline {
1533 break;
1534 }
1535 tokio::time::sleep(Duration::from_millis(100)).await;
1536 }
1537 }
1538
1539 pub fn resize(&self, new_max: usize) {
1547 self.set_max_size(new_max as u32);
1548 }
1549
1550 pub fn set_max_size(&self, new_max: u32) {
1552 self.dynamic_max_size.store(new_max, Ordering::SeqCst);
1553 }
1554
1555 pub fn max_size(&self) -> u32 {
1557 self.dynamic_max_size.load(Ordering::Acquire)
1558 }
1559
1560 pub async fn warmup(&self, min_idle: usize) -> Result<(), PoolError> {
1564 for _ in 0..min_idle {
1565 let current_max = self.dynamic_max_size.load(Ordering::Acquire);
1566 let current = self.total_count.load(Ordering::Acquire);
1567 if current >= current_max {
1568 break;
1569 }
1570 match self.total_count.compare_exchange(
1572 current,
1573 current + 1,
1574 Ordering::SeqCst,
1575 Ordering::Acquire,
1576 ) {
1577 Ok(_) => {}
1578 Err(_) => continue, }
1580 match self.factory.create().await {
1581 Ok(conn) => {
1582 let now = Instant::now();
1583 let pooled = PooledConnection {
1584 conn,
1585 created_at: now,
1586 last_used_at: now,
1587 pool: None,
1588 };
1589 if let Err(rejected) = self.idle.push(pooled) {
1590 self.close_connection(rejected).await;
1592 self.total_count.fetch_sub(1, Ordering::SeqCst);
1593 }
1594 self.emit_event(PoolEvent::ConnectionCreated);
1595 }
1596 Err(_) => {
1597 self.total_count.fetch_sub(1, Ordering::SeqCst);
1599 break;
1600 }
1601 }
1602 }
1603 Ok(())
1604 }
1605
1606 pub async fn query_with_timeout(&self, sql: &str) -> Result<QueryRows, crate::DbError> {
1611 let timeout = self.config.query_timeout.unwrap_or(Duration::from_secs(30));
1612 let mut conn = self.acquire().await.map_err(crate::DbError::PoolError)?;
1613 tokio::time::timeout(timeout, conn.query(sql))
1614 .await
1615 .map_err(|_| crate::DbError::QueryError(format!("Query timeout after {:?}", timeout)))?
1616 }
1617}
1618
1619#[cfg(test)]
1620mod tests {
1621 use super::*;
1622
1623 struct MockConnection {
1625 connected: bool,
1626 }
1627
1628 impl MockConnection {
1629 fn new() -> Self {
1630 Self { connected: true }
1631 }
1632 }
1633
1634 impl Connection for MockConnection {
1635 fn execute<'a>(
1636 &'a mut self,
1637 _sql: &'a str,
1638 ) -> Pin<Box<dyn Future<Output = Result<u64, crate::DbError>> + Send + 'a>> {
1639 Box::pin(async move { Ok(1) })
1640 }
1641
1642 fn query<'a>(
1643 &'a mut self,
1644 _sql: &'a str,
1645 ) -> Pin<
1646 Box<
1647 dyn Future<
1648 Output = Result<
1649 Vec<std::collections::HashMap<String, crate::value::Value>>,
1650 crate::DbError,
1651 >,
1652 > + Send
1653 + 'a,
1654 >,
1655 > {
1656 Box::pin(async move { Ok(vec![]) })
1657 }
1658
1659 fn begin_transaction<'a>(
1660 &'a mut self,
1661 ) -> Pin<Box<dyn Future<Output = Result<(), crate::DbError>> + Send + 'a>> {
1662 Box::pin(async move { Ok(()) })
1663 }
1664
1665 fn commit<'a>(
1666 &'a mut self,
1667 ) -> Pin<Box<dyn Future<Output = Result<(), crate::DbError>> + Send + 'a>> {
1668 Box::pin(async move { Ok(()) })
1669 }
1670
1671 fn rollback<'a>(
1672 &'a mut self,
1673 ) -> Pin<Box<dyn Future<Output = Result<(), crate::DbError>> + Send + 'a>> {
1674 Box::pin(async move { Ok(()) })
1675 }
1676
1677 fn is_connected(&self) -> bool {
1678 self.connected
1679 }
1680
1681 fn ping<'a>(&'a mut self) -> Pin<Box<dyn Future<Output = bool> + Send + 'a>> {
1682 Box::pin(async move { true })
1683 }
1684
1685 fn close<'a>(
1686 &'a mut self,
1687 ) -> Pin<Box<dyn Future<Output = Result<(), crate::DbError>> + Send + 'a>> {
1688 Box::pin(async move {
1689 self.connected = false;
1690 Ok(())
1691 })
1692 }
1693 }
1694
1695 struct MockConnectionFactory;
1696
1697 #[async_trait]
1698 impl ConnectionFactory for MockConnectionFactory {
1699 async fn create(&self) -> Result<Box<dyn Connection>, crate::DbError> {
1700 Ok(Box::new(MockConnection::new()))
1701 }
1702 }
1703
1704 #[tokio::test]
1705 async fn test_pool_config_builder() -> Result<(), Box<dyn std::error::Error>> {
1706 let config = PoolConfigBuilder::new().max_size(50).min_idle(10).build()?;
1707
1708 assert_eq!(config.max_size, 50);
1709 assert_eq!(config.min_idle, 10);
1710 Ok(())
1711 }
1712
1713 #[test]
1714 fn test_pool_status_display() {
1715 let status = PoolStatus {
1716 idle: 5,
1717 active: 10,
1718 max: 100,
1719 min: 5,
1720 waiters: 0,
1721 };
1722
1723 let display = format!("{:?}", status);
1724 assert!(display.contains("idle"));
1725 assert!(display.contains("active"));
1726 }
1727
1728 #[test]
1729 fn test_default_pool_config() {
1730 let config = PoolConfig::default();
1731 assert_eq!(config.max_size, 100);
1732 assert_eq!(config.min_idle, 0);
1733 assert_eq!(config.acquire_timeout.as_secs(), 30);
1734 assert_eq!(config.idle_timeout.as_secs(), 600);
1735 assert_eq!(config.max_lifetime.as_secs(), 1800);
1736 }
1737
1738 #[tokio::test]
1739 async fn test_pool_config_clone() {
1740 let config = PoolConfig::default();
1741 let cloned = config.clone();
1742 assert_eq!(cloned.max_size, config.max_size);
1743 assert_eq!(cloned.min_idle, config.min_idle);
1744 }
1745
1746 #[test]
1747 fn test_pool_config_builder_default() -> Result<(), Box<dyn std::error::Error>> {
1748 let builder = PoolConfigBuilder::new();
1749 let config = builder.build()?;
1750 assert_eq!(config.max_size, 100);
1751 Ok(())
1752 }
1753
1754 #[test]
1755 fn test_pool_config_validate() {
1756 let result = PoolConfigBuilder::new().max_size(0).build();
1757 assert!(result.is_err());
1758
1759 let result = PoolConfigBuilder::new().max_size(10).min_idle(20).build();
1760 assert!(result.is_err());
1761 }
1762
1763 #[test]
1764 fn test_pool_config_validate_duration_upper_bound() {
1765 use std::time::Duration;
1766
1767 let config = PoolConfig {
1769 max_size: 10,
1770 min_idle: 1,
1771 acquire_timeout: Duration::from_secs(u64::MAX),
1772 idle_timeout: Duration::from_secs(1),
1773 max_lifetime: Duration::from_secs(1),
1774 connection_timeout: Duration::from_secs(5),
1775 tls: None,
1776 query_timeout: None,
1777 max_rows: None,
1778 memory_limit: None,
1779 on_event: None,
1780 test_before_acquire: false,
1781 prewarm: false,
1782 };
1783 assert!(config.validate().is_err());
1784
1785 let config = PoolConfig {
1787 max_size: 10,
1788 min_idle: 1,
1789 acquire_timeout: Duration::from_secs(u32::MAX as u64),
1790 idle_timeout: Duration::from_secs(1),
1791 max_lifetime: Duration::from_secs(1),
1792 connection_timeout: Duration::from_secs(5),
1793 tls: None,
1794 query_timeout: None,
1795 max_rows: None,
1796 memory_limit: None,
1797 on_event: None,
1798 test_before_acquire: false,
1799 prewarm: false,
1800 };
1801 assert!(config.validate().is_ok());
1802
1803 let config = PoolConfig {
1805 max_size: 10,
1806 min_idle: 1,
1807 acquire_timeout: Duration::from_secs(u32::MAX as u64 + 1),
1808 idle_timeout: Duration::from_secs(1),
1809 max_lifetime: Duration::from_secs(1),
1810 connection_timeout: Duration::from_secs(5),
1811 tls: None,
1812 query_timeout: None,
1813 max_rows: None,
1814 memory_limit: None,
1815 on_event: None,
1816 test_before_acquire: false,
1817 prewarm: false,
1818 };
1819 assert!(config.validate().is_err());
1820 }
1821
1822 #[test]
1823 fn test_pool_config_test_before_acquire_default() {
1824 let config = PoolConfig::default();
1826 assert!(!config.test_before_acquire);
1827 }
1828
1829 #[test]
1830 fn test_pool_config_builder_test_before_acquire() {
1831 let config = PoolConfigBuilder::new()
1833 .test_before_acquire(true)
1834 .build()
1835 .unwrap();
1836 assert!(config.test_before_acquire);
1837 }
1838
1839 #[tokio::test]
1840 async fn test_pool_acquire_and_release() -> Result<(), Box<dyn std::error::Error>> {
1841 let config = PoolConfigBuilder::new().max_size(5).min_idle(1).build()?;
1842 let factory = Arc::new(MockConnectionFactory);
1843 let pool = Pool::new(config, factory)?;
1844
1845 let conn = pool.acquire().await?;
1846 let status = pool.status().await;
1847 assert_eq!(status.active, 1);
1848 assert_eq!(status.idle, 0);
1849
1850 pool.release(conn).await;
1851 let status = pool.status().await;
1852 assert_eq!(status.idle, 1);
1853
1854 let _conn2 = pool.acquire().await?;
1856 let status = pool.status().await;
1857 assert_eq!(status.idle, 0);
1858 Ok(())
1859 }
1860
1861 #[tokio::test]
1862 async fn test_pool_status() -> Result<(), Box<dyn std::error::Error>> {
1863 let config = PoolConfigBuilder::new().max_size(10).min_idle(2).build()?;
1864 let factory = Arc::new(MockConnectionFactory);
1865 let pool = Pool::new(config, factory)?;
1866
1867 let status = pool.status().await;
1868 assert_eq!(status.max, 10);
1869 assert_eq!(status.min, 2);
1870 assert_eq!(status.active, 0);
1871 Ok(())
1872 }
1873
1874 #[tokio::test]
1875 async fn test_pool_close_all() -> Result<(), Box<dyn std::error::Error>> {
1876 let config = PoolConfigBuilder::new().max_size(5).build()?;
1877 let factory = Arc::new(MockConnectionFactory);
1878 let pool = Pool::new(config, factory)?;
1879
1880 let conn1 = pool.acquire().await?;
1882 let conn2 = pool.acquire().await?;
1883 pool.release(conn1).await;
1884 pool.release(conn2).await;
1885
1886 pool.close_all().await;
1887 let status = pool.status().await;
1888 assert_eq!(status.idle, 0);
1889 assert_eq!(status.active, 0);
1890 Ok(())
1891 }
1892
1893 #[tokio::test]
1894 async fn test_pool_reap_idle() -> Result<(), Box<dyn std::error::Error>> {
1895 let config = PoolConfigBuilder::new()
1896 .max_size(5)
1897 .idle_timeout(0) .build()?;
1899 let factory = Arc::new(MockConnectionFactory);
1900 let pool = Pool::new(config, factory)?;
1901
1902 let conn = pool.acquire().await?;
1903 pool.release(conn).await;
1904
1905 tokio::time::sleep(tokio::time::Duration::from_millis(10)).await;
1907
1908 pool.reap_idle().await;
1909 let status = pool.status().await;
1910 assert_eq!(status.idle, 0);
1911 Ok(())
1912 }
1913
1914 #[tokio::test]
1920 async fn test_h7_acquire_timeout_default_30s() {
1921 let config = PoolConfig::default();
1922 assert_eq!(
1923 config.acquire_timeout,
1924 Duration::from_secs(30),
1925 "H-7: acquire_timeout 默认应为 30s"
1926 );
1927 }
1928
1929 #[tokio::test]
1931 async fn test_h7_acquire_timeout_configurable() -> Result<(), Box<dyn std::error::Error>> {
1932 let config = PoolConfigBuilder::new()
1933 .max_size(1)
1934 .acquire_timeout(5) .build()?;
1936 assert_eq!(config.acquire_timeout, Duration::from_secs(5));
1937
1938 let factory = Arc::new(MockConnectionFactory);
1940 let pool = Pool::new(config, factory)?;
1941 let _conn1 = pool.acquire().await?;
1942
1943 let fast_config = PoolConfigBuilder::new()
1945 .max_size(1)
1946 .acquire_timeout(0) .build()?;
1948 let fast_pool = Pool::new(fast_config, Arc::new(MockConnectionFactory))?;
1951 let _fast_conn = fast_pool.acquire().await?; let result = fast_pool.acquire().await;
1953 assert!(
1954 matches!(result, Err(PoolError::Timeout)),
1955 "H-7: 应返回 Timeout"
1956 );
1957 Ok(())
1958 }
1959
1960 #[tokio::test]
1963 async fn test_m7_health_check_removes_nothing_when_all_healthy(
1964 ) -> Result<(), Box<dyn std::error::Error>> {
1965 let config = PoolConfigBuilder::new().max_size(5).build()?;
1967 let factory = Arc::new(MockConnectionFactory);
1968 let pool = Pool::new(config, factory)?;
1969
1970 let conn1 = pool.acquire().await?;
1972 let conn2 = pool.acquire().await?;
1973 let conn3 = pool.acquire().await?;
1974 pool.release(conn1).await;
1975 pool.release(conn2).await;
1976 pool.release(conn3).await;
1977
1978 let removed = pool.health_check().await;
1979 assert_eq!(removed, 0, "Healthy connections should not be removed");
1980
1981 let status = pool.status().await;
1982 assert_eq!(status.idle, 3);
1983 assert_eq!(status.active, 3);
1984 Ok(())
1985 }
1986
1987 #[tokio::test]
1988 async fn test_m7_health_check_returns_zero_for_empty_pool(
1989 ) -> Result<(), Box<dyn std::error::Error>> {
1990 let config = PoolConfigBuilder::new().max_size(5).build()?;
1991 let factory = Arc::new(MockConnectionFactory);
1992 let pool = Pool::new(config, factory)?;
1993
1994 let removed = pool.health_check().await;
1995 assert_eq!(removed, 0);
1996 Ok(())
1997 }
1998
1999 struct CountingFactory {
2003 count: AtomicU32,
2004 }
2005
2006 impl CountingFactory {
2007 fn new() -> Self {
2008 Self {
2009 count: AtomicU32::new(0),
2010 }
2011 }
2012 fn created_count(&self) -> u32 {
2013 self.count.load(Ordering::SeqCst)
2014 }
2015 }
2016
2017 #[async_trait]
2018 impl ConnectionFactory for CountingFactory {
2019 async fn create(&self) -> Result<Box<dyn Connection>, crate::DbError> {
2020 self.count.fetch_add(1, Ordering::SeqCst);
2021 Ok(Box::new(MockConnection::new()))
2022 }
2023 }
2024
2025 #[tokio::test]
2031 async fn test_production_bug_max_lifetime_never_expires(
2032 ) -> Result<(), Box<dyn std::error::Error>> {
2033 let config = PoolConfig {
2036 max_size: 5,
2037 min_idle: 0,
2038 acquire_timeout: Duration::from_secs(30),
2039 idle_timeout: Duration::from_secs(600),
2040 max_lifetime: Duration::from_millis(100), connection_timeout: Duration::from_secs(10),
2042 tls: None,
2043 query_timeout: None,
2044 max_rows: None,
2045 memory_limit: None,
2046 on_event: None,
2047 test_before_acquire: false,
2048 prewarm: false,
2049 };
2050 let factory = Arc::new(CountingFactory::new());
2051 let pool = Pool::new(config, factory.clone())?;
2052
2053 let conn = pool.acquire().await?;
2055 assert_eq!(factory.created_count(), 1, "应创建 1 个连接");
2056
2057 pool.release(conn).await;
2059
2060 tokio::time::sleep(Duration::from_millis(150)).await;
2062
2063 let conn2 = pool.acquire().await?;
2065
2066 assert_eq!(
2069 factory.created_count(),
2070 2,
2071 "超过 max_lifetime 后应创建新连接(旧连接应被回收)"
2072 );
2073
2074 pool.release(conn2).await;
2075 Ok(())
2076 }
2077
2078 #[tokio::test]
2085 async fn test_drop_auto_release_connection() -> Result<(), Box<dyn std::error::Error>> {
2086 let config = PoolConfigBuilder::new().max_size(2).build()?;
2087 let factory = Arc::new(CountingFactory::new());
2088 let pool = Pool::new(config, factory.clone())?;
2089
2090 {
2092 let _conn = pool.acquire().await?;
2093 assert_eq!(factory.created_count(), 1, "应创建 1 个连接");
2094 let status = pool.status().await;
2095 assert_eq!(status.active, 1, "active 应为 1");
2096 assert_eq!(status.idle, 0, "idle 应为 0");
2097 }
2099
2100 tokio::time::sleep(Duration::from_millis(50)).await;
2102
2103 let status = pool.status().await;
2105 assert_eq!(status.idle, 1, "Drop 后连接应自动归还,idle 应为 1");
2106 assert_eq!(status.active, 1, "total_count 应为 1");
2107 assert_eq!(factory.created_count(), 1, "应复用归还的连接,不创建新连接");
2108 Ok(())
2109 }
2110
2111 #[tokio::test]
2113 async fn test_drop_auto_release_then_reuse() -> Result<(), Box<dyn std::error::Error>> {
2114 let config = PoolConfigBuilder::new().max_size(1).build()?;
2115 let factory = Arc::new(CountingFactory::new());
2116 let pool = Pool::new(config, factory.clone())?;
2117
2118 {
2120 let _conn = pool.acquire().await?;
2121 }
2122
2123 tokio::time::sleep(Duration::from_millis(50)).await;
2125
2126 let conn = pool.acquire().await?;
2128 assert_eq!(factory.created_count(), 1, "应复用归还的连接,不创建新连接");
2129
2130 pool.release(conn).await;
2131 Ok(())
2132 }
2133
2134 #[tokio::test]
2136 async fn test_into_inner_does_not_return_to_pool() -> Result<(), Box<dyn std::error::Error>> {
2137 let config = PoolConfigBuilder::new().max_size(2).build()?;
2138 let factory = Arc::new(CountingFactory::new());
2139 let pool = Pool::new(config, factory.clone())?;
2140
2141 let conn = pool.acquire().await?;
2142 assert_eq!(factory.created_count(), 1);
2143
2144 let _raw_conn = conn.into_inner();
2146
2147 tokio::time::sleep(Duration::from_millis(50)).await;
2149
2150 let status = pool.status().await;
2151 assert_eq!(status.idle, 0, "into_inner 后连接不应归还");
2152 assert_eq!(status.active, 1, "total_count 仍为 1(连接被外部持有)");
2153 Ok(())
2154 }
2155
2156 #[tokio::test]
2158 async fn test_explicit_release_no_double_return() -> Result<(), Box<dyn std::error::Error>> {
2159 let config = PoolConfigBuilder::new().max_size(2).build()?;
2160 let factory = Arc::new(CountingFactory::new());
2161 let pool = Pool::new(config, factory.clone())?;
2162
2163 let conn = pool.acquire().await?;
2164 pool.release(conn).await;
2165
2166 let status = pool.status().await;
2167 assert_eq!(status.idle, 1, "release 后 idle 应为 1");
2168
2169 let conn = pool.acquire().await?;
2171 pool.release(conn).await;
2172
2173 let status = pool.status().await;
2174 assert_eq!(status.idle, 1, "再次 release 后 idle 仍应为 1(不重复)");
2175 assert_eq!(status.active, 1, "total_count 应为 1");
2176 Ok(())
2177 }
2178
2179 struct CursorMockConn {
2185 rows: QueryRows,
2186 call_count: usize,
2187 }
2188
2189 impl CursorMockConn {
2190 fn new(rows: QueryRows) -> Self {
2191 Self {
2192 rows,
2193 call_count: 0,
2194 }
2195 }
2196 }
2197
2198 impl Connection for CursorMockConn {
2199 fn execute<'a>(
2200 &'a mut self,
2201 _sql: &'a str,
2202 ) -> Pin<Box<dyn Future<Output = Result<u64, crate::DbError>> + Send + 'a>> {
2203 Box::pin(async move { Ok(1) })
2204 }
2205
2206 fn query<'a>(
2207 &'a mut self,
2208 _sql: &'a str,
2209 ) -> Pin<Box<dyn Future<Output = Result<QueryRows, crate::DbError>> + Send + 'a>> {
2210 Box::pin(async move {
2211 self.call_count += 1;
2212 Ok(self.rows.clone())
2213 })
2214 }
2215
2216 fn begin_transaction<'a>(
2217 &'a mut self,
2218 ) -> Pin<Box<dyn Future<Output = Result<(), crate::DbError>> + Send + 'a>> {
2219 Box::pin(async move { Ok(()) })
2220 }
2221
2222 fn commit<'a>(
2223 &'a mut self,
2224 ) -> Pin<Box<dyn Future<Output = Result<(), crate::DbError>> + Send + 'a>> {
2225 Box::pin(async move { Ok(()) })
2226 }
2227
2228 fn rollback<'a>(
2229 &'a mut self,
2230 ) -> Pin<Box<dyn Future<Output = Result<(), crate::DbError>> + Send + 'a>> {
2231 Box::pin(async move { Ok(()) })
2232 }
2233
2234 fn is_connected(&self) -> bool {
2235 true
2236 }
2237
2238 fn ping<'a>(&'a mut self) -> Pin<Box<dyn Future<Output = bool> + Send + 'a>> {
2239 Box::pin(async move { true })
2240 }
2241
2242 fn close<'a>(
2243 &'a mut self,
2244 ) -> Pin<Box<dyn Future<Output = Result<(), crate::DbError>> + Send + 'a>> {
2245 Box::pin(async move { Ok(()) })
2246 }
2247 }
2248
2249 struct CursorOverrideMockConn {
2251 rows: Vec<crate::value::Value>,
2252 yielded: usize,
2253 }
2254
2255 impl CursorOverrideMockConn {
2256 fn new(rows: Vec<crate::value::Value>) -> Self {
2257 Self { rows, yielded: 0 }
2258 }
2259 }
2260
2261 impl Connection for CursorOverrideMockConn {
2262 fn execute<'a>(
2263 &'a mut self,
2264 _sql: &'a str,
2265 ) -> Pin<Box<dyn Future<Output = Result<u64, crate::DbError>> + Send + 'a>> {
2266 Box::pin(async move { Ok(1) })
2267 }
2268
2269 fn query<'a>(
2270 &'a mut self,
2271 _sql: &'a str,
2272 ) -> Pin<Box<dyn Future<Output = Result<QueryRows, crate::DbError>> + Send + 'a>> {
2273 Box::pin(async move {
2275 Ok(self
2276 .rows
2277 .iter()
2278 .map(|v| {
2279 let mut m = std::collections::HashMap::new();
2280 m.insert("v".to_string(), v.clone());
2281 m
2282 })
2283 .collect())
2284 })
2285 }
2286
2287 fn query_stream<'a>(
2289 &'a mut self,
2290 _sql: &'a str,
2291 ) -> Pin<Box<dyn futures::Stream<Item = QueryStreamItem> + Send + 'a>> {
2292 Box::pin(futures::stream::iter(
2293 self.rows
2294 .iter()
2295 .enumerate()
2296 .map(|(i, v)| {
2297 self.yielded = i + 1;
2298 let mut m = std::collections::HashMap::new();
2299 m.insert("v".to_string(), v.clone());
2300 Ok(m)
2301 })
2302 .collect::<Vec<_>>(),
2303 ))
2304 }
2305
2306 fn begin_transaction<'a>(
2307 &'a mut self,
2308 ) -> Pin<Box<dyn Future<Output = Result<(), crate::DbError>> + Send + 'a>> {
2309 Box::pin(async move { Ok(()) })
2310 }
2311
2312 fn commit<'a>(
2313 &'a mut self,
2314 ) -> Pin<Box<dyn Future<Output = Result<(), crate::DbError>> + Send + 'a>> {
2315 Box::pin(async move { Ok(()) })
2316 }
2317
2318 fn rollback<'a>(
2319 &'a mut self,
2320 ) -> Pin<Box<dyn Future<Output = Result<(), crate::DbError>> + Send + 'a>> {
2321 Box::pin(async move { Ok(()) })
2322 }
2323
2324 fn is_connected(&self) -> bool {
2325 true
2326 }
2327
2328 fn ping<'a>(&'a mut self) -> Pin<Box<dyn Future<Output = bool> + Send + 'a>> {
2329 Box::pin(async move { true })
2330 }
2331
2332 fn close<'a>(
2333 &'a mut self,
2334 ) -> Pin<Box<dyn Future<Output = Result<(), crate::DbError>> + Send + 'a>> {
2335 Box::pin(async move { Ok(()) })
2336 }
2337 }
2338
2339 #[tokio::test]
2341 async fn test_query_stream_default_impl_yields_all_rows() {
2342 use futures::StreamExt;
2343 let rows: QueryRows = vec![
2344 std::collections::HashMap::from([
2345 ("id".to_string(), crate::value::Value::I64(1)),
2346 (
2347 "name".to_string(),
2348 crate::value::Value::String("alice".to_string()),
2349 ),
2350 ]),
2351 std::collections::HashMap::from([
2352 ("id".to_string(), crate::value::Value::I64(2)),
2353 (
2354 "name".to_string(),
2355 crate::value::Value::String("bob".to_string()),
2356 ),
2357 ]),
2358 std::collections::HashMap::from([
2359 ("id".to_string(), crate::value::Value::I64(3)),
2360 (
2361 "name".to_string(),
2362 crate::value::Value::String("carol".to_string()),
2363 ),
2364 ]),
2365 ];
2366 let mut conn = CursorMockConn::new(rows);
2367 let mut stream = conn.query_stream("SELECT id, name FROM users");
2368 let mut received: Vec<QueryStreamItem> = Vec::new();
2369 while let Some(item) = stream.next().await {
2370 received.push(item);
2371 }
2372 assert_eq!(received.len(), 3, "应收到 3 行");
2373 assert!(received.iter().all(|r| r.is_ok()), "所有项应为 Ok");
2374 drop(stream);
2375 assert_eq!(conn.call_count, 1, "默认实现应调用 query() 一次");
2376 }
2377
2378 #[tokio::test]
2380 async fn test_query_stream_default_empty_result() {
2381 use futures::StreamExt;
2382 let mut conn = CursorMockConn::new(Vec::new());
2383 let mut stream = conn.query_stream("SELECT * FROM empty_table");
2384 let mut count = 0;
2385 while let Some(_item) = stream.next().await {
2386 count += 1;
2387 }
2388 assert_eq!(count, 0, "空结果集应产生 0 项");
2389 }
2390
2391 #[tokio::test]
2393 async fn test_query_stream_default_error_propagation() {
2394 use futures::StreamExt;
2395 struct ErrorMockConn;
2397 impl Connection for ErrorMockConn {
2398 fn execute<'a>(
2399 &'a mut self,
2400 _sql: &'a str,
2401 ) -> Pin<Box<dyn Future<Output = Result<u64, crate::DbError>> + Send + 'a>>
2402 {
2403 Box::pin(async move { Ok(1) })
2404 }
2405 fn query<'a>(
2406 &'a mut self,
2407 _sql: &'a str,
2408 ) -> Pin<Box<dyn Future<Output = Result<QueryRows, crate::DbError>> + Send + 'a>>
2409 {
2410 Box::pin(async move { Err(crate::DbError::Internal("query failed".to_string())) })
2411 }
2412 fn begin_transaction<'a>(
2413 &'a mut self,
2414 ) -> Pin<Box<dyn Future<Output = Result<(), crate::DbError>> + Send + 'a>> {
2415 Box::pin(async move { Ok(()) })
2416 }
2417 fn commit<'a>(
2418 &'a mut self,
2419 ) -> Pin<Box<dyn Future<Output = Result<(), crate::DbError>> + Send + 'a>> {
2420 Box::pin(async move { Ok(()) })
2421 }
2422 fn rollback<'a>(
2423 &'a mut self,
2424 ) -> Pin<Box<dyn Future<Output = Result<(), crate::DbError>> + Send + 'a>> {
2425 Box::pin(async move { Ok(()) })
2426 }
2427 fn is_connected(&self) -> bool {
2428 true
2429 }
2430 fn ping<'a>(&'a mut self) -> Pin<Box<dyn Future<Output = bool> + Send + 'a>> {
2431 Box::pin(async move { true })
2432 }
2433 fn close<'a>(
2434 &'a mut self,
2435 ) -> Pin<Box<dyn Future<Output = Result<(), crate::DbError>> + Send + 'a>> {
2436 Box::pin(async move { Ok(()) })
2437 }
2438 }
2439 let mut conn = ErrorMockConn;
2440 let mut stream = conn.query_stream("SELECT * FROM bad_table");
2441 let item = stream.next().await;
2442 assert!(item.is_some(), "应产生一项");
2443 assert!(item.unwrap().is_err(), "该项应为 Err");
2444 }
2445
2446 #[tokio::test]
2448 async fn test_query_stream_override_yields_rows_one_by_one() {
2449 use futures::StreamExt;
2450 let rows = vec![
2451 crate::value::Value::I64(10),
2452 crate::value::Value::I64(20),
2453 crate::value::Value::I64(30),
2454 crate::value::Value::I64(40),
2455 crate::value::Value::I64(50),
2456 ];
2457 let mut conn = CursorOverrideMockConn::new(rows);
2458 let values: Vec<i64> = {
2459 let mut stream = conn.query_stream("SELECT v FROM seq");
2460 let mut vals: Vec<i64> = Vec::new();
2461 while let Some(Ok(row)) = stream.next().await {
2462 if let crate::value::Value::I64(v) = row.get("v").unwrap() {
2463 vals.push(*v);
2464 }
2465 }
2466 vals
2467 };
2468 assert_eq!(values, vec![10, 20, 30, 40, 50], "应按顺序收到全部 5 行");
2469 assert_eq!(conn.yielded, 5, "应逐行 yield 5 次(真游标覆盖)");
2470 }
2471
2472 #[tokio::test]
2474 async fn test_query_stream_override_early_drop() {
2475 use futures::StreamExt;
2476 let rows = vec![
2477 crate::value::Value::I64(1),
2478 crate::value::Value::I64(2),
2479 crate::value::Value::I64(3),
2480 ];
2481 let mut conn = CursorOverrideMockConn::new(rows);
2482 {
2483 let mut stream = conn.query_stream("SELECT v FROM seq");
2484 let first = stream.next().await;
2485 assert!(first.is_some(), "第一项应存在");
2486 drop(stream);
2488 }
2489 assert!(conn.is_connected(), "提前 drop 流后连接仍应可用");
2491 }
2492
2493 #[tokio::test]
2495 async fn test_pool_prewarm() -> Result<(), Box<dyn std::error::Error>> {
2496 use std::sync::atomic::AtomicU32;
2497
2498 let create_count = Arc::new(AtomicU32::new(0));
2500 let create_count_clone = create_count.clone();
2501
2502 struct CountingFactory {
2503 count: Arc<AtomicU32>,
2504 }
2505
2506 #[async_trait]
2507 impl ConnectionFactory for CountingFactory {
2508 async fn create(&self) -> Result<Box<dyn Connection>, crate::DbError> {
2509 self.count.fetch_add(1, Ordering::SeqCst);
2510 Ok(Box::new(MockConnection::new()))
2511 }
2512 }
2513
2514 let config = PoolConfigBuilder::new()
2516 .max_size(10)
2517 .min_idle(5)
2518 .prewarm(true)
2519 .build()?;
2520
2521 let factory = Arc::new(CountingFactory {
2522 count: create_count_clone,
2523 });
2524
2525 let pool = Pool::new(config, factory)?;
2526
2527 let status_before = pool.status().await;
2529 assert_eq!(status_before.idle, 0, "预热前 idle 应为 0");
2530
2531 pool.prewarm().await;
2533
2534 let status_after = pool.status().await;
2536 assert!(
2537 status_after.idle >= 5,
2538 "预热后 idle 应 >= 5,实际: {}",
2539 status_after.idle
2540 );
2541
2542 assert_eq!(
2544 create_count.load(Ordering::SeqCst),
2545 5,
2546 "工厂应被调用 5 次(min_idle)"
2547 );
2548
2549 Ok(())
2550 }
2551
2552 #[tokio::test]
2554 async fn test_pool_prewarm_failure_non_blocking() -> Result<(), Box<dyn std::error::Error>> {
2555 use std::sync::atomic::AtomicBool;
2556
2557 struct FailingFactory {
2558 failed: Arc<AtomicBool>,
2559 }
2560
2561 #[async_trait]
2562 impl ConnectionFactory for FailingFactory {
2563 async fn create(&self) -> Result<Box<dyn Connection>, crate::DbError> {
2564 self.failed.store(true, Ordering::SeqCst);
2565 Err(crate::DbError::Internal(
2567 "simulated connection failure".to_string(),
2568 ))
2569 }
2570 }
2571
2572 let failed = Arc::new(AtomicBool::new(false));
2573 let mut config = PoolConfigBuilder::new()
2574 .max_size(10)
2575 .min_idle(3)
2576 .prewarm(true)
2577 .build()?;
2578 config.connection_timeout = std::time::Duration::from_secs(1); let factory = Arc::new(FailingFactory {
2581 failed: failed.clone(),
2582 });
2583
2584 let pool = Pool::new(config, factory)?;
2586 pool.prewarm().await; assert!(failed.load(Ordering::SeqCst), "工厂应被调用且失败");
2590
2591 let status = pool.status().await;
2593 assert_eq!(status.max, 10, "池配置应正常");
2594
2595 Ok(())
2596 }
2597
2598 #[tokio::test]
2600 async fn test_pool_prewarm_disabled() -> Result<(), Box<dyn std::error::Error>> {
2601 use std::sync::atomic::AtomicU32;
2602
2603 let create_count = Arc::new(AtomicU32::new(0));
2604 let create_count_clone = create_count.clone();
2605
2606 struct CountingFactory {
2607 count: Arc<AtomicU32>,
2608 }
2609
2610 #[async_trait]
2611 impl ConnectionFactory for CountingFactory {
2612 async fn create(&self) -> Result<Box<dyn Connection>, crate::DbError> {
2613 self.count.fetch_add(1, Ordering::SeqCst);
2614 Ok(Box::new(MockConnection::new()))
2615 }
2616 }
2617
2618 let config = PoolConfigBuilder::new()
2620 .max_size(10)
2621 .min_idle(5)
2622 .prewarm(false) .build()?;
2624
2625 let factory = Arc::new(CountingFactory {
2626 count: create_count_clone,
2627 });
2628
2629 let pool = Pool::new(config, factory)?;
2630 pool.prewarm().await; assert_eq!(
2634 create_count.load(Ordering::SeqCst),
2635 0,
2636 "prewarm=false 时工厂不应被调用"
2637 );
2638
2639 let status = pool.status().await;
2640 assert_eq!(status.idle, 0, "idle 应为 0");
2641
2642 Ok(())
2643 }
2644
2645 #[tokio::test]
2647 async fn test_pool_metrics_acquire_release() -> Result<(), Box<dyn std::error::Error>> {
2648 let config = PoolConfigBuilder::new().max_size(10).build()?;
2649 let pool = Pool::new(config, Arc::new(MockConnectionFactory))?;
2650
2651 let metrics = pool.pool_metrics();
2652 assert_eq!(metrics.acquire_count, 0);
2653 assert_eq!(metrics.release_count, 0);
2654 assert_eq!(metrics.connection_created_count, 0);
2655
2656 let conn = pool.acquire().await?;
2657 let metrics = pool.pool_metrics();
2658 assert_eq!(metrics.acquire_count, 1);
2659 assert_eq!(metrics.connection_created_count, 1);
2660 assert_eq!(metrics.acquire_failed_count, 0);
2661
2662 pool.release(conn).await;
2663 let metrics = pool.pool_metrics();
2664 assert_eq!(metrics.release_count, 1);
2665 assert_eq!(metrics.connection_closed_count, 0);
2667
2668 Ok(())
2669 }
2670
2671 #[tokio::test]
2673 async fn test_pool_metrics_acquire_failed() -> Result<(), Box<dyn std::error::Error>> {
2674 struct FailingFactory;
2675
2676 #[async_trait]
2677 impl ConnectionFactory for FailingFactory {
2678 async fn create(&self) -> Result<Box<dyn Connection>, crate::DbError> {
2679 Err(crate::DbError::Internal("simulated failure".to_string()))
2680 }
2681 }
2682
2683 let config = PoolConfigBuilder::new().max_size(10).build()?;
2684 let pool = Pool::new(config, Arc::new(FailingFactory))?;
2685
2686 let result = pool.acquire().await;
2687 assert!(result.is_err());
2688
2689 let metrics = pool.pool_metrics();
2690 assert_eq!(metrics.acquire_failed_count, 1);
2691 assert_eq!(metrics.acquire_count, 0);
2692
2693 Ok(())
2694 }
2695
2696 #[tokio::test]
2698 async fn test_pool_metrics_connection_closed() -> Result<(), Box<dyn std::error::Error>> {
2699 let config = PoolConfigBuilder::new().max_size(10).build()?;
2700 let pool = Pool::new(config, Arc::new(MockConnectionFactory))?;
2701
2702 let conn = pool.acquire().await?;
2703 pool.release(conn).await;
2704
2705 let status = pool.status().await;
2706 assert_eq!(status.idle, 1);
2707
2708 pool.close_all().await;
2709
2710 let metrics = pool.pool_metrics();
2711 assert_eq!(metrics.connection_closed_count, 1);
2712 assert_eq!(metrics.connection_created_count, 1);
2713
2714 Ok(())
2715 }
2716
2717 #[test]
2719 fn test_pool_metrics_average_wait_time() {
2720 let metrics = PoolMetrics {
2721 acquire_count: 4,
2722 acquire_failed_count: 1,
2723 acquire_wait_time: Duration::from_millis(200),
2724 release_count: 4,
2725 connection_created_count: 2,
2726 connection_closed_count: 0,
2727 };
2728 assert_eq!(
2729 metrics.average_acquire_wait_time(),
2730 Duration::from_millis(50)
2731 );
2732
2733 let empty = PoolMetrics::default();
2735 assert_eq!(empty.average_acquire_wait_time(), Duration::ZERO);
2736 }
2737}