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 mut deadline: Option<Instant> = None;
1420 let mut backoff = Duration::from_millis(1);
1422 const MAX_BACKOFF: Duration = Duration::from_millis(100);
1424 let mut to_close: Vec<PooledConnection> = Vec::with_capacity(4);
1426
1427 loop {
1428 let acquired: Option<PooledConnection> = {
1435 let mut found: Option<PooledConnection> = None;
1436 while let Some(pooled) = self.idle.pop() {
1437 if pooled.is_expired(self.config.max_lifetime) {
1439 to_close.push(pooled);
1440 continue;
1441 }
1442 if pooled.is_idle_too_long(self.config.idle_timeout) {
1444 to_close.push(pooled);
1445 continue;
1446 }
1447 if !pooled.conn.is_connected() {
1450 to_close.push(pooled);
1451 continue;
1452 }
1453 found = Some(pooled);
1454 break;
1455 }
1456 found
1457 };
1458
1459 for pooled in to_close.drain(..) {
1461 self.close_connection(pooled).await;
1462 self.total_count.fetch_sub(1, Ordering::SeqCst);
1464 }
1465
1466 if let Some(mut pooled) = acquired {
1467 if self.config.test_before_acquire {
1469 let ping_timeout = self.config.connection_timeout / 2;
1470 let alive = match tokio::time::timeout(ping_timeout, pooled.conn.ping()).await {
1471 Ok(true) => true,
1472 Ok(false) => false,
1473 Err(_) => false, };
1475 if !alive {
1476 self.close_connection(pooled).await;
1478 self.total_count.fetch_sub(1, Ordering::SeqCst);
1479 continue;
1480 }
1481 }
1482 pooled.pool = Some(self.clone());
1485 self.acquire_count.fetch_add(1, Ordering::Relaxed);
1486 return Ok(pooled);
1487 }
1488
1489 let current_max = self.dynamic_max_size.load(Ordering::Acquire);
1494 let created = loop {
1495 let current = self.total_count.load(Ordering::Acquire);
1496 if current >= current_max {
1497 break None; }
1499 match self.total_count.compare_exchange(
1500 current,
1501 current + 1,
1502 Ordering::SeqCst,
1503 Ordering::Acquire,
1504 ) {
1505 Ok(_) => break Some(()), Err(_) => continue, }
1508 };
1509
1510 if created.is_some() {
1511 match tokio::time::timeout(self.config.connection_timeout, self.factory.create())
1512 .await
1513 {
1514 Ok(Ok(conn)) => {
1515 #[cfg(feature = "circuit-breaker")]
1518 {
1519 self.circuit_breaker.lock().record_success();
1520 }
1521 self.emit_event(PoolEvent::ConnectionCreated);
1522 self.emit_event(PoolEvent::ConnectionAcquired);
1523 self.acquire_count.fetch_add(1, Ordering::Relaxed);
1524 return Ok(PooledConnection::new(conn, self.clone()));
1525 }
1526 Ok(Err(e)) => {
1527 self.total_count.fetch_sub(1, Ordering::SeqCst);
1529 #[cfg(feature = "circuit-breaker")]
1532 {
1533 self.circuit_breaker.lock().record_failure();
1534 }
1535 self.acquire_failed_count.fetch_add(1, Ordering::Relaxed);
1536 return Err(PoolError::ConnectionFailed(e.to_string()));
1537 }
1538 Err(_) => {
1539 self.total_count.fetch_sub(1, Ordering::SeqCst);
1541 #[cfg(feature = "circuit-breaker")]
1544 {
1545 self.circuit_breaker.lock().record_failure();
1546 }
1547 self.acquire_failed_count.fetch_add(1, Ordering::Relaxed);
1548 return Err(PoolError::Timeout);
1549 }
1550 }
1551 }
1552
1553 let now = Instant::now();
1555 let dl = deadline.get_or_insert_with(|| now + self.config.acquire_timeout);
1556 if now >= *dl {
1557 self.emit_event(PoolEvent::AcquireTimeout);
1558 self.acquire_failed_count.fetch_add(1, Ordering::Relaxed);
1559 return Err(PoolError::Timeout);
1560 }
1561 self.waiters_count.fetch_add(1, Ordering::SeqCst);
1563 let wait = std::cmp::min(backoff, *dl - now);
1564 match tokio::time::timeout(wait, self.notify.notified()).await {
1565 Ok(()) => {
1566 backoff = Duration::from_millis(1);
1568 }
1569 Err(_) => {
1570 backoff = std::cmp::min(backoff * 2, MAX_BACKOFF);
1572 }
1573 }
1574 self.waiters_count.fetch_sub(1, Ordering::SeqCst);
1576 self.acquire_wait_time_ns
1578 .fetch_add(wait.as_nanos() as u64, Ordering::Relaxed);
1579 }
1580 }
1581
1582 #[tracing::instrument(skip(self, pooled))]
1590 pub async fn release(&self, mut pooled: PooledConnection) {
1591 pooled.pool = None;
1593 self.release_count.fetch_add(1, Ordering::Relaxed);
1595
1596 if self.closed.load(Ordering::Acquire) {
1598 self.close_connection(pooled).await;
1599 self.total_count.fetch_sub(1, Ordering::SeqCst);
1601 self.emit_event(PoolEvent::ConnectionClosed);
1602 return;
1603 }
1604
1605 if !pooled.conn.is_connected() {
1607 self.close_connection(pooled).await;
1608 self.total_count.fetch_sub(1, Ordering::SeqCst);
1609 self.emit_event(PoolEvent::ConnectionClosed);
1610 return;
1611 }
1612
1613 pooled.last_used_at = Instant::now();
1615
1616 if let Err(rejected) = self.idle.push(pooled) {
1622 self.close_connection(rejected).await;
1624 self.total_count.fetch_sub(1, Ordering::SeqCst);
1625 self.emit_event(PoolEvent::ConnectionClosed);
1626 } else {
1627 self.emit_event(PoolEvent::ConnectionReleased);
1628 }
1629 self.notify.notify_one();
1630 }
1631
1632 pub async fn status(&self) -> PoolStatus {
1637 let idle_count = self.idle.len() as u32;
1638 let active = self.total_count.load(Ordering::Acquire);
1640 let waiters = self.waiters_count.load(Ordering::Acquire);
1641 PoolStatus {
1642 idle: idle_count,
1643 active,
1644 max: self.dynamic_max_size.load(Ordering::Acquire),
1645 min: self.config.min_idle,
1646 waiters,
1647 }
1648 }
1649
1650 pub fn pool_metrics(&self) -> PoolMetrics {
1661 PoolMetrics {
1662 acquire_count: self.acquire_count.load(Ordering::Acquire),
1663 acquire_failed_count: self.acquire_failed_count.load(Ordering::Acquire),
1664 acquire_wait_time: Duration::from_nanos(
1665 self.acquire_wait_time_ns.load(Ordering::Acquire),
1666 ),
1667 release_count: self.release_count.load(Ordering::Acquire),
1668 connection_created_count: self.connection_created_count.load(Ordering::Acquire),
1669 connection_closed_count: self.connection_closed_count.load(Ordering::Acquire),
1670 }
1671 }
1672
1673 #[must_use]
1682 pub fn suggest_tuning(&self) -> PoolTuningAdvice {
1683 let metrics = self.pool_metrics();
1684 let current_max = self.dynamic_max_size.load(Ordering::Acquire);
1685
1686 if metrics.acquire_count == 0 {
1687 return PoolTuningAdvice {
1688 suggested_max_size: None,
1689 suggested_min_idle: None,
1690 suggested_idle_timeout: None,
1691 reason: "数据不足".to_string(),
1692 };
1693 }
1694
1695 let reuse_rate = metrics.connection_reuse_rate();
1696 let mut advice = PoolTuningAdvice {
1697 suggested_max_size: None,
1698 suggested_min_idle: None,
1699 suggested_idle_timeout: None,
1700 reason: String::new(),
1701 };
1702
1703 if reuse_rate < 0.5 {
1704 advice.suggested_max_size = Some(current_max.saturating_mul(2));
1705 advice.reason = "复用率过低,池过小或回收过激".to_string();
1706 } else if reuse_rate < 0.9 {
1707 advice.suggested_min_idle = Some(current_max / 4);
1708 advice.reason = "复用率偏低,预热不足".to_string();
1709 }
1710
1711 let avg_wait = metrics.average_acquire_wait_time();
1712 if avg_wait > Duration::from_millis(100) {
1713 advice.suggested_max_size = Some(current_max.saturating_mul(2));
1714 if !advice.reason.is_empty() {
1715 advice.reason.push(';');
1716 }
1717 advice.reason.push_str("等待时长过高,池容量不足");
1718 }
1719
1720 if metrics.connection_created_count > 0
1721 && metrics.connection_closed_count as f64
1722 > metrics.connection_created_count as f64 * 0.5
1723 {
1724 advice.suggested_idle_timeout = Some(self.config.idle_timeout * 2);
1725 if !advice.reason.is_empty() {
1726 advice.reason.push(';');
1727 }
1728 advice.reason.push_str("连接关闭过快,空闲回收过激");
1729 }
1730
1731 if advice.reason.is_empty() {
1732 advice.reason = "池配置合理".to_string();
1733 }
1734
1735 advice
1736 }
1737
1738 pub fn metrics_snapshot_json(&self) -> String {
1743 serde_json::to_string(&self.pool_metrics()).unwrap_or_else(|_| "{}".to_string())
1744 }
1745
1746 #[tracing::instrument(skip(self))]
1748 pub async fn reap_idle(&self) {
1749 let mut all: Vec<PooledConnection> = Vec::new();
1753 while let Some(pooled) = self.idle.pop() {
1754 all.push(pooled);
1755 }
1756
1757 let mut to_close = Vec::new();
1759 for pooled in all {
1760 if pooled.is_idle_too_long(self.config.idle_timeout)
1761 || pooled.is_expired(self.config.max_lifetime)
1762 {
1763 to_close.push(pooled);
1764 } else {
1765 if let Err(rejected) = self.idle.push(pooled) {
1767 self.close_connection(rejected).await;
1768 self.total_count.fetch_sub(1, Ordering::SeqCst);
1769 }
1770 }
1771 }
1772
1773 for pooled in to_close {
1775 self.close_connection(pooled).await;
1776 self.total_count.fetch_sub(1, Ordering::SeqCst);
1778 }
1779 }
1780
1781 pub async fn close_all(&self) {
1785 self.closed.store(true, Ordering::Release);
1787 let mut to_close: Vec<PooledConnection> = Vec::new();
1790 while let Some(pooled) = self.idle.pop() {
1791 to_close.push(pooled);
1792 }
1793 let closed_count: u32 = to_close.len() as u32;
1795 for pooled in to_close {
1796 self.close_connection(pooled).await;
1797 }
1798 self.total_count.fetch_sub(closed_count, Ordering::SeqCst);
1801 }
1802
1803 pub async fn health_check(&self) -> u32 {
1819 let mut to_check: Vec<PooledConnection> = Vec::new();
1821 while let Some(pooled) = self.idle.pop() {
1822 to_check.push(pooled);
1823 }
1824
1825 let mut removed: u32 = 0;
1826 let mut alive: Vec<PooledConnection> = Vec::with_capacity(to_check.len());
1827 for mut pooled in to_check.drain(..) {
1828 if !pooled.conn.is_connected() {
1830 self.close_connection(pooled).await;
1831 removed += 1;
1832 continue;
1833 }
1834 let ping_timeout = self.config.connection_timeout / 2;
1836 match tokio::time::timeout(ping_timeout, pooled.conn.ping()).await {
1837 Ok(true) => alive.push(pooled),
1838 Ok(false) => {
1839 self.close_connection(pooled).await;
1841 removed += 1;
1842 }
1843 Err(_) => {
1844 self.close_connection(pooled).await;
1846 removed += 1;
1847 }
1848 }
1849 }
1850
1851 let alive_count: u32 = alive.len() as u32;
1853 for pooled in alive {
1854 if let Err(rejected) = self.idle.push(pooled) {
1856 self.close_connection(rejected).await;
1857 removed += 1;
1858 }
1859 }
1860
1861 if removed > 0 {
1863 self.total_count.fetch_sub(removed, Ordering::SeqCst);
1864 }
1865
1866 if alive_count > 0 {
1868 self.notify.notify_one();
1869 }
1870
1871 removed
1872 }
1873
1874 pub async fn shutdown(&self) {
1881 self.shutdown_with_timeout(Duration::from_secs(30)).await;
1882 }
1883
1884 pub async fn shutdown_with_timeout(&self, timeout: Duration) {
1892 if self.closed.swap(true, Ordering::SeqCst) {
1894 return;
1895 }
1896 self.notify.notify_waiters();
1898 self.close_all().await;
1900 let deadline = Instant::now() + timeout;
1902 while self.total_count.load(Ordering::SeqCst) > 0 {
1903 if Instant::now() >= deadline {
1904 let remaining = self.total_count.load(Ordering::SeqCst);
1905 if remaining > 0 {
1906 eprintln!(
1907 "graceful shutdown timeout, {} connections force closed",
1908 remaining
1909 );
1910 }
1911 break;
1912 }
1913 tokio::time::sleep(Duration::from_millis(100)).await;
1914 }
1915 }
1916
1917 pub fn resize(&self, new_max: usize) {
1925 self.set_max_size(new_max as u32);
1926 }
1927
1928 pub fn set_max_size(&self, new_max: u32) {
1930 self.dynamic_max_size.store(new_max, Ordering::SeqCst);
1931 }
1932
1933 pub fn max_size(&self) -> u32 {
1935 self.dynamic_max_size.load(Ordering::Acquire)
1936 }
1937
1938 pub async fn warmup(&self, min_idle: usize) -> Result<(), PoolError> {
1942 for _ in 0..min_idle {
1943 let current_max = self.dynamic_max_size.load(Ordering::Acquire);
1944 let current = self.total_count.load(Ordering::Acquire);
1945 if current >= current_max {
1946 break;
1947 }
1948 match self.total_count.compare_exchange(
1950 current,
1951 current + 1,
1952 Ordering::SeqCst,
1953 Ordering::Acquire,
1954 ) {
1955 Ok(_) => {}
1956 Err(_) => continue, }
1958 match self.factory.create().await {
1959 Ok(conn) => {
1960 let now = Instant::now();
1961 let pooled = PooledConnection {
1962 conn,
1963 created_at: now,
1964 last_used_at: now,
1965 pool: None,
1966 };
1967 if let Err(rejected) = self.idle.push(pooled) {
1968 self.close_connection(rejected).await;
1970 self.total_count.fetch_sub(1, Ordering::SeqCst);
1971 }
1972 self.emit_event(PoolEvent::ConnectionCreated);
1973 }
1974 Err(_) => {
1975 self.total_count.fetch_sub(1, Ordering::SeqCst);
1977 break;
1978 }
1979 }
1980 }
1981 Ok(())
1982 }
1983
1984 pub async fn query_with_timeout(&self, sql: &str) -> Result<QueryRows, crate::DbError> {
1989 let timeout = self.config.query_timeout.unwrap_or(Duration::from_secs(30));
1990 let mut conn = self.acquire().await.map_err(crate::DbError::PoolError)?;
1991 tokio::time::timeout(timeout, conn.query(sql))
1992 .await
1993 .map_err(|_| crate::DbError::QueryError(format!("Query timeout after {:?}", timeout)))?
1994 }
1995}
1996
1997#[cfg(feature = "prod-pool-tuning")]
2002mod pool_prod {
2003 use super::PoolConfig;
2004 use serde::{Deserialize, Serialize};
2005 use std::time::Duration;
2006
2007 #[derive(Debug, Clone, PartialEq, Eq, thiserror::Error)]
2009 pub enum PoolProdError {
2010 #[error("pool max_size must be positive")]
2012 MaxSizeNotPositive,
2013 #[error("pool acquire_timeout must be positive")]
2015 AcquireTimeoutNotPositive,
2016 #[error("pool min_idle cannot exceed max_size")]
2018 MinIdleExceedsMaxSize,
2019 }
2020
2021 #[derive(Debug, Clone, Serialize, Deserialize)]
2023 pub struct PoolProdConfig {
2024 pub max_size: u32,
2026 pub acquire_timeout: Duration,
2028 pub idle_timeout: Duration,
2030 pub connection_timeout: Duration,
2032 pub query_timeout: Duration,
2034 pub min_idle: u32,
2036 pub prewarm: bool,
2038 }
2039
2040 impl Default for PoolProdConfig {
2041 fn default() -> Self {
2042 Self {
2043 max_size: 100,
2044 acquire_timeout: Duration::from_secs(30),
2045 idle_timeout: Duration::from_secs(600),
2046 connection_timeout: Duration::from_secs(10),
2047 query_timeout: Duration::from_secs(30),
2048 min_idle: 0,
2049 prewarm: false,
2050 }
2051 }
2052 }
2053
2054 impl PoolProdConfig {
2055 pub fn new(
2057 max_size: u32,
2058 acquire_timeout: Duration,
2059 idle_timeout: Duration,
2060 connection_timeout: Duration,
2061 query_timeout: Duration,
2062 min_idle: u32,
2063 prewarm: bool,
2064 ) -> Self {
2065 Self {
2066 max_size,
2067 acquire_timeout,
2068 idle_timeout,
2069 connection_timeout,
2070 query_timeout,
2071 min_idle,
2072 prewarm,
2073 }
2074 }
2075
2076 pub fn validate(&self) -> Result<(), PoolProdError> {
2078 if self.max_size == 0 {
2079 return Err(PoolProdError::MaxSizeNotPositive);
2080 }
2081 if self.acquire_timeout.is_zero() {
2082 return Err(PoolProdError::AcquireTimeoutNotPositive);
2083 }
2084 if self.min_idle > self.max_size {
2085 return Err(PoolProdError::MinIdleExceedsMaxSize);
2086 }
2087 Ok(())
2088 }
2089
2090 pub fn to_pool_config(&self) -> PoolConfig {
2092 PoolConfig {
2093 max_size: self.max_size,
2094 min_idle: self.min_idle,
2095 acquire_timeout: self.acquire_timeout,
2096 idle_timeout: self.idle_timeout,
2097 max_lifetime: Duration::from_secs(1800),
2098 connection_timeout: self.connection_timeout,
2099 tls: None,
2100 query_timeout: Some(self.query_timeout),
2101 max_rows: None,
2102 memory_limit: None,
2103 on_event: None,
2104 test_before_acquire: false,
2105 prewarm: self.prewarm,
2106 }
2107 }
2108 }
2109}
2110
2111#[cfg(feature = "prod-pool-tuning")]
2112pub use pool_prod::{PoolProdConfig, PoolProdError};
2113
2114#[cfg(feature = "prod-leak-detection")]
2119mod leak_detection {
2120 use serde::{Deserialize, Serialize};
2121 use std::time::Duration;
2122
2123 #[derive(Debug, Clone, Serialize, Deserialize)]
2125 pub struct LeakDetectionConfig {
2126 pub enabled: bool,
2128 pub interval: Duration,
2130 pub threshold: u32,
2132 pub borrow_timeout: Duration,
2134 }
2135
2136 impl Default for LeakDetectionConfig {
2137 fn default() -> Self {
2138 Self {
2139 enabled: false,
2140 interval: Duration::from_secs(60),
2141 threshold: 5,
2142 borrow_timeout: Duration::from_secs(60),
2143 }
2144 }
2145 }
2146
2147 impl LeakDetectionConfig {
2148 pub fn new(
2150 enabled: bool,
2151 interval: Duration,
2152 threshold: u32,
2153 borrow_timeout: Duration,
2154 ) -> Self {
2155 Self {
2156 enabled,
2157 interval,
2158 threshold,
2159 borrow_timeout,
2160 }
2161 }
2162
2163 pub fn validate(&self) -> Result<(), LeakDetectionError> {
2165 if self.interval.is_zero() {
2166 return Err(LeakDetectionError::IntervalNotPositive);
2167 }
2168 if self.borrow_timeout.is_zero() {
2169 return Err(LeakDetectionError::BorrowTimeoutNotPositive);
2170 }
2171 Ok(())
2172 }
2173 }
2174
2175 #[derive(Debug, Clone, PartialEq, Eq, thiserror::Error)]
2177 pub enum LeakDetectionError {
2178 #[error("leak detection interval must be positive")]
2180 IntervalNotPositive,
2181 #[error("leak detection borrow_timeout must be positive")]
2183 BorrowTimeoutNotPositive,
2184 }
2185
2186 #[derive(Debug, Clone, Serialize, Deserialize)]
2188 pub struct LeakEntry {
2189 pub conn_id: u64,
2191 pub borrowed_at: String,
2193 pub borrow_duration: Duration,
2195 }
2196
2197 #[derive(Debug, Clone, Serialize, Deserialize)]
2199 pub struct LeakReport {
2200 pub borrowed_count: u32,
2202 pub max_borrow_duration: Duration,
2204 pub suspected_leaks: Vec<LeakEntry>,
2206 }
2207
2208 impl LeakReport {
2209 pub fn empty() -> Self {
2211 Self {
2212 borrowed_count: 0,
2213 max_borrow_duration: Duration::ZERO,
2214 suspected_leaks: vec![],
2215 }
2216 }
2217 }
2218}
2219
2220#[cfg(feature = "prod-leak-detection")]
2221pub use leak_detection::{LeakDetectionConfig, LeakDetectionError, LeakEntry, LeakReport};
2222
2223#[cfg(test)]
2224mod tests {
2225 use super::*;
2226
2227 struct MockConnection {
2229 connected: bool,
2230 }
2231
2232 impl MockConnection {
2233 fn new() -> Self {
2234 Self { connected: true }
2235 }
2236 }
2237
2238 impl Connection for MockConnection {
2239 fn execute<'a>(
2240 &'a mut self,
2241 _sql: &'a str,
2242 ) -> Pin<Box<dyn Future<Output = Result<u64, crate::DbError>> + Send + 'a>> {
2243 Box::pin(async move { Ok(1) })
2244 }
2245
2246 fn query<'a>(
2247 &'a mut self,
2248 _sql: &'a str,
2249 ) -> Pin<
2250 Box<
2251 dyn Future<
2252 Output = Result<
2253 Vec<std::collections::HashMap<String, crate::value::Value>>,
2254 crate::DbError,
2255 >,
2256 > + Send
2257 + 'a,
2258 >,
2259 > {
2260 Box::pin(async move { Ok(vec![]) })
2261 }
2262
2263 fn begin_transaction<'a>(
2264 &'a mut self,
2265 ) -> Pin<Box<dyn Future<Output = Result<(), crate::DbError>> + Send + 'a>> {
2266 Box::pin(async move { Ok(()) })
2267 }
2268
2269 fn commit<'a>(
2270 &'a mut self,
2271 ) -> Pin<Box<dyn Future<Output = Result<(), crate::DbError>> + Send + 'a>> {
2272 Box::pin(async move { Ok(()) })
2273 }
2274
2275 fn rollback<'a>(
2276 &'a mut self,
2277 ) -> Pin<Box<dyn Future<Output = Result<(), crate::DbError>> + Send + 'a>> {
2278 Box::pin(async move { Ok(()) })
2279 }
2280
2281 fn is_connected(&self) -> bool {
2282 self.connected
2283 }
2284
2285 fn ping<'a>(&'a mut self) -> Pin<Box<dyn Future<Output = bool> + Send + 'a>> {
2286 Box::pin(async move { true })
2287 }
2288
2289 fn close<'a>(
2290 &'a mut self,
2291 ) -> Pin<Box<dyn Future<Output = Result<(), crate::DbError>> + Send + 'a>> {
2292 Box::pin(async move {
2293 self.connected = false;
2294 Ok(())
2295 })
2296 }
2297 }
2298
2299 struct MockConnectionFactory;
2300
2301 #[async_trait]
2302 impl ConnectionFactory for MockConnectionFactory {
2303 async fn create(&self) -> Result<Box<dyn Connection>, crate::DbError> {
2304 Ok(Box::new(MockConnection::new()))
2305 }
2306 }
2307
2308 #[tokio::test]
2309 async fn test_pool_config_builder() -> Result<(), Box<dyn std::error::Error>> {
2310 let config = PoolConfigBuilder::new().max_size(50).min_idle(10).build()?;
2311
2312 assert_eq!(config.max_size, 50);
2313 assert_eq!(config.min_idle, 10);
2314 Ok(())
2315 }
2316
2317 #[test]
2318 fn test_pool_status_display() {
2319 let status = PoolStatus {
2320 idle: 5,
2321 active: 10,
2322 max: 100,
2323 min: 5,
2324 waiters: 0,
2325 };
2326
2327 let display = format!("{:?}", status);
2328 assert!(display.contains("idle"));
2329 assert!(display.contains("active"));
2330 }
2331
2332 #[test]
2333 fn test_default_pool_config() {
2334 let config = PoolConfig::default();
2335 assert_eq!(config.max_size, 100);
2336 assert_eq!(config.min_idle, 0);
2337 assert_eq!(config.acquire_timeout.as_secs(), 30);
2338 assert_eq!(config.idle_timeout.as_secs(), 600);
2339 assert_eq!(config.max_lifetime.as_secs(), 1800);
2340 }
2341
2342 #[tokio::test]
2343 async fn test_pool_config_clone() {
2344 let config = PoolConfig::default();
2345 let cloned = config.clone();
2346 assert_eq!(cloned.max_size, config.max_size);
2347 assert_eq!(cloned.min_idle, config.min_idle);
2348 }
2349
2350 #[test]
2351 fn test_pool_config_builder_default() -> Result<(), Box<dyn std::error::Error>> {
2352 let builder = PoolConfigBuilder::new();
2353 let config = builder.build()?;
2354 assert_eq!(config.max_size, 100);
2355 Ok(())
2356 }
2357
2358 #[test]
2359 fn test_pool_config_validate() {
2360 let result = PoolConfigBuilder::new().max_size(0).build();
2361 assert!(result.is_err());
2362
2363 let result = PoolConfigBuilder::new().max_size(10).min_idle(20).build();
2364 assert!(result.is_err());
2365 }
2366
2367 #[test]
2368 fn test_pool_config_validate_duration_upper_bound() {
2369 use std::time::Duration;
2370
2371 let config = PoolConfig {
2373 max_size: 10,
2374 min_idle: 1,
2375 acquire_timeout: Duration::from_secs(u64::MAX),
2376 idle_timeout: Duration::from_secs(1),
2377 max_lifetime: Duration::from_secs(1),
2378 connection_timeout: Duration::from_secs(5),
2379 tls: None,
2380 query_timeout: None,
2381 max_rows: None,
2382 memory_limit: None,
2383 on_event: None,
2384 test_before_acquire: false,
2385 prewarm: false,
2386 };
2387 assert!(config.validate().is_err());
2388
2389 let config = PoolConfig {
2391 max_size: 10,
2392 min_idle: 1,
2393 acquire_timeout: Duration::from_secs(u32::MAX as u64),
2394 idle_timeout: Duration::from_secs(1),
2395 max_lifetime: Duration::from_secs(1),
2396 connection_timeout: Duration::from_secs(5),
2397 tls: None,
2398 query_timeout: None,
2399 max_rows: None,
2400 memory_limit: None,
2401 on_event: None,
2402 test_before_acquire: false,
2403 prewarm: false,
2404 };
2405 assert!(config.validate().is_ok());
2406
2407 let config = PoolConfig {
2409 max_size: 10,
2410 min_idle: 1,
2411 acquire_timeout: Duration::from_secs(u32::MAX as u64 + 1),
2412 idle_timeout: Duration::from_secs(1),
2413 max_lifetime: Duration::from_secs(1),
2414 connection_timeout: Duration::from_secs(5),
2415 tls: None,
2416 query_timeout: None,
2417 max_rows: None,
2418 memory_limit: None,
2419 on_event: None,
2420 test_before_acquire: false,
2421 prewarm: false,
2422 };
2423 assert!(config.validate().is_err());
2424 }
2425
2426 #[test]
2427 fn test_pool_config_test_before_acquire_default() {
2428 let config = PoolConfig::default();
2430 assert!(!config.test_before_acquire);
2431 }
2432
2433 #[test]
2434 fn test_pool_config_builder_test_before_acquire() {
2435 let config = PoolConfigBuilder::new()
2437 .test_before_acquire(true)
2438 .build()
2439 .unwrap();
2440 assert!(config.test_before_acquire);
2441 }
2442
2443 #[tokio::test]
2444 async fn test_pool_acquire_and_release() -> Result<(), Box<dyn std::error::Error>> {
2445 let config = PoolConfigBuilder::new().max_size(5).min_idle(1).build()?;
2446 let factory = Arc::new(MockConnectionFactory);
2447 let pool = Pool::new(config, factory)?;
2448
2449 let conn = pool.acquire().await?;
2450 let status = pool.status().await;
2451 assert_eq!(status.active, 1);
2452 assert_eq!(status.idle, 0);
2453
2454 pool.release(conn).await;
2455 let status = pool.status().await;
2456 assert_eq!(status.idle, 1);
2457
2458 let _conn2 = pool.acquire().await?;
2460 let status = pool.status().await;
2461 assert_eq!(status.idle, 0);
2462 Ok(())
2463 }
2464
2465 #[tokio::test]
2466 async fn test_pool_status() -> Result<(), Box<dyn std::error::Error>> {
2467 let config = PoolConfigBuilder::new().max_size(10).min_idle(2).build()?;
2468 let factory = Arc::new(MockConnectionFactory);
2469 let pool = Pool::new(config, factory)?;
2470
2471 let status = pool.status().await;
2472 assert_eq!(status.max, 10);
2473 assert_eq!(status.min, 2);
2474 assert_eq!(status.active, 0);
2475 Ok(())
2476 }
2477
2478 #[tokio::test]
2479 async fn test_pool_close_all() -> Result<(), Box<dyn std::error::Error>> {
2480 let config = PoolConfigBuilder::new().max_size(5).build()?;
2481 let factory = Arc::new(MockConnectionFactory);
2482 let pool = Pool::new(config, factory)?;
2483
2484 let conn1 = pool.acquire().await?;
2486 let conn2 = pool.acquire().await?;
2487 pool.release(conn1).await;
2488 pool.release(conn2).await;
2489
2490 pool.close_all().await;
2491 let status = pool.status().await;
2492 assert_eq!(status.idle, 0);
2493 assert_eq!(status.active, 0);
2494 Ok(())
2495 }
2496
2497 #[tokio::test]
2498 async fn test_pool_reap_idle() -> Result<(), Box<dyn std::error::Error>> {
2499 let config = PoolConfigBuilder::new()
2500 .max_size(5)
2501 .idle_timeout(0) .build()?;
2503 let factory = Arc::new(MockConnectionFactory);
2504 let pool = Pool::new(config, factory)?;
2505
2506 let conn = pool.acquire().await?;
2507 pool.release(conn).await;
2508
2509 tokio::time::sleep(tokio::time::Duration::from_millis(10)).await;
2511
2512 pool.reap_idle().await;
2513 let status = pool.status().await;
2514 assert_eq!(status.idle, 0);
2515 Ok(())
2516 }
2517
2518 #[tokio::test]
2524 async fn test_h7_acquire_timeout_default_30s() {
2525 let config = PoolConfig::default();
2526 assert_eq!(
2527 config.acquire_timeout,
2528 Duration::from_secs(30),
2529 "H-7: acquire_timeout 默认应为 30s"
2530 );
2531 }
2532
2533 #[tokio::test]
2535 async fn test_h7_acquire_timeout_configurable() -> Result<(), Box<dyn std::error::Error>> {
2536 let config = PoolConfigBuilder::new()
2537 .max_size(1)
2538 .acquire_timeout(5) .build()?;
2540 assert_eq!(config.acquire_timeout, Duration::from_secs(5));
2541
2542 let factory = Arc::new(MockConnectionFactory);
2544 let pool = Pool::new(config, factory)?;
2545 let _conn1 = pool.acquire().await?;
2546
2547 let fast_config = PoolConfigBuilder::new()
2549 .max_size(1)
2550 .acquire_timeout(0) .build()?;
2552 let fast_pool = Pool::new(fast_config, Arc::new(MockConnectionFactory))?;
2555 let _fast_conn = fast_pool.acquire().await?; let result = fast_pool.acquire().await;
2557 assert!(
2558 matches!(result, Err(PoolError::Timeout)),
2559 "H-7: 应返回 Timeout"
2560 );
2561 Ok(())
2562 }
2563
2564 #[tokio::test]
2567 async fn test_m7_health_check_removes_nothing_when_all_healthy(
2568 ) -> Result<(), Box<dyn std::error::Error>> {
2569 let config = PoolConfigBuilder::new().max_size(5).build()?;
2571 let factory = Arc::new(MockConnectionFactory);
2572 let pool = Pool::new(config, factory)?;
2573
2574 let conn1 = pool.acquire().await?;
2576 let conn2 = pool.acquire().await?;
2577 let conn3 = pool.acquire().await?;
2578 pool.release(conn1).await;
2579 pool.release(conn2).await;
2580 pool.release(conn3).await;
2581
2582 let removed = pool.health_check().await;
2583 assert_eq!(removed, 0, "Healthy connections should not be removed");
2584
2585 let status = pool.status().await;
2586 assert_eq!(status.idle, 3);
2587 assert_eq!(status.active, 3);
2588 Ok(())
2589 }
2590
2591 #[tokio::test]
2592 async fn test_m7_health_check_returns_zero_for_empty_pool(
2593 ) -> Result<(), Box<dyn std::error::Error>> {
2594 let config = PoolConfigBuilder::new().max_size(5).build()?;
2595 let factory = Arc::new(MockConnectionFactory);
2596 let pool = Pool::new(config, factory)?;
2597
2598 let removed = pool.health_check().await;
2599 assert_eq!(removed, 0);
2600 Ok(())
2601 }
2602
2603 struct CountingFactory {
2607 count: AtomicU32,
2608 }
2609
2610 impl CountingFactory {
2611 fn new() -> Self {
2612 Self {
2613 count: AtomicU32::new(0),
2614 }
2615 }
2616 fn created_count(&self) -> u32 {
2617 self.count.load(Ordering::SeqCst)
2618 }
2619 }
2620
2621 #[async_trait]
2622 impl ConnectionFactory for CountingFactory {
2623 async fn create(&self) -> Result<Box<dyn Connection>, crate::DbError> {
2624 self.count.fetch_add(1, Ordering::SeqCst);
2625 Ok(Box::new(MockConnection::new()))
2626 }
2627 }
2628
2629 #[tokio::test]
2635 async fn test_production_bug_max_lifetime_never_expires(
2636 ) -> Result<(), Box<dyn std::error::Error>> {
2637 let config = PoolConfig {
2640 max_size: 5,
2641 min_idle: 0,
2642 acquire_timeout: Duration::from_secs(30),
2643 idle_timeout: Duration::from_secs(600),
2644 max_lifetime: Duration::from_millis(100), connection_timeout: Duration::from_secs(10),
2646 tls: None,
2647 query_timeout: None,
2648 max_rows: None,
2649 memory_limit: None,
2650 on_event: None,
2651 test_before_acquire: false,
2652 prewarm: false,
2653 };
2654 let factory = Arc::new(CountingFactory::new());
2655 let pool = Pool::new(config, factory.clone())?;
2656
2657 let conn = pool.acquire().await?;
2659 assert_eq!(factory.created_count(), 1, "应创建 1 个连接");
2660
2661 pool.release(conn).await;
2663
2664 tokio::time::sleep(Duration::from_millis(150)).await;
2666
2667 let conn2 = pool.acquire().await?;
2669
2670 assert_eq!(
2673 factory.created_count(),
2674 2,
2675 "超过 max_lifetime 后应创建新连接(旧连接应被回收)"
2676 );
2677
2678 pool.release(conn2).await;
2679 Ok(())
2680 }
2681
2682 #[tokio::test]
2689 async fn test_drop_auto_release_connection() -> Result<(), Box<dyn std::error::Error>> {
2690 let config = PoolConfigBuilder::new().max_size(2).build()?;
2691 let factory = Arc::new(CountingFactory::new());
2692 let pool = Pool::new(config, factory.clone())?;
2693
2694 {
2696 let _conn = pool.acquire().await?;
2697 assert_eq!(factory.created_count(), 1, "应创建 1 个连接");
2698 let status = pool.status().await;
2699 assert_eq!(status.active, 1, "active 应为 1");
2700 assert_eq!(status.idle, 0, "idle 应为 0");
2701 }
2703
2704 tokio::time::sleep(Duration::from_millis(50)).await;
2706
2707 let status = pool.status().await;
2709 assert_eq!(status.idle, 1, "Drop 后连接应自动归还,idle 应为 1");
2710 assert_eq!(status.active, 1, "total_count 应为 1");
2711 assert_eq!(factory.created_count(), 1, "应复用归还的连接,不创建新连接");
2712 Ok(())
2713 }
2714
2715 #[tokio::test]
2717 async fn test_drop_auto_release_then_reuse() -> Result<(), Box<dyn std::error::Error>> {
2718 let config = PoolConfigBuilder::new().max_size(1).build()?;
2719 let factory = Arc::new(CountingFactory::new());
2720 let pool = Pool::new(config, factory.clone())?;
2721
2722 {
2724 let _conn = pool.acquire().await?;
2725 }
2726
2727 tokio::time::sleep(Duration::from_millis(50)).await;
2729
2730 let conn = pool.acquire().await?;
2732 assert_eq!(factory.created_count(), 1, "应复用归还的连接,不创建新连接");
2733
2734 pool.release(conn).await;
2735 Ok(())
2736 }
2737
2738 #[tokio::test]
2740 async fn test_into_inner_does_not_return_to_pool() -> Result<(), Box<dyn std::error::Error>> {
2741 let config = PoolConfigBuilder::new().max_size(2).build()?;
2742 let factory = Arc::new(CountingFactory::new());
2743 let pool = Pool::new(config, factory.clone())?;
2744
2745 let conn = pool.acquire().await?;
2746 assert_eq!(factory.created_count(), 1);
2747
2748 let _raw_conn = conn.into_inner();
2750
2751 tokio::time::sleep(Duration::from_millis(50)).await;
2753
2754 let status = pool.status().await;
2755 assert_eq!(status.idle, 0, "into_inner 后连接不应归还");
2756 assert_eq!(status.active, 1, "total_count 仍为 1(连接被外部持有)");
2757 Ok(())
2758 }
2759
2760 #[tokio::test]
2762 async fn test_explicit_release_no_double_return() -> Result<(), Box<dyn std::error::Error>> {
2763 let config = PoolConfigBuilder::new().max_size(2).build()?;
2764 let factory = Arc::new(CountingFactory::new());
2765 let pool = Pool::new(config, factory.clone())?;
2766
2767 let conn = pool.acquire().await?;
2768 pool.release(conn).await;
2769
2770 let status = pool.status().await;
2771 assert_eq!(status.idle, 1, "release 后 idle 应为 1");
2772
2773 let conn = pool.acquire().await?;
2775 pool.release(conn).await;
2776
2777 let status = pool.status().await;
2778 assert_eq!(status.idle, 1, "再次 release 后 idle 仍应为 1(不重复)");
2779 assert_eq!(status.active, 1, "total_count 应为 1");
2780 Ok(())
2781 }
2782
2783 struct CursorMockConn {
2789 rows: QueryRows,
2790 call_count: usize,
2791 }
2792
2793 impl CursorMockConn {
2794 fn new(rows: QueryRows) -> Self {
2795 Self {
2796 rows,
2797 call_count: 0,
2798 }
2799 }
2800 }
2801
2802 impl Connection for CursorMockConn {
2803 fn execute<'a>(
2804 &'a mut self,
2805 _sql: &'a str,
2806 ) -> Pin<Box<dyn Future<Output = Result<u64, crate::DbError>> + Send + 'a>> {
2807 Box::pin(async move { Ok(1) })
2808 }
2809
2810 fn query<'a>(
2811 &'a mut self,
2812 _sql: &'a str,
2813 ) -> Pin<Box<dyn Future<Output = Result<QueryRows, crate::DbError>> + Send + 'a>> {
2814 Box::pin(async move {
2815 self.call_count += 1;
2816 Ok(self.rows.clone())
2817 })
2818 }
2819
2820 fn begin_transaction<'a>(
2821 &'a mut self,
2822 ) -> Pin<Box<dyn Future<Output = Result<(), crate::DbError>> + Send + 'a>> {
2823 Box::pin(async move { Ok(()) })
2824 }
2825
2826 fn commit<'a>(
2827 &'a mut self,
2828 ) -> Pin<Box<dyn Future<Output = Result<(), crate::DbError>> + Send + 'a>> {
2829 Box::pin(async move { Ok(()) })
2830 }
2831
2832 fn rollback<'a>(
2833 &'a mut self,
2834 ) -> Pin<Box<dyn Future<Output = Result<(), crate::DbError>> + Send + 'a>> {
2835 Box::pin(async move { Ok(()) })
2836 }
2837
2838 fn is_connected(&self) -> bool {
2839 true
2840 }
2841
2842 fn ping<'a>(&'a mut self) -> Pin<Box<dyn Future<Output = bool> + Send + 'a>> {
2843 Box::pin(async move { true })
2844 }
2845
2846 fn close<'a>(
2847 &'a mut self,
2848 ) -> Pin<Box<dyn Future<Output = Result<(), crate::DbError>> + Send + 'a>> {
2849 Box::pin(async move { Ok(()) })
2850 }
2851 }
2852
2853 struct CursorOverrideMockConn {
2855 rows: Vec<crate::value::Value>,
2856 yielded: usize,
2857 }
2858
2859 impl CursorOverrideMockConn {
2860 fn new(rows: Vec<crate::value::Value>) -> Self {
2861 Self { rows, yielded: 0 }
2862 }
2863 }
2864
2865 impl Connection for CursorOverrideMockConn {
2866 fn execute<'a>(
2867 &'a mut self,
2868 _sql: &'a str,
2869 ) -> Pin<Box<dyn Future<Output = Result<u64, crate::DbError>> + Send + 'a>> {
2870 Box::pin(async move { Ok(1) })
2871 }
2872
2873 fn query<'a>(
2874 &'a mut self,
2875 _sql: &'a str,
2876 ) -> Pin<Box<dyn Future<Output = Result<QueryRows, crate::DbError>> + Send + 'a>> {
2877 Box::pin(async move {
2879 Ok(self
2880 .rows
2881 .iter()
2882 .map(|v| {
2883 let mut m = std::collections::HashMap::new();
2884 m.insert("v".to_string(), v.clone());
2885 m
2886 })
2887 .collect())
2888 })
2889 }
2890
2891 fn query_stream<'a>(
2893 &'a mut self,
2894 _sql: &'a str,
2895 ) -> Pin<Box<dyn futures::Stream<Item = QueryStreamItem> + Send + 'a>> {
2896 Box::pin(futures::stream::iter(
2897 self.rows
2898 .iter()
2899 .enumerate()
2900 .map(|(i, v)| {
2901 self.yielded = i + 1;
2902 let mut m = std::collections::HashMap::new();
2903 m.insert("v".to_string(), v.clone());
2904 Ok(m)
2905 })
2906 .collect::<Vec<_>>(),
2907 ))
2908 }
2909
2910 fn begin_transaction<'a>(
2911 &'a mut self,
2912 ) -> Pin<Box<dyn Future<Output = Result<(), crate::DbError>> + Send + 'a>> {
2913 Box::pin(async move { Ok(()) })
2914 }
2915
2916 fn commit<'a>(
2917 &'a mut self,
2918 ) -> Pin<Box<dyn Future<Output = Result<(), crate::DbError>> + Send + 'a>> {
2919 Box::pin(async move { Ok(()) })
2920 }
2921
2922 fn rollback<'a>(
2923 &'a mut self,
2924 ) -> Pin<Box<dyn Future<Output = Result<(), crate::DbError>> + Send + 'a>> {
2925 Box::pin(async move { Ok(()) })
2926 }
2927
2928 fn is_connected(&self) -> bool {
2929 true
2930 }
2931
2932 fn ping<'a>(&'a mut self) -> Pin<Box<dyn Future<Output = bool> + Send + 'a>> {
2933 Box::pin(async move { true })
2934 }
2935
2936 fn close<'a>(
2937 &'a mut self,
2938 ) -> Pin<Box<dyn Future<Output = Result<(), crate::DbError>> + Send + 'a>> {
2939 Box::pin(async move { Ok(()) })
2940 }
2941 }
2942
2943 #[tokio::test]
2945 async fn test_query_stream_default_impl_yields_all_rows() {
2946 use futures::StreamExt;
2947 let rows: QueryRows = vec![
2948 std::collections::HashMap::from([
2949 ("id".to_string(), crate::value::Value::I64(1)),
2950 (
2951 "name".to_string(),
2952 crate::value::Value::String("alice".to_string()),
2953 ),
2954 ]),
2955 std::collections::HashMap::from([
2956 ("id".to_string(), crate::value::Value::I64(2)),
2957 (
2958 "name".to_string(),
2959 crate::value::Value::String("bob".to_string()),
2960 ),
2961 ]),
2962 std::collections::HashMap::from([
2963 ("id".to_string(), crate::value::Value::I64(3)),
2964 (
2965 "name".to_string(),
2966 crate::value::Value::String("carol".to_string()),
2967 ),
2968 ]),
2969 ];
2970 let mut conn = CursorMockConn::new(rows);
2971 let mut stream = conn.query_stream("SELECT id, name FROM users");
2972 let mut received: Vec<QueryStreamItem> = Vec::new();
2973 while let Some(item) = stream.next().await {
2974 received.push(item);
2975 }
2976 assert_eq!(received.len(), 3, "应收到 3 行");
2977 assert!(received.iter().all(|r| r.is_ok()), "所有项应为 Ok");
2978 drop(stream);
2979 assert_eq!(conn.call_count, 1, "默认实现应调用 query() 一次");
2980 }
2981
2982 #[tokio::test]
2984 async fn test_query_stream_default_empty_result() {
2985 use futures::StreamExt;
2986 let mut conn = CursorMockConn::new(Vec::new());
2987 let mut stream = conn.query_stream("SELECT * FROM empty_table");
2988 let mut count = 0;
2989 while let Some(_item) = stream.next().await {
2990 count += 1;
2991 }
2992 assert_eq!(count, 0, "空结果集应产生 0 项");
2993 }
2994
2995 #[tokio::test]
2997 async fn test_query_stream_default_error_propagation() {
2998 use futures::StreamExt;
2999 struct ErrorMockConn;
3001 impl Connection for ErrorMockConn {
3002 fn execute<'a>(
3003 &'a mut self,
3004 _sql: &'a str,
3005 ) -> Pin<Box<dyn Future<Output = Result<u64, crate::DbError>> + Send + 'a>>
3006 {
3007 Box::pin(async move { Ok(1) })
3008 }
3009 fn query<'a>(
3010 &'a mut self,
3011 _sql: &'a str,
3012 ) -> Pin<Box<dyn Future<Output = Result<QueryRows, crate::DbError>> + Send + 'a>>
3013 {
3014 Box::pin(async move { Err(crate::DbError::Internal("query failed".to_string())) })
3015 }
3016 fn begin_transaction<'a>(
3017 &'a mut self,
3018 ) -> Pin<Box<dyn Future<Output = Result<(), crate::DbError>> + Send + 'a>> {
3019 Box::pin(async move { Ok(()) })
3020 }
3021 fn commit<'a>(
3022 &'a mut self,
3023 ) -> Pin<Box<dyn Future<Output = Result<(), crate::DbError>> + Send + 'a>> {
3024 Box::pin(async move { Ok(()) })
3025 }
3026 fn rollback<'a>(
3027 &'a mut self,
3028 ) -> Pin<Box<dyn Future<Output = Result<(), crate::DbError>> + Send + 'a>> {
3029 Box::pin(async move { Ok(()) })
3030 }
3031 fn is_connected(&self) -> bool {
3032 true
3033 }
3034 fn ping<'a>(&'a mut self) -> Pin<Box<dyn Future<Output = bool> + Send + 'a>> {
3035 Box::pin(async move { true })
3036 }
3037 fn close<'a>(
3038 &'a mut self,
3039 ) -> Pin<Box<dyn Future<Output = Result<(), crate::DbError>> + Send + 'a>> {
3040 Box::pin(async move { Ok(()) })
3041 }
3042 }
3043 let mut conn = ErrorMockConn;
3044 let mut stream = conn.query_stream("SELECT * FROM bad_table");
3045 let item = stream.next().await;
3046 assert!(item.is_some(), "应产生一项");
3047 assert!(item.unwrap().is_err(), "该项应为 Err");
3048 }
3049
3050 #[tokio::test]
3052 async fn test_query_stream_override_yields_rows_one_by_one() {
3053 use futures::StreamExt;
3054 let rows = vec![
3055 crate::value::Value::I64(10),
3056 crate::value::Value::I64(20),
3057 crate::value::Value::I64(30),
3058 crate::value::Value::I64(40),
3059 crate::value::Value::I64(50),
3060 ];
3061 let mut conn = CursorOverrideMockConn::new(rows);
3062 let values: Vec<i64> = {
3063 let mut stream = conn.query_stream("SELECT v FROM seq");
3064 let mut vals: Vec<i64> = Vec::new();
3065 while let Some(Ok(row)) = stream.next().await {
3066 if let crate::value::Value::I64(v) = row.get("v").unwrap() {
3067 vals.push(*v);
3068 }
3069 }
3070 vals
3071 };
3072 assert_eq!(values, vec![10, 20, 30, 40, 50], "应按顺序收到全部 5 行");
3073 assert_eq!(conn.yielded, 5, "应逐行 yield 5 次(真游标覆盖)");
3074 }
3075
3076 #[tokio::test]
3078 async fn test_query_stream_override_early_drop() {
3079 use futures::StreamExt;
3080 let rows = vec![
3081 crate::value::Value::I64(1),
3082 crate::value::Value::I64(2),
3083 crate::value::Value::I64(3),
3084 ];
3085 let mut conn = CursorOverrideMockConn::new(rows);
3086 {
3087 let mut stream = conn.query_stream("SELECT v FROM seq");
3088 let first = stream.next().await;
3089 assert!(first.is_some(), "第一项应存在");
3090 drop(stream);
3092 }
3093 assert!(conn.is_connected(), "提前 drop 流后连接仍应可用");
3095 }
3096
3097 #[tokio::test]
3099 async fn test_pool_prewarm() -> Result<(), Box<dyn std::error::Error>> {
3100 use std::sync::atomic::AtomicU32;
3101
3102 let create_count = Arc::new(AtomicU32::new(0));
3104 let create_count_clone = create_count.clone();
3105
3106 struct CountingFactory {
3107 count: Arc<AtomicU32>,
3108 }
3109
3110 #[async_trait]
3111 impl ConnectionFactory for CountingFactory {
3112 async fn create(&self) -> Result<Box<dyn Connection>, crate::DbError> {
3113 self.count.fetch_add(1, Ordering::SeqCst);
3114 Ok(Box::new(MockConnection::new()))
3115 }
3116 }
3117
3118 let config = PoolConfigBuilder::new()
3120 .max_size(10)
3121 .min_idle(5)
3122 .prewarm(true)
3123 .build()?;
3124
3125 let factory = Arc::new(CountingFactory {
3126 count: create_count_clone,
3127 });
3128
3129 let pool = Pool::new(config, factory)?;
3130
3131 let status_before = pool.status().await;
3133 assert_eq!(status_before.idle, 0, "预热前 idle 应为 0");
3134
3135 pool.prewarm().await;
3137
3138 let status_after = pool.status().await;
3140 assert!(
3141 status_after.idle >= 5,
3142 "预热后 idle 应 >= 5,实际: {}",
3143 status_after.idle
3144 );
3145
3146 assert_eq!(
3148 create_count.load(Ordering::SeqCst),
3149 5,
3150 "工厂应被调用 5 次(min_idle)"
3151 );
3152
3153 Ok(())
3154 }
3155
3156 #[tokio::test]
3158 async fn test_pool_prewarm_failure_non_blocking() -> Result<(), Box<dyn std::error::Error>> {
3159 use std::sync::atomic::AtomicBool;
3160
3161 struct FailingFactory {
3162 failed: Arc<AtomicBool>,
3163 }
3164
3165 #[async_trait]
3166 impl ConnectionFactory for FailingFactory {
3167 async fn create(&self) -> Result<Box<dyn Connection>, crate::DbError> {
3168 self.failed.store(true, Ordering::SeqCst);
3169 Err(crate::DbError::Internal(
3171 "simulated connection failure".to_string(),
3172 ))
3173 }
3174 }
3175
3176 let failed = Arc::new(AtomicBool::new(false));
3177 let mut config = PoolConfigBuilder::new()
3178 .max_size(10)
3179 .min_idle(3)
3180 .prewarm(true)
3181 .build()?;
3182 config.connection_timeout = std::time::Duration::from_secs(1); let factory = Arc::new(FailingFactory {
3185 failed: failed.clone(),
3186 });
3187
3188 let pool = Pool::new(config, factory)?;
3190 pool.prewarm().await; assert!(failed.load(Ordering::SeqCst), "工厂应被调用且失败");
3194
3195 let status = pool.status().await;
3197 assert_eq!(status.max, 10, "池配置应正常");
3198
3199 Ok(())
3200 }
3201
3202 #[tokio::test]
3204 async fn test_pool_prewarm_disabled() -> Result<(), Box<dyn std::error::Error>> {
3205 use std::sync::atomic::AtomicU32;
3206
3207 let create_count = Arc::new(AtomicU32::new(0));
3208 let create_count_clone = create_count.clone();
3209
3210 struct CountingFactory {
3211 count: Arc<AtomicU32>,
3212 }
3213
3214 #[async_trait]
3215 impl ConnectionFactory for CountingFactory {
3216 async fn create(&self) -> Result<Box<dyn Connection>, crate::DbError> {
3217 self.count.fetch_add(1, Ordering::SeqCst);
3218 Ok(Box::new(MockConnection::new()))
3219 }
3220 }
3221
3222 let config = PoolConfigBuilder::new()
3224 .max_size(10)
3225 .min_idle(5)
3226 .prewarm(false) .build()?;
3228
3229 let factory = Arc::new(CountingFactory {
3230 count: create_count_clone,
3231 });
3232
3233 let pool = Pool::new(config, factory)?;
3234 pool.prewarm().await; assert_eq!(
3238 create_count.load(Ordering::SeqCst),
3239 0,
3240 "prewarm=false 时工厂不应被调用"
3241 );
3242
3243 let status = pool.status().await;
3244 assert_eq!(status.idle, 0, "idle 应为 0");
3245
3246 Ok(())
3247 }
3248
3249 #[tokio::test]
3251 async fn test_pool_new_async_with_prewarm() -> Result<(), Box<dyn std::error::Error>> {
3252 use std::sync::atomic::AtomicU32;
3253
3254 let create_count = Arc::new(AtomicU32::new(0));
3255 let create_count_clone = create_count.clone();
3256
3257 struct CountingFactory {
3258 count: Arc<AtomicU32>,
3259 }
3260
3261 #[async_trait]
3262 impl ConnectionFactory for CountingFactory {
3263 async fn create(&self) -> Result<Box<dyn Connection>, crate::DbError> {
3264 self.count.fetch_add(1, Ordering::SeqCst);
3265 Ok(Box::new(MockConnection::new()))
3266 }
3267 }
3268
3269 let config = PoolConfigBuilder::new()
3270 .max_size(10)
3271 .min_idle(5)
3272 .prewarm(true)
3273 .build()?;
3274
3275 let factory = Arc::new(CountingFactory {
3276 count: create_count_clone,
3277 });
3278
3279 let pool = Pool::new_async(config, factory).await?;
3280
3281 let status = pool.status().await;
3282 assert!(
3283 status.idle >= 5,
3284 "new_async prewarm=true 后 idle 应 >= 5,实际: {}",
3285 status.idle
3286 );
3287 assert_eq!(create_count.load(Ordering::SeqCst), 5, "工厂应被调用 5 次");
3288
3289 Ok(())
3290 }
3291
3292 #[tokio::test]
3294 async fn test_pool_new_async_without_prewarm() -> Result<(), Box<dyn std::error::Error>> {
3295 use std::sync::atomic::AtomicU32;
3296
3297 let create_count = Arc::new(AtomicU32::new(0));
3298 let create_count_clone = create_count.clone();
3299
3300 struct CountingFactory {
3301 count: Arc<AtomicU32>,
3302 }
3303
3304 #[async_trait]
3305 impl ConnectionFactory for CountingFactory {
3306 async fn create(&self) -> Result<Box<dyn Connection>, crate::DbError> {
3307 self.count.fetch_add(1, Ordering::SeqCst);
3308 Ok(Box::new(MockConnection::new()))
3309 }
3310 }
3311
3312 let config = PoolConfigBuilder::new()
3313 .max_size(10)
3314 .min_idle(5)
3315 .prewarm(false)
3316 .build()?;
3317
3318 let factory = Arc::new(CountingFactory {
3319 count: create_count_clone,
3320 });
3321
3322 let pool = Pool::new_async(config, factory).await?;
3323
3324 let status = pool.status().await;
3325 assert_eq!(status.idle, 0, "prewarm=false 时 idle 应为 0");
3326 assert_eq!(create_count.load(Ordering::SeqCst), 0, "工厂不应被调用");
3327
3328 Ok(())
3329 }
3330
3331 #[tokio::test]
3333 async fn test_pool_new_async_failure_non_blocking() -> Result<(), Box<dyn std::error::Error>> {
3334 struct FailingFactory;
3335
3336 #[async_trait]
3337 impl ConnectionFactory for FailingFactory {
3338 async fn create(&self) -> Result<Box<dyn Connection>, crate::DbError> {
3339 Err(crate::DbError::Internal("simulated failure".to_string()))
3340 }
3341 }
3342
3343 let mut config = PoolConfigBuilder::new()
3344 .max_size(10)
3345 .min_idle(3)
3346 .prewarm(true)
3347 .build()?;
3348 config.connection_timeout = std::time::Duration::from_secs(1);
3349
3350 let pool = Pool::new_async(config, Arc::new(FailingFactory)).await?;
3351
3352 let status = pool.status().await;
3353 assert_eq!(status.max, 10, "池配置应正常");
3354
3355 Ok(())
3356 }
3357
3358 #[cfg(feature = "auto-prewarm")]
3360 #[tokio::test]
3361 async fn test_pool_progressive_prewarm() -> Result<(), Box<dyn std::error::Error>> {
3362 use std::sync::atomic::AtomicU32;
3363
3364 let create_count = Arc::new(AtomicU32::new(0));
3365 let create_count_clone = create_count.clone();
3366
3367 struct CountingFactory {
3368 count: Arc<AtomicU32>,
3369 }
3370
3371 #[async_trait]
3372 impl ConnectionFactory for CountingFactory {
3373 async fn create(&self) -> Result<Box<dyn Connection>, crate::DbError> {
3374 self.count.fetch_add(1, Ordering::SeqCst);
3375 Ok(Box::new(MockConnection::new()))
3376 }
3377 }
3378
3379 let config = PoolConfigBuilder::new()
3380 .max_size(20)
3381 .min_idle(6)
3382 .prewarm(true)
3383 .build()?;
3384
3385 let factory = Arc::new(CountingFactory {
3386 count: create_count_clone,
3387 });
3388
3389 let pool = Pool::new(config, factory)?;
3390
3391 let progress = crate::prewarm::PrewarmProgress::new(6);
3392 pool.progressive_prewarm(
3393 2,
3394 std::time::Duration::from_millis(5),
3395 std::time::Duration::from_secs(10),
3396 &progress,
3397 )
3398 .await;
3399
3400 let snap = progress.snapshot();
3401 assert!(
3402 snap.warmed >= 6,
3403 "progressive_prewarm 后 warmed 应 >= 6,实际: {}",
3404 snap.warmed
3405 );
3406 assert!(snap.is_completed, "应标记完成");
3407 assert_eq!(create_count.load(Ordering::SeqCst), 6, "工厂应被调用 6 次");
3408
3409 let status = pool.status().await;
3410 assert!(status.idle >= 6, "池中 idle 应 >= 6");
3411
3412 Ok(())
3413 }
3414
3415 #[cfg(feature = "auto-prewarm")]
3417 #[tokio::test]
3418 async fn test_pool_progressive_prewarm_timeout_zero() -> Result<(), Box<dyn std::error::Error>>
3419 {
3420 use std::sync::atomic::AtomicU32;
3421
3422 let create_count = Arc::new(AtomicU32::new(0));
3423 let create_count_clone = create_count.clone();
3424
3425 struct CountingFactory {
3426 count: Arc<AtomicU32>,
3427 }
3428
3429 #[async_trait]
3430 impl ConnectionFactory for CountingFactory {
3431 async fn create(&self) -> Result<Box<dyn Connection>, crate::DbError> {
3432 self.count.fetch_add(1, Ordering::SeqCst);
3433 Ok(Box::new(MockConnection::new()))
3434 }
3435 }
3436
3437 let config = PoolConfigBuilder::new()
3438 .max_size(20)
3439 .min_idle(10)
3440 .prewarm(true)
3441 .build()?;
3442
3443 let factory = Arc::new(CountingFactory {
3444 count: create_count_clone,
3445 });
3446
3447 let pool = Pool::new(config, factory)?;
3448
3449 let progress = crate::prewarm::PrewarmProgress::new(10);
3450 pool.progressive_prewarm(
3451 2,
3452 std::time::Duration::from_millis(5),
3453 std::time::Duration::ZERO,
3454 &progress,
3455 )
3456 .await;
3457
3458 let snap = progress.snapshot();
3459 assert!(snap.is_completed, "应标记完成");
3460 assert!(
3461 snap.warmed <= 2,
3462 "total_timeout=0 时最多建一批(batch_size=2),实际: {}",
3463 snap.warmed
3464 );
3465
3466 Ok(())
3467 }
3468
3469 #[cfg(feature = "auto-prewarm")]
3471 #[tokio::test]
3472 async fn test_pool_progressive_prewarm_disabled() -> Result<(), Box<dyn std::error::Error>> {
3473 use std::sync::atomic::AtomicU32;
3474
3475 let create_count = Arc::new(AtomicU32::new(0));
3476 let create_count_clone = create_count.clone();
3477
3478 struct CountingFactory {
3479 count: Arc<AtomicU32>,
3480 }
3481
3482 #[async_trait]
3483 impl ConnectionFactory for CountingFactory {
3484 async fn create(&self) -> Result<Box<dyn Connection>, crate::DbError> {
3485 self.count.fetch_add(1, Ordering::SeqCst);
3486 Ok(Box::new(MockConnection::new()))
3487 }
3488 }
3489
3490 let config = PoolConfigBuilder::new()
3491 .max_size(20)
3492 .min_idle(10)
3493 .prewarm(false)
3494 .build()?;
3495
3496 let factory = Arc::new(CountingFactory {
3497 count: create_count_clone,
3498 });
3499
3500 let pool = Pool::new(config, factory)?;
3501
3502 let progress = crate::prewarm::PrewarmProgress::new(10);
3503 pool.progressive_prewarm(
3504 2,
3505 std::time::Duration::from_millis(5),
3506 std::time::Duration::from_secs(10),
3507 &progress,
3508 )
3509 .await;
3510
3511 let snap = progress.snapshot();
3512 assert!(snap.is_completed, "应标记完成");
3513 assert_eq!(snap.warmed, 0, "prewarm=false 时不应建连");
3514 assert_eq!(create_count.load(Ordering::SeqCst), 0, "工厂不应被调用");
3515
3516 Ok(())
3517 }
3518
3519 #[cfg(feature = "auto-prewarm")]
3521 #[tokio::test]
3522 async fn test_pool_progressive_prewarm_failure_non_blocking(
3523 ) -> Result<(), Box<dyn std::error::Error>> {
3524 struct FailingFactory;
3525
3526 #[async_trait]
3527 impl ConnectionFactory for FailingFactory {
3528 async fn create(&self) -> Result<Box<dyn Connection>, crate::DbError> {
3529 Err(crate::DbError::Internal("simulated failure".to_string()))
3530 }
3531 }
3532
3533 let mut config = PoolConfigBuilder::new()
3534 .max_size(20)
3535 .min_idle(5)
3536 .prewarm(true)
3537 .build()?;
3538 config.connection_timeout = std::time::Duration::from_secs(1);
3539
3540 let pool = Pool::new(config, Arc::new(FailingFactory))?;
3541
3542 let progress = crate::prewarm::PrewarmProgress::new(5);
3543 pool.progressive_prewarm(
3544 2,
3545 std::time::Duration::from_millis(5),
3546 std::time::Duration::from_secs(5),
3547 &progress,
3548 )
3549 .await;
3550
3551 let snap = progress.snapshot();
3552 assert!(snap.is_completed, "应标记完成");
3553 assert_eq!(snap.warmed, 0, "全部失败时 warmed=0");
3554 assert!(snap.failed > 0, "应有失败记录");
3555
3556 Ok(())
3557 }
3558
3559 #[tokio::test]
3561 async fn test_pool_metrics_acquire_release() -> Result<(), Box<dyn std::error::Error>> {
3562 let config = PoolConfigBuilder::new().max_size(10).build()?;
3563 let pool = Pool::new(config, Arc::new(MockConnectionFactory))?;
3564
3565 let metrics = pool.pool_metrics();
3566 assert_eq!(metrics.acquire_count, 0);
3567 assert_eq!(metrics.release_count, 0);
3568 assert_eq!(metrics.connection_created_count, 0);
3569
3570 let conn = pool.acquire().await?;
3571 let metrics = pool.pool_metrics();
3572 assert_eq!(metrics.acquire_count, 1);
3573 assert_eq!(metrics.connection_created_count, 1);
3574 assert_eq!(metrics.acquire_failed_count, 0);
3575
3576 pool.release(conn).await;
3577 let metrics = pool.pool_metrics();
3578 assert_eq!(metrics.release_count, 1);
3579 assert_eq!(metrics.connection_closed_count, 0);
3581
3582 Ok(())
3583 }
3584
3585 #[tokio::test]
3587 async fn test_pool_metrics_acquire_failed() -> Result<(), Box<dyn std::error::Error>> {
3588 struct FailingFactory;
3589
3590 #[async_trait]
3591 impl ConnectionFactory for FailingFactory {
3592 async fn create(&self) -> Result<Box<dyn Connection>, crate::DbError> {
3593 Err(crate::DbError::Internal("simulated failure".to_string()))
3594 }
3595 }
3596
3597 let config = PoolConfigBuilder::new().max_size(10).build()?;
3598 let pool = Pool::new(config, Arc::new(FailingFactory))?;
3599
3600 let result = pool.acquire().await;
3601 assert!(result.is_err());
3602
3603 let metrics = pool.pool_metrics();
3604 assert_eq!(metrics.acquire_failed_count, 1);
3605 assert_eq!(metrics.acquire_count, 0);
3606
3607 Ok(())
3608 }
3609
3610 #[tokio::test]
3612 async fn test_pool_metrics_connection_closed() -> Result<(), Box<dyn std::error::Error>> {
3613 let config = PoolConfigBuilder::new().max_size(10).build()?;
3614 let pool = Pool::new(config, Arc::new(MockConnectionFactory))?;
3615
3616 let conn = pool.acquire().await?;
3617 pool.release(conn).await;
3618
3619 let status = pool.status().await;
3620 assert_eq!(status.idle, 1);
3621
3622 pool.close_all().await;
3623
3624 let metrics = pool.pool_metrics();
3625 assert_eq!(metrics.connection_closed_count, 1);
3626 assert_eq!(metrics.connection_created_count, 1);
3627
3628 Ok(())
3629 }
3630
3631 #[test]
3633 fn test_pool_metrics_average_wait_time() {
3634 let metrics = PoolMetrics {
3635 acquire_count: 4,
3636 acquire_failed_count: 1,
3637 acquire_wait_time: Duration::from_millis(200),
3638 release_count: 4,
3639 connection_created_count: 2,
3640 connection_closed_count: 0,
3641 };
3642 assert_eq!(
3643 metrics.average_acquire_wait_time(),
3644 Duration::from_millis(50)
3645 );
3646
3647 let empty = PoolMetrics::default();
3649 assert_eq!(empty.average_acquire_wait_time(), Duration::ZERO);
3650 }
3651
3652 #[tokio::test]
3653 async fn test_shutdown_with_timeout_fast_return_when_empty() {
3654 let factory = Arc::new(MockConnectionFactory) as Arc<dyn ConnectionFactory>;
3655 let config = PoolConfigBuilder::new().max_size(4).build().unwrap();
3656 let pool = Pool::new(config, factory).unwrap();
3657 pool.shutdown_with_timeout(Duration::from_secs(1)).await;
3658 assert!(pool.closed.load(Ordering::SeqCst));
3659 assert_eq!(pool.total_count.load(Ordering::SeqCst), 0);
3660 }
3661
3662 #[tokio::test]
3663 async fn test_shutdown_delegates_to_shutdown_with_timeout() {
3664 let factory = Arc::new(MockConnectionFactory) as Arc<dyn ConnectionFactory>;
3665 let config = PoolConfigBuilder::new().max_size(4).build().unwrap();
3666 let pool = Pool::new(config, factory).unwrap();
3667 pool.shutdown().await;
3668 assert!(pool.closed.load(Ordering::SeqCst));
3669 }
3670
3671 #[tokio::test]
3672 async fn test_shutdown_with_timeout_idempotent() {
3673 let factory = Arc::new(MockConnectionFactory) as Arc<dyn ConnectionFactory>;
3674 let config = PoolConfigBuilder::new().max_size(4).build().unwrap();
3675 let pool = Pool::new(config, factory).unwrap();
3676 pool.shutdown_with_timeout(Duration::from_secs(1)).await;
3677 let count_after_first = pool.total_count.load(Ordering::SeqCst);
3678 pool.shutdown_with_timeout(Duration::from_secs(1)).await;
3679 let count_after_second = pool.total_count.load(Ordering::SeqCst);
3680 assert_eq!(count_after_first, count_after_second);
3681 }
3682
3683 #[tokio::test]
3684 async fn test_shutdown_with_timeout_rejects_new_acquire() {
3685 let factory = Arc::new(MockConnectionFactory) as Arc<dyn ConnectionFactory>;
3686 let config = PoolConfigBuilder::new().max_size(4).build().unwrap();
3687 let pool = Pool::new(config, factory).unwrap();
3688 pool.shutdown_with_timeout(Duration::from_secs(1)).await;
3689 let result = pool.acquire().await;
3690 assert!(result.is_err());
3691 }
3692}
3693
3694#[cfg(all(test, feature = "prod-pool-tuning"))]
3695mod pool_prod_tests {
3696 use super::*;
3697
3698 struct MockFactory;
3699
3700 #[async_trait]
3701 impl ConnectionFactory for MockFactory {
3702 async fn create(&self) -> Result<Box<dyn Connection>, crate::DbError> {
3703 Ok(Box::new(MockConn))
3704 }
3705 }
3706
3707 struct MockConn;
3708
3709 impl Connection for MockConn {
3710 fn execute<'a>(
3711 &'a mut self,
3712 _sql: &'a str,
3713 ) -> Pin<Box<dyn Future<Output = Result<u64, crate::DbError>> + Send + 'a>> {
3714 Box::pin(async move { Ok(1) })
3715 }
3716 fn query<'a>(
3717 &'a mut self,
3718 _sql: &'a str,
3719 ) -> Pin<
3720 Box<
3721 dyn Future<
3722 Output = Result<
3723 Vec<std::collections::HashMap<String, crate::value::Value>>,
3724 crate::DbError,
3725 >,
3726 > + Send
3727 + 'a,
3728 >,
3729 > {
3730 Box::pin(async move { Ok(vec![]) })
3731 }
3732 fn begin_transaction<'a>(
3733 &'a mut self,
3734 ) -> Pin<Box<dyn Future<Output = Result<(), crate::DbError>> + Send + 'a>> {
3735 Box::pin(async move { Ok(()) })
3736 }
3737 fn commit<'a>(
3738 &'a mut self,
3739 ) -> Pin<Box<dyn Future<Output = Result<(), crate::DbError>> + Send + 'a>> {
3740 Box::pin(async move { Ok(()) })
3741 }
3742 fn rollback<'a>(
3743 &'a mut self,
3744 ) -> Pin<Box<dyn Future<Output = Result<(), crate::DbError>> + Send + 'a>> {
3745 Box::pin(async move { Ok(()) })
3746 }
3747 fn is_connected(&self) -> bool {
3748 true
3749 }
3750 fn ping<'a>(&'a mut self) -> Pin<Box<dyn Future<Output = bool> + Send + 'a>> {
3751 Box::pin(async move { true })
3752 }
3753 fn close<'a>(
3754 &'a mut self,
3755 ) -> Pin<Box<dyn Future<Output = Result<(), crate::DbError>> + Send + 'a>> {
3756 Box::pin(async move { Ok(()) })
3757 }
3758 }
3759
3760 #[test]
3761 fn test_pool_prod_config_validate_ok() {
3762 let config = PoolProdConfig::new(
3763 50,
3764 Duration::from_secs(10),
3765 Duration::from_secs(600),
3766 Duration::from_secs(5),
3767 Duration::from_secs(30),
3768 5,
3769 true,
3770 );
3771 assert!(config.validate().is_ok());
3772 }
3773
3774 #[test]
3775 fn test_pool_prod_config_max_size_zero_rejected() {
3776 let config = PoolProdConfig::default();
3777 let mut c = config;
3778 c.max_size = 0;
3779 let err = c.validate().unwrap_err();
3780 assert!(err.to_string().contains("max_size must be positive"));
3781 }
3782
3783 #[test]
3784 fn test_pool_prod_config_min_idle_exceeds_max_size() {
3785 let config = PoolProdConfig::new(
3786 10,
3787 Duration::from_secs(10),
3788 Duration::from_secs(600),
3789 Duration::from_secs(5),
3790 Duration::from_secs(30),
3791 20,
3792 false,
3793 );
3794 let err = config.validate().unwrap_err();
3795 assert!(err.to_string().contains("min_idle cannot exceed max_size"));
3796 }
3797
3798 #[test]
3799 fn test_pool_prod_config_to_pool_config() {
3800 let config = PoolProdConfig::new(
3801 50,
3802 Duration::from_secs(10),
3803 Duration::from_secs(600),
3804 Duration::from_secs(5),
3805 Duration::from_secs(30),
3806 5,
3807 true,
3808 );
3809 let pool_config = config.to_pool_config();
3810 assert_eq!(pool_config.max_size, 50);
3811 assert_eq!(pool_config.min_idle, 5);
3812 assert_eq!(pool_config.acquire_timeout, Duration::from_secs(10));
3813 assert!(pool_config.prewarm);
3814 }
3815
3816 #[tokio::test]
3817 async fn test_pool_prod_config_runtime_resize() {
3818 let factory = Arc::new(MockFactory) as Arc<dyn ConnectionFactory>;
3819 let config = PoolProdConfig::default();
3820 let pool = Pool::new(config.to_pool_config(), factory).unwrap();
3821 assert_eq!(pool.max_size(), 100);
3822 pool.resize(50);
3823 assert_eq!(pool.max_size(), 50);
3824 }
3825}
3826
3827#[cfg(all(test, feature = "prod-leak-detection"))]
3828mod leak_prod_tests {
3829 use super::*;
3830
3831 #[test]
3832 fn test_leak_config_default() {
3833 let config = LeakDetectionConfig::default();
3834 assert!(!config.enabled);
3835 assert_eq!(config.interval, Duration::from_secs(60));
3836 assert_eq!(config.threshold, 5);
3837 }
3838
3839 #[test]
3840 fn test_leak_config_validate_ok() {
3841 let config =
3842 LeakDetectionConfig::new(true, Duration::from_secs(30), 10, Duration::from_secs(60));
3843 assert!(config.validate().is_ok());
3844 }
3845
3846 #[test]
3847 fn test_leak_config_interval_zero_rejected() {
3848 let config = LeakDetectionConfig::new(true, Duration::ZERO, 10, Duration::from_secs(60));
3849 assert!(config.validate().is_err());
3850 }
3851
3852 #[test]
3853 fn test_leak_report_empty() {
3854 let report = LeakReport::empty();
3855 assert_eq!(report.borrowed_count, 0);
3856 assert!(report.suspected_leaks.is_empty());
3857 }
3858
3859 #[test]
3860 fn connection_reuse_rate_zero() {
3861 let metrics = PoolMetrics::default();
3862 assert_eq!(metrics.connection_reuse_rate(), 0.0);
3863 }
3864
3865 #[test]
3866 fn connection_reuse_rate_full() {
3867 let metrics = PoolMetrics {
3868 acquire_count: 100,
3869 connection_created_count: 1,
3870 ..Default::default()
3871 };
3872 let rate = metrics.connection_reuse_rate();
3873 assert!(
3874 (rate - 0.99).abs() < 0.001,
3875 "复用率应接近 0.99,实际 {rate}"
3876 );
3877 }
3878
3879 #[test]
3880 fn connection_reuse_rate_partial() {
3881 let metrics = PoolMetrics {
3882 acquire_count: 10,
3883 connection_created_count: 2,
3884 ..Default::default()
3885 };
3886 assert!((metrics.connection_reuse_rate() - 0.8).abs() < 0.001);
3887 }
3888
3889 #[test]
3890 fn pool_tuning_advice_is_optimal() {
3891 let advice = PoolTuningAdvice {
3892 suggested_max_size: None,
3893 suggested_min_idle: None,
3894 suggested_idle_timeout: None,
3895 reason: "池配置合理".to_string(),
3896 };
3897 assert!(advice.is_optimal());
3898
3899 let not_optimal = PoolTuningAdvice {
3900 suggested_max_size: Some(20),
3901 suggested_min_idle: None,
3902 suggested_idle_timeout: None,
3903 reason: "test".to_string(),
3904 };
3905 assert!(!not_optimal.is_optimal());
3906 }
3907
3908 #[test]
3909 fn suggest_tuning_low_reuse() {
3910 let metrics = PoolMetrics {
3911 acquire_count: 100,
3912 connection_created_count: 60,
3913 ..Default::default()
3914 };
3915 let reuse = metrics.connection_reuse_rate();
3916 assert!(reuse < 0.5, "复用率 {reuse} 应 < 0.5");
3917 }
3918
3919 #[test]
3920 fn suggest_tuning_optimal() {
3921 let metrics = PoolMetrics {
3922 acquire_count: 1000,
3923 connection_created_count: 10,
3924 acquire_wait_time: Duration::from_millis(10),
3925 ..Default::default()
3926 };
3927 let reuse = metrics.connection_reuse_rate();
3928 assert!(reuse >= 0.9, "复用率 {reuse} 应 >= 0.9");
3929 let avg_wait = metrics.average_acquire_wait_time();
3930 assert!(avg_wait <= Duration::from_millis(100));
3931 }
3932
3933 #[test]
3934 fn suggest_tuning_high_wait() {
3935 let metrics = PoolMetrics {
3936 acquire_count: 100,
3937 acquire_wait_time: Duration::from_millis(200 * 100),
3938 ..Default::default()
3939 };
3940 let avg_wait = metrics.average_acquire_wait_time();
3941 assert!(
3942 avg_wait > Duration::from_millis(100),
3943 "平均等待 {avg_wait:?} 应 > 100ms"
3944 );
3945 }
3946}