1use async_trait::async_trait;
6use crossbeam_queue::ArrayQueue;
7use futures::StreamExt;
8use serde::{Deserialize, Serialize};
10#[cfg(any(feature = "circuit-breaker", feature = "tenant-quota-rls-enhanced"))]
15use parking_lot::Mutex as PlMutex;
16#[cfg(feature = "rate-limit")]
17use parking_lot::RwLock as PlRwLock;
18use std::future::Future;
19use std::ops::{Deref, DerefMut};
20use std::pin::Pin;
21use std::sync::atomic::{AtomicBool, AtomicU32, AtomicU64, Ordering};
22use std::sync::Arc;
23use std::time::{Duration, Instant};
24use tokio::sync::Notify;
25
26#[cfg(feature = "circuit-breaker")]
30use crate::circuit_breaker::{CircuitBreaker, CircuitState, DefaultCircuitBreaker};
31use crate::error::PoolError;
32#[cfg(feature = "rate-limit")]
33use crate::rate_limiter::RateLimiter;
34#[cfg(feature = "tenant-quota-rls-enhanced")]
35use crate::tenant_quota_rls::{QuotaEnforcer, QuotaResource};
36
37pub type QueryRows = Vec<std::collections::HashMap<String, crate::value::Value>>;
39
40pub type QueryStreamItem =
42 Result<std::collections::HashMap<String, crate::value::Value>, crate::DbError>;
43
44pub trait Connection: Send + Sync {
51 fn execute<'a>(
53 &'a mut self,
54 sql: &'a str,
55 ) -> Pin<Box<dyn Future<Output = Result<u64, crate::DbError>> + Send + 'a>>;
56 fn query<'a>(
58 &'a mut self,
59 sql: &'a str,
60 ) -> Pin<Box<dyn Future<Output = Result<QueryRows, crate::DbError>> + Send + 'a>>;
61 fn begin_transaction<'a>(
63 &'a mut self,
64 ) -> Pin<Box<dyn Future<Output = Result<(), crate::DbError>> + Send + 'a>>;
65 fn commit<'a>(
67 &'a mut self,
68 ) -> Pin<Box<dyn Future<Output = Result<(), crate::DbError>> + Send + 'a>>;
69 fn rollback<'a>(
71 &'a mut self,
72 ) -> Pin<Box<dyn Future<Output = Result<(), crate::DbError>> + Send + 'a>>;
73 fn is_connected(&self) -> bool;
75 fn ping<'a>(&'a mut self) -> Pin<Box<dyn Future<Output = bool> + Send + 'a>>;
77 fn close<'a>(
79 &'a mut self,
80 ) -> Pin<Box<dyn Future<Output = Result<(), crate::DbError>> + Send + 'a>>;
81
82 fn execute_with_params<'a>(
88 &'a mut self,
89 sql: &'a str,
90 params: &'a [crate::value::Value],
91 ) -> Pin<Box<dyn Future<Output = Result<u64, crate::DbError>> + Send + 'a>> {
92 let _ = (sql, params);
93 Box::pin(async move {
94 Err(crate::DbError::Internal(
95 "execute_with_params not implemented for this adapter".to_string(),
96 ))
97 })
98 }
99
100 fn query_with_params<'a>(
106 &'a mut self,
107 sql: &'a str,
108 params: &'a [crate::value::Value],
109 ) -> Pin<Box<dyn Future<Output = Result<QueryRows, crate::DbError>> + Send + 'a>> {
110 let _ = (sql, params);
111 Box::pin(async move {
112 Err(crate::DbError::Internal(
113 "query_with_params not implemented for this adapter".to_string(),
114 ))
115 })
116 }
117
118 fn query_values<'a>(
123 &'a mut self,
124 sql: &'a str,
125 ) -> Pin<Box<dyn Future<Output = Result<crate::value::QueryValues, crate::DbError>> + Send + 'a>>
126 {
127 let _ = sql;
128 Box::pin(async move {
129 Err(crate::DbError::Internal(
130 "query_values not implemented for this adapter".to_string(),
131 ))
132 })
133 }
134
135 fn query_values_with_params<'a>(
139 &'a mut self,
140 sql: &'a str,
141 params: &'a [crate::value::Value],
142 ) -> Pin<Box<dyn Future<Output = Result<crate::value::QueryValues, crate::DbError>> + Send + 'a>>
143 {
144 let _ = (sql, params);
145 Box::pin(async move {
146 Err(crate::DbError::Internal(
147 "query_values_with_params not implemented for this adapter".to_string(),
148 ))
149 })
150 }
151
152 fn query_stream<'a>(
164 &'a mut self,
165 sql: &'a str,
166 ) -> Pin<Box<dyn futures::Stream<Item = QueryStreamItem> + Send + 'a>> {
167 let sql_owned = sql.to_string();
169 let stream = futures::stream::once(async move { self.query(&sql_owned).await })
171 .map(|result| {
173 let items: Vec<QueryStreamItem> = match result {
174 Ok(rows) => rows.into_iter().map(Ok).collect(),
175 Err(e) => vec![Err(e)],
176 };
177 futures::stream::iter(items)
178 })
179 .flatten();
180 Box::pin(stream)
181 }
182
183 fn query_stream_cursor<'a>(
194 &'a mut self,
195 sql: &'a str,
196 _batch_size: usize,
197 ) -> Pin<Box<dyn futures::Stream<Item = QueryStreamItem> + Send + 'a>> {
198 self.query_stream(sql)
199 }
200
201 fn execute_batch<'a>(
206 &'a mut self,
207 sqls: &'a [String],
208 ) -> Pin<Box<dyn Future<Output = Result<u64, crate::DbError>> + Send + 'a>> {
209 Box::pin(async move {
210 let mut total = 0u64;
211 for sql in sqls {
212 total += self.execute(sql).await?;
213 }
214 Ok(total)
215 })
216 }
217
218 fn execute_batch_params<'a>(
223 &'a mut self,
224 sql: &'a str,
225 params_batch: &'a [Vec<crate::value::Value>],
226 ) -> Pin<Box<dyn Future<Output = Result<u64, crate::DbError>> + Send + 'a>> {
227 Box::pin(async move {
228 let mut total = 0u64;
229 for params in params_batch {
230 total += self.execute_with_params(sql, params).await?;
231 }
232 Ok(total)
233 })
234 }
235}
236
237pub struct PooledConnection {
245 conn: Box<dyn Connection>,
246 created_at: Instant,
247 last_used_at: Instant,
248 pool: Option<Pool>,
249}
250
251impl PooledConnection {
252 fn new(conn: Box<dyn Connection>, pool: Pool) -> Self {
253 let now = Instant::now();
254 Self {
255 conn,
256 created_at: now,
257 last_used_at: now,
258 pool: Some(pool),
259 }
260 }
261
262 fn is_expired(&self, max_lifetime: Duration) -> bool {
263 self.created_at.elapsed() >= max_lifetime
264 }
265
266 fn is_idle_too_long(&self, idle_timeout: Duration) -> bool {
267 self.last_used_at.elapsed() >= idle_timeout
268 }
269
270 pub fn created_at(&self) -> Instant {
272 self.created_at
273 }
274
275 pub fn into_inner(mut self) -> Box<dyn Connection> {
280 self.pool = None; std::mem::replace(&mut self.conn, Box::new(ClosedConnection))
284 }
285}
286
287impl Drop for PooledConnection {
299 fn drop(&mut self) {
300 if let Some(pool) = self.pool.take() {
301 let conn = std::mem::replace(&mut self.conn, Box::new(ClosedConnection));
303 let pooled = PooledConnection {
304 conn,
305 created_at: self.created_at,
306 last_used_at: self.last_used_at,
307 pool: None,
308 };
309 if let Ok(handle) = tokio::runtime::Handle::try_current() {
311 handle.spawn(async move {
312 pool.release(pooled).await;
313 });
314 } else {
315 drop(pooled);
319 pool.total_count.fetch_sub(1, Ordering::SeqCst);
320 }
321 }
322 }
323}
324
325struct ClosedConnection;
329
330impl Connection for ClosedConnection {
331 fn execute<'a>(
332 &'a mut self,
333 _sql: &'a str,
334 ) -> Pin<Box<dyn Future<Output = Result<u64, crate::DbError>> + Send + 'a>> {
335 Box::pin(async {
336 Err(crate::DbError::ConnectionError(
337 "connection already returned to pool".to_string(),
338 ))
339 })
340 }
341
342 fn query<'a>(
343 &'a mut self,
344 _sql: &'a str,
345 ) -> Pin<Box<dyn Future<Output = Result<QueryRows, crate::DbError>> + Send + 'a>> {
346 Box::pin(async {
347 Err(crate::DbError::ConnectionError(
348 "connection already returned to pool".to_string(),
349 ))
350 })
351 }
352
353 fn begin_transaction<'a>(
354 &'a mut self,
355 ) -> Pin<Box<dyn Future<Output = Result<(), crate::DbError>> + Send + 'a>> {
356 Box::pin(async {
357 Err(crate::DbError::ConnectionError(
358 "connection already returned to pool".to_string(),
359 ))
360 })
361 }
362
363 fn commit<'a>(
364 &'a mut self,
365 ) -> Pin<Box<dyn Future<Output = Result<(), crate::DbError>> + Send + 'a>> {
366 Box::pin(async { Ok(()) })
367 }
368
369 fn rollback<'a>(
370 &'a mut self,
371 ) -> Pin<Box<dyn Future<Output = Result<(), crate::DbError>> + Send + 'a>> {
372 Box::pin(async { Ok(()) })
373 }
374
375 fn is_connected(&self) -> bool {
376 false
377 }
378
379 fn ping<'a>(&'a mut self) -> Pin<Box<dyn Future<Output = bool> + Send + 'a>> {
380 Box::pin(async { false })
381 }
382
383 fn close<'a>(
384 &'a mut self,
385 ) -> Pin<Box<dyn Future<Output = Result<(), crate::DbError>> + Send + 'a>> {
386 Box::pin(async { Ok(()) })
387 }
388}
389
390impl Deref for PooledConnection {
391 type Target = dyn Connection;
392
393 fn deref(&self) -> &Self::Target {
394 self.conn.as_ref()
395 }
396}
397
398impl DerefMut for PooledConnection {
399 fn deref_mut(&mut self) -> &mut Self::Target {
400 self.conn.as_mut()
401 }
402}
403
404#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
406pub enum TlsVersion {
407 #[default]
409 Tls12,
410 Tls13,
412}
413
414#[derive(Debug, Clone, Default)]
416pub struct TlsConfig {
417 pub enabled: bool,
419 pub ca_cert_path: Option<String>,
421 pub client_cert_path: Option<String>,
423 pub client_key_path: Option<String>,
425 pub min_version: TlsVersion,
427}
428
429#[derive(Debug, Clone)]
431pub enum PoolEvent {
432 ConnectionCreated,
434 ConnectionClosed,
436 ConnectionAcquired,
438 ConnectionReleased,
440 AcquireTimeout,
442}
443
444pub type PoolEventCallback = Arc<dyn Fn(PoolEvent) + Send + Sync>;
446
447pub struct PoolConfig {
449 pub max_size: u32,
451 pub min_idle: u32,
453 pub acquire_timeout: Duration,
455 pub idle_timeout: Duration,
457 pub max_lifetime: Duration,
459 pub connection_timeout: Duration,
461 pub tls: Option<TlsConfig>,
463 pub query_timeout: Option<Duration>,
465 pub max_rows: Option<usize>,
467 pub memory_limit: Option<usize>,
469 pub on_event: Option<PoolEventCallback>,
471 pub test_before_acquire: bool,
479 pub prewarm: bool,
491}
492
493impl Default for PoolConfig {
494 fn default() -> Self {
495 Self {
496 max_size: 100,
497 min_idle: 0,
498 acquire_timeout: Duration::from_secs(30),
499 idle_timeout: Duration::from_secs(600),
500 max_lifetime: Duration::from_secs(1800),
501 connection_timeout: Duration::from_secs(10),
502 tls: None,
503 query_timeout: Some(Duration::from_secs(30)),
504 max_rows: None,
505 memory_limit: None,
506 on_event: None,
507 test_before_acquire: false,
508 prewarm: false,
509 }
510 }
511}
512
513impl Clone for PoolConfig {
514 fn clone(&self) -> Self {
515 Self {
516 max_size: self.max_size,
517 min_idle: self.min_idle,
518 acquire_timeout: self.acquire_timeout,
519 idle_timeout: self.idle_timeout,
520 max_lifetime: self.max_lifetime,
521 connection_timeout: self.connection_timeout,
522 tls: self.tls.clone(),
523 query_timeout: self.query_timeout,
524 max_rows: self.max_rows,
525 memory_limit: self.memory_limit,
526 on_event: self.on_event.clone(),
527 test_before_acquire: self.test_before_acquire,
528 prewarm: self.prewarm,
529 }
530 }
531}
532
533impl PoolConfig {
534 pub fn validate(&self) -> Result<(), PoolError> {
536 if self.max_size == 0 {
537 return Err(PoolError::InvalidConfig("max_size cannot be 0".to_string()));
538 }
539 if self.min_idle > self.max_size {
540 return Err(PoolError::InvalidConfig(
541 "min_idle cannot exceed max_size".to_string(),
542 ));
543 }
544 const MAX_DURATION_SECS: u64 = u32::MAX as u64; for (name, dur) in [
550 ("acquire_timeout", self.acquire_timeout),
551 ("idle_timeout", self.idle_timeout),
552 ("max_lifetime", self.max_lifetime),
553 ("connection_timeout", self.connection_timeout),
554 ] {
555 if dur.as_secs() > MAX_DURATION_SECS {
556 return Err(PoolError::InvalidConfig(format!(
557 "{name} ({:?}) exceeds maximum allowed duration ({} seconds)",
558 dur, MAX_DURATION_SECS
559 )));
560 }
561 }
562 Ok(())
563 }
564
565 #[must_use]
567 pub fn with_prewarm(mut self, prewarm: bool) -> Self {
568 self.prewarm = prewarm;
569 self
570 }
571}
572
573pub struct PoolStatus {
575 pub idle: u32,
577 pub active: u32,
579 pub max: u32,
581 pub min: u32,
583 pub waiters: u32,
585}
586
587impl std::fmt::Debug for PoolStatus {
588 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
589 f.debug_struct("PoolStatus")
590 .field("idle", &self.idle)
591 .field("active", &self.active)
592 .field("max", &self.max)
593 .field("min", &self.min)
594 .field("waiters", &self.waiters)
595 .finish()
596 }
597}
598
599#[derive(Debug, Clone, Copy, Default, PartialEq, Eq, Serialize, Deserialize)]
606pub struct PoolMetrics {
607 pub acquire_count: u64,
609 pub acquire_failed_count: u64,
611 pub acquire_wait_time: Duration,
613 pub release_count: u64,
615 pub connection_created_count: u64,
617 pub connection_closed_count: u64,
619}
620
621impl PoolMetrics {
622 #[must_use]
624 pub fn average_acquire_wait_time(&self) -> Duration {
625 if self.acquire_count == 0 {
626 Duration::ZERO
627 } else {
628 self.acquire_wait_time / self.acquire_count as u32
629 }
630 }
631
632 #[must_use]
638 pub fn connection_reuse_rate(&self) -> f64 {
639 if self.acquire_count == 0 || self.connection_created_count > self.acquire_count {
640 0.0
641 } else {
642 let reused = self.acquire_count - self.connection_created_count;
643 reused as f64 / self.acquire_count as f64
644 }
645 }
646}
647
648#[derive(Debug, Clone)]
652pub struct PoolTuningAdvice {
653 pub suggested_max_size: Option<u32>,
655 pub suggested_min_idle: Option<u32>,
657 pub suggested_idle_timeout: Option<Duration>,
659 pub reason: String,
661}
662
663impl PoolTuningAdvice {
664 #[must_use]
666 pub fn is_optimal(&self) -> bool {
667 self.suggested_max_size.is_none()
668 && self.suggested_min_idle.is_none()
669 && self.suggested_idle_timeout.is_none()
670 }
671}
672
673pub struct PoolConfigBuilder {
675 config: PoolConfig,
676}
677
678impl PoolConfigBuilder {
679 pub fn new() -> Self {
681 Self {
682 config: PoolConfig::default(),
683 }
684 }
685
686 pub fn max_size(mut self, size: u32) -> Self {
688 self.config.max_size = size;
689 self
690 }
691
692 pub fn min_idle(mut self, count: u32) -> Self {
694 self.config.min_idle = count;
695 self
696 }
697
698 pub fn acquire_timeout(mut self, timeout_secs: u64) -> Self {
700 self.config.acquire_timeout = Duration::from_secs(timeout_secs);
701 self
702 }
703
704 pub fn idle_timeout(mut self, timeout_secs: u64) -> Self {
706 self.config.idle_timeout = Duration::from_secs(timeout_secs);
707 self
708 }
709
710 pub fn max_lifetime(mut self, lifetime_secs: u64) -> Self {
712 self.config.max_lifetime = Duration::from_secs(lifetime_secs);
713 self
714 }
715
716 pub fn tls(mut self, tls: TlsConfig) -> Self {
718 self.config.tls = Some(tls);
719 self
720 }
721
722 pub fn query_timeout(mut self, timeout: Duration) -> Self {
724 self.config.query_timeout = Some(timeout);
725 self
726 }
727
728 pub fn max_rows(mut self, max_rows: usize) -> Self {
730 self.config.max_rows = Some(max_rows);
731 self
732 }
733
734 pub fn memory_limit(mut self, memory_limit: usize) -> Self {
736 self.config.memory_limit = Some(memory_limit);
737 self
738 }
739
740 pub fn on_event(mut self, callback: PoolEventCallback) -> Self {
742 self.config.on_event = Some(callback);
743 self
744 }
745
746 pub fn test_before_acquire(mut self, enabled: bool) -> Self {
751 self.config.test_before_acquire = enabled;
752 self
753 }
754
755 pub fn prewarm(mut self, enabled: bool) -> Self {
760 self.config.prewarm = enabled;
761 self
762 }
763
764 pub fn build(self) -> Result<PoolConfig, PoolError> {
766 self.config.validate()?;
767 Ok(self.config)
768 }
769}
770
771impl Default for PoolConfigBuilder {
772 fn default() -> Self {
773 Self::new()
774 }
775}
776
777#[async_trait]
779pub trait ConnectionFactory: Send + Sync {
780 async fn create(&self) -> Result<Box<dyn Connection>, crate::DbError>;
782}
783
784pub struct Pool {
790 config: PoolConfig,
791 factory: Arc<dyn ConnectionFactory>,
792 idle: Arc<ArrayQueue<PooledConnection>>,
798 total_count: Arc<AtomicU32>,
808 closed: Arc<AtomicBool>,
810 notify: Arc<Notify>,
811 waiters_count: Arc<AtomicU32>,
813 dynamic_max_size: Arc<AtomicU32>,
815 #[cfg(feature = "circuit-breaker")]
821 circuit_breaker: Arc<PlMutex<DefaultCircuitBreaker>>,
822 #[cfg(feature = "rate-limit")]
831 rate_limiter: Arc<PlRwLock<Option<Arc<dyn RateLimiter>>>>,
832 #[cfg(feature = "rate-limit")]
834 rate_limit_key: String,
835 #[cfg(feature = "tenant-quota-rls-enhanced")]
840 quota_enforcer: Arc<PlMutex<Option<Arc<QuotaEnforcer>>>>,
841 acquire_count: Arc<AtomicU64>,
843 acquire_failed_count: Arc<AtomicU64>,
845 acquire_wait_time_ns: Arc<AtomicU64>,
847 release_count: Arc<AtomicU64>,
849 connection_created_count: Arc<AtomicU64>,
851 connection_closed_count: Arc<AtomicU64>,
853}
854
855impl Clone for Pool {
859 fn clone(&self) -> Self {
860 Self {
861 config: self.config.clone(),
862 factory: self.factory.clone(),
863 idle: self.idle.clone(),
864 total_count: self.total_count.clone(),
865 closed: self.closed.clone(),
866 notify: Arc::clone(&self.notify),
867 waiters_count: self.waiters_count.clone(),
868 dynamic_max_size: self.dynamic_max_size.clone(),
869 #[cfg(feature = "circuit-breaker")]
870 circuit_breaker: Arc::clone(&self.circuit_breaker),
871 #[cfg(feature = "rate-limit")]
872 rate_limiter: Arc::clone(&self.rate_limiter),
873 #[cfg(feature = "rate-limit")]
874 rate_limit_key: self.rate_limit_key.clone(),
875 #[cfg(feature = "tenant-quota-rls-enhanced")]
876 quota_enforcer: Arc::clone(&self.quota_enforcer),
877 acquire_count: self.acquire_count.clone(),
878 acquire_failed_count: self.acquire_failed_count.clone(),
879 acquire_wait_time_ns: self.acquire_wait_time_ns.clone(),
880 release_count: self.release_count.clone(),
881 connection_created_count: self.connection_created_count.clone(),
882 connection_closed_count: self.connection_closed_count.clone(),
883 }
884 }
885}
886
887impl Pool {
888 pub fn new(config: PoolConfig, factory: Arc<dyn ConnectionFactory>) -> Result<Self, PoolError> {
912 config.validate()?;
913 let max_size = config.max_size as usize;
916 let dynamic_max = config.max_size;
917 Ok(Self {
918 config,
919 factory,
920 idle: Arc::new(ArrayQueue::new(max_size)),
921 total_count: Arc::new(AtomicU32::new(0)),
922 closed: Arc::new(AtomicBool::new(false)),
923 notify: Arc::new(Notify::new()),
924 waiters_count: Arc::new(AtomicU32::new(0)),
925 dynamic_max_size: Arc::new(AtomicU32::new(dynamic_max)),
926 #[cfg(feature = "circuit-breaker")]
929 circuit_breaker: Arc::new(PlMutex::new(DefaultCircuitBreaker::new(
930 5,
931 std::time::Duration::from_secs(30),
932 ))),
933 #[cfg(feature = "rate-limit")]
936 rate_limiter: Arc::new(PlRwLock::new(None)),
937 #[cfg(feature = "rate-limit")]
938 rate_limit_key: "pool".to_string(),
939 #[cfg(feature = "tenant-quota-rls-enhanced")]
940 quota_enforcer: Arc::new(PlMutex::new(None)),
941 acquire_count: Arc::new(AtomicU64::new(0)),
942 acquire_failed_count: Arc::new(AtomicU64::new(0)),
943 acquire_wait_time_ns: Arc::new(AtomicU64::new(0)),
944 release_count: Arc::new(AtomicU64::new(0)),
945 connection_created_count: Arc::new(AtomicU64::new(0)),
946 connection_closed_count: Arc::new(AtomicU64::new(0)),
947 })
948 }
949
950 pub async fn new_async(
957 config: PoolConfig,
958 factory: Arc<dyn ConnectionFactory>,
959 ) -> Result<Self, PoolError> {
960 let pool = Self::new(config, factory)?;
961 if pool.config.prewarm {
962 pool.prewarm().await;
963 }
964 Ok(pool)
965 }
966
967 pub async fn prewarm(&self) {
984 if !self.config.prewarm {
985 return;
986 }
987
988 let min_idle = self.config.min_idle as usize;
989 let mut warmed = 0;
990
991 for i in 0..min_idle {
992 if self.closed.load(Ordering::Acquire) {
994 break;
995 }
996
997 let current_max = self.dynamic_max_size.load(Ordering::Acquire);
999 let current = self.total_count.load(Ordering::Acquire);
1000 if current >= current_max {
1001 break;
1002 }
1003
1004 let created = loop {
1006 let current = self.total_count.load(Ordering::Acquire);
1007 if current >= current_max {
1008 break None;
1009 }
1010 match self.total_count.compare_exchange(
1011 current,
1012 current + 1,
1013 Ordering::SeqCst,
1014 Ordering::Acquire,
1015 ) {
1016 Ok(_) => break Some(()),
1017 Err(_) => continue,
1018 }
1019 };
1020
1021 if created.is_some() {
1022 match tokio::time::timeout(self.config.connection_timeout, self.factory.create())
1023 .await
1024 {
1025 Ok(Ok(conn)) => {
1026 #[cfg(feature = "circuit-breaker")]
1027 {
1028 self.circuit_breaker.lock().record_success();
1029 }
1030 self.emit_event(PoolEvent::ConnectionCreated);
1031 let pooled = PooledConnection::new(conn, self.clone());
1032 if self.idle.push(pooled).is_err() {
1034 let _ = self.total_count.fetch_sub(1, Ordering::SeqCst);
1036 tracing::warn!(
1037 target: "sz_orm::pool::prewarm",
1038 "prewarm connection {} failed: idle queue full",
1039 i
1040 );
1041 } else {
1042 warmed += 1;
1043 self.notify.notify_one();
1044 }
1045 }
1046 Ok(Err(e)) => {
1047 let _ = self.total_count.fetch_sub(1, Ordering::SeqCst);
1048 #[cfg(feature = "circuit-breaker")]
1049 {
1050 self.circuit_breaker.lock().record_failure();
1051 }
1052 tracing::warn!(
1053 target: "sz_orm::pool::prewarm",
1054 "prewarm connection {} failed: {}",
1055 i,
1056 e
1057 );
1058 }
1059 Err(_) => {
1060 let _ = self.total_count.fetch_sub(1, Ordering::SeqCst);
1061 #[cfg(feature = "circuit-breaker")]
1062 {
1063 self.circuit_breaker.lock().record_failure();
1064 }
1065 tracing::warn!(
1066 target: "sz_orm::pool::prewarm",
1067 "prewarm connection {} timeout",
1068 i
1069 );
1070 }
1071 }
1072 }
1073 }
1074
1075 if warmed > 0 {
1076 tracing::info!(
1077 target: "sz_orm::pool::prewarm",
1078 "pool prewarm completed: {}/{} connections established",
1079 warmed,
1080 min_idle
1081 );
1082 }
1083 }
1084
1085 #[cfg(feature = "auto-prewarm")]
1090 pub async fn progressive_prewarm(
1091 &self,
1092 batch_size: u32,
1093 interval: std::time::Duration,
1094 total_timeout: std::time::Duration,
1095 progress: &crate::prewarm::PrewarmProgress,
1096 ) {
1097 use std::time::Instant;
1098
1099 let min_idle = self.config.min_idle;
1100 if min_idle == 0 || !self.config.prewarm {
1101 progress.mark_completed();
1102 return;
1103 }
1104
1105 let start = Instant::now();
1106 let batch = batch_size.max(1);
1107 let mut warmed_total: u32 = 0;
1108
1109 while warmed_total < min_idle {
1110 if start.elapsed() >= total_timeout {
1111 tracing::warn!(
1112 target: "sz_orm::pool::prewarm",
1113 "progressive prewarm timeout: {}/{} connections established",
1114 warmed_total,
1115 min_idle
1116 );
1117 break;
1118 }
1119
1120 if self.closed.load(Ordering::Acquire) {
1121 break;
1122 }
1123
1124 let remaining = min_idle - warmed_total;
1125 let this_batch = batch.min(remaining);
1126
1127 for _ in 0..this_batch {
1128 let current_max = self.dynamic_max_size.load(Ordering::Acquire);
1129 let current = self.total_count.load(Ordering::Acquire);
1130 if current >= current_max {
1131 break;
1132 }
1133
1134 let created = loop {
1135 let current = self.total_count.load(Ordering::Acquire);
1136 if current >= current_max {
1137 break None;
1138 }
1139 match self.total_count.compare_exchange(
1140 current,
1141 current + 1,
1142 Ordering::SeqCst,
1143 Ordering::Acquire,
1144 ) {
1145 Ok(_) => break Some(()),
1146 Err(_) => continue,
1147 }
1148 };
1149
1150 if created.is_some() {
1151 match tokio::time::timeout(
1152 self.config.connection_timeout,
1153 self.factory.create(),
1154 )
1155 .await
1156 {
1157 Ok(Ok(conn)) => {
1158 #[cfg(feature = "circuit-breaker")]
1159 {
1160 self.circuit_breaker.lock().record_success();
1161 }
1162 self.emit_event(PoolEvent::ConnectionCreated);
1163 let pooled = PooledConnection::new(conn, self.clone());
1164 if self.idle.push(pooled).is_err() {
1165 let _ = self.total_count.fetch_sub(1, Ordering::SeqCst);
1166 progress.record_failure();
1167 } else {
1168 progress.record_success();
1169 warmed_total += 1;
1170 self.notify.notify_one();
1171 }
1172 }
1173 Ok(Err(_)) => {
1174 let _ = self.total_count.fetch_sub(1, Ordering::SeqCst);
1175 progress.record_failure();
1176 #[cfg(feature = "circuit-breaker")]
1177 {
1178 self.circuit_breaker.lock().record_failure();
1179 }
1180 }
1181 Err(_) => {
1182 let _ = self.total_count.fetch_sub(1, Ordering::SeqCst);
1183 progress.record_failure();
1184 #[cfg(feature = "circuit-breaker")]
1185 {
1186 self.circuit_breaker.lock().record_failure();
1187 }
1188 }
1189 }
1190 }
1191 }
1192
1193 if warmed_total < min_idle && interval > std::time::Duration::ZERO {
1194 tokio::time::sleep(interval).await;
1195 }
1196 }
1197
1198 progress.set_elapsed(start.elapsed());
1199 progress.mark_completed();
1200
1201 tracing::info!(
1202 target: "sz_orm::pool::prewarm",
1203 "progressive prewarm completed: {} warmed, {} failed, elapsed {:?}",
1204 progress.snapshot().warmed,
1205 progress.snapshot().failed,
1206 start.elapsed()
1207 );
1208 }
1209
1210 pub fn config(&self) -> &PoolConfig {
1212 &self.config
1213 }
1214
1215 #[cfg(feature = "circuit-breaker")]
1229 pub fn configure_circuit_breaker(
1230 &self,
1231 failure_threshold: usize,
1232 reset_timeout: std::time::Duration,
1233 ) {
1234 let new_cb = DefaultCircuitBreaker::new(failure_threshold, reset_timeout);
1235 let mut guard = self.circuit_breaker.lock();
1237 *guard = new_cb;
1238 }
1239
1240 #[cfg(feature = "circuit-breaker")]
1245 pub fn reset_circuit_breaker(&self) -> bool {
1246 let mut guard = self.circuit_breaker.lock();
1248 guard.reset()
1249 }
1250
1251 #[cfg(feature = "circuit-breaker")]
1253 pub fn circuit_state(&self) -> CircuitState {
1254 let guard = self.circuit_breaker.lock();
1256 guard.state()
1257 }
1258
1259 #[cfg(feature = "rate-limit")]
1268 pub fn set_rate_limiter(&self, limiter: Option<Arc<dyn RateLimiter>>) {
1269 let mut guard = self.rate_limiter.write();
1271 *guard = limiter;
1272 }
1273
1274 #[cfg(feature = "rate-limit")]
1276 pub fn with_rate_limit_key(mut self, key: impl Into<String>) -> Self {
1277 self.rate_limit_key = key.into();
1278 self
1279 }
1280
1281 #[cfg(feature = "tenant-quota-rls-enhanced")]
1286 pub fn set_quota_enforcer(&self, enforcer: Option<Arc<QuotaEnforcer>>) {
1287 let mut guard = self.quota_enforcer.lock();
1288 *guard = enforcer;
1289 }
1290
1291 #[cfg(feature = "tenant-quota-rls-enhanced")]
1300 pub async fn acquire_with_tenant(
1301 &self,
1302 tenant_id: &str,
1303 ) -> Result<PooledConnection, PoolError> {
1304 {
1305 let guard = self.quota_enforcer.lock();
1306 if let Some(ref enforcer) = *guard {
1307 let current = enforcer.current_usage(tenant_id, QuotaResource::Connection);
1308 enforcer
1309 .check_and_record(tenant_id, QuotaResource::Connection, 1)
1310 .map_err(|e| PoolError::Internal(e.to_string()))?;
1311 let _ = current;
1312 }
1313 }
1314 self.acquire().await
1315 }
1316
1317 #[cfg(feature = "tenant-quota-rls-enhanced")]
1322 pub async fn release_with_tenant(&self, tenant_id: &str, pooled: PooledConnection) {
1323 {
1324 let guard = self.quota_enforcer.lock();
1325 if let Some(ref enforcer) = *guard {
1326 enforcer.release_usage(tenant_id, QuotaResource::Connection, 1);
1329 }
1330 }
1331 self.release(pooled).await;
1332 }
1333
1334 fn emit_event(&self, event: PoolEvent) {
1336 if matches!(event, PoolEvent::ConnectionCreated) {
1339 self.connection_created_count
1340 .fetch_add(1, Ordering::Relaxed);
1341 }
1342 if let Some(ref callback) = self.config.on_event {
1343 callback(event);
1344 }
1345 }
1346
1347 async fn close_connection(&self, pooled: PooledConnection) {
1352 let mut pooled = pooled;
1353 let _ = pooled.conn.close().await;
1354 self.connection_closed_count.fetch_add(1, Ordering::Relaxed);
1355 }
1356
1357 #[tracing::instrument(skip(self), fields(max_size = self.config.max_size, acquire_timeout = ?self.config.acquire_timeout))]
1377 pub async fn acquire(&self) -> Result<PooledConnection, PoolError> {
1378 if self.closed.load(Ordering::Acquire) {
1380 self.acquire_failed_count.fetch_add(1, Ordering::Relaxed);
1381 return Err(PoolError::Closed);
1382 }
1383
1384 #[cfg(feature = "circuit-breaker")]
1388 {
1389 let mut guard = self.circuit_breaker.lock();
1390 if !guard.can_execute() {
1391 self.acquire_failed_count.fetch_add(1, Ordering::Relaxed);
1392 return Err(PoolError::CircuitOpen);
1393 }
1394 }
1395
1396 #[cfg(feature = "rate-limit")]
1400 {
1401 let guard = self.rate_limiter.read();
1402 if let Some(ref limiter) = *guard {
1403 match limiter.try_acquire(&self.rate_limit_key) {
1404 Ok(result) if !result.allowed => {
1405 self.acquire_failed_count.fetch_add(1, Ordering::Relaxed);
1406 return Err(PoolError::RateLimited {
1407 remaining: result.remaining,
1408 reset_at: result.reset_at,
1409 });
1410 }
1411 Ok(_) => {} Err(_) => {
1413 }
1415 }
1416 }
1417 }
1418
1419 let deadline = Instant::now() + self.config.acquire_timeout;
1420 let mut backoff = Duration::from_millis(1);
1422 const MAX_BACKOFF: Duration = Duration::from_millis(100);
1424
1425 loop {
1426 let mut to_close: Vec<PooledConnection> = Vec::new();
1432 let acquired: Option<PooledConnection> = {
1433 let mut found: Option<PooledConnection> = None;
1434 while let Some(pooled) = self.idle.pop() {
1435 if pooled.is_expired(self.config.max_lifetime) {
1437 to_close.push(pooled);
1438 continue;
1439 }
1440 if pooled.is_idle_too_long(self.config.idle_timeout) {
1442 to_close.push(pooled);
1443 continue;
1444 }
1445 if !pooled.conn.is_connected() {
1448 to_close.push(pooled);
1449 continue;
1450 }
1451 found = Some(pooled);
1452 break;
1453 }
1454 found
1455 };
1456
1457 for pooled in to_close {
1459 self.close_connection(pooled).await;
1460 self.total_count.fetch_sub(1, Ordering::SeqCst);
1462 }
1463
1464 if let Some(mut pooled) = acquired {
1465 if self.config.test_before_acquire {
1467 let ping_timeout = self.config.connection_timeout / 2;
1468 let alive = match tokio::time::timeout(ping_timeout, pooled.conn.ping()).await {
1469 Ok(true) => true,
1470 Ok(false) => false,
1471 Err(_) => false, };
1473 if !alive {
1474 self.close_connection(pooled).await;
1476 self.total_count.fetch_sub(1, Ordering::SeqCst);
1477 continue;
1478 }
1479 }
1480 pooled.pool = Some(self.clone());
1483 self.acquire_count.fetch_add(1, Ordering::Relaxed);
1484 return Ok(pooled);
1485 }
1486
1487 let current_max = self.dynamic_max_size.load(Ordering::Acquire);
1492 let created = loop {
1493 let current = self.total_count.load(Ordering::Acquire);
1494 if current >= current_max {
1495 break None; }
1497 match self.total_count.compare_exchange(
1498 current,
1499 current + 1,
1500 Ordering::SeqCst,
1501 Ordering::Acquire,
1502 ) {
1503 Ok(_) => break Some(()), Err(_) => continue, }
1506 };
1507
1508 if created.is_some() {
1509 match tokio::time::timeout(self.config.connection_timeout, self.factory.create())
1510 .await
1511 {
1512 Ok(Ok(conn)) => {
1513 #[cfg(feature = "circuit-breaker")]
1516 {
1517 self.circuit_breaker.lock().record_success();
1518 }
1519 self.emit_event(PoolEvent::ConnectionCreated);
1520 self.emit_event(PoolEvent::ConnectionAcquired);
1521 self.acquire_count.fetch_add(1, Ordering::Relaxed);
1522 return Ok(PooledConnection::new(conn, self.clone()));
1523 }
1524 Ok(Err(e)) => {
1525 self.total_count.fetch_sub(1, Ordering::SeqCst);
1527 #[cfg(feature = "circuit-breaker")]
1530 {
1531 self.circuit_breaker.lock().record_failure();
1532 }
1533 self.acquire_failed_count.fetch_add(1, Ordering::Relaxed);
1534 return Err(PoolError::ConnectionFailed(e.to_string()));
1535 }
1536 Err(_) => {
1537 self.total_count.fetch_sub(1, Ordering::SeqCst);
1539 #[cfg(feature = "circuit-breaker")]
1542 {
1543 self.circuit_breaker.lock().record_failure();
1544 }
1545 self.acquire_failed_count.fetch_add(1, Ordering::Relaxed);
1546 return Err(PoolError::Timeout);
1547 }
1548 }
1549 }
1550
1551 let now = Instant::now();
1553 if now >= deadline {
1554 self.emit_event(PoolEvent::AcquireTimeout);
1555 self.acquire_failed_count.fetch_add(1, Ordering::Relaxed);
1556 return Err(PoolError::Timeout);
1557 }
1558 self.waiters_count.fetch_add(1, Ordering::SeqCst);
1560 let wait = std::cmp::min(backoff, deadline - now);
1561 match tokio::time::timeout(wait, self.notify.notified()).await {
1562 Ok(()) => {
1563 backoff = Duration::from_millis(1);
1565 }
1566 Err(_) => {
1567 backoff = std::cmp::min(backoff * 2, MAX_BACKOFF);
1569 }
1570 }
1571 self.waiters_count.fetch_sub(1, Ordering::SeqCst);
1573 self.acquire_wait_time_ns
1575 .fetch_add(wait.as_nanos() as u64, Ordering::Relaxed);
1576 }
1577 }
1578
1579 #[tracing::instrument(skip(self, pooled))]
1587 pub async fn release(&self, mut pooled: PooledConnection) {
1588 pooled.pool = None;
1590 self.release_count.fetch_add(1, Ordering::Relaxed);
1592
1593 if self.closed.load(Ordering::Acquire) {
1595 self.close_connection(pooled).await;
1596 self.total_count.fetch_sub(1, Ordering::SeqCst);
1598 self.emit_event(PoolEvent::ConnectionClosed);
1599 return;
1600 }
1601
1602 if !pooled.conn.is_connected() {
1604 self.close_connection(pooled).await;
1605 self.total_count.fetch_sub(1, Ordering::SeqCst);
1606 self.emit_event(PoolEvent::ConnectionClosed);
1607 return;
1608 }
1609
1610 pooled.last_used_at = Instant::now();
1612
1613 if let Err(rejected) = self.idle.push(pooled) {
1619 self.close_connection(rejected).await;
1621 self.total_count.fetch_sub(1, Ordering::SeqCst);
1622 self.emit_event(PoolEvent::ConnectionClosed);
1623 } else {
1624 self.emit_event(PoolEvent::ConnectionReleased);
1625 }
1626 self.notify.notify_one();
1627 }
1628
1629 pub async fn status(&self) -> PoolStatus {
1634 let idle_count = self.idle.len() as u32;
1635 let active = self.total_count.load(Ordering::Acquire);
1637 let waiters = self.waiters_count.load(Ordering::Acquire);
1638 PoolStatus {
1639 idle: idle_count,
1640 active,
1641 max: self.dynamic_max_size.load(Ordering::Acquire),
1642 min: self.config.min_idle,
1643 waiters,
1644 }
1645 }
1646
1647 pub fn pool_metrics(&self) -> PoolMetrics {
1658 PoolMetrics {
1659 acquire_count: self.acquire_count.load(Ordering::Acquire),
1660 acquire_failed_count: self.acquire_failed_count.load(Ordering::Acquire),
1661 acquire_wait_time: Duration::from_nanos(
1662 self.acquire_wait_time_ns.load(Ordering::Acquire),
1663 ),
1664 release_count: self.release_count.load(Ordering::Acquire),
1665 connection_created_count: self.connection_created_count.load(Ordering::Acquire),
1666 connection_closed_count: self.connection_closed_count.load(Ordering::Acquire),
1667 }
1668 }
1669
1670 #[must_use]
1679 pub fn suggest_tuning(&self) -> PoolTuningAdvice {
1680 let metrics = self.pool_metrics();
1681 let current_max = self.dynamic_max_size.load(Ordering::Acquire);
1682
1683 if metrics.acquire_count == 0 {
1684 return PoolTuningAdvice {
1685 suggested_max_size: None,
1686 suggested_min_idle: None,
1687 suggested_idle_timeout: None,
1688 reason: "数据不足".to_string(),
1689 };
1690 }
1691
1692 let reuse_rate = metrics.connection_reuse_rate();
1693 let mut advice = PoolTuningAdvice {
1694 suggested_max_size: None,
1695 suggested_min_idle: None,
1696 suggested_idle_timeout: None,
1697 reason: String::new(),
1698 };
1699
1700 if reuse_rate < 0.5 {
1701 advice.suggested_max_size = Some(current_max.saturating_mul(2));
1702 advice.reason = "复用率过低,池过小或回收过激".to_string();
1703 } else if reuse_rate < 0.9 {
1704 advice.suggested_min_idle = Some(current_max / 4);
1705 advice.reason = "复用率偏低,预热不足".to_string();
1706 }
1707
1708 let avg_wait = metrics.average_acquire_wait_time();
1709 if avg_wait > Duration::from_millis(100) {
1710 advice.suggested_max_size = Some(current_max.saturating_mul(2));
1711 if !advice.reason.is_empty() {
1712 advice.reason.push(';');
1713 }
1714 advice.reason.push_str("等待时长过高,池容量不足");
1715 }
1716
1717 if metrics.connection_created_count > 0
1718 && metrics.connection_closed_count as f64
1719 > metrics.connection_created_count as f64 * 0.5
1720 {
1721 advice.suggested_idle_timeout = Some(self.config.idle_timeout * 2);
1722 if !advice.reason.is_empty() {
1723 advice.reason.push(';');
1724 }
1725 advice.reason.push_str("连接关闭过快,空闲回收过激");
1726 }
1727
1728 if advice.reason.is_empty() {
1729 advice.reason = "池配置合理".to_string();
1730 }
1731
1732 advice
1733 }
1734
1735 pub fn metrics_snapshot_json(&self) -> String {
1740 serde_json::to_string(&self.pool_metrics()).unwrap_or_else(|_| "{}".to_string())
1741 }
1742
1743 #[tracing::instrument(skip(self))]
1745 pub async fn reap_idle(&self) {
1746 let mut all: Vec<PooledConnection> = Vec::new();
1750 while let Some(pooled) = self.idle.pop() {
1751 all.push(pooled);
1752 }
1753
1754 let mut to_close = Vec::new();
1756 for pooled in all {
1757 if pooled.is_idle_too_long(self.config.idle_timeout)
1758 || pooled.is_expired(self.config.max_lifetime)
1759 {
1760 to_close.push(pooled);
1761 } else {
1762 if let Err(rejected) = self.idle.push(pooled) {
1764 self.close_connection(rejected).await;
1765 self.total_count.fetch_sub(1, Ordering::SeqCst);
1766 }
1767 }
1768 }
1769
1770 for pooled in to_close {
1772 self.close_connection(pooled).await;
1773 self.total_count.fetch_sub(1, Ordering::SeqCst);
1775 }
1776 }
1777
1778 pub async fn close_all(&self) {
1782 self.closed.store(true, Ordering::Release);
1784 let mut to_close: Vec<PooledConnection> = Vec::new();
1787 while let Some(pooled) = self.idle.pop() {
1788 to_close.push(pooled);
1789 }
1790 let closed_count: u32 = to_close.len() as u32;
1792 for pooled in to_close {
1793 self.close_connection(pooled).await;
1794 }
1795 self.total_count.fetch_sub(closed_count, Ordering::SeqCst);
1798 }
1799
1800 pub async fn health_check(&self) -> u32 {
1816 let mut to_check: Vec<PooledConnection> = Vec::new();
1818 while let Some(pooled) = self.idle.pop() {
1819 to_check.push(pooled);
1820 }
1821
1822 let mut removed: u32 = 0;
1823 let mut alive: Vec<PooledConnection> = Vec::with_capacity(to_check.len());
1824 for mut pooled in to_check.drain(..) {
1825 if !pooled.conn.is_connected() {
1827 self.close_connection(pooled).await;
1828 removed += 1;
1829 continue;
1830 }
1831 let ping_timeout = self.config.connection_timeout / 2;
1833 match tokio::time::timeout(ping_timeout, pooled.conn.ping()).await {
1834 Ok(true) => alive.push(pooled),
1835 Ok(false) => {
1836 self.close_connection(pooled).await;
1838 removed += 1;
1839 }
1840 Err(_) => {
1841 self.close_connection(pooled).await;
1843 removed += 1;
1844 }
1845 }
1846 }
1847
1848 let alive_count: u32 = alive.len() as u32;
1850 for pooled in alive {
1851 if let Err(rejected) = self.idle.push(pooled) {
1853 self.close_connection(rejected).await;
1854 removed += 1;
1855 }
1856 }
1857
1858 if removed > 0 {
1860 self.total_count.fetch_sub(removed, Ordering::SeqCst);
1861 }
1862
1863 if alive_count > 0 {
1865 self.notify.notify_one();
1866 }
1867
1868 removed
1869 }
1870
1871 pub async fn shutdown(&self) {
1878 self.shutdown_with_timeout(Duration::from_secs(30)).await;
1879 }
1880
1881 pub async fn shutdown_with_timeout(&self, timeout: Duration) {
1889 if self.closed.swap(true, Ordering::SeqCst) {
1891 return;
1892 }
1893 self.notify.notify_waiters();
1895 self.close_all().await;
1897 let deadline = Instant::now() + timeout;
1899 while self.total_count.load(Ordering::SeqCst) > 0 {
1900 if Instant::now() >= deadline {
1901 let remaining = self.total_count.load(Ordering::SeqCst);
1902 if remaining > 0 {
1903 eprintln!(
1904 "graceful shutdown timeout, {} connections force closed",
1905 remaining
1906 );
1907 }
1908 break;
1909 }
1910 tokio::time::sleep(Duration::from_millis(100)).await;
1911 }
1912 }
1913
1914 pub fn resize(&self, new_max: usize) {
1922 self.set_max_size(new_max as u32);
1923 }
1924
1925 pub fn set_max_size(&self, new_max: u32) {
1927 self.dynamic_max_size.store(new_max, Ordering::SeqCst);
1928 }
1929
1930 pub fn max_size(&self) -> u32 {
1932 self.dynamic_max_size.load(Ordering::Acquire)
1933 }
1934
1935 pub async fn warmup(&self, min_idle: usize) -> Result<(), PoolError> {
1939 for _ in 0..min_idle {
1940 let current_max = self.dynamic_max_size.load(Ordering::Acquire);
1941 let current = self.total_count.load(Ordering::Acquire);
1942 if current >= current_max {
1943 break;
1944 }
1945 match self.total_count.compare_exchange(
1947 current,
1948 current + 1,
1949 Ordering::SeqCst,
1950 Ordering::Acquire,
1951 ) {
1952 Ok(_) => {}
1953 Err(_) => continue, }
1955 match self.factory.create().await {
1956 Ok(conn) => {
1957 let now = Instant::now();
1958 let pooled = PooledConnection {
1959 conn,
1960 created_at: now,
1961 last_used_at: now,
1962 pool: None,
1963 };
1964 if let Err(rejected) = self.idle.push(pooled) {
1965 self.close_connection(rejected).await;
1967 self.total_count.fetch_sub(1, Ordering::SeqCst);
1968 }
1969 self.emit_event(PoolEvent::ConnectionCreated);
1970 }
1971 Err(_) => {
1972 self.total_count.fetch_sub(1, Ordering::SeqCst);
1974 break;
1975 }
1976 }
1977 }
1978 Ok(())
1979 }
1980
1981 pub async fn query_with_timeout(&self, sql: &str) -> Result<QueryRows, crate::DbError> {
1986 let timeout = self.config.query_timeout.unwrap_or(Duration::from_secs(30));
1987 let mut conn = self.acquire().await.map_err(crate::DbError::PoolError)?;
1988 tokio::time::timeout(timeout, conn.query(sql))
1989 .await
1990 .map_err(|_| crate::DbError::QueryError(format!("Query timeout after {:?}", timeout)))?
1991 }
1992}
1993
1994#[cfg(feature = "prod-pool-tuning")]
1999mod pool_prod {
2000 use super::PoolConfig;
2001 use serde::{Deserialize, Serialize};
2002 use std::time::Duration;
2003
2004 #[derive(Debug, Clone, PartialEq, Eq, thiserror::Error)]
2006 pub enum PoolProdError {
2007 #[error("pool max_size must be positive")]
2009 MaxSizeNotPositive,
2010 #[error("pool acquire_timeout must be positive")]
2012 AcquireTimeoutNotPositive,
2013 #[error("pool min_idle cannot exceed max_size")]
2015 MinIdleExceedsMaxSize,
2016 }
2017
2018 #[derive(Debug, Clone, Serialize, Deserialize)]
2020 pub struct PoolProdConfig {
2021 pub max_size: u32,
2023 pub acquire_timeout: Duration,
2025 pub idle_timeout: Duration,
2027 pub connection_timeout: Duration,
2029 pub query_timeout: Duration,
2031 pub min_idle: u32,
2033 pub prewarm: bool,
2035 }
2036
2037 impl Default for PoolProdConfig {
2038 fn default() -> Self {
2039 Self {
2040 max_size: 100,
2041 acquire_timeout: Duration::from_secs(30),
2042 idle_timeout: Duration::from_secs(600),
2043 connection_timeout: Duration::from_secs(10),
2044 query_timeout: Duration::from_secs(30),
2045 min_idle: 0,
2046 prewarm: false,
2047 }
2048 }
2049 }
2050
2051 impl PoolProdConfig {
2052 pub fn new(
2054 max_size: u32,
2055 acquire_timeout: Duration,
2056 idle_timeout: Duration,
2057 connection_timeout: Duration,
2058 query_timeout: Duration,
2059 min_idle: u32,
2060 prewarm: bool,
2061 ) -> Self {
2062 Self {
2063 max_size,
2064 acquire_timeout,
2065 idle_timeout,
2066 connection_timeout,
2067 query_timeout,
2068 min_idle,
2069 prewarm,
2070 }
2071 }
2072
2073 pub fn validate(&self) -> Result<(), PoolProdError> {
2075 if self.max_size == 0 {
2076 return Err(PoolProdError::MaxSizeNotPositive);
2077 }
2078 if self.acquire_timeout.is_zero() {
2079 return Err(PoolProdError::AcquireTimeoutNotPositive);
2080 }
2081 if self.min_idle > self.max_size {
2082 return Err(PoolProdError::MinIdleExceedsMaxSize);
2083 }
2084 Ok(())
2085 }
2086
2087 pub fn to_pool_config(&self) -> PoolConfig {
2089 PoolConfig {
2090 max_size: self.max_size,
2091 min_idle: self.min_idle,
2092 acquire_timeout: self.acquire_timeout,
2093 idle_timeout: self.idle_timeout,
2094 max_lifetime: Duration::from_secs(1800),
2095 connection_timeout: self.connection_timeout,
2096 tls: None,
2097 query_timeout: Some(self.query_timeout),
2098 max_rows: None,
2099 memory_limit: None,
2100 on_event: None,
2101 test_before_acquire: false,
2102 prewarm: self.prewarm,
2103 }
2104 }
2105 }
2106}
2107
2108#[cfg(feature = "prod-pool-tuning")]
2109pub use pool_prod::{PoolProdConfig, PoolProdError};
2110
2111#[cfg(feature = "prod-leak-detection")]
2116mod leak_detection {
2117 use serde::{Deserialize, Serialize};
2118 use std::time::Duration;
2119
2120 #[derive(Debug, Clone, Serialize, Deserialize)]
2122 pub struct LeakDetectionConfig {
2123 pub enabled: bool,
2125 pub interval: Duration,
2127 pub threshold: u32,
2129 pub borrow_timeout: Duration,
2131 }
2132
2133 impl Default for LeakDetectionConfig {
2134 fn default() -> Self {
2135 Self {
2136 enabled: false,
2137 interval: Duration::from_secs(60),
2138 threshold: 5,
2139 borrow_timeout: Duration::from_secs(60),
2140 }
2141 }
2142 }
2143
2144 impl LeakDetectionConfig {
2145 pub fn new(
2147 enabled: bool,
2148 interval: Duration,
2149 threshold: u32,
2150 borrow_timeout: Duration,
2151 ) -> Self {
2152 Self {
2153 enabled,
2154 interval,
2155 threshold,
2156 borrow_timeout,
2157 }
2158 }
2159
2160 pub fn validate(&self) -> Result<(), LeakDetectionError> {
2162 if self.interval.is_zero() {
2163 return Err(LeakDetectionError::IntervalNotPositive);
2164 }
2165 if self.borrow_timeout.is_zero() {
2166 return Err(LeakDetectionError::BorrowTimeoutNotPositive);
2167 }
2168 Ok(())
2169 }
2170 }
2171
2172 #[derive(Debug, Clone, PartialEq, Eq, thiserror::Error)]
2174 pub enum LeakDetectionError {
2175 #[error("leak detection interval must be positive")]
2177 IntervalNotPositive,
2178 #[error("leak detection borrow_timeout must be positive")]
2180 BorrowTimeoutNotPositive,
2181 }
2182
2183 #[derive(Debug, Clone, Serialize, Deserialize)]
2185 pub struct LeakEntry {
2186 pub conn_id: u64,
2188 pub borrowed_at: String,
2190 pub borrow_duration: Duration,
2192 }
2193
2194 #[derive(Debug, Clone, Serialize, Deserialize)]
2196 pub struct LeakReport {
2197 pub borrowed_count: u32,
2199 pub max_borrow_duration: Duration,
2201 pub suspected_leaks: Vec<LeakEntry>,
2203 }
2204
2205 impl LeakReport {
2206 pub fn empty() -> Self {
2208 Self {
2209 borrowed_count: 0,
2210 max_borrow_duration: Duration::ZERO,
2211 suspected_leaks: vec![],
2212 }
2213 }
2214 }
2215}
2216
2217#[cfg(feature = "prod-leak-detection")]
2218pub use leak_detection::{LeakDetectionConfig, LeakDetectionError, LeakEntry, LeakReport};
2219
2220#[cfg(test)]
2221mod tests {
2222 use super::*;
2223
2224 struct MockConnection {
2226 connected: bool,
2227 }
2228
2229 impl MockConnection {
2230 fn new() -> Self {
2231 Self { connected: true }
2232 }
2233 }
2234
2235 impl Connection for MockConnection {
2236 fn execute<'a>(
2237 &'a mut self,
2238 _sql: &'a str,
2239 ) -> Pin<Box<dyn Future<Output = Result<u64, crate::DbError>> + Send + 'a>> {
2240 Box::pin(async move { Ok(1) })
2241 }
2242
2243 fn query<'a>(
2244 &'a mut self,
2245 _sql: &'a str,
2246 ) -> Pin<
2247 Box<
2248 dyn Future<
2249 Output = Result<
2250 Vec<std::collections::HashMap<String, crate::value::Value>>,
2251 crate::DbError,
2252 >,
2253 > + Send
2254 + 'a,
2255 >,
2256 > {
2257 Box::pin(async move { Ok(vec![]) })
2258 }
2259
2260 fn begin_transaction<'a>(
2261 &'a mut self,
2262 ) -> Pin<Box<dyn Future<Output = Result<(), crate::DbError>> + Send + 'a>> {
2263 Box::pin(async move { Ok(()) })
2264 }
2265
2266 fn commit<'a>(
2267 &'a mut self,
2268 ) -> Pin<Box<dyn Future<Output = Result<(), crate::DbError>> + Send + 'a>> {
2269 Box::pin(async move { Ok(()) })
2270 }
2271
2272 fn rollback<'a>(
2273 &'a mut self,
2274 ) -> Pin<Box<dyn Future<Output = Result<(), crate::DbError>> + Send + 'a>> {
2275 Box::pin(async move { Ok(()) })
2276 }
2277
2278 fn is_connected(&self) -> bool {
2279 self.connected
2280 }
2281
2282 fn ping<'a>(&'a mut self) -> Pin<Box<dyn Future<Output = bool> + Send + 'a>> {
2283 Box::pin(async move { true })
2284 }
2285
2286 fn close<'a>(
2287 &'a mut self,
2288 ) -> Pin<Box<dyn Future<Output = Result<(), crate::DbError>> + Send + 'a>> {
2289 Box::pin(async move {
2290 self.connected = false;
2291 Ok(())
2292 })
2293 }
2294 }
2295
2296 struct MockConnectionFactory;
2297
2298 #[async_trait]
2299 impl ConnectionFactory for MockConnectionFactory {
2300 async fn create(&self) -> Result<Box<dyn Connection>, crate::DbError> {
2301 Ok(Box::new(MockConnection::new()))
2302 }
2303 }
2304
2305 #[tokio::test]
2306 async fn test_pool_config_builder() -> Result<(), Box<dyn std::error::Error>> {
2307 let config = PoolConfigBuilder::new().max_size(50).min_idle(10).build()?;
2308
2309 assert_eq!(config.max_size, 50);
2310 assert_eq!(config.min_idle, 10);
2311 Ok(())
2312 }
2313
2314 #[test]
2315 fn test_pool_status_display() {
2316 let status = PoolStatus {
2317 idle: 5,
2318 active: 10,
2319 max: 100,
2320 min: 5,
2321 waiters: 0,
2322 };
2323
2324 let display = format!("{:?}", status);
2325 assert!(display.contains("idle"));
2326 assert!(display.contains("active"));
2327 }
2328
2329 #[test]
2330 fn test_default_pool_config() {
2331 let config = PoolConfig::default();
2332 assert_eq!(config.max_size, 100);
2333 assert_eq!(config.min_idle, 0);
2334 assert_eq!(config.acquire_timeout.as_secs(), 30);
2335 assert_eq!(config.idle_timeout.as_secs(), 600);
2336 assert_eq!(config.max_lifetime.as_secs(), 1800);
2337 }
2338
2339 #[tokio::test]
2340 async fn test_pool_config_clone() {
2341 let config = PoolConfig::default();
2342 let cloned = config.clone();
2343 assert_eq!(cloned.max_size, config.max_size);
2344 assert_eq!(cloned.min_idle, config.min_idle);
2345 }
2346
2347 #[test]
2348 fn test_pool_config_builder_default() -> Result<(), Box<dyn std::error::Error>> {
2349 let builder = PoolConfigBuilder::new();
2350 let config = builder.build()?;
2351 assert_eq!(config.max_size, 100);
2352 Ok(())
2353 }
2354
2355 #[test]
2356 fn test_pool_config_validate() {
2357 let result = PoolConfigBuilder::new().max_size(0).build();
2358 assert!(result.is_err());
2359
2360 let result = PoolConfigBuilder::new().max_size(10).min_idle(20).build();
2361 assert!(result.is_err());
2362 }
2363
2364 #[test]
2365 fn test_pool_config_validate_duration_upper_bound() {
2366 use std::time::Duration;
2367
2368 let config = PoolConfig {
2370 max_size: 10,
2371 min_idle: 1,
2372 acquire_timeout: Duration::from_secs(u64::MAX),
2373 idle_timeout: Duration::from_secs(1),
2374 max_lifetime: Duration::from_secs(1),
2375 connection_timeout: Duration::from_secs(5),
2376 tls: None,
2377 query_timeout: None,
2378 max_rows: None,
2379 memory_limit: None,
2380 on_event: None,
2381 test_before_acquire: false,
2382 prewarm: false,
2383 };
2384 assert!(config.validate().is_err());
2385
2386 let config = PoolConfig {
2388 max_size: 10,
2389 min_idle: 1,
2390 acquire_timeout: Duration::from_secs(u32::MAX as u64),
2391 idle_timeout: Duration::from_secs(1),
2392 max_lifetime: Duration::from_secs(1),
2393 connection_timeout: Duration::from_secs(5),
2394 tls: None,
2395 query_timeout: None,
2396 max_rows: None,
2397 memory_limit: None,
2398 on_event: None,
2399 test_before_acquire: false,
2400 prewarm: false,
2401 };
2402 assert!(config.validate().is_ok());
2403
2404 let config = PoolConfig {
2406 max_size: 10,
2407 min_idle: 1,
2408 acquire_timeout: Duration::from_secs(u32::MAX as u64 + 1),
2409 idle_timeout: Duration::from_secs(1),
2410 max_lifetime: Duration::from_secs(1),
2411 connection_timeout: Duration::from_secs(5),
2412 tls: None,
2413 query_timeout: None,
2414 max_rows: None,
2415 memory_limit: None,
2416 on_event: None,
2417 test_before_acquire: false,
2418 prewarm: false,
2419 };
2420 assert!(config.validate().is_err());
2421 }
2422
2423 #[test]
2424 fn test_pool_config_test_before_acquire_default() {
2425 let config = PoolConfig::default();
2427 assert!(!config.test_before_acquire);
2428 }
2429
2430 #[test]
2431 fn test_pool_config_builder_test_before_acquire() {
2432 let config = PoolConfigBuilder::new()
2434 .test_before_acquire(true)
2435 .build()
2436 .unwrap();
2437 assert!(config.test_before_acquire);
2438 }
2439
2440 #[tokio::test]
2441 async fn test_pool_acquire_and_release() -> Result<(), Box<dyn std::error::Error>> {
2442 let config = PoolConfigBuilder::new().max_size(5).min_idle(1).build()?;
2443 let factory = Arc::new(MockConnectionFactory);
2444 let pool = Pool::new(config, factory)?;
2445
2446 let conn = pool.acquire().await?;
2447 let status = pool.status().await;
2448 assert_eq!(status.active, 1);
2449 assert_eq!(status.idle, 0);
2450
2451 pool.release(conn).await;
2452 let status = pool.status().await;
2453 assert_eq!(status.idle, 1);
2454
2455 let _conn2 = pool.acquire().await?;
2457 let status = pool.status().await;
2458 assert_eq!(status.idle, 0);
2459 Ok(())
2460 }
2461
2462 #[tokio::test]
2463 async fn test_pool_status() -> Result<(), Box<dyn std::error::Error>> {
2464 let config = PoolConfigBuilder::new().max_size(10).min_idle(2).build()?;
2465 let factory = Arc::new(MockConnectionFactory);
2466 let pool = Pool::new(config, factory)?;
2467
2468 let status = pool.status().await;
2469 assert_eq!(status.max, 10);
2470 assert_eq!(status.min, 2);
2471 assert_eq!(status.active, 0);
2472 Ok(())
2473 }
2474
2475 #[tokio::test]
2476 async fn test_pool_close_all() -> Result<(), Box<dyn std::error::Error>> {
2477 let config = PoolConfigBuilder::new().max_size(5).build()?;
2478 let factory = Arc::new(MockConnectionFactory);
2479 let pool = Pool::new(config, factory)?;
2480
2481 let conn1 = pool.acquire().await?;
2483 let conn2 = pool.acquire().await?;
2484 pool.release(conn1).await;
2485 pool.release(conn2).await;
2486
2487 pool.close_all().await;
2488 let status = pool.status().await;
2489 assert_eq!(status.idle, 0);
2490 assert_eq!(status.active, 0);
2491 Ok(())
2492 }
2493
2494 #[tokio::test]
2495 async fn test_pool_reap_idle() -> Result<(), Box<dyn std::error::Error>> {
2496 let config = PoolConfigBuilder::new()
2497 .max_size(5)
2498 .idle_timeout(0) .build()?;
2500 let factory = Arc::new(MockConnectionFactory);
2501 let pool = Pool::new(config, factory)?;
2502
2503 let conn = pool.acquire().await?;
2504 pool.release(conn).await;
2505
2506 tokio::time::sleep(tokio::time::Duration::from_millis(10)).await;
2508
2509 pool.reap_idle().await;
2510 let status = pool.status().await;
2511 assert_eq!(status.idle, 0);
2512 Ok(())
2513 }
2514
2515 #[tokio::test]
2521 async fn test_h7_acquire_timeout_default_30s() {
2522 let config = PoolConfig::default();
2523 assert_eq!(
2524 config.acquire_timeout,
2525 Duration::from_secs(30),
2526 "H-7: acquire_timeout 默认应为 30s"
2527 );
2528 }
2529
2530 #[tokio::test]
2532 async fn test_h7_acquire_timeout_configurable() -> Result<(), Box<dyn std::error::Error>> {
2533 let config = PoolConfigBuilder::new()
2534 .max_size(1)
2535 .acquire_timeout(5) .build()?;
2537 assert_eq!(config.acquire_timeout, Duration::from_secs(5));
2538
2539 let factory = Arc::new(MockConnectionFactory);
2541 let pool = Pool::new(config, factory)?;
2542 let _conn1 = pool.acquire().await?;
2543
2544 let fast_config = PoolConfigBuilder::new()
2546 .max_size(1)
2547 .acquire_timeout(0) .build()?;
2549 let fast_pool = Pool::new(fast_config, Arc::new(MockConnectionFactory))?;
2552 let _fast_conn = fast_pool.acquire().await?; let result = fast_pool.acquire().await;
2554 assert!(
2555 matches!(result, Err(PoolError::Timeout)),
2556 "H-7: 应返回 Timeout"
2557 );
2558 Ok(())
2559 }
2560
2561 #[tokio::test]
2564 async fn test_m7_health_check_removes_nothing_when_all_healthy(
2565 ) -> Result<(), Box<dyn std::error::Error>> {
2566 let config = PoolConfigBuilder::new().max_size(5).build()?;
2568 let factory = Arc::new(MockConnectionFactory);
2569 let pool = Pool::new(config, factory)?;
2570
2571 let conn1 = pool.acquire().await?;
2573 let conn2 = pool.acquire().await?;
2574 let conn3 = pool.acquire().await?;
2575 pool.release(conn1).await;
2576 pool.release(conn2).await;
2577 pool.release(conn3).await;
2578
2579 let removed = pool.health_check().await;
2580 assert_eq!(removed, 0, "Healthy connections should not be removed");
2581
2582 let status = pool.status().await;
2583 assert_eq!(status.idle, 3);
2584 assert_eq!(status.active, 3);
2585 Ok(())
2586 }
2587
2588 #[tokio::test]
2589 async fn test_m7_health_check_returns_zero_for_empty_pool(
2590 ) -> Result<(), Box<dyn std::error::Error>> {
2591 let config = PoolConfigBuilder::new().max_size(5).build()?;
2592 let factory = Arc::new(MockConnectionFactory);
2593 let pool = Pool::new(config, factory)?;
2594
2595 let removed = pool.health_check().await;
2596 assert_eq!(removed, 0);
2597 Ok(())
2598 }
2599
2600 struct CountingFactory {
2604 count: AtomicU32,
2605 }
2606
2607 impl CountingFactory {
2608 fn new() -> Self {
2609 Self {
2610 count: AtomicU32::new(0),
2611 }
2612 }
2613 fn created_count(&self) -> u32 {
2614 self.count.load(Ordering::SeqCst)
2615 }
2616 }
2617
2618 #[async_trait]
2619 impl ConnectionFactory for CountingFactory {
2620 async fn create(&self) -> Result<Box<dyn Connection>, crate::DbError> {
2621 self.count.fetch_add(1, Ordering::SeqCst);
2622 Ok(Box::new(MockConnection::new()))
2623 }
2624 }
2625
2626 #[tokio::test]
2632 async fn test_production_bug_max_lifetime_never_expires(
2633 ) -> Result<(), Box<dyn std::error::Error>> {
2634 let config = PoolConfig {
2637 max_size: 5,
2638 min_idle: 0,
2639 acquire_timeout: Duration::from_secs(30),
2640 idle_timeout: Duration::from_secs(600),
2641 max_lifetime: Duration::from_millis(100), connection_timeout: Duration::from_secs(10),
2643 tls: None,
2644 query_timeout: None,
2645 max_rows: None,
2646 memory_limit: None,
2647 on_event: None,
2648 test_before_acquire: false,
2649 prewarm: false,
2650 };
2651 let factory = Arc::new(CountingFactory::new());
2652 let pool = Pool::new(config, factory.clone())?;
2653
2654 let conn = pool.acquire().await?;
2656 assert_eq!(factory.created_count(), 1, "应创建 1 个连接");
2657
2658 pool.release(conn).await;
2660
2661 tokio::time::sleep(Duration::from_millis(150)).await;
2663
2664 let conn2 = pool.acquire().await?;
2666
2667 assert_eq!(
2670 factory.created_count(),
2671 2,
2672 "超过 max_lifetime 后应创建新连接(旧连接应被回收)"
2673 );
2674
2675 pool.release(conn2).await;
2676 Ok(())
2677 }
2678
2679 #[tokio::test]
2686 async fn test_drop_auto_release_connection() -> Result<(), Box<dyn std::error::Error>> {
2687 let config = PoolConfigBuilder::new().max_size(2).build()?;
2688 let factory = Arc::new(CountingFactory::new());
2689 let pool = Pool::new(config, factory.clone())?;
2690
2691 {
2693 let _conn = pool.acquire().await?;
2694 assert_eq!(factory.created_count(), 1, "应创建 1 个连接");
2695 let status = pool.status().await;
2696 assert_eq!(status.active, 1, "active 应为 1");
2697 assert_eq!(status.idle, 0, "idle 应为 0");
2698 }
2700
2701 tokio::time::sleep(Duration::from_millis(50)).await;
2703
2704 let status = pool.status().await;
2706 assert_eq!(status.idle, 1, "Drop 后连接应自动归还,idle 应为 1");
2707 assert_eq!(status.active, 1, "total_count 应为 1");
2708 assert_eq!(factory.created_count(), 1, "应复用归还的连接,不创建新连接");
2709 Ok(())
2710 }
2711
2712 #[tokio::test]
2714 async fn test_drop_auto_release_then_reuse() -> Result<(), Box<dyn std::error::Error>> {
2715 let config = PoolConfigBuilder::new().max_size(1).build()?;
2716 let factory = Arc::new(CountingFactory::new());
2717 let pool = Pool::new(config, factory.clone())?;
2718
2719 {
2721 let _conn = pool.acquire().await?;
2722 }
2723
2724 tokio::time::sleep(Duration::from_millis(50)).await;
2726
2727 let conn = pool.acquire().await?;
2729 assert_eq!(factory.created_count(), 1, "应复用归还的连接,不创建新连接");
2730
2731 pool.release(conn).await;
2732 Ok(())
2733 }
2734
2735 #[tokio::test]
2737 async fn test_into_inner_does_not_return_to_pool() -> Result<(), Box<dyn std::error::Error>> {
2738 let config = PoolConfigBuilder::new().max_size(2).build()?;
2739 let factory = Arc::new(CountingFactory::new());
2740 let pool = Pool::new(config, factory.clone())?;
2741
2742 let conn = pool.acquire().await?;
2743 assert_eq!(factory.created_count(), 1);
2744
2745 let _raw_conn = conn.into_inner();
2747
2748 tokio::time::sleep(Duration::from_millis(50)).await;
2750
2751 let status = pool.status().await;
2752 assert_eq!(status.idle, 0, "into_inner 后连接不应归还");
2753 assert_eq!(status.active, 1, "total_count 仍为 1(连接被外部持有)");
2754 Ok(())
2755 }
2756
2757 #[tokio::test]
2759 async fn test_explicit_release_no_double_return() -> Result<(), Box<dyn std::error::Error>> {
2760 let config = PoolConfigBuilder::new().max_size(2).build()?;
2761 let factory = Arc::new(CountingFactory::new());
2762 let pool = Pool::new(config, factory.clone())?;
2763
2764 let conn = pool.acquire().await?;
2765 pool.release(conn).await;
2766
2767 let status = pool.status().await;
2768 assert_eq!(status.idle, 1, "release 后 idle 应为 1");
2769
2770 let conn = pool.acquire().await?;
2772 pool.release(conn).await;
2773
2774 let status = pool.status().await;
2775 assert_eq!(status.idle, 1, "再次 release 后 idle 仍应为 1(不重复)");
2776 assert_eq!(status.active, 1, "total_count 应为 1");
2777 Ok(())
2778 }
2779
2780 struct CursorMockConn {
2786 rows: QueryRows,
2787 call_count: usize,
2788 }
2789
2790 impl CursorMockConn {
2791 fn new(rows: QueryRows) -> Self {
2792 Self {
2793 rows,
2794 call_count: 0,
2795 }
2796 }
2797 }
2798
2799 impl Connection for CursorMockConn {
2800 fn execute<'a>(
2801 &'a mut self,
2802 _sql: &'a str,
2803 ) -> Pin<Box<dyn Future<Output = Result<u64, crate::DbError>> + Send + 'a>> {
2804 Box::pin(async move { Ok(1) })
2805 }
2806
2807 fn query<'a>(
2808 &'a mut self,
2809 _sql: &'a str,
2810 ) -> Pin<Box<dyn Future<Output = Result<QueryRows, crate::DbError>> + Send + 'a>> {
2811 Box::pin(async move {
2812 self.call_count += 1;
2813 Ok(self.rows.clone())
2814 })
2815 }
2816
2817 fn begin_transaction<'a>(
2818 &'a mut self,
2819 ) -> Pin<Box<dyn Future<Output = Result<(), crate::DbError>> + Send + 'a>> {
2820 Box::pin(async move { Ok(()) })
2821 }
2822
2823 fn commit<'a>(
2824 &'a mut self,
2825 ) -> Pin<Box<dyn Future<Output = Result<(), crate::DbError>> + Send + 'a>> {
2826 Box::pin(async move { Ok(()) })
2827 }
2828
2829 fn rollback<'a>(
2830 &'a mut self,
2831 ) -> Pin<Box<dyn Future<Output = Result<(), crate::DbError>> + Send + 'a>> {
2832 Box::pin(async move { Ok(()) })
2833 }
2834
2835 fn is_connected(&self) -> bool {
2836 true
2837 }
2838
2839 fn ping<'a>(&'a mut self) -> Pin<Box<dyn Future<Output = bool> + Send + 'a>> {
2840 Box::pin(async move { true })
2841 }
2842
2843 fn close<'a>(
2844 &'a mut self,
2845 ) -> Pin<Box<dyn Future<Output = Result<(), crate::DbError>> + Send + 'a>> {
2846 Box::pin(async move { Ok(()) })
2847 }
2848 }
2849
2850 struct CursorOverrideMockConn {
2852 rows: Vec<crate::value::Value>,
2853 yielded: usize,
2854 }
2855
2856 impl CursorOverrideMockConn {
2857 fn new(rows: Vec<crate::value::Value>) -> Self {
2858 Self { rows, yielded: 0 }
2859 }
2860 }
2861
2862 impl Connection for CursorOverrideMockConn {
2863 fn execute<'a>(
2864 &'a mut self,
2865 _sql: &'a str,
2866 ) -> Pin<Box<dyn Future<Output = Result<u64, crate::DbError>> + Send + 'a>> {
2867 Box::pin(async move { Ok(1) })
2868 }
2869
2870 fn query<'a>(
2871 &'a mut self,
2872 _sql: &'a str,
2873 ) -> Pin<Box<dyn Future<Output = Result<QueryRows, crate::DbError>> + Send + 'a>> {
2874 Box::pin(async move {
2876 Ok(self
2877 .rows
2878 .iter()
2879 .map(|v| {
2880 let mut m = std::collections::HashMap::new();
2881 m.insert("v".to_string(), v.clone());
2882 m
2883 })
2884 .collect())
2885 })
2886 }
2887
2888 fn query_stream<'a>(
2890 &'a mut self,
2891 _sql: &'a str,
2892 ) -> Pin<Box<dyn futures::Stream<Item = QueryStreamItem> + Send + 'a>> {
2893 Box::pin(futures::stream::iter(
2894 self.rows
2895 .iter()
2896 .enumerate()
2897 .map(|(i, v)| {
2898 self.yielded = i + 1;
2899 let mut m = std::collections::HashMap::new();
2900 m.insert("v".to_string(), v.clone());
2901 Ok(m)
2902 })
2903 .collect::<Vec<_>>(),
2904 ))
2905 }
2906
2907 fn begin_transaction<'a>(
2908 &'a mut self,
2909 ) -> Pin<Box<dyn Future<Output = Result<(), crate::DbError>> + Send + 'a>> {
2910 Box::pin(async move { Ok(()) })
2911 }
2912
2913 fn commit<'a>(
2914 &'a mut self,
2915 ) -> Pin<Box<dyn Future<Output = Result<(), crate::DbError>> + Send + 'a>> {
2916 Box::pin(async move { Ok(()) })
2917 }
2918
2919 fn rollback<'a>(
2920 &'a mut self,
2921 ) -> Pin<Box<dyn Future<Output = Result<(), crate::DbError>> + Send + 'a>> {
2922 Box::pin(async move { Ok(()) })
2923 }
2924
2925 fn is_connected(&self) -> bool {
2926 true
2927 }
2928
2929 fn ping<'a>(&'a mut self) -> Pin<Box<dyn Future<Output = bool> + Send + 'a>> {
2930 Box::pin(async move { true })
2931 }
2932
2933 fn close<'a>(
2934 &'a mut self,
2935 ) -> Pin<Box<dyn Future<Output = Result<(), crate::DbError>> + Send + 'a>> {
2936 Box::pin(async move { Ok(()) })
2937 }
2938 }
2939
2940 #[tokio::test]
2942 async fn test_query_stream_default_impl_yields_all_rows() {
2943 use futures::StreamExt;
2944 let rows: QueryRows = vec![
2945 std::collections::HashMap::from([
2946 ("id".to_string(), crate::value::Value::I64(1)),
2947 (
2948 "name".to_string(),
2949 crate::value::Value::String("alice".to_string()),
2950 ),
2951 ]),
2952 std::collections::HashMap::from([
2953 ("id".to_string(), crate::value::Value::I64(2)),
2954 (
2955 "name".to_string(),
2956 crate::value::Value::String("bob".to_string()),
2957 ),
2958 ]),
2959 std::collections::HashMap::from([
2960 ("id".to_string(), crate::value::Value::I64(3)),
2961 (
2962 "name".to_string(),
2963 crate::value::Value::String("carol".to_string()),
2964 ),
2965 ]),
2966 ];
2967 let mut conn = CursorMockConn::new(rows);
2968 let mut stream = conn.query_stream("SELECT id, name FROM users");
2969 let mut received: Vec<QueryStreamItem> = Vec::new();
2970 while let Some(item) = stream.next().await {
2971 received.push(item);
2972 }
2973 assert_eq!(received.len(), 3, "应收到 3 行");
2974 assert!(received.iter().all(|r| r.is_ok()), "所有项应为 Ok");
2975 drop(stream);
2976 assert_eq!(conn.call_count, 1, "默认实现应调用 query() 一次");
2977 }
2978
2979 #[tokio::test]
2981 async fn test_query_stream_default_empty_result() {
2982 use futures::StreamExt;
2983 let mut conn = CursorMockConn::new(Vec::new());
2984 let mut stream = conn.query_stream("SELECT * FROM empty_table");
2985 let mut count = 0;
2986 while let Some(_item) = stream.next().await {
2987 count += 1;
2988 }
2989 assert_eq!(count, 0, "空结果集应产生 0 项");
2990 }
2991
2992 #[tokio::test]
2994 async fn test_query_stream_default_error_propagation() {
2995 use futures::StreamExt;
2996 struct ErrorMockConn;
2998 impl Connection for ErrorMockConn {
2999 fn execute<'a>(
3000 &'a mut self,
3001 _sql: &'a str,
3002 ) -> Pin<Box<dyn Future<Output = Result<u64, crate::DbError>> + Send + 'a>>
3003 {
3004 Box::pin(async move { Ok(1) })
3005 }
3006 fn query<'a>(
3007 &'a mut self,
3008 _sql: &'a str,
3009 ) -> Pin<Box<dyn Future<Output = Result<QueryRows, crate::DbError>> + Send + 'a>>
3010 {
3011 Box::pin(async move { Err(crate::DbError::Internal("query failed".to_string())) })
3012 }
3013 fn begin_transaction<'a>(
3014 &'a mut self,
3015 ) -> Pin<Box<dyn Future<Output = Result<(), crate::DbError>> + Send + 'a>> {
3016 Box::pin(async move { Ok(()) })
3017 }
3018 fn commit<'a>(
3019 &'a mut self,
3020 ) -> Pin<Box<dyn Future<Output = Result<(), crate::DbError>> + Send + 'a>> {
3021 Box::pin(async move { Ok(()) })
3022 }
3023 fn rollback<'a>(
3024 &'a mut self,
3025 ) -> Pin<Box<dyn Future<Output = Result<(), crate::DbError>> + Send + 'a>> {
3026 Box::pin(async move { Ok(()) })
3027 }
3028 fn is_connected(&self) -> bool {
3029 true
3030 }
3031 fn ping<'a>(&'a mut self) -> Pin<Box<dyn Future<Output = bool> + Send + 'a>> {
3032 Box::pin(async move { true })
3033 }
3034 fn close<'a>(
3035 &'a mut self,
3036 ) -> Pin<Box<dyn Future<Output = Result<(), crate::DbError>> + Send + 'a>> {
3037 Box::pin(async move { Ok(()) })
3038 }
3039 }
3040 let mut conn = ErrorMockConn;
3041 let mut stream = conn.query_stream("SELECT * FROM bad_table");
3042 let item = stream.next().await;
3043 assert!(item.is_some(), "应产生一项");
3044 assert!(item.unwrap().is_err(), "该项应为 Err");
3045 }
3046
3047 #[tokio::test]
3049 async fn test_query_stream_override_yields_rows_one_by_one() {
3050 use futures::StreamExt;
3051 let rows = vec![
3052 crate::value::Value::I64(10),
3053 crate::value::Value::I64(20),
3054 crate::value::Value::I64(30),
3055 crate::value::Value::I64(40),
3056 crate::value::Value::I64(50),
3057 ];
3058 let mut conn = CursorOverrideMockConn::new(rows);
3059 let values: Vec<i64> = {
3060 let mut stream = conn.query_stream("SELECT v FROM seq");
3061 let mut vals: Vec<i64> = Vec::new();
3062 while let Some(Ok(row)) = stream.next().await {
3063 if let crate::value::Value::I64(v) = row.get("v").unwrap() {
3064 vals.push(*v);
3065 }
3066 }
3067 vals
3068 };
3069 assert_eq!(values, vec![10, 20, 30, 40, 50], "应按顺序收到全部 5 行");
3070 assert_eq!(conn.yielded, 5, "应逐行 yield 5 次(真游标覆盖)");
3071 }
3072
3073 #[tokio::test]
3075 async fn test_query_stream_override_early_drop() {
3076 use futures::StreamExt;
3077 let rows = vec![
3078 crate::value::Value::I64(1),
3079 crate::value::Value::I64(2),
3080 crate::value::Value::I64(3),
3081 ];
3082 let mut conn = CursorOverrideMockConn::new(rows);
3083 {
3084 let mut stream = conn.query_stream("SELECT v FROM seq");
3085 let first = stream.next().await;
3086 assert!(first.is_some(), "第一项应存在");
3087 drop(stream);
3089 }
3090 assert!(conn.is_connected(), "提前 drop 流后连接仍应可用");
3092 }
3093
3094 #[tokio::test]
3096 async fn test_pool_prewarm() -> Result<(), Box<dyn std::error::Error>> {
3097 use std::sync::atomic::AtomicU32;
3098
3099 let create_count = Arc::new(AtomicU32::new(0));
3101 let create_count_clone = create_count.clone();
3102
3103 struct CountingFactory {
3104 count: Arc<AtomicU32>,
3105 }
3106
3107 #[async_trait]
3108 impl ConnectionFactory for CountingFactory {
3109 async fn create(&self) -> Result<Box<dyn Connection>, crate::DbError> {
3110 self.count.fetch_add(1, Ordering::SeqCst);
3111 Ok(Box::new(MockConnection::new()))
3112 }
3113 }
3114
3115 let config = PoolConfigBuilder::new()
3117 .max_size(10)
3118 .min_idle(5)
3119 .prewarm(true)
3120 .build()?;
3121
3122 let factory = Arc::new(CountingFactory {
3123 count: create_count_clone,
3124 });
3125
3126 let pool = Pool::new(config, factory)?;
3127
3128 let status_before = pool.status().await;
3130 assert_eq!(status_before.idle, 0, "预热前 idle 应为 0");
3131
3132 pool.prewarm().await;
3134
3135 let status_after = pool.status().await;
3137 assert!(
3138 status_after.idle >= 5,
3139 "预热后 idle 应 >= 5,实际: {}",
3140 status_after.idle
3141 );
3142
3143 assert_eq!(
3145 create_count.load(Ordering::SeqCst),
3146 5,
3147 "工厂应被调用 5 次(min_idle)"
3148 );
3149
3150 Ok(())
3151 }
3152
3153 #[tokio::test]
3155 async fn test_pool_prewarm_failure_non_blocking() -> Result<(), Box<dyn std::error::Error>> {
3156 use std::sync::atomic::AtomicBool;
3157
3158 struct FailingFactory {
3159 failed: Arc<AtomicBool>,
3160 }
3161
3162 #[async_trait]
3163 impl ConnectionFactory for FailingFactory {
3164 async fn create(&self) -> Result<Box<dyn Connection>, crate::DbError> {
3165 self.failed.store(true, Ordering::SeqCst);
3166 Err(crate::DbError::Internal(
3168 "simulated connection failure".to_string(),
3169 ))
3170 }
3171 }
3172
3173 let failed = Arc::new(AtomicBool::new(false));
3174 let mut config = PoolConfigBuilder::new()
3175 .max_size(10)
3176 .min_idle(3)
3177 .prewarm(true)
3178 .build()?;
3179 config.connection_timeout = std::time::Duration::from_secs(1); let factory = Arc::new(FailingFactory {
3182 failed: failed.clone(),
3183 });
3184
3185 let pool = Pool::new(config, factory)?;
3187 pool.prewarm().await; assert!(failed.load(Ordering::SeqCst), "工厂应被调用且失败");
3191
3192 let status = pool.status().await;
3194 assert_eq!(status.max, 10, "池配置应正常");
3195
3196 Ok(())
3197 }
3198
3199 #[tokio::test]
3201 async fn test_pool_prewarm_disabled() -> Result<(), Box<dyn std::error::Error>> {
3202 use std::sync::atomic::AtomicU32;
3203
3204 let create_count = Arc::new(AtomicU32::new(0));
3205 let create_count_clone = create_count.clone();
3206
3207 struct CountingFactory {
3208 count: Arc<AtomicU32>,
3209 }
3210
3211 #[async_trait]
3212 impl ConnectionFactory for CountingFactory {
3213 async fn create(&self) -> Result<Box<dyn Connection>, crate::DbError> {
3214 self.count.fetch_add(1, Ordering::SeqCst);
3215 Ok(Box::new(MockConnection::new()))
3216 }
3217 }
3218
3219 let config = PoolConfigBuilder::new()
3221 .max_size(10)
3222 .min_idle(5)
3223 .prewarm(false) .build()?;
3225
3226 let factory = Arc::new(CountingFactory {
3227 count: create_count_clone,
3228 });
3229
3230 let pool = Pool::new(config, factory)?;
3231 pool.prewarm().await; assert_eq!(
3235 create_count.load(Ordering::SeqCst),
3236 0,
3237 "prewarm=false 时工厂不应被调用"
3238 );
3239
3240 let status = pool.status().await;
3241 assert_eq!(status.idle, 0, "idle 应为 0");
3242
3243 Ok(())
3244 }
3245
3246 #[tokio::test]
3248 async fn test_pool_new_async_with_prewarm() -> Result<(), Box<dyn std::error::Error>> {
3249 use std::sync::atomic::AtomicU32;
3250
3251 let create_count = Arc::new(AtomicU32::new(0));
3252 let create_count_clone = create_count.clone();
3253
3254 struct CountingFactory {
3255 count: Arc<AtomicU32>,
3256 }
3257
3258 #[async_trait]
3259 impl ConnectionFactory for CountingFactory {
3260 async fn create(&self) -> Result<Box<dyn Connection>, crate::DbError> {
3261 self.count.fetch_add(1, Ordering::SeqCst);
3262 Ok(Box::new(MockConnection::new()))
3263 }
3264 }
3265
3266 let config = PoolConfigBuilder::new()
3267 .max_size(10)
3268 .min_idle(5)
3269 .prewarm(true)
3270 .build()?;
3271
3272 let factory = Arc::new(CountingFactory {
3273 count: create_count_clone,
3274 });
3275
3276 let pool = Pool::new_async(config, factory).await?;
3277
3278 let status = pool.status().await;
3279 assert!(
3280 status.idle >= 5,
3281 "new_async prewarm=true 后 idle 应 >= 5,实际: {}",
3282 status.idle
3283 );
3284 assert_eq!(create_count.load(Ordering::SeqCst), 5, "工厂应被调用 5 次");
3285
3286 Ok(())
3287 }
3288
3289 #[tokio::test]
3291 async fn test_pool_new_async_without_prewarm() -> Result<(), Box<dyn std::error::Error>> {
3292 use std::sync::atomic::AtomicU32;
3293
3294 let create_count = Arc::new(AtomicU32::new(0));
3295 let create_count_clone = create_count.clone();
3296
3297 struct CountingFactory {
3298 count: Arc<AtomicU32>,
3299 }
3300
3301 #[async_trait]
3302 impl ConnectionFactory for CountingFactory {
3303 async fn create(&self) -> Result<Box<dyn Connection>, crate::DbError> {
3304 self.count.fetch_add(1, Ordering::SeqCst);
3305 Ok(Box::new(MockConnection::new()))
3306 }
3307 }
3308
3309 let config = PoolConfigBuilder::new()
3310 .max_size(10)
3311 .min_idle(5)
3312 .prewarm(false)
3313 .build()?;
3314
3315 let factory = Arc::new(CountingFactory {
3316 count: create_count_clone,
3317 });
3318
3319 let pool = Pool::new_async(config, factory).await?;
3320
3321 let status = pool.status().await;
3322 assert_eq!(status.idle, 0, "prewarm=false 时 idle 应为 0");
3323 assert_eq!(create_count.load(Ordering::SeqCst), 0, "工厂不应被调用");
3324
3325 Ok(())
3326 }
3327
3328 #[tokio::test]
3330 async fn test_pool_new_async_failure_non_blocking() -> Result<(), Box<dyn std::error::Error>> {
3331 struct FailingFactory;
3332
3333 #[async_trait]
3334 impl ConnectionFactory for FailingFactory {
3335 async fn create(&self) -> Result<Box<dyn Connection>, crate::DbError> {
3336 Err(crate::DbError::Internal("simulated failure".to_string()))
3337 }
3338 }
3339
3340 let mut config = PoolConfigBuilder::new()
3341 .max_size(10)
3342 .min_idle(3)
3343 .prewarm(true)
3344 .build()?;
3345 config.connection_timeout = std::time::Duration::from_secs(1);
3346
3347 let pool = Pool::new_async(config, Arc::new(FailingFactory)).await?;
3348
3349 let status = pool.status().await;
3350 assert_eq!(status.max, 10, "池配置应正常");
3351
3352 Ok(())
3353 }
3354
3355 #[cfg(feature = "auto-prewarm")]
3357 #[tokio::test]
3358 async fn test_pool_progressive_prewarm() -> Result<(), Box<dyn std::error::Error>> {
3359 use std::sync::atomic::AtomicU32;
3360
3361 let create_count = Arc::new(AtomicU32::new(0));
3362 let create_count_clone = create_count.clone();
3363
3364 struct CountingFactory {
3365 count: Arc<AtomicU32>,
3366 }
3367
3368 #[async_trait]
3369 impl ConnectionFactory for CountingFactory {
3370 async fn create(&self) -> Result<Box<dyn Connection>, crate::DbError> {
3371 self.count.fetch_add(1, Ordering::SeqCst);
3372 Ok(Box::new(MockConnection::new()))
3373 }
3374 }
3375
3376 let config = PoolConfigBuilder::new()
3377 .max_size(20)
3378 .min_idle(6)
3379 .prewarm(true)
3380 .build()?;
3381
3382 let factory = Arc::new(CountingFactory {
3383 count: create_count_clone,
3384 });
3385
3386 let pool = Pool::new(config, factory)?;
3387
3388 let progress = crate::prewarm::PrewarmProgress::new(6);
3389 pool.progressive_prewarm(
3390 2,
3391 std::time::Duration::from_millis(5),
3392 std::time::Duration::from_secs(10),
3393 &progress,
3394 )
3395 .await;
3396
3397 let snap = progress.snapshot();
3398 assert!(
3399 snap.warmed >= 6,
3400 "progressive_prewarm 后 warmed 应 >= 6,实际: {}",
3401 snap.warmed
3402 );
3403 assert!(snap.is_completed, "应标记完成");
3404 assert_eq!(create_count.load(Ordering::SeqCst), 6, "工厂应被调用 6 次");
3405
3406 let status = pool.status().await;
3407 assert!(status.idle >= 6, "池中 idle 应 >= 6");
3408
3409 Ok(())
3410 }
3411
3412 #[cfg(feature = "auto-prewarm")]
3414 #[tokio::test]
3415 async fn test_pool_progressive_prewarm_timeout_zero() -> Result<(), Box<dyn std::error::Error>>
3416 {
3417 use std::sync::atomic::AtomicU32;
3418
3419 let create_count = Arc::new(AtomicU32::new(0));
3420 let create_count_clone = create_count.clone();
3421
3422 struct CountingFactory {
3423 count: Arc<AtomicU32>,
3424 }
3425
3426 #[async_trait]
3427 impl ConnectionFactory for CountingFactory {
3428 async fn create(&self) -> Result<Box<dyn Connection>, crate::DbError> {
3429 self.count.fetch_add(1, Ordering::SeqCst);
3430 Ok(Box::new(MockConnection::new()))
3431 }
3432 }
3433
3434 let config = PoolConfigBuilder::new()
3435 .max_size(20)
3436 .min_idle(10)
3437 .prewarm(true)
3438 .build()?;
3439
3440 let factory = Arc::new(CountingFactory {
3441 count: create_count_clone,
3442 });
3443
3444 let pool = Pool::new(config, factory)?;
3445
3446 let progress = crate::prewarm::PrewarmProgress::new(10);
3447 pool.progressive_prewarm(
3448 2,
3449 std::time::Duration::from_millis(5),
3450 std::time::Duration::ZERO,
3451 &progress,
3452 )
3453 .await;
3454
3455 let snap = progress.snapshot();
3456 assert!(snap.is_completed, "应标记完成");
3457 assert!(
3458 snap.warmed <= 2,
3459 "total_timeout=0 时最多建一批(batch_size=2),实际: {}",
3460 snap.warmed
3461 );
3462
3463 Ok(())
3464 }
3465
3466 #[cfg(feature = "auto-prewarm")]
3468 #[tokio::test]
3469 async fn test_pool_progressive_prewarm_disabled() -> Result<(), Box<dyn std::error::Error>> {
3470 use std::sync::atomic::AtomicU32;
3471
3472 let create_count = Arc::new(AtomicU32::new(0));
3473 let create_count_clone = create_count.clone();
3474
3475 struct CountingFactory {
3476 count: Arc<AtomicU32>,
3477 }
3478
3479 #[async_trait]
3480 impl ConnectionFactory for CountingFactory {
3481 async fn create(&self) -> Result<Box<dyn Connection>, crate::DbError> {
3482 self.count.fetch_add(1, Ordering::SeqCst);
3483 Ok(Box::new(MockConnection::new()))
3484 }
3485 }
3486
3487 let config = PoolConfigBuilder::new()
3488 .max_size(20)
3489 .min_idle(10)
3490 .prewarm(false)
3491 .build()?;
3492
3493 let factory = Arc::new(CountingFactory {
3494 count: create_count_clone,
3495 });
3496
3497 let pool = Pool::new(config, factory)?;
3498
3499 let progress = crate::prewarm::PrewarmProgress::new(10);
3500 pool.progressive_prewarm(
3501 2,
3502 std::time::Duration::from_millis(5),
3503 std::time::Duration::from_secs(10),
3504 &progress,
3505 )
3506 .await;
3507
3508 let snap = progress.snapshot();
3509 assert!(snap.is_completed, "应标记完成");
3510 assert_eq!(snap.warmed, 0, "prewarm=false 时不应建连");
3511 assert_eq!(create_count.load(Ordering::SeqCst), 0, "工厂不应被调用");
3512
3513 Ok(())
3514 }
3515
3516 #[cfg(feature = "auto-prewarm")]
3518 #[tokio::test]
3519 async fn test_pool_progressive_prewarm_failure_non_blocking(
3520 ) -> Result<(), Box<dyn std::error::Error>> {
3521 struct FailingFactory;
3522
3523 #[async_trait]
3524 impl ConnectionFactory for FailingFactory {
3525 async fn create(&self) -> Result<Box<dyn Connection>, crate::DbError> {
3526 Err(crate::DbError::Internal("simulated failure".to_string()))
3527 }
3528 }
3529
3530 let mut config = PoolConfigBuilder::new()
3531 .max_size(20)
3532 .min_idle(5)
3533 .prewarm(true)
3534 .build()?;
3535 config.connection_timeout = std::time::Duration::from_secs(1);
3536
3537 let pool = Pool::new(config, Arc::new(FailingFactory))?;
3538
3539 let progress = crate::prewarm::PrewarmProgress::new(5);
3540 pool.progressive_prewarm(
3541 2,
3542 std::time::Duration::from_millis(5),
3543 std::time::Duration::from_secs(5),
3544 &progress,
3545 )
3546 .await;
3547
3548 let snap = progress.snapshot();
3549 assert!(snap.is_completed, "应标记完成");
3550 assert_eq!(snap.warmed, 0, "全部失败时 warmed=0");
3551 assert!(snap.failed > 0, "应有失败记录");
3552
3553 Ok(())
3554 }
3555
3556 #[tokio::test]
3558 async fn test_pool_metrics_acquire_release() -> Result<(), Box<dyn std::error::Error>> {
3559 let config = PoolConfigBuilder::new().max_size(10).build()?;
3560 let pool = Pool::new(config, Arc::new(MockConnectionFactory))?;
3561
3562 let metrics = pool.pool_metrics();
3563 assert_eq!(metrics.acquire_count, 0);
3564 assert_eq!(metrics.release_count, 0);
3565 assert_eq!(metrics.connection_created_count, 0);
3566
3567 let conn = pool.acquire().await?;
3568 let metrics = pool.pool_metrics();
3569 assert_eq!(metrics.acquire_count, 1);
3570 assert_eq!(metrics.connection_created_count, 1);
3571 assert_eq!(metrics.acquire_failed_count, 0);
3572
3573 pool.release(conn).await;
3574 let metrics = pool.pool_metrics();
3575 assert_eq!(metrics.release_count, 1);
3576 assert_eq!(metrics.connection_closed_count, 0);
3578
3579 Ok(())
3580 }
3581
3582 #[tokio::test]
3584 async fn test_pool_metrics_acquire_failed() -> Result<(), Box<dyn std::error::Error>> {
3585 struct FailingFactory;
3586
3587 #[async_trait]
3588 impl ConnectionFactory for FailingFactory {
3589 async fn create(&self) -> Result<Box<dyn Connection>, crate::DbError> {
3590 Err(crate::DbError::Internal("simulated failure".to_string()))
3591 }
3592 }
3593
3594 let config = PoolConfigBuilder::new().max_size(10).build()?;
3595 let pool = Pool::new(config, Arc::new(FailingFactory))?;
3596
3597 let result = pool.acquire().await;
3598 assert!(result.is_err());
3599
3600 let metrics = pool.pool_metrics();
3601 assert_eq!(metrics.acquire_failed_count, 1);
3602 assert_eq!(metrics.acquire_count, 0);
3603
3604 Ok(())
3605 }
3606
3607 #[tokio::test]
3609 async fn test_pool_metrics_connection_closed() -> Result<(), Box<dyn std::error::Error>> {
3610 let config = PoolConfigBuilder::new().max_size(10).build()?;
3611 let pool = Pool::new(config, Arc::new(MockConnectionFactory))?;
3612
3613 let conn = pool.acquire().await?;
3614 pool.release(conn).await;
3615
3616 let status = pool.status().await;
3617 assert_eq!(status.idle, 1);
3618
3619 pool.close_all().await;
3620
3621 let metrics = pool.pool_metrics();
3622 assert_eq!(metrics.connection_closed_count, 1);
3623 assert_eq!(metrics.connection_created_count, 1);
3624
3625 Ok(())
3626 }
3627
3628 #[test]
3630 fn test_pool_metrics_average_wait_time() {
3631 let metrics = PoolMetrics {
3632 acquire_count: 4,
3633 acquire_failed_count: 1,
3634 acquire_wait_time: Duration::from_millis(200),
3635 release_count: 4,
3636 connection_created_count: 2,
3637 connection_closed_count: 0,
3638 };
3639 assert_eq!(
3640 metrics.average_acquire_wait_time(),
3641 Duration::from_millis(50)
3642 );
3643
3644 let empty = PoolMetrics::default();
3646 assert_eq!(empty.average_acquire_wait_time(), Duration::ZERO);
3647 }
3648
3649 #[tokio::test]
3650 async fn test_shutdown_with_timeout_fast_return_when_empty() {
3651 let factory = Arc::new(MockConnectionFactory) as Arc<dyn ConnectionFactory>;
3652 let config = PoolConfigBuilder::new().max_size(4).build().unwrap();
3653 let pool = Pool::new(config, factory).unwrap();
3654 pool.shutdown_with_timeout(Duration::from_secs(1)).await;
3655 assert!(pool.closed.load(Ordering::SeqCst));
3656 assert_eq!(pool.total_count.load(Ordering::SeqCst), 0);
3657 }
3658
3659 #[tokio::test]
3660 async fn test_shutdown_delegates_to_shutdown_with_timeout() {
3661 let factory = Arc::new(MockConnectionFactory) as Arc<dyn ConnectionFactory>;
3662 let config = PoolConfigBuilder::new().max_size(4).build().unwrap();
3663 let pool = Pool::new(config, factory).unwrap();
3664 pool.shutdown().await;
3665 assert!(pool.closed.load(Ordering::SeqCst));
3666 }
3667
3668 #[tokio::test]
3669 async fn test_shutdown_with_timeout_idempotent() {
3670 let factory = Arc::new(MockConnectionFactory) as Arc<dyn ConnectionFactory>;
3671 let config = PoolConfigBuilder::new().max_size(4).build().unwrap();
3672 let pool = Pool::new(config, factory).unwrap();
3673 pool.shutdown_with_timeout(Duration::from_secs(1)).await;
3674 let count_after_first = pool.total_count.load(Ordering::SeqCst);
3675 pool.shutdown_with_timeout(Duration::from_secs(1)).await;
3676 let count_after_second = pool.total_count.load(Ordering::SeqCst);
3677 assert_eq!(count_after_first, count_after_second);
3678 }
3679
3680 #[tokio::test]
3681 async fn test_shutdown_with_timeout_rejects_new_acquire() {
3682 let factory = Arc::new(MockConnectionFactory) as Arc<dyn ConnectionFactory>;
3683 let config = PoolConfigBuilder::new().max_size(4).build().unwrap();
3684 let pool = Pool::new(config, factory).unwrap();
3685 pool.shutdown_with_timeout(Duration::from_secs(1)).await;
3686 let result = pool.acquire().await;
3687 assert!(result.is_err());
3688 }
3689}
3690
3691#[cfg(all(test, feature = "prod-pool-tuning"))]
3692mod pool_prod_tests {
3693 use super::*;
3694
3695 struct MockFactory;
3696
3697 #[async_trait]
3698 impl ConnectionFactory for MockFactory {
3699 async fn create(&self) -> Result<Box<dyn Connection>, crate::DbError> {
3700 Ok(Box::new(MockConn))
3701 }
3702 }
3703
3704 struct MockConn;
3705
3706 impl Connection for MockConn {
3707 fn execute<'a>(
3708 &'a mut self,
3709 _sql: &'a str,
3710 ) -> Pin<Box<dyn Future<Output = Result<u64, crate::DbError>> + Send + 'a>> {
3711 Box::pin(async move { Ok(1) })
3712 }
3713 fn query<'a>(
3714 &'a mut self,
3715 _sql: &'a str,
3716 ) -> Pin<
3717 Box<
3718 dyn Future<
3719 Output = Result<
3720 Vec<std::collections::HashMap<String, crate::value::Value>>,
3721 crate::DbError,
3722 >,
3723 > + Send
3724 + 'a,
3725 >,
3726 > {
3727 Box::pin(async move { Ok(vec![]) })
3728 }
3729 fn begin_transaction<'a>(
3730 &'a mut self,
3731 ) -> Pin<Box<dyn Future<Output = Result<(), crate::DbError>> + Send + 'a>> {
3732 Box::pin(async move { Ok(()) })
3733 }
3734 fn commit<'a>(
3735 &'a mut self,
3736 ) -> Pin<Box<dyn Future<Output = Result<(), crate::DbError>> + Send + 'a>> {
3737 Box::pin(async move { Ok(()) })
3738 }
3739 fn rollback<'a>(
3740 &'a mut self,
3741 ) -> Pin<Box<dyn Future<Output = Result<(), crate::DbError>> + Send + 'a>> {
3742 Box::pin(async move { Ok(()) })
3743 }
3744 fn is_connected(&self) -> bool {
3745 true
3746 }
3747 fn ping<'a>(&'a mut self) -> Pin<Box<dyn Future<Output = bool> + Send + 'a>> {
3748 Box::pin(async move { true })
3749 }
3750 fn close<'a>(
3751 &'a mut self,
3752 ) -> Pin<Box<dyn Future<Output = Result<(), crate::DbError>> + Send + 'a>> {
3753 Box::pin(async move { Ok(()) })
3754 }
3755 }
3756
3757 #[test]
3758 fn test_pool_prod_config_validate_ok() {
3759 let config = PoolProdConfig::new(
3760 50,
3761 Duration::from_secs(10),
3762 Duration::from_secs(600),
3763 Duration::from_secs(5),
3764 Duration::from_secs(30),
3765 5,
3766 true,
3767 );
3768 assert!(config.validate().is_ok());
3769 }
3770
3771 #[test]
3772 fn test_pool_prod_config_max_size_zero_rejected() {
3773 let config = PoolProdConfig::default();
3774 let mut c = config;
3775 c.max_size = 0;
3776 let err = c.validate().unwrap_err();
3777 assert!(err.to_string().contains("max_size must be positive"));
3778 }
3779
3780 #[test]
3781 fn test_pool_prod_config_min_idle_exceeds_max_size() {
3782 let config = PoolProdConfig::new(
3783 10,
3784 Duration::from_secs(10),
3785 Duration::from_secs(600),
3786 Duration::from_secs(5),
3787 Duration::from_secs(30),
3788 20,
3789 false,
3790 );
3791 let err = config.validate().unwrap_err();
3792 assert!(err.to_string().contains("min_idle cannot exceed max_size"));
3793 }
3794
3795 #[test]
3796 fn test_pool_prod_config_to_pool_config() {
3797 let config = PoolProdConfig::new(
3798 50,
3799 Duration::from_secs(10),
3800 Duration::from_secs(600),
3801 Duration::from_secs(5),
3802 Duration::from_secs(30),
3803 5,
3804 true,
3805 );
3806 let pool_config = config.to_pool_config();
3807 assert_eq!(pool_config.max_size, 50);
3808 assert_eq!(pool_config.min_idle, 5);
3809 assert_eq!(pool_config.acquire_timeout, Duration::from_secs(10));
3810 assert!(pool_config.prewarm);
3811 }
3812
3813 #[tokio::test]
3814 async fn test_pool_prod_config_runtime_resize() {
3815 let factory = Arc::new(MockFactory) as Arc<dyn ConnectionFactory>;
3816 let config = PoolProdConfig::default();
3817 let pool = Pool::new(config.to_pool_config(), factory).unwrap();
3818 assert_eq!(pool.max_size(), 100);
3819 pool.resize(50);
3820 assert_eq!(pool.max_size(), 50);
3821 }
3822}
3823
3824#[cfg(all(test, feature = "prod-leak-detection"))]
3825mod leak_prod_tests {
3826 use super::*;
3827
3828 #[test]
3829 fn test_leak_config_default() {
3830 let config = LeakDetectionConfig::default();
3831 assert!(!config.enabled);
3832 assert_eq!(config.interval, Duration::from_secs(60));
3833 assert_eq!(config.threshold, 5);
3834 }
3835
3836 #[test]
3837 fn test_leak_config_validate_ok() {
3838 let config =
3839 LeakDetectionConfig::new(true, Duration::from_secs(30), 10, Duration::from_secs(60));
3840 assert!(config.validate().is_ok());
3841 }
3842
3843 #[test]
3844 fn test_leak_config_interval_zero_rejected() {
3845 let config = LeakDetectionConfig::new(true, Duration::ZERO, 10, Duration::from_secs(60));
3846 assert!(config.validate().is_err());
3847 }
3848
3849 #[test]
3850 fn test_leak_report_empty() {
3851 let report = LeakReport::empty();
3852 assert_eq!(report.borrowed_count, 0);
3853 assert!(report.suspected_leaks.is_empty());
3854 }
3855
3856 #[test]
3857 fn connection_reuse_rate_zero() {
3858 let metrics = PoolMetrics::default();
3859 assert_eq!(metrics.connection_reuse_rate(), 0.0);
3860 }
3861
3862 #[test]
3863 fn connection_reuse_rate_full() {
3864 let metrics = PoolMetrics {
3865 acquire_count: 100,
3866 connection_created_count: 1,
3867 ..Default::default()
3868 };
3869 let rate = metrics.connection_reuse_rate();
3870 assert!((rate - 0.99).abs() < 0.001, "复用率应接近 0.99,实际 {rate}");
3871 }
3872
3873 #[test]
3874 fn connection_reuse_rate_partial() {
3875 let metrics = PoolMetrics {
3876 acquire_count: 10,
3877 connection_created_count: 2,
3878 ..Default::default()
3879 };
3880 assert!((metrics.connection_reuse_rate() - 0.8).abs() < 0.001);
3881 }
3882
3883 #[test]
3884 fn pool_tuning_advice_is_optimal() {
3885 let advice = PoolTuningAdvice {
3886 suggested_max_size: None,
3887 suggested_min_idle: None,
3888 suggested_idle_timeout: None,
3889 reason: "池配置合理".to_string(),
3890 };
3891 assert!(advice.is_optimal());
3892
3893 let not_optimal = PoolTuningAdvice {
3894 suggested_max_size: Some(20),
3895 suggested_min_idle: None,
3896 suggested_idle_timeout: None,
3897 reason: "test".to_string(),
3898 };
3899 assert!(!not_optimal.is_optimal());
3900 }
3901
3902 #[test]
3903 fn suggest_tuning_low_reuse() {
3904 let metrics = PoolMetrics {
3905 acquire_count: 100,
3906 connection_created_count: 60,
3907 ..Default::default()
3908 };
3909 let reuse = metrics.connection_reuse_rate();
3910 assert!(reuse < 0.5, "复用率 {reuse} 应 < 0.5");
3911 }
3912
3913 #[test]
3914 fn suggest_tuning_optimal() {
3915 let metrics = PoolMetrics {
3916 acquire_count: 1000,
3917 connection_created_count: 10,
3918 acquire_wait_time: Duration::from_millis(10),
3919 ..Default::default()
3920 };
3921 let reuse = metrics.connection_reuse_rate();
3922 assert!(reuse >= 0.9, "复用率 {reuse} 应 >= 0.9");
3923 let avg_wait = metrics.average_acquire_wait_time();
3924 assert!(avg_wait <= Duration::from_millis(100));
3925 }
3926
3927 #[test]
3928 fn suggest_tuning_high_wait() {
3929 let metrics = PoolMetrics {
3930 acquire_count: 100,
3931 acquire_wait_time: Duration::from_millis(200 * 100),
3932 ..Default::default()
3933 };
3934 let avg_wait = metrics.average_acquire_wait_time();
3935 assert!(avg_wait > Duration::from_millis(100), "平均等待 {avg_wait:?} 应 > 100ms");
3936 }
3937}