1use async_trait::async_trait;
6use crossbeam_queue::ArrayQueue;
7use futures::StreamExt;
8#[cfg(feature = "circuit-breaker")]
12use parking_lot::Mutex as PlMutex;
13#[cfg(feature = "rate-limit")]
14use parking_lot::RwLock as PlRwLock;
15use std::future::Future;
16use std::ops::{Deref, DerefMut};
17use std::pin::Pin;
18use std::sync::atomic::{AtomicBool, AtomicU32, Ordering};
19use std::sync::Arc;
20use std::time::{Duration, Instant};
21use tokio::sync::Notify;
22
23#[cfg(feature = "circuit-breaker")]
27use crate::circuit_breaker::{CircuitBreaker, CircuitState, DefaultCircuitBreaker};
28use crate::error::PoolError;
29#[cfg(feature = "rate-limit")]
30use crate::rate_limiter::RateLimiter;
31
32pub type QueryRows = Vec<std::collections::HashMap<String, crate::value::Value>>;
34
35pub type QueryStreamItem =
37 Result<std::collections::HashMap<String, crate::value::Value>, crate::DbError>;
38
39pub trait Connection: Send + Sync {
46 fn execute<'a>(
47 &'a mut self,
48 sql: &'a str,
49 ) -> Pin<Box<dyn Future<Output = Result<u64, crate::DbError>> + Send + 'a>>;
50 fn query<'a>(
51 &'a mut self,
52 sql: &'a str,
53 ) -> Pin<Box<dyn Future<Output = Result<QueryRows, crate::DbError>> + Send + 'a>>;
54 fn begin_transaction<'a>(
55 &'a mut self,
56 ) -> Pin<Box<dyn Future<Output = Result<(), crate::DbError>> + Send + 'a>>;
57 fn commit<'a>(
58 &'a mut self,
59 ) -> Pin<Box<dyn Future<Output = Result<(), crate::DbError>> + Send + 'a>>;
60 fn rollback<'a>(
61 &'a mut self,
62 ) -> Pin<Box<dyn Future<Output = Result<(), crate::DbError>> + Send + 'a>>;
63 fn is_connected(&self) -> bool;
64 fn ping<'a>(&'a mut self) -> Pin<Box<dyn Future<Output = bool> + Send + 'a>>;
65 fn close<'a>(
66 &'a mut self,
67 ) -> Pin<Box<dyn Future<Output = Result<(), crate::DbError>> + Send + 'a>>;
68
69 fn execute_with_params<'a>(
75 &'a mut self,
76 sql: &'a str,
77 params: &'a [crate::value::Value],
78 ) -> Pin<Box<dyn Future<Output = Result<u64, crate::DbError>> + Send + 'a>> {
79 let _ = (sql, params);
80 Box::pin(async move {
81 Err(crate::DbError::Internal(
82 "execute_with_params not implemented for this adapter".to_string(),
83 ))
84 })
85 }
86
87 fn query_with_params<'a>(
93 &'a mut self,
94 sql: &'a str,
95 params: &'a [crate::value::Value],
96 ) -> Pin<Box<dyn Future<Output = Result<QueryRows, crate::DbError>> + Send + 'a>> {
97 let _ = (sql, params);
98 Box::pin(async move {
99 Err(crate::DbError::Internal(
100 "query_with_params not implemented for this adapter".to_string(),
101 ))
102 })
103 }
104
105 fn query_values<'a>(
110 &'a mut self,
111 sql: &'a str,
112 ) -> Pin<Box<dyn Future<Output = Result<crate::value::QueryValues, crate::DbError>> + Send + 'a>>
113 {
114 let _ = sql;
115 Box::pin(async move {
116 Err(crate::DbError::Internal(
117 "query_values not implemented for this adapter".to_string(),
118 ))
119 })
120 }
121
122 fn query_values_with_params<'a>(
126 &'a mut self,
127 sql: &'a str,
128 params: &'a [crate::value::Value],
129 ) -> Pin<Box<dyn Future<Output = Result<crate::value::QueryValues, crate::DbError>> + Send + 'a>>
130 {
131 let _ = (sql, params);
132 Box::pin(async move {
133 Err(crate::DbError::Internal(
134 "query_values_with_params not implemented for this adapter".to_string(),
135 ))
136 })
137 }
138
139 fn query_stream<'a>(
151 &'a mut self,
152 sql: &'a str,
153 ) -> Pin<Box<dyn futures::Stream<Item = QueryStreamItem> + Send + 'a>> {
154 let sql_owned = sql.to_string();
156 let stream = futures::stream::once(async move { self.query(&sql_owned).await })
158 .map(|result| {
160 let items: Vec<QueryStreamItem> = match result {
161 Ok(rows) => rows.into_iter().map(Ok).collect(),
162 Err(e) => vec![Err(e)],
163 };
164 futures::stream::iter(items)
165 })
166 .flatten();
167 Box::pin(stream)
168 }
169
170 fn query_stream_cursor<'a>(
181 &'a mut self,
182 sql: &'a str,
183 _batch_size: usize,
184 ) -> Pin<Box<dyn futures::Stream<Item = QueryStreamItem> + Send + 'a>> {
185 self.query_stream(sql)
186 }
187
188 fn execute_batch<'a>(
193 &'a mut self,
194 sqls: &'a [String],
195 ) -> Pin<Box<dyn Future<Output = Result<u64, crate::DbError>> + Send + 'a>> {
196 Box::pin(async move {
197 let mut total = 0u64;
198 for sql in sqls {
199 total += self.execute(sql).await?;
200 }
201 Ok(total)
202 })
203 }
204
205 fn execute_batch_params<'a>(
210 &'a mut self,
211 sql: &'a str,
212 params_batch: &'a [Vec<crate::value::Value>],
213 ) -> Pin<Box<dyn Future<Output = Result<u64, crate::DbError>> + Send + 'a>> {
214 Box::pin(async move {
215 let mut total = 0u64;
216 for params in params_batch {
217 total += self.execute_with_params(sql, params).await?;
218 }
219 Ok(total)
220 })
221 }
222}
223
224pub struct PooledConnection {
232 conn: Box<dyn Connection>,
233 created_at: Instant,
234 last_used_at: Instant,
235 pool: Option<Pool>,
236}
237
238impl PooledConnection {
239 fn new(conn: Box<dyn Connection>, pool: Pool) -> Self {
240 let now = Instant::now();
241 Self {
242 conn,
243 created_at: now,
244 last_used_at: now,
245 pool: Some(pool),
246 }
247 }
248
249 fn is_expired(&self, max_lifetime: Duration) -> bool {
250 self.created_at.elapsed() >= max_lifetime
251 }
252
253 fn is_idle_too_long(&self, idle_timeout: Duration) -> bool {
254 self.last_used_at.elapsed() >= idle_timeout
255 }
256
257 pub fn created_at(&self) -> Instant {
259 self.created_at
260 }
261
262 pub fn into_inner(mut self) -> Box<dyn Connection> {
267 self.pool = None; std::mem::replace(&mut self.conn, Box::new(ClosedConnection))
271 }
272}
273
274impl Drop for PooledConnection {
286 fn drop(&mut self) {
287 if let Some(pool) = self.pool.take() {
288 let conn = std::mem::replace(&mut self.conn, Box::new(ClosedConnection));
290 let pooled = PooledConnection {
291 conn,
292 created_at: self.created_at,
293 last_used_at: self.last_used_at,
294 pool: None,
295 };
296 if let Ok(handle) = tokio::runtime::Handle::try_current() {
298 handle.spawn(async move {
299 pool.release(pooled).await;
300 });
301 } else {
302 drop(pooled);
306 pool.total_count.fetch_sub(1, Ordering::SeqCst);
307 }
308 }
309 }
310}
311
312struct ClosedConnection;
316
317impl Connection for ClosedConnection {
318 fn execute<'a>(
319 &'a mut self,
320 _sql: &'a str,
321 ) -> Pin<Box<dyn Future<Output = Result<u64, crate::DbError>> + Send + 'a>> {
322 Box::pin(async {
323 Err(crate::DbError::ConnectionError(
324 "connection already returned to pool".to_string(),
325 ))
326 })
327 }
328
329 fn query<'a>(
330 &'a mut self,
331 _sql: &'a str,
332 ) -> Pin<Box<dyn Future<Output = Result<QueryRows, crate::DbError>> + Send + 'a>> {
333 Box::pin(async {
334 Err(crate::DbError::ConnectionError(
335 "connection already returned to pool".to_string(),
336 ))
337 })
338 }
339
340 fn begin_transaction<'a>(
341 &'a mut self,
342 ) -> Pin<Box<dyn Future<Output = Result<(), crate::DbError>> + Send + 'a>> {
343 Box::pin(async {
344 Err(crate::DbError::ConnectionError(
345 "connection already returned to pool".to_string(),
346 ))
347 })
348 }
349
350 fn commit<'a>(
351 &'a mut self,
352 ) -> Pin<Box<dyn Future<Output = Result<(), crate::DbError>> + Send + 'a>> {
353 Box::pin(async { Ok(()) })
354 }
355
356 fn rollback<'a>(
357 &'a mut self,
358 ) -> Pin<Box<dyn Future<Output = Result<(), crate::DbError>> + Send + 'a>> {
359 Box::pin(async { Ok(()) })
360 }
361
362 fn is_connected(&self) -> bool {
363 false
364 }
365
366 fn ping<'a>(&'a mut self) -> Pin<Box<dyn Future<Output = bool> + Send + 'a>> {
367 Box::pin(async { false })
368 }
369
370 fn close<'a>(
371 &'a mut self,
372 ) -> Pin<Box<dyn Future<Output = Result<(), crate::DbError>> + Send + 'a>> {
373 Box::pin(async { Ok(()) })
374 }
375}
376
377impl Deref for PooledConnection {
378 type Target = dyn Connection;
379
380 fn deref(&self) -> &Self::Target {
381 self.conn.as_ref()
382 }
383}
384
385impl DerefMut for PooledConnection {
386 fn deref_mut(&mut self) -> &mut Self::Target {
387 self.conn.as_mut()
388 }
389}
390
391#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
393pub enum TlsVersion {
394 #[default]
395 Tls12,
396 Tls13,
397}
398
399#[derive(Debug, Clone, Default)]
401pub struct TlsConfig {
402 pub enabled: bool,
404 pub ca_cert_path: Option<String>,
406 pub client_cert_path: Option<String>,
408 pub client_key_path: Option<String>,
410 pub min_version: TlsVersion,
412}
413
414#[derive(Debug, Clone)]
416pub enum PoolEvent {
417 ConnectionCreated,
419 ConnectionClosed,
421 ConnectionAcquired,
423 ConnectionReleased,
425 AcquireTimeout,
427}
428
429pub type PoolEventCallback = Arc<dyn Fn(PoolEvent) + Send + Sync>;
431
432pub struct PoolConfig {
433 pub max_size: u32,
434 pub min_idle: u32,
435 pub acquire_timeout: Duration,
436 pub idle_timeout: Duration,
437 pub max_lifetime: Duration,
438 pub connection_timeout: Duration,
439 pub tls: Option<TlsConfig>,
441 pub query_timeout: Option<Duration>,
443 pub max_rows: Option<usize>,
445 pub memory_limit: Option<usize>,
447 pub on_event: Option<PoolEventCallback>,
449 pub test_before_acquire: bool,
457 pub prewarm: bool,
469}
470
471impl Default for PoolConfig {
472 fn default() -> Self {
473 Self {
474 max_size: 100,
475 min_idle: 0,
476 acquire_timeout: Duration::from_secs(30),
477 idle_timeout: Duration::from_secs(600),
478 max_lifetime: Duration::from_secs(1800),
479 connection_timeout: Duration::from_secs(10),
480 tls: None,
481 query_timeout: Some(Duration::from_secs(30)),
482 max_rows: None,
483 memory_limit: None,
484 on_event: None,
485 test_before_acquire: false,
486 prewarm: false,
487 }
488 }
489}
490
491impl Clone for PoolConfig {
492 fn clone(&self) -> Self {
493 Self {
494 max_size: self.max_size,
495 min_idle: self.min_idle,
496 acquire_timeout: self.acquire_timeout,
497 idle_timeout: self.idle_timeout,
498 max_lifetime: self.max_lifetime,
499 connection_timeout: self.connection_timeout,
500 tls: self.tls.clone(),
501 query_timeout: self.query_timeout,
502 max_rows: self.max_rows,
503 memory_limit: self.memory_limit,
504 on_event: self.on_event.clone(),
505 test_before_acquire: self.test_before_acquire,
506 prewarm: self.prewarm,
507 }
508 }
509}
510
511impl PoolConfig {
512 pub fn validate(&self) -> Result<(), PoolError> {
514 if self.max_size == 0 {
515 return Err(PoolError::InvalidConfig("max_size cannot be 0".to_string()));
516 }
517 if self.min_idle > self.max_size {
518 return Err(PoolError::InvalidConfig(
519 "min_idle cannot exceed max_size".to_string(),
520 ));
521 }
522 const MAX_DURATION_SECS: u64 = u32::MAX as u64; for (name, dur) in [
528 ("acquire_timeout", self.acquire_timeout),
529 ("idle_timeout", self.idle_timeout),
530 ("max_lifetime", self.max_lifetime),
531 ("connection_timeout", self.connection_timeout),
532 ] {
533 if dur.as_secs() > MAX_DURATION_SECS {
534 return Err(PoolError::InvalidConfig(format!(
535 "{name} ({:?}) exceeds maximum allowed duration ({} seconds)",
536 dur, MAX_DURATION_SECS
537 )));
538 }
539 }
540 Ok(())
541 }
542
543 #[must_use]
545 pub fn with_prewarm(mut self, prewarm: bool) -> Self {
546 self.prewarm = prewarm;
547 self
548 }
549}
550
551pub struct PoolStatus {
552 pub idle: u32,
553 pub active: u32,
554 pub max: u32,
555 pub min: u32,
556 pub waiters: u32,
558}
559
560impl std::fmt::Debug for PoolStatus {
561 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
562 f.debug_struct("PoolStatus")
563 .field("idle", &self.idle)
564 .field("active", &self.active)
565 .field("max", &self.max)
566 .field("min", &self.min)
567 .field("waiters", &self.waiters)
568 .finish()
569 }
570}
571
572pub struct PoolConfigBuilder {
573 config: PoolConfig,
574}
575
576impl PoolConfigBuilder {
577 pub fn new() -> Self {
578 Self {
579 config: PoolConfig::default(),
580 }
581 }
582
583 pub fn max_size(mut self, size: u32) -> Self {
584 self.config.max_size = size;
585 self
586 }
587
588 pub fn min_idle(mut self, count: u32) -> Self {
589 self.config.min_idle = count;
590 self
591 }
592
593 pub fn acquire_timeout(mut self, timeout_secs: u64) -> Self {
594 self.config.acquire_timeout = Duration::from_secs(timeout_secs);
595 self
596 }
597
598 pub fn idle_timeout(mut self, timeout_secs: u64) -> Self {
599 self.config.idle_timeout = Duration::from_secs(timeout_secs);
600 self
601 }
602
603 pub fn max_lifetime(mut self, lifetime_secs: u64) -> Self {
604 self.config.max_lifetime = Duration::from_secs(lifetime_secs);
605 self
606 }
607
608 pub fn tls(mut self, tls: TlsConfig) -> Self {
610 self.config.tls = Some(tls);
611 self
612 }
613
614 pub fn query_timeout(mut self, timeout: Duration) -> Self {
616 self.config.query_timeout = Some(timeout);
617 self
618 }
619
620 pub fn max_rows(mut self, max_rows: usize) -> Self {
622 self.config.max_rows = Some(max_rows);
623 self
624 }
625
626 pub fn memory_limit(mut self, memory_limit: usize) -> Self {
628 self.config.memory_limit = Some(memory_limit);
629 self
630 }
631
632 pub fn on_event(mut self, callback: PoolEventCallback) -> Self {
634 self.config.on_event = Some(callback);
635 self
636 }
637
638 pub fn test_before_acquire(mut self, enabled: bool) -> Self {
643 self.config.test_before_acquire = enabled;
644 self
645 }
646
647 pub fn prewarm(mut self, enabled: bool) -> Self {
652 self.config.prewarm = enabled;
653 self
654 }
655
656 pub fn build(self) -> Result<PoolConfig, PoolError> {
657 self.config.validate()?;
658 Ok(self.config)
659 }
660}
661
662impl Default for PoolConfigBuilder {
663 fn default() -> Self {
664 Self::new()
665 }
666}
667
668#[async_trait]
670pub trait ConnectionFactory: Send + Sync {
671 async fn create(&self) -> Result<Box<dyn Connection>, crate::DbError>;
672}
673
674pub struct Pool {
680 config: PoolConfig,
681 factory: Arc<dyn ConnectionFactory>,
682 idle: Arc<ArrayQueue<PooledConnection>>,
688 total_count: Arc<AtomicU32>,
698 closed: Arc<AtomicBool>,
700 notify: Arc<Notify>,
701 waiters_count: Arc<AtomicU32>,
703 dynamic_max_size: Arc<AtomicU32>,
705 #[cfg(feature = "circuit-breaker")]
711 circuit_breaker: Arc<PlMutex<DefaultCircuitBreaker>>,
712 #[cfg(feature = "rate-limit")]
721 rate_limiter: Arc<PlRwLock<Option<Arc<dyn RateLimiter>>>>,
722 #[cfg(feature = "rate-limit")]
724 rate_limit_key: String,
725}
726
727impl Clone for Pool {
731 fn clone(&self) -> Self {
732 Self {
733 config: self.config.clone(),
734 factory: self.factory.clone(),
735 idle: self.idle.clone(),
736 total_count: self.total_count.clone(),
737 closed: self.closed.clone(),
738 notify: Arc::clone(&self.notify),
739 waiters_count: self.waiters_count.clone(),
740 dynamic_max_size: self.dynamic_max_size.clone(),
741 #[cfg(feature = "circuit-breaker")]
742 circuit_breaker: Arc::clone(&self.circuit_breaker),
743 #[cfg(feature = "rate-limit")]
744 rate_limiter: Arc::clone(&self.rate_limiter),
745 #[cfg(feature = "rate-limit")]
746 rate_limit_key: self.rate_limit_key.clone(),
747 }
748 }
749}
750
751impl Pool {
752 pub fn new(config: PoolConfig, factory: Arc<dyn ConnectionFactory>) -> Result<Self, PoolError> {
776 config.validate()?;
777 let max_size = config.max_size as usize;
780 let dynamic_max = config.max_size;
781 Ok(Self {
782 config,
783 factory,
784 idle: Arc::new(ArrayQueue::new(max_size)),
785 total_count: Arc::new(AtomicU32::new(0)),
786 closed: Arc::new(AtomicBool::new(false)),
787 notify: Arc::new(Notify::new()),
788 waiters_count: Arc::new(AtomicU32::new(0)),
789 dynamic_max_size: Arc::new(AtomicU32::new(dynamic_max)),
790 #[cfg(feature = "circuit-breaker")]
793 circuit_breaker: Arc::new(PlMutex::new(DefaultCircuitBreaker::new(
794 5,
795 std::time::Duration::from_secs(30),
796 ))),
797 #[cfg(feature = "rate-limit")]
800 rate_limiter: Arc::new(PlRwLock::new(None)),
801 #[cfg(feature = "rate-limit")]
802 rate_limit_key: "pool".to_string(),
803 })
804 }
805
806 pub async fn prewarm(&self) {
823 if !self.config.prewarm {
824 return;
825 }
826
827 let min_idle = self.config.min_idle as usize;
828 let mut warmed = 0;
829
830 for i in 0..min_idle {
831 if self.closed.load(Ordering::Acquire) {
833 break;
834 }
835
836 let current_max = self.dynamic_max_size.load(Ordering::Acquire);
838 let current = self.total_count.load(Ordering::Acquire);
839 if current >= current_max {
840 break;
841 }
842
843 let created = loop {
845 let current = self.total_count.load(Ordering::Acquire);
846 if current >= current_max {
847 break None;
848 }
849 match self.total_count.compare_exchange(
850 current,
851 current + 1,
852 Ordering::SeqCst,
853 Ordering::Acquire,
854 ) {
855 Ok(_) => break Some(()),
856 Err(_) => continue,
857 }
858 };
859
860 if created.is_some() {
861 match tokio::time::timeout(self.config.connection_timeout, self.factory.create())
862 .await
863 {
864 Ok(Ok(conn)) => {
865 #[cfg(feature = "circuit-breaker")]
866 {
867 self.circuit_breaker.lock().record_success();
868 }
869 self.emit_event(PoolEvent::ConnectionCreated);
870 let pooled = PooledConnection::new(conn, self.clone());
871 if let Err(_) = self.idle.push(pooled) {
873 let _ = self.total_count.fetch_sub(1, Ordering::SeqCst);
875 tracing::warn!(
876 target: "sz_orm::pool::prewarm",
877 "prewarm connection {} failed: idle queue full",
878 i
879 );
880 } else {
881 warmed += 1;
882 self.notify.notify_one();
883 }
884 }
885 Ok(Err(e)) => {
886 let _ = self.total_count.fetch_sub(1, Ordering::SeqCst);
887 #[cfg(feature = "circuit-breaker")]
888 {
889 self.circuit_breaker.lock().record_failure();
890 }
891 tracing::warn!(
892 target: "sz_orm::pool::prewarm",
893 "prewarm connection {} failed: {}",
894 i,
895 e
896 );
897 }
898 Err(_) => {
899 let _ = self.total_count.fetch_sub(1, Ordering::SeqCst);
900 #[cfg(feature = "circuit-breaker")]
901 {
902 self.circuit_breaker.lock().record_failure();
903 }
904 tracing::warn!(
905 target: "sz_orm::pool::prewarm",
906 "prewarm connection {} timeout",
907 i
908 );
909 }
910 }
911 }
912 }
913
914 if warmed > 0 {
915 tracing::info!(
916 target: "sz_orm::pool::prewarm",
917 "pool prewarm completed: {}/{} connections established",
918 warmed,
919 min_idle
920 );
921 }
922 }
923
924 pub fn config(&self) -> &PoolConfig {
926 &self.config
927 }
928
929 #[cfg(feature = "circuit-breaker")]
943 pub fn configure_circuit_breaker(
944 &self,
945 failure_threshold: usize,
946 reset_timeout: std::time::Duration,
947 ) {
948 let new_cb = DefaultCircuitBreaker::new(failure_threshold, reset_timeout);
949 let mut guard = self.circuit_breaker.lock();
951 *guard = new_cb;
952 }
953
954 #[cfg(feature = "circuit-breaker")]
959 pub fn reset_circuit_breaker(&self) -> bool {
960 let mut guard = self.circuit_breaker.lock();
962 guard.reset()
963 }
964
965 #[cfg(feature = "circuit-breaker")]
967 pub fn circuit_state(&self) -> CircuitState {
968 let guard = self.circuit_breaker.lock();
970 guard.state()
971 }
972
973 #[cfg(feature = "rate-limit")]
982 pub fn set_rate_limiter(&self, limiter: Option<Arc<dyn RateLimiter>>) {
983 let mut guard = self.rate_limiter.write();
985 *guard = limiter;
986 }
987
988 #[cfg(feature = "rate-limit")]
990 pub fn with_rate_limit_key(mut self, key: impl Into<String>) -> Self {
991 self.rate_limit_key = key.into();
992 self
993 }
994
995 fn emit_event(&self, event: PoolEvent) {
997 if let Some(ref callback) = self.config.on_event {
998 callback(event);
999 }
1000 }
1001
1002 #[tracing::instrument(skip(self), fields(max_size = self.config.max_size, acquire_timeout = ?self.config.acquire_timeout))]
1022 pub async fn acquire(&self) -> Result<PooledConnection, PoolError> {
1023 if self.closed.load(Ordering::Acquire) {
1025 return Err(PoolError::Closed);
1026 }
1027
1028 #[cfg(feature = "circuit-breaker")]
1032 {
1033 let mut guard = self.circuit_breaker.lock();
1034 if !guard.can_execute() {
1035 return Err(PoolError::CircuitOpen);
1036 }
1037 }
1038
1039 #[cfg(feature = "rate-limit")]
1043 {
1044 let guard = self.rate_limiter.read();
1045 if let Some(ref limiter) = *guard {
1046 match limiter.try_acquire(&self.rate_limit_key) {
1047 Ok(result) if !result.allowed => {
1048 return Err(PoolError::RateLimited {
1049 remaining: result.remaining,
1050 reset_at: result.reset_at,
1051 });
1052 }
1053 Ok(_) => {} Err(_) => {
1055 }
1057 }
1058 }
1059 }
1060
1061 let deadline = Instant::now() + self.config.acquire_timeout;
1062 let mut backoff = Duration::from_millis(1);
1064 const MAX_BACKOFF: Duration = Duration::from_millis(100);
1066
1067 loop {
1068 let mut to_close: Vec<PooledConnection> = Vec::new();
1074 let acquired: Option<PooledConnection> = {
1075 let mut found: Option<PooledConnection> = None;
1076 while let Some(pooled) = self.idle.pop() {
1077 if pooled.is_expired(self.config.max_lifetime) {
1079 to_close.push(pooled);
1080 continue;
1081 }
1082 if pooled.is_idle_too_long(self.config.idle_timeout) {
1084 to_close.push(pooled);
1085 continue;
1086 }
1087 if !pooled.conn.is_connected() {
1090 to_close.push(pooled);
1091 continue;
1092 }
1093 found = Some(pooled);
1094 break;
1095 }
1096 found
1097 };
1098
1099 for mut pooled in to_close {
1101 let _ = pooled.conn.close().await;
1102 self.total_count.fetch_sub(1, Ordering::SeqCst);
1104 }
1105
1106 if let Some(mut pooled) = acquired {
1107 if self.config.test_before_acquire {
1109 let ping_timeout = self.config.connection_timeout / 2;
1110 let alive = match tokio::time::timeout(ping_timeout, pooled.conn.ping()).await {
1111 Ok(true) => true,
1112 Ok(false) => false,
1113 Err(_) => false, };
1115 if !alive {
1116 let _ = pooled.conn.close().await;
1118 self.total_count.fetch_sub(1, Ordering::SeqCst);
1119 continue;
1120 }
1121 }
1122 pooled.pool = Some(self.clone());
1125 return Ok(pooled);
1126 }
1127
1128 let current_max = self.dynamic_max_size.load(Ordering::Acquire);
1133 let created = loop {
1134 let current = self.total_count.load(Ordering::Acquire);
1135 if current >= current_max {
1136 break None; }
1138 match self.total_count.compare_exchange(
1139 current,
1140 current + 1,
1141 Ordering::SeqCst,
1142 Ordering::Acquire,
1143 ) {
1144 Ok(_) => break Some(()), Err(_) => continue, }
1147 };
1148
1149 if created.is_some() {
1150 match tokio::time::timeout(self.config.connection_timeout, self.factory.create())
1151 .await
1152 {
1153 Ok(Ok(conn)) => {
1154 #[cfg(feature = "circuit-breaker")]
1157 {
1158 self.circuit_breaker.lock().record_success();
1159 }
1160 self.emit_event(PoolEvent::ConnectionCreated);
1161 self.emit_event(PoolEvent::ConnectionAcquired);
1162 return Ok(PooledConnection::new(conn, self.clone()));
1163 }
1164 Ok(Err(e)) => {
1165 self.total_count.fetch_sub(1, Ordering::SeqCst);
1167 #[cfg(feature = "circuit-breaker")]
1170 {
1171 self.circuit_breaker.lock().record_failure();
1172 }
1173 return Err(PoolError::ConnectionFailed(e.to_string()));
1174 }
1175 Err(_) => {
1176 self.total_count.fetch_sub(1, Ordering::SeqCst);
1178 #[cfg(feature = "circuit-breaker")]
1181 {
1182 self.circuit_breaker.lock().record_failure();
1183 }
1184 return Err(PoolError::Timeout);
1185 }
1186 }
1187 }
1188
1189 let now = Instant::now();
1191 if now >= deadline {
1192 self.emit_event(PoolEvent::AcquireTimeout);
1193 return Err(PoolError::Timeout);
1194 }
1195 self.waiters_count.fetch_add(1, Ordering::SeqCst);
1197 let wait = std::cmp::min(backoff, deadline - now);
1198 match tokio::time::timeout(wait, self.notify.notified()).await {
1199 Ok(()) => {
1200 backoff = Duration::from_millis(1);
1202 }
1203 Err(_) => {
1204 backoff = std::cmp::min(backoff * 2, MAX_BACKOFF);
1206 }
1207 }
1208 self.waiters_count.fetch_sub(1, Ordering::SeqCst);
1210 }
1211 }
1212
1213 #[tracing::instrument(skip(self, pooled))]
1221 pub async fn release(&self, mut pooled: PooledConnection) {
1222 pooled.pool = None;
1224
1225 if self.closed.load(Ordering::Acquire) {
1227 let _ = pooled.conn.close().await;
1228 self.total_count.fetch_sub(1, Ordering::SeqCst);
1230 self.emit_event(PoolEvent::ConnectionClosed);
1231 return;
1232 }
1233
1234 if !pooled.conn.is_connected() {
1236 let _ = pooled.conn.close().await;
1237 self.total_count.fetch_sub(1, Ordering::SeqCst);
1238 self.emit_event(PoolEvent::ConnectionClosed);
1239 return;
1240 }
1241
1242 pooled.last_used_at = Instant::now();
1244
1245 if let Err(mut rejected) = self.idle.push(pooled) {
1251 let _ = rejected.conn.close().await;
1253 self.total_count.fetch_sub(1, Ordering::SeqCst);
1254 self.emit_event(PoolEvent::ConnectionClosed);
1255 } else {
1256 self.emit_event(PoolEvent::ConnectionReleased);
1257 }
1258 self.notify.notify_one();
1259 }
1260
1261 pub async fn status(&self) -> PoolStatus {
1266 let idle_count = self.idle.len() as u32;
1267 let active = self.total_count.load(Ordering::Acquire);
1269 let waiters = self.waiters_count.load(Ordering::Acquire);
1270 PoolStatus {
1271 idle: idle_count,
1272 active,
1273 max: self.dynamic_max_size.load(Ordering::Acquire),
1274 min: self.config.min_idle,
1275 waiters,
1276 }
1277 }
1278
1279 #[tracing::instrument(skip(self))]
1281 pub async fn reap_idle(&self) {
1282 let mut all: Vec<PooledConnection> = Vec::new();
1286 while let Some(pooled) = self.idle.pop() {
1287 all.push(pooled);
1288 }
1289
1290 let mut to_close = Vec::new();
1292 for pooled in all {
1293 if pooled.is_idle_too_long(self.config.idle_timeout)
1294 || pooled.is_expired(self.config.max_lifetime)
1295 {
1296 to_close.push(pooled);
1297 } else {
1298 if let Err(mut rejected) = self.idle.push(pooled) {
1300 let _ = rejected.conn.close().await;
1301 self.total_count.fetch_sub(1, Ordering::SeqCst);
1302 }
1303 }
1304 }
1305
1306 for mut pooled in to_close {
1308 let _ = pooled.conn.close().await;
1309 self.total_count.fetch_sub(1, Ordering::SeqCst);
1311 }
1312 }
1313
1314 pub async fn close_all(&self) {
1318 self.closed.store(true, Ordering::Release);
1320 let mut to_close: Vec<PooledConnection> = Vec::new();
1323 while let Some(pooled) = self.idle.pop() {
1324 to_close.push(pooled);
1325 }
1326 let closed_count: u32 = to_close.len() as u32;
1328 for mut pooled in to_close {
1329 let _ = pooled.conn.close().await;
1330 }
1331 self.total_count.fetch_sub(closed_count, Ordering::SeqCst);
1334 }
1335
1336 pub async fn health_check(&self) -> u32 {
1352 let mut to_check: Vec<PooledConnection> = Vec::new();
1354 while let Some(pooled) = self.idle.pop() {
1355 to_check.push(pooled);
1356 }
1357
1358 let mut removed: u32 = 0;
1359 let mut alive: Vec<PooledConnection> = Vec::with_capacity(to_check.len());
1360 for mut pooled in to_check.drain(..) {
1361 if !pooled.conn.is_connected() {
1363 let _ = pooled.conn.close().await;
1364 removed += 1;
1365 continue;
1366 }
1367 let ping_timeout = self.config.connection_timeout / 2;
1369 match tokio::time::timeout(ping_timeout, pooled.conn.ping()).await {
1370 Ok(true) => alive.push(pooled),
1371 Ok(false) => {
1372 let _ = pooled.conn.close().await;
1374 removed += 1;
1375 }
1376 Err(_) => {
1377 let _ = pooled.conn.close().await;
1379 removed += 1;
1380 }
1381 }
1382 }
1383
1384 let alive_count: u32 = alive.len() as u32;
1386 for pooled in alive {
1387 if let Err(mut rejected) = self.idle.push(pooled) {
1389 let _ = rejected.conn.close().await;
1390 removed += 1;
1391 }
1392 }
1393
1394 if removed > 0 {
1396 self.total_count.fetch_sub(removed, Ordering::SeqCst);
1397 }
1398
1399 if alive_count > 0 {
1401 self.notify.notify_one();
1402 }
1403
1404 removed
1405 }
1406
1407 pub async fn shutdown(&self) {
1414 self.closed.store(true, Ordering::SeqCst);
1416 self.notify.notify_waiters();
1418 self.close_all().await;
1420 let deadline = Instant::now() + Duration::from_secs(30);
1422 while self.total_count.load(Ordering::SeqCst) > 0 {
1423 if Instant::now() >= deadline {
1424 break;
1425 }
1426 tokio::time::sleep(Duration::from_millis(100)).await;
1427 }
1428 }
1429
1430 pub fn resize(&self, new_max: usize) {
1438 self.set_max_size(new_max as u32);
1439 }
1440
1441 pub fn set_max_size(&self, new_max: u32) {
1443 self.dynamic_max_size.store(new_max, Ordering::SeqCst);
1444 }
1445
1446 pub fn max_size(&self) -> u32 {
1448 self.dynamic_max_size.load(Ordering::Acquire)
1449 }
1450
1451 pub async fn warmup(&self, min_idle: usize) -> Result<(), PoolError> {
1455 for _ in 0..min_idle {
1456 let current_max = self.dynamic_max_size.load(Ordering::Acquire);
1457 let current = self.total_count.load(Ordering::Acquire);
1458 if current >= current_max {
1459 break;
1460 }
1461 match self.total_count.compare_exchange(
1463 current,
1464 current + 1,
1465 Ordering::SeqCst,
1466 Ordering::Acquire,
1467 ) {
1468 Ok(_) => {}
1469 Err(_) => continue, }
1471 match self.factory.create().await {
1472 Ok(conn) => {
1473 let now = Instant::now();
1474 let pooled = PooledConnection {
1475 conn,
1476 created_at: now,
1477 last_used_at: now,
1478 pool: None,
1479 };
1480 if let Err(mut rejected) = self.idle.push(pooled) {
1481 let _ = rejected.conn.close().await;
1483 self.total_count.fetch_sub(1, Ordering::SeqCst);
1484 }
1485 self.emit_event(PoolEvent::ConnectionCreated);
1486 }
1487 Err(_) => {
1488 self.total_count.fetch_sub(1, Ordering::SeqCst);
1490 break;
1491 }
1492 }
1493 }
1494 Ok(())
1495 }
1496
1497 pub async fn query_with_timeout(&self, sql: &str) -> Result<QueryRows, crate::DbError> {
1502 let timeout = self.config.query_timeout.unwrap_or(Duration::from_secs(30));
1503 let mut conn = self.acquire().await.map_err(crate::DbError::PoolError)?;
1504 tokio::time::timeout(timeout, conn.query(sql))
1505 .await
1506 .map_err(|_| crate::DbError::QueryError(format!("Query timeout after {:?}", timeout)))?
1507 }
1508}
1509
1510#[cfg(test)]
1511mod tests {
1512 use super::*;
1513
1514 struct MockConnection {
1516 connected: bool,
1517 }
1518
1519 impl MockConnection {
1520 fn new() -> Self {
1521 Self { connected: true }
1522 }
1523 }
1524
1525 impl Connection for MockConnection {
1526 fn execute<'a>(
1527 &'a mut self,
1528 _sql: &'a str,
1529 ) -> Pin<Box<dyn Future<Output = Result<u64, crate::DbError>> + Send + 'a>> {
1530 Box::pin(async move { Ok(1) })
1531 }
1532
1533 fn query<'a>(
1534 &'a mut self,
1535 _sql: &'a str,
1536 ) -> Pin<
1537 Box<
1538 dyn Future<
1539 Output = Result<
1540 Vec<std::collections::HashMap<String, crate::value::Value>>,
1541 crate::DbError,
1542 >,
1543 > + Send
1544 + 'a,
1545 >,
1546 > {
1547 Box::pin(async move { Ok(vec![]) })
1548 }
1549
1550 fn begin_transaction<'a>(
1551 &'a mut self,
1552 ) -> Pin<Box<dyn Future<Output = Result<(), crate::DbError>> + Send + 'a>> {
1553 Box::pin(async move { Ok(()) })
1554 }
1555
1556 fn commit<'a>(
1557 &'a mut self,
1558 ) -> Pin<Box<dyn Future<Output = Result<(), crate::DbError>> + Send + 'a>> {
1559 Box::pin(async move { Ok(()) })
1560 }
1561
1562 fn rollback<'a>(
1563 &'a mut self,
1564 ) -> Pin<Box<dyn Future<Output = Result<(), crate::DbError>> + Send + 'a>> {
1565 Box::pin(async move { Ok(()) })
1566 }
1567
1568 fn is_connected(&self) -> bool {
1569 self.connected
1570 }
1571
1572 fn ping<'a>(&'a mut self) -> Pin<Box<dyn Future<Output = bool> + Send + 'a>> {
1573 Box::pin(async move { true })
1574 }
1575
1576 fn close<'a>(
1577 &'a mut self,
1578 ) -> Pin<Box<dyn Future<Output = Result<(), crate::DbError>> + Send + 'a>> {
1579 Box::pin(async move {
1580 self.connected = false;
1581 Ok(())
1582 })
1583 }
1584 }
1585
1586 struct MockConnectionFactory;
1587
1588 #[async_trait]
1589 impl ConnectionFactory for MockConnectionFactory {
1590 async fn create(&self) -> Result<Box<dyn Connection>, crate::DbError> {
1591 Ok(Box::new(MockConnection::new()))
1592 }
1593 }
1594
1595 #[tokio::test]
1596 async fn test_pool_config_builder() -> Result<(), Box<dyn std::error::Error>> {
1597 let config = PoolConfigBuilder::new().max_size(50).min_idle(10).build()?;
1598
1599 assert_eq!(config.max_size, 50);
1600 assert_eq!(config.min_idle, 10);
1601 Ok(())
1602 }
1603
1604 #[test]
1605 fn test_pool_status_display() {
1606 let status = PoolStatus {
1607 idle: 5,
1608 active: 10,
1609 max: 100,
1610 min: 5,
1611 waiters: 0,
1612 };
1613
1614 let display = format!("{:?}", status);
1615 assert!(display.contains("idle"));
1616 assert!(display.contains("active"));
1617 }
1618
1619 #[test]
1620 fn test_default_pool_config() {
1621 let config = PoolConfig::default();
1622 assert_eq!(config.max_size, 100);
1623 assert_eq!(config.min_idle, 0);
1624 assert_eq!(config.acquire_timeout.as_secs(), 30);
1625 assert_eq!(config.idle_timeout.as_secs(), 600);
1626 assert_eq!(config.max_lifetime.as_secs(), 1800);
1627 }
1628
1629 #[tokio::test]
1630 async fn test_pool_config_clone() {
1631 let config = PoolConfig::default();
1632 let cloned = config.clone();
1633 assert_eq!(cloned.max_size, config.max_size);
1634 assert_eq!(cloned.min_idle, config.min_idle);
1635 }
1636
1637 #[test]
1638 fn test_pool_config_builder_default() -> Result<(), Box<dyn std::error::Error>> {
1639 let builder = PoolConfigBuilder::new();
1640 let config = builder.build()?;
1641 assert_eq!(config.max_size, 100);
1642 Ok(())
1643 }
1644
1645 #[test]
1646 fn test_pool_config_validate() {
1647 let result = PoolConfigBuilder::new().max_size(0).build();
1648 assert!(result.is_err());
1649
1650 let result = PoolConfigBuilder::new().max_size(10).min_idle(20).build();
1651 assert!(result.is_err());
1652 }
1653
1654 #[test]
1655 fn test_pool_config_validate_duration_upper_bound() {
1656 use std::time::Duration;
1657
1658 let config = PoolConfig {
1660 max_size: 10,
1661 min_idle: 1,
1662 acquire_timeout: Duration::from_secs(u64::MAX),
1663 idle_timeout: Duration::from_secs(1),
1664 max_lifetime: Duration::from_secs(1),
1665 connection_timeout: Duration::from_secs(5),
1666 tls: None,
1667 query_timeout: None,
1668 max_rows: None,
1669 memory_limit: None,
1670 on_event: None,
1671 test_before_acquire: false,
1672 prewarm: false,
1673 };
1674 assert!(config.validate().is_err());
1675
1676 let config = PoolConfig {
1678 max_size: 10,
1679 min_idle: 1,
1680 acquire_timeout: Duration::from_secs(u32::MAX as u64),
1681 idle_timeout: Duration::from_secs(1),
1682 max_lifetime: Duration::from_secs(1),
1683 connection_timeout: Duration::from_secs(5),
1684 tls: None,
1685 query_timeout: None,
1686 max_rows: None,
1687 memory_limit: None,
1688 on_event: None,
1689 test_before_acquire: false,
1690 prewarm: false,
1691 };
1692 assert!(config.validate().is_ok());
1693
1694 let config = PoolConfig {
1696 max_size: 10,
1697 min_idle: 1,
1698 acquire_timeout: Duration::from_secs(u32::MAX as u64 + 1),
1699 idle_timeout: Duration::from_secs(1),
1700 max_lifetime: Duration::from_secs(1),
1701 connection_timeout: Duration::from_secs(5),
1702 tls: None,
1703 query_timeout: None,
1704 max_rows: None,
1705 memory_limit: None,
1706 on_event: None,
1707 test_before_acquire: false,
1708 prewarm: false,
1709 };
1710 assert!(config.validate().is_err());
1711 }
1712
1713 #[test]
1714 fn test_pool_config_test_before_acquire_default() {
1715 let config = PoolConfig::default();
1717 assert!(!config.test_before_acquire);
1718 }
1719
1720 #[test]
1721 fn test_pool_config_builder_test_before_acquire() {
1722 let config = PoolConfigBuilder::new()
1724 .test_before_acquire(true)
1725 .build()
1726 .unwrap();
1727 assert!(config.test_before_acquire);
1728 }
1729
1730 #[tokio::test]
1731 async fn test_pool_acquire_and_release() -> Result<(), Box<dyn std::error::Error>> {
1732 let config = PoolConfigBuilder::new().max_size(5).min_idle(1).build()?;
1733 let factory = Arc::new(MockConnectionFactory);
1734 let pool = Pool::new(config, factory)?;
1735
1736 let conn = pool.acquire().await?;
1737 let status = pool.status().await;
1738 assert_eq!(status.active, 1);
1739 assert_eq!(status.idle, 0);
1740
1741 pool.release(conn).await;
1742 let status = pool.status().await;
1743 assert_eq!(status.idle, 1);
1744
1745 let _conn2 = pool.acquire().await?;
1747 let status = pool.status().await;
1748 assert_eq!(status.idle, 0);
1749 Ok(())
1750 }
1751
1752 #[tokio::test]
1753 async fn test_pool_status() -> Result<(), Box<dyn std::error::Error>> {
1754 let config = PoolConfigBuilder::new().max_size(10).min_idle(2).build()?;
1755 let factory = Arc::new(MockConnectionFactory);
1756 let pool = Pool::new(config, factory)?;
1757
1758 let status = pool.status().await;
1759 assert_eq!(status.max, 10);
1760 assert_eq!(status.min, 2);
1761 assert_eq!(status.active, 0);
1762 Ok(())
1763 }
1764
1765 #[tokio::test]
1766 async fn test_pool_close_all() -> Result<(), Box<dyn std::error::Error>> {
1767 let config = PoolConfigBuilder::new().max_size(5).build()?;
1768 let factory = Arc::new(MockConnectionFactory);
1769 let pool = Pool::new(config, factory)?;
1770
1771 let conn1 = pool.acquire().await?;
1773 let conn2 = pool.acquire().await?;
1774 pool.release(conn1).await;
1775 pool.release(conn2).await;
1776
1777 pool.close_all().await;
1778 let status = pool.status().await;
1779 assert_eq!(status.idle, 0);
1780 assert_eq!(status.active, 0);
1781 Ok(())
1782 }
1783
1784 #[tokio::test]
1785 async fn test_pool_reap_idle() -> Result<(), Box<dyn std::error::Error>> {
1786 let config = PoolConfigBuilder::new()
1787 .max_size(5)
1788 .idle_timeout(0) .build()?;
1790 let factory = Arc::new(MockConnectionFactory);
1791 let pool = Pool::new(config, factory)?;
1792
1793 let conn = pool.acquire().await?;
1794 pool.release(conn).await;
1795
1796 tokio::time::sleep(tokio::time::Duration::from_millis(10)).await;
1798
1799 pool.reap_idle().await;
1800 let status = pool.status().await;
1801 assert_eq!(status.idle, 0);
1802 Ok(())
1803 }
1804
1805 #[tokio::test]
1811 async fn test_h7_acquire_timeout_default_30s() {
1812 let config = PoolConfig::default();
1813 assert_eq!(
1814 config.acquire_timeout,
1815 Duration::from_secs(30),
1816 "H-7: acquire_timeout 默认应为 30s"
1817 );
1818 }
1819
1820 #[tokio::test]
1822 async fn test_h7_acquire_timeout_configurable() -> Result<(), Box<dyn std::error::Error>> {
1823 let config = PoolConfigBuilder::new()
1824 .max_size(1)
1825 .acquire_timeout(5) .build()?;
1827 assert_eq!(config.acquire_timeout, Duration::from_secs(5));
1828
1829 let factory = Arc::new(MockConnectionFactory);
1831 let pool = Pool::new(config, factory)?;
1832 let _conn1 = pool.acquire().await?;
1833
1834 let fast_config = PoolConfigBuilder::new()
1836 .max_size(1)
1837 .acquire_timeout(0) .build()?;
1839 let fast_pool = Pool::new(fast_config, Arc::new(MockConnectionFactory))?;
1842 let _fast_conn = fast_pool.acquire().await?; let result = fast_pool.acquire().await;
1844 assert!(
1845 matches!(result, Err(PoolError::Timeout)),
1846 "H-7: 应返回 Timeout"
1847 );
1848 Ok(())
1849 }
1850
1851 #[tokio::test]
1854 async fn test_m7_health_check_removes_nothing_when_all_healthy(
1855 ) -> Result<(), Box<dyn std::error::Error>> {
1856 let config = PoolConfigBuilder::new().max_size(5).build()?;
1858 let factory = Arc::new(MockConnectionFactory);
1859 let pool = Pool::new(config, factory)?;
1860
1861 let conn1 = pool.acquire().await?;
1863 let conn2 = pool.acquire().await?;
1864 let conn3 = pool.acquire().await?;
1865 pool.release(conn1).await;
1866 pool.release(conn2).await;
1867 pool.release(conn3).await;
1868
1869 let removed = pool.health_check().await;
1870 assert_eq!(removed, 0, "Healthy connections should not be removed");
1871
1872 let status = pool.status().await;
1873 assert_eq!(status.idle, 3);
1874 assert_eq!(status.active, 3);
1875 Ok(())
1876 }
1877
1878 #[tokio::test]
1879 async fn test_m7_health_check_returns_zero_for_empty_pool(
1880 ) -> Result<(), Box<dyn std::error::Error>> {
1881 let config = PoolConfigBuilder::new().max_size(5).build()?;
1882 let factory = Arc::new(MockConnectionFactory);
1883 let pool = Pool::new(config, factory)?;
1884
1885 let removed = pool.health_check().await;
1886 assert_eq!(removed, 0);
1887 Ok(())
1888 }
1889
1890 struct CountingFactory {
1894 count: AtomicU32,
1895 }
1896
1897 impl CountingFactory {
1898 fn new() -> Self {
1899 Self {
1900 count: AtomicU32::new(0),
1901 }
1902 }
1903 fn created_count(&self) -> u32 {
1904 self.count.load(Ordering::SeqCst)
1905 }
1906 }
1907
1908 #[async_trait]
1909 impl ConnectionFactory for CountingFactory {
1910 async fn create(&self) -> Result<Box<dyn Connection>, crate::DbError> {
1911 self.count.fetch_add(1, Ordering::SeqCst);
1912 Ok(Box::new(MockConnection::new()))
1913 }
1914 }
1915
1916 #[tokio::test]
1922 async fn test_production_bug_max_lifetime_never_expires(
1923 ) -> Result<(), Box<dyn std::error::Error>> {
1924 let config = PoolConfig {
1927 max_size: 5,
1928 min_idle: 0,
1929 acquire_timeout: Duration::from_secs(30),
1930 idle_timeout: Duration::from_secs(600),
1931 max_lifetime: Duration::from_millis(100), connection_timeout: Duration::from_secs(10),
1933 tls: None,
1934 query_timeout: None,
1935 max_rows: None,
1936 memory_limit: None,
1937 on_event: None,
1938 test_before_acquire: false,
1939 prewarm: false,
1940 };
1941 let factory = Arc::new(CountingFactory::new());
1942 let pool = Pool::new(config, factory.clone())?;
1943
1944 let conn = pool.acquire().await?;
1946 assert_eq!(factory.created_count(), 1, "应创建 1 个连接");
1947
1948 pool.release(conn).await;
1950
1951 tokio::time::sleep(Duration::from_millis(150)).await;
1953
1954 let conn2 = pool.acquire().await?;
1956
1957 assert_eq!(
1960 factory.created_count(),
1961 2,
1962 "超过 max_lifetime 后应创建新连接(旧连接应被回收)"
1963 );
1964
1965 pool.release(conn2).await;
1966 Ok(())
1967 }
1968
1969 #[tokio::test]
1976 async fn test_drop_auto_release_connection() -> Result<(), Box<dyn std::error::Error>> {
1977 let config = PoolConfigBuilder::new().max_size(2).build()?;
1978 let factory = Arc::new(CountingFactory::new());
1979 let pool = Pool::new(config, factory.clone())?;
1980
1981 {
1983 let _conn = pool.acquire().await?;
1984 assert_eq!(factory.created_count(), 1, "应创建 1 个连接");
1985 let status = pool.status().await;
1986 assert_eq!(status.active, 1, "active 应为 1");
1987 assert_eq!(status.idle, 0, "idle 应为 0");
1988 }
1990
1991 tokio::time::sleep(Duration::from_millis(50)).await;
1993
1994 let status = pool.status().await;
1996 assert_eq!(status.idle, 1, "Drop 后连接应自动归还,idle 应为 1");
1997 assert_eq!(status.active, 1, "total_count 应为 1");
1998 assert_eq!(factory.created_count(), 1, "应复用归还的连接,不创建新连接");
1999 Ok(())
2000 }
2001
2002 #[tokio::test]
2004 async fn test_drop_auto_release_then_reuse() -> Result<(), Box<dyn std::error::Error>> {
2005 let config = PoolConfigBuilder::new().max_size(1).build()?;
2006 let factory = Arc::new(CountingFactory::new());
2007 let pool = Pool::new(config, factory.clone())?;
2008
2009 {
2011 let _conn = pool.acquire().await?;
2012 }
2013
2014 tokio::time::sleep(Duration::from_millis(50)).await;
2016
2017 let conn = pool.acquire().await?;
2019 assert_eq!(factory.created_count(), 1, "应复用归还的连接,不创建新连接");
2020
2021 pool.release(conn).await;
2022 Ok(())
2023 }
2024
2025 #[tokio::test]
2027 async fn test_into_inner_does_not_return_to_pool() -> Result<(), Box<dyn std::error::Error>> {
2028 let config = PoolConfigBuilder::new().max_size(2).build()?;
2029 let factory = Arc::new(CountingFactory::new());
2030 let pool = Pool::new(config, factory.clone())?;
2031
2032 let conn = pool.acquire().await?;
2033 assert_eq!(factory.created_count(), 1);
2034
2035 let _raw_conn = conn.into_inner();
2037
2038 tokio::time::sleep(Duration::from_millis(50)).await;
2040
2041 let status = pool.status().await;
2042 assert_eq!(status.idle, 0, "into_inner 后连接不应归还");
2043 assert_eq!(status.active, 1, "total_count 仍为 1(连接被外部持有)");
2044 Ok(())
2045 }
2046
2047 #[tokio::test]
2049 async fn test_explicit_release_no_double_return() -> Result<(), Box<dyn std::error::Error>> {
2050 let config = PoolConfigBuilder::new().max_size(2).build()?;
2051 let factory = Arc::new(CountingFactory::new());
2052 let pool = Pool::new(config, factory.clone())?;
2053
2054 let conn = pool.acquire().await?;
2055 pool.release(conn).await;
2056
2057 let status = pool.status().await;
2058 assert_eq!(status.idle, 1, "release 后 idle 应为 1");
2059
2060 let conn = pool.acquire().await?;
2062 pool.release(conn).await;
2063
2064 let status = pool.status().await;
2065 assert_eq!(status.idle, 1, "再次 release 后 idle 仍应为 1(不重复)");
2066 assert_eq!(status.active, 1, "total_count 应为 1");
2067 Ok(())
2068 }
2069
2070 struct CursorMockConn {
2076 rows: QueryRows,
2077 call_count: usize,
2078 }
2079
2080 impl CursorMockConn {
2081 fn new(rows: QueryRows) -> Self {
2082 Self {
2083 rows,
2084 call_count: 0,
2085 }
2086 }
2087 }
2088
2089 impl Connection for CursorMockConn {
2090 fn execute<'a>(
2091 &'a mut self,
2092 _sql: &'a str,
2093 ) -> Pin<Box<dyn Future<Output = Result<u64, crate::DbError>> + Send + 'a>> {
2094 Box::pin(async move { Ok(1) })
2095 }
2096
2097 fn query<'a>(
2098 &'a mut self,
2099 _sql: &'a str,
2100 ) -> Pin<Box<dyn Future<Output = Result<QueryRows, crate::DbError>> + Send + 'a>> {
2101 Box::pin(async move {
2102 self.call_count += 1;
2103 Ok(self.rows.clone())
2104 })
2105 }
2106
2107 fn begin_transaction<'a>(
2108 &'a mut self,
2109 ) -> Pin<Box<dyn Future<Output = Result<(), crate::DbError>> + Send + 'a>> {
2110 Box::pin(async move { Ok(()) })
2111 }
2112
2113 fn commit<'a>(
2114 &'a mut self,
2115 ) -> Pin<Box<dyn Future<Output = Result<(), crate::DbError>> + Send + 'a>> {
2116 Box::pin(async move { Ok(()) })
2117 }
2118
2119 fn rollback<'a>(
2120 &'a mut self,
2121 ) -> Pin<Box<dyn Future<Output = Result<(), crate::DbError>> + Send + 'a>> {
2122 Box::pin(async move { Ok(()) })
2123 }
2124
2125 fn is_connected(&self) -> bool {
2126 true
2127 }
2128
2129 fn ping<'a>(&'a mut self) -> Pin<Box<dyn Future<Output = bool> + Send + 'a>> {
2130 Box::pin(async move { true })
2131 }
2132
2133 fn close<'a>(
2134 &'a mut self,
2135 ) -> Pin<Box<dyn Future<Output = Result<(), crate::DbError>> + Send + 'a>> {
2136 Box::pin(async move { Ok(()) })
2137 }
2138 }
2139
2140 struct CursorOverrideMockConn {
2142 rows: Vec<crate::value::Value>,
2143 yielded: usize,
2144 }
2145
2146 impl CursorOverrideMockConn {
2147 fn new(rows: Vec<crate::value::Value>) -> Self {
2148 Self { rows, yielded: 0 }
2149 }
2150 }
2151
2152 impl Connection for CursorOverrideMockConn {
2153 fn execute<'a>(
2154 &'a mut self,
2155 _sql: &'a str,
2156 ) -> Pin<Box<dyn Future<Output = Result<u64, crate::DbError>> + Send + 'a>> {
2157 Box::pin(async move { Ok(1) })
2158 }
2159
2160 fn query<'a>(
2161 &'a mut self,
2162 _sql: &'a str,
2163 ) -> Pin<Box<dyn Future<Output = Result<QueryRows, crate::DbError>> + Send + 'a>> {
2164 Box::pin(async move {
2166 Ok(self
2167 .rows
2168 .iter()
2169 .map(|v| {
2170 let mut m = std::collections::HashMap::new();
2171 m.insert("v".to_string(), v.clone());
2172 m
2173 })
2174 .collect())
2175 })
2176 }
2177
2178 fn query_stream<'a>(
2180 &'a mut self,
2181 _sql: &'a str,
2182 ) -> Pin<Box<dyn futures::Stream<Item = QueryStreamItem> + Send + 'a>> {
2183 Box::pin(futures::stream::iter(
2184 self.rows
2185 .iter()
2186 .enumerate()
2187 .map(|(i, v)| {
2188 self.yielded = i + 1;
2189 let mut m = std::collections::HashMap::new();
2190 m.insert("v".to_string(), v.clone());
2191 Ok(m)
2192 })
2193 .collect::<Vec<_>>(),
2194 ))
2195 }
2196
2197 fn begin_transaction<'a>(
2198 &'a mut self,
2199 ) -> Pin<Box<dyn Future<Output = Result<(), crate::DbError>> + Send + 'a>> {
2200 Box::pin(async move { Ok(()) })
2201 }
2202
2203 fn commit<'a>(
2204 &'a mut self,
2205 ) -> Pin<Box<dyn Future<Output = Result<(), crate::DbError>> + Send + 'a>> {
2206 Box::pin(async move { Ok(()) })
2207 }
2208
2209 fn rollback<'a>(
2210 &'a mut self,
2211 ) -> Pin<Box<dyn Future<Output = Result<(), crate::DbError>> + Send + 'a>> {
2212 Box::pin(async move { Ok(()) })
2213 }
2214
2215 fn is_connected(&self) -> bool {
2216 true
2217 }
2218
2219 fn ping<'a>(&'a mut self) -> Pin<Box<dyn Future<Output = bool> + Send + 'a>> {
2220 Box::pin(async move { true })
2221 }
2222
2223 fn close<'a>(
2224 &'a mut self,
2225 ) -> Pin<Box<dyn Future<Output = Result<(), crate::DbError>> + Send + 'a>> {
2226 Box::pin(async move { Ok(()) })
2227 }
2228 }
2229
2230 #[tokio::test]
2232 async fn test_query_stream_default_impl_yields_all_rows() {
2233 use futures::StreamExt;
2234 let rows: QueryRows = vec![
2235 std::collections::HashMap::from([
2236 ("id".to_string(), crate::value::Value::I64(1)),
2237 (
2238 "name".to_string(),
2239 crate::value::Value::String("alice".to_string()),
2240 ),
2241 ]),
2242 std::collections::HashMap::from([
2243 ("id".to_string(), crate::value::Value::I64(2)),
2244 (
2245 "name".to_string(),
2246 crate::value::Value::String("bob".to_string()),
2247 ),
2248 ]),
2249 std::collections::HashMap::from([
2250 ("id".to_string(), crate::value::Value::I64(3)),
2251 (
2252 "name".to_string(),
2253 crate::value::Value::String("carol".to_string()),
2254 ),
2255 ]),
2256 ];
2257 let mut conn = CursorMockConn::new(rows);
2258 let mut stream = conn.query_stream("SELECT id, name FROM users");
2259 let mut received: Vec<QueryStreamItem> = Vec::new();
2260 while let Some(item) = stream.next().await {
2261 received.push(item);
2262 }
2263 assert_eq!(received.len(), 3, "应收到 3 行");
2264 assert!(received.iter().all(|r| r.is_ok()), "所有项应为 Ok");
2265 drop(stream);
2266 assert_eq!(conn.call_count, 1, "默认实现应调用 query() 一次");
2267 }
2268
2269 #[tokio::test]
2271 async fn test_query_stream_default_empty_result() {
2272 use futures::StreamExt;
2273 let mut conn = CursorMockConn::new(Vec::new());
2274 let mut stream = conn.query_stream("SELECT * FROM empty_table");
2275 let mut count = 0;
2276 while let Some(_item) = stream.next().await {
2277 count += 1;
2278 }
2279 assert_eq!(count, 0, "空结果集应产生 0 项");
2280 }
2281
2282 #[tokio::test]
2284 async fn test_query_stream_default_error_propagation() {
2285 use futures::StreamExt;
2286 struct ErrorMockConn;
2288 impl Connection for ErrorMockConn {
2289 fn execute<'a>(
2290 &'a mut self,
2291 _sql: &'a str,
2292 ) -> Pin<Box<dyn Future<Output = Result<u64, crate::DbError>> + Send + 'a>>
2293 {
2294 Box::pin(async move { Ok(1) })
2295 }
2296 fn query<'a>(
2297 &'a mut self,
2298 _sql: &'a str,
2299 ) -> Pin<Box<dyn Future<Output = Result<QueryRows, crate::DbError>> + Send + 'a>>
2300 {
2301 Box::pin(async move { Err(crate::DbError::Internal("query failed".to_string())) })
2302 }
2303 fn begin_transaction<'a>(
2304 &'a mut self,
2305 ) -> Pin<Box<dyn Future<Output = Result<(), crate::DbError>> + Send + 'a>> {
2306 Box::pin(async move { Ok(()) })
2307 }
2308 fn commit<'a>(
2309 &'a mut self,
2310 ) -> Pin<Box<dyn Future<Output = Result<(), crate::DbError>> + Send + 'a>> {
2311 Box::pin(async move { Ok(()) })
2312 }
2313 fn rollback<'a>(
2314 &'a mut self,
2315 ) -> Pin<Box<dyn Future<Output = Result<(), crate::DbError>> + Send + 'a>> {
2316 Box::pin(async move { Ok(()) })
2317 }
2318 fn is_connected(&self) -> bool {
2319 true
2320 }
2321 fn ping<'a>(&'a mut self) -> Pin<Box<dyn Future<Output = bool> + Send + 'a>> {
2322 Box::pin(async move { true })
2323 }
2324 fn close<'a>(
2325 &'a mut self,
2326 ) -> Pin<Box<dyn Future<Output = Result<(), crate::DbError>> + Send + 'a>> {
2327 Box::pin(async move { Ok(()) })
2328 }
2329 }
2330 let mut conn = ErrorMockConn;
2331 let mut stream = conn.query_stream("SELECT * FROM bad_table");
2332 let item = stream.next().await;
2333 assert!(item.is_some(), "应产生一项");
2334 assert!(item.unwrap().is_err(), "该项应为 Err");
2335 }
2336
2337 #[tokio::test]
2339 async fn test_query_stream_override_yields_rows_one_by_one() {
2340 use futures::StreamExt;
2341 let rows = vec![
2342 crate::value::Value::I64(10),
2343 crate::value::Value::I64(20),
2344 crate::value::Value::I64(30),
2345 crate::value::Value::I64(40),
2346 crate::value::Value::I64(50),
2347 ];
2348 let mut conn = CursorOverrideMockConn::new(rows);
2349 let values: Vec<i64> = {
2350 let mut stream = conn.query_stream("SELECT v FROM seq");
2351 let mut vals: Vec<i64> = Vec::new();
2352 while let Some(Ok(row)) = stream.next().await {
2353 if let crate::value::Value::I64(v) = row.get("v").unwrap() {
2354 vals.push(*v);
2355 }
2356 }
2357 vals
2358 };
2359 assert_eq!(values, vec![10, 20, 30, 40, 50], "应按顺序收到全部 5 行");
2360 assert_eq!(conn.yielded, 5, "应逐行 yield 5 次(真游标覆盖)");
2361 }
2362
2363 #[tokio::test]
2365 async fn test_query_stream_override_early_drop() {
2366 use futures::StreamExt;
2367 let rows = vec![
2368 crate::value::Value::I64(1),
2369 crate::value::Value::I64(2),
2370 crate::value::Value::I64(3),
2371 ];
2372 let mut conn = CursorOverrideMockConn::new(rows);
2373 {
2374 let mut stream = conn.query_stream("SELECT v FROM seq");
2375 let first = stream.next().await;
2376 assert!(first.is_some(), "第一项应存在");
2377 drop(stream);
2379 }
2380 assert!(conn.is_connected(), "提前 drop 流后连接仍应可用");
2382 }
2383
2384 #[tokio::test]
2386 async fn test_pool_prewarm() -> Result<(), Box<dyn std::error::Error>> {
2387 use std::sync::atomic::AtomicU32;
2388
2389 let create_count = Arc::new(AtomicU32::new(0));
2391 let create_count_clone = create_count.clone();
2392
2393 struct CountingFactory {
2394 count: Arc<AtomicU32>,
2395 }
2396
2397 #[async_trait]
2398 impl ConnectionFactory for CountingFactory {
2399 async fn create(&self) -> Result<Box<dyn Connection>, crate::DbError> {
2400 self.count.fetch_add(1, Ordering::SeqCst);
2401 Ok(Box::new(MockConnection::new()))
2402 }
2403 }
2404
2405 let config = PoolConfigBuilder::new()
2407 .max_size(10)
2408 .min_idle(5)
2409 .prewarm(true)
2410 .build()?;
2411
2412 let factory = Arc::new(CountingFactory {
2413 count: create_count_clone,
2414 });
2415
2416 let pool = Pool::new(config, factory)?;
2417
2418 let status_before = pool.status().await;
2420 assert_eq!(status_before.idle, 0, "预热前 idle 应为 0");
2421
2422 pool.prewarm().await;
2424
2425 let status_after = pool.status().await;
2427 assert!(
2428 status_after.idle >= 5,
2429 "预热后 idle 应 >= 5,实际: {}",
2430 status_after.idle
2431 );
2432
2433 assert_eq!(
2435 create_count.load(Ordering::SeqCst),
2436 5,
2437 "工厂应被调用 5 次(min_idle)"
2438 );
2439
2440 Ok(())
2441 }
2442
2443 #[tokio::test]
2445 async fn test_pool_prewarm_failure_non_blocking() -> Result<(), Box<dyn std::error::Error>> {
2446 use std::sync::atomic::AtomicBool;
2447
2448 struct FailingFactory {
2449 failed: Arc<AtomicBool>,
2450 }
2451
2452 #[async_trait]
2453 impl ConnectionFactory for FailingFactory {
2454 async fn create(&self) -> Result<Box<dyn Connection>, crate::DbError> {
2455 self.failed.store(true, Ordering::SeqCst);
2456 Err(crate::DbError::Internal(
2458 "simulated connection failure".to_string(),
2459 ))
2460 }
2461 }
2462
2463 let failed = Arc::new(AtomicBool::new(false));
2464 let mut config = PoolConfigBuilder::new()
2465 .max_size(10)
2466 .min_idle(3)
2467 .prewarm(true)
2468 .build()?;
2469 config.connection_timeout = std::time::Duration::from_secs(1); let factory = Arc::new(FailingFactory {
2472 failed: failed.clone(),
2473 });
2474
2475 let pool = Pool::new(config, factory)?;
2477 pool.prewarm().await; assert!(failed.load(Ordering::SeqCst), "工厂应被调用且失败");
2481
2482 let status = pool.status().await;
2484 assert_eq!(status.max, 10, "池配置应正常");
2485
2486 Ok(())
2487 }
2488
2489 #[tokio::test]
2491 async fn test_pool_prewarm_disabled() -> Result<(), Box<dyn std::error::Error>> {
2492 use std::sync::atomic::AtomicU32;
2493
2494 let create_count = Arc::new(AtomicU32::new(0));
2495 let create_count_clone = create_count.clone();
2496
2497 struct CountingFactory {
2498 count: Arc<AtomicU32>,
2499 }
2500
2501 #[async_trait]
2502 impl ConnectionFactory for CountingFactory {
2503 async fn create(&self) -> Result<Box<dyn Connection>, crate::DbError> {
2504 self.count.fetch_add(1, Ordering::SeqCst);
2505 Ok(Box::new(MockConnection::new()))
2506 }
2507 }
2508
2509 let config = PoolConfigBuilder::new()
2511 .max_size(10)
2512 .min_idle(5)
2513 .prewarm(false) .build()?;
2515
2516 let factory = Arc::new(CountingFactory {
2517 count: create_count_clone,
2518 });
2519
2520 let pool = Pool::new(config, factory)?;
2521 pool.prewarm().await; assert_eq!(
2525 create_count.load(Ordering::SeqCst),
2526 0,
2527 "prewarm=false 时工厂不应被调用"
2528 );
2529
2530 let status = pool.status().await;
2531 assert_eq!(status.idle, 0, "idle 应为 0");
2532
2533 Ok(())
2534 }
2535}