1use async_trait::async_trait;
6use crossbeam_queue::ArrayQueue;
7#[cfg(feature = "circuit-breaker")]
11use parking_lot::Mutex as PlMutex;
12#[cfg(feature = "rate-limit")]
13use parking_lot::RwLock as PlRwLock;
14use std::future::Future;
15use std::ops::{Deref, DerefMut};
16use std::pin::Pin;
17use std::sync::atomic::{AtomicBool, AtomicU32, Ordering};
18use std::sync::Arc;
19use std::time::{Duration, Instant};
20use tokio::sync::Notify;
21
22#[cfg(feature = "circuit-breaker")]
26use crate::circuit_breaker::{CircuitBreaker, CircuitState, DefaultCircuitBreaker};
27#[cfg(feature = "rate-limit")]
28use crate::rate_limiter::RateLimiter;
29use sz_orm_model::PoolError;
30
31pub use sz_orm_model::QueryRows;
33
34pub type QueryStreamItem =
36 Result<std::collections::HashMap<String, sz_orm_model::Value>, sz_orm_model::DbError>;
37
38pub trait Connection: Send + Sync {
45 fn execute<'a>(
46 &'a mut self,
47 sql: &'a str,
48 ) -> Pin<Box<dyn Future<Output = Result<u64, sz_orm_model::DbError>> + Send + 'a>>;
49 fn query<'a>(
50 &'a mut self,
51 sql: &'a str,
52 ) -> Pin<Box<dyn Future<Output = Result<QueryRows, sz_orm_model::DbError>> + Send + 'a>>;
53 fn begin_transaction<'a>(
54 &'a mut self,
55 ) -> Pin<Box<dyn Future<Output = Result<(), sz_orm_model::DbError>> + Send + 'a>>;
56 fn commit<'a>(
57 &'a mut self,
58 ) -> Pin<Box<dyn Future<Output = Result<(), sz_orm_model::DbError>> + Send + 'a>>;
59 fn rollback<'a>(
60 &'a mut self,
61 ) -> Pin<Box<dyn Future<Output = Result<(), sz_orm_model::DbError>> + Send + 'a>>;
62 fn is_connected(&self) -> bool;
63
64 fn in_transaction(&self) -> bool {
70 false
71 }
72
73 fn ping<'a>(&'a mut self) -> Pin<Box<dyn Future<Output = bool> + Send + 'a>>;
74 fn close<'a>(
75 &'a mut self,
76 ) -> Pin<Box<dyn Future<Output = Result<(), sz_orm_model::DbError>> + Send + 'a>>;
77
78 fn execute_with_params<'a>(
84 &'a mut self,
85 sql: &'a str,
86 params: &'a [sz_orm_model::Value],
87 ) -> Pin<Box<dyn Future<Output = Result<u64, sz_orm_model::DbError>> + Send + 'a>> {
88 let _ = (sql, params);
89 Box::pin(async move {
90 Err(sz_orm_model::DbError::Internal(
91 "execute_with_params not implemented for this adapter".to_string(),
92 ))
93 })
94 }
95
96 fn query_with_params<'a>(
102 &'a mut self,
103 sql: &'a str,
104 params: &'a [sz_orm_model::Value],
105 ) -> Pin<Box<dyn Future<Output = Result<QueryRows, sz_orm_model::DbError>> + Send + 'a>> {
106 let _ = (sql, params);
107 Box::pin(async move {
108 Err(sz_orm_model::DbError::Internal(
109 "query_with_params not implemented for this adapter".to_string(),
110 ))
111 })
112 }
113
114 fn query_values<'a>(
119 &'a mut self,
120 sql: &'a str,
121 ) -> Pin<
122 Box<
123 dyn Future<Output = Result<sz_orm_model::QueryValues, sz_orm_model::DbError>>
124 + Send
125 + 'a,
126 >,
127 > {
128 let _ = sql;
129 Box::pin(async move {
130 Err(sz_orm_model::DbError::Internal(
131 "query_values not implemented for this adapter".to_string(),
132 ))
133 })
134 }
135
136 fn query_values_with_params<'a>(
140 &'a mut self,
141 sql: &'a str,
142 params: &'a [sz_orm_model::Value],
143 ) -> Pin<
144 Box<
145 dyn Future<Output = Result<sz_orm_model::QueryValues, sz_orm_model::DbError>>
146 + Send
147 + 'a,
148 >,
149 > {
150 let _ = (sql, params);
151 Box::pin(async move {
152 Err(sz_orm_model::DbError::Internal(
153 "query_values_with_params not implemented for this adapter".to_string(),
154 ))
155 })
156 }
157
158 fn query_stream<'a>(
163 &'a mut self,
164 sql: &'a str,
165 ) -> Pin<Box<dyn futures::Stream<Item = QueryStreamItem> + Send + 'a>> {
166 let _ = sql;
167 Box::pin(futures::stream::empty())
168 }
169
170 fn execute_batch<'a>(
175 &'a mut self,
176 sqls: &'a [String],
177 ) -> Pin<Box<dyn Future<Output = Result<u64, sz_orm_model::DbError>> + Send + 'a>> {
178 Box::pin(async move {
179 let mut total = 0u64;
180 for sql in sqls {
181 total += self.execute(sql).await?;
182 }
183 Ok(total)
184 })
185 }
186
187 fn execute_batch_params<'a>(
192 &'a mut self,
193 sql: &'a str,
194 params_batch: &'a [Vec<sz_orm_model::Value>],
195 ) -> Pin<Box<dyn Future<Output = Result<u64, sz_orm_model::DbError>> + Send + 'a>> {
196 Box::pin(async move {
197 let mut total = 0u64;
198 for params in params_batch {
199 total += self.execute_with_params(sql, params).await?;
200 }
201 Ok(total)
202 })
203 }
204}
205
206pub struct PooledConnection {
214 conn: Box<dyn Connection>,
215 created_at: Instant,
216 last_used_at: Instant,
217 pool: Option<Pool>,
218}
219
220impl PooledConnection {
221 fn new(conn: Box<dyn Connection>, pool: Pool) -> Self {
222 let now = Instant::now();
223 Self {
224 conn,
225 created_at: now,
226 last_used_at: now,
227 pool: Some(pool),
228 }
229 }
230
231 fn is_expired(&self, max_lifetime: Duration) -> bool {
232 self.created_at.elapsed() >= max_lifetime
233 }
234
235 fn is_idle_too_long(&self, idle_timeout: Duration) -> bool {
236 self.last_used_at.elapsed() >= idle_timeout
237 }
238
239 pub fn created_at(&self) -> Instant {
241 self.created_at
242 }
243
244 pub fn into_inner(mut self) -> Box<dyn Connection> {
249 self.pool = None; std::mem::replace(&mut self.conn, Box::new(ClosedConnection))
253 }
254}
255
256impl Drop for PooledConnection {
268 fn drop(&mut self) {
269 if let Some(pool) = self.pool.take() {
270 let conn = std::mem::replace(&mut self.conn, Box::new(ClosedConnection));
272 let pooled = PooledConnection {
273 conn,
274 created_at: self.created_at,
275 last_used_at: self.last_used_at,
276 pool: None,
277 };
278 if let Ok(handle) = tokio::runtime::Handle::try_current() {
280 handle.spawn(async move {
281 pool.release(pooled).await;
282 });
283 } else {
284 drop(pooled);
288 pool.total_count.fetch_sub(1, Ordering::SeqCst);
289 }
290 }
291 }
292}
293
294struct ClosedConnection;
298
299impl Connection for ClosedConnection {
300 fn execute<'a>(
301 &'a mut self,
302 _sql: &'a str,
303 ) -> Pin<Box<dyn Future<Output = Result<u64, sz_orm_model::DbError>> + Send + 'a>> {
304 Box::pin(async {
305 Err(sz_orm_model::DbError::ConnectionError(
306 "connection already returned to pool".to_string(),
307 ))
308 })
309 }
310
311 fn query<'a>(
312 &'a mut self,
313 _sql: &'a str,
314 ) -> Pin<Box<dyn Future<Output = Result<QueryRows, sz_orm_model::DbError>> + Send + 'a>> {
315 Box::pin(async {
316 Err(sz_orm_model::DbError::ConnectionError(
317 "connection already returned to pool".to_string(),
318 ))
319 })
320 }
321
322 fn begin_transaction<'a>(
323 &'a mut self,
324 ) -> Pin<Box<dyn Future<Output = Result<(), sz_orm_model::DbError>> + Send + 'a>> {
325 Box::pin(async {
326 Err(sz_orm_model::DbError::ConnectionError(
327 "connection already returned to pool".to_string(),
328 ))
329 })
330 }
331
332 fn commit<'a>(
333 &'a mut self,
334 ) -> Pin<Box<dyn Future<Output = Result<(), sz_orm_model::DbError>> + Send + 'a>> {
335 Box::pin(async { Ok(()) })
336 }
337
338 fn rollback<'a>(
339 &'a mut self,
340 ) -> Pin<Box<dyn Future<Output = Result<(), sz_orm_model::DbError>> + Send + 'a>> {
341 Box::pin(async { Ok(()) })
342 }
343
344 fn is_connected(&self) -> bool {
345 false
346 }
347
348 fn ping<'a>(&'a mut self) -> Pin<Box<dyn Future<Output = bool> + Send + 'a>> {
349 Box::pin(async { false })
350 }
351
352 fn close<'a>(
353 &'a mut self,
354 ) -> Pin<Box<dyn Future<Output = Result<(), sz_orm_model::DbError>> + Send + 'a>> {
355 Box::pin(async { Ok(()) })
356 }
357}
358
359impl Deref for PooledConnection {
360 type Target = dyn Connection;
361
362 fn deref(&self) -> &Self::Target {
363 self.conn.as_ref()
364 }
365}
366
367impl DerefMut for PooledConnection {
368 fn deref_mut(&mut self) -> &mut Self::Target {
369 self.conn.as_mut()
370 }
371}
372
373#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
375pub enum TlsVersion {
376 #[default]
377 Tls12,
378 Tls13,
379}
380
381#[derive(Debug, Clone, Default)]
383pub struct TlsConfig {
384 pub enabled: bool,
386 pub ca_cert_path: Option<String>,
388 pub client_cert_path: Option<String>,
390 pub client_key_path: Option<String>,
392 pub min_version: TlsVersion,
394}
395
396#[derive(Debug, Clone)]
398pub enum PoolEvent {
399 ConnectionCreated,
401 ConnectionClosed,
403 ConnectionAcquired,
405 ConnectionReleased,
407 AcquireTimeout,
409}
410
411pub type PoolEventCallback = Arc<dyn Fn(PoolEvent) + Send + Sync>;
413
414pub struct PoolConfig {
415 pub max_size: u32,
416 pub min_idle: u32,
417 pub acquire_timeout: Duration,
418 pub idle_timeout: Duration,
419 pub max_lifetime: Duration,
420 pub connection_timeout: Duration,
421 pub tls: Option<TlsConfig>,
423 pub query_timeout: Option<Duration>,
425 pub max_rows: Option<usize>,
427 pub memory_limit: Option<usize>,
429 pub on_event: Option<PoolEventCallback>,
431}
432
433impl Default for PoolConfig {
434 fn default() -> Self {
435 Self {
436 max_size: 100,
437 min_idle: 0,
438 acquire_timeout: Duration::from_secs(30),
439 idle_timeout: Duration::from_secs(600),
440 max_lifetime: Duration::from_secs(1800),
441 connection_timeout: Duration::from_secs(10),
442 tls: None,
443 query_timeout: Some(Duration::from_secs(30)),
444 max_rows: None,
445 memory_limit: None,
446 on_event: None,
447 }
448 }
449}
450
451impl Clone for PoolConfig {
452 fn clone(&self) -> Self {
453 Self {
454 max_size: self.max_size,
455 min_idle: self.min_idle,
456 acquire_timeout: self.acquire_timeout,
457 idle_timeout: self.idle_timeout,
458 max_lifetime: self.max_lifetime,
459 connection_timeout: self.connection_timeout,
460 tls: self.tls.clone(),
461 query_timeout: self.query_timeout,
462 max_rows: self.max_rows,
463 memory_limit: self.memory_limit,
464 on_event: self.on_event.clone(),
465 }
466 }
467}
468
469impl PoolConfig {
470 pub fn validate(&self) -> Result<(), PoolError> {
472 if self.max_size == 0 {
473 return Err(PoolError::InvalidConfig("max_size cannot be 0".to_string()));
474 }
475 if self.min_idle > self.max_size {
476 return Err(PoolError::InvalidConfig(
477 "min_idle cannot exceed max_size".to_string(),
478 ));
479 }
480 Ok(())
481 }
482}
483
484pub struct PoolStatus {
485 pub idle: u32,
486 pub active: u32,
487 pub max: u32,
488 pub min: u32,
489 pub waiters: u32,
491}
492
493impl std::fmt::Debug for PoolStatus {
494 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
495 f.debug_struct("PoolStatus")
496 .field("idle", &self.idle)
497 .field("active", &self.active)
498 .field("max", &self.max)
499 .field("min", &self.min)
500 .field("waiters", &self.waiters)
501 .finish()
502 }
503}
504
505pub struct PoolConfigBuilder {
506 config: PoolConfig,
507}
508
509impl PoolConfigBuilder {
510 pub fn new() -> Self {
511 Self {
512 config: PoolConfig::default(),
513 }
514 }
515
516 pub fn max_size(mut self, size: u32) -> Self {
517 self.config.max_size = size;
518 self
519 }
520
521 pub fn min_idle(mut self, count: u32) -> Self {
522 self.config.min_idle = count;
523 self
524 }
525
526 pub fn acquire_timeout(mut self, timeout_secs: u64) -> Self {
527 self.config.acquire_timeout = Duration::from_secs(timeout_secs);
528 self
529 }
530
531 pub fn idle_timeout(mut self, timeout_secs: u64) -> Self {
532 self.config.idle_timeout = Duration::from_secs(timeout_secs);
533 self
534 }
535
536 pub fn max_lifetime(mut self, lifetime_secs: u64) -> Self {
537 self.config.max_lifetime = Duration::from_secs(lifetime_secs);
538 self
539 }
540
541 pub fn tls(mut self, tls: TlsConfig) -> Self {
543 self.config.tls = Some(tls);
544 self
545 }
546
547 pub fn query_timeout(mut self, timeout: Duration) -> Self {
549 self.config.query_timeout = Some(timeout);
550 self
551 }
552
553 pub fn max_rows(mut self, max_rows: usize) -> Self {
555 self.config.max_rows = Some(max_rows);
556 self
557 }
558
559 pub fn memory_limit(mut self, memory_limit: usize) -> Self {
561 self.config.memory_limit = Some(memory_limit);
562 self
563 }
564
565 pub fn on_event(mut self, callback: PoolEventCallback) -> Self {
567 self.config.on_event = Some(callback);
568 self
569 }
570
571 pub fn build(self) -> Result<PoolConfig, PoolError> {
572 self.config.validate()?;
573 Ok(self.config)
574 }
575}
576
577impl Default for PoolConfigBuilder {
578 fn default() -> Self {
579 Self::new()
580 }
581}
582
583#[async_trait]
585pub trait ConnectionFactory: Send + Sync {
586 async fn create(&self) -> Result<Box<dyn Connection>, sz_orm_model::DbError>;
587}
588
589pub struct Pool {
595 config: PoolConfig,
596 factory: Arc<dyn ConnectionFactory>,
597 idle: Arc<ArrayQueue<PooledConnection>>,
603 total_count: Arc<AtomicU32>,
613 closed: Arc<AtomicBool>,
615 notify: Arc<Notify>,
616 waiters_count: Arc<AtomicU32>,
618 dynamic_max_size: Arc<AtomicU32>,
620 #[cfg(feature = "circuit-breaker")]
626 circuit_breaker: Arc<PlMutex<DefaultCircuitBreaker>>,
627 #[cfg(feature = "rate-limit")]
636 rate_limiter: Arc<PlRwLock<Option<Arc<dyn RateLimiter>>>>,
637 #[cfg(feature = "rate-limit")]
639 rate_limit_key: String,
640}
641
642impl Clone for Pool {
646 fn clone(&self) -> Self {
647 Self {
648 config: self.config.clone(),
649 factory: self.factory.clone(),
650 idle: self.idle.clone(),
651 total_count: self.total_count.clone(),
652 closed: self.closed.clone(),
653 notify: Arc::clone(&self.notify),
654 waiters_count: self.waiters_count.clone(),
655 dynamic_max_size: self.dynamic_max_size.clone(),
656 #[cfg(feature = "circuit-breaker")]
657 circuit_breaker: Arc::clone(&self.circuit_breaker),
658 #[cfg(feature = "rate-limit")]
659 rate_limiter: Arc::clone(&self.rate_limiter),
660 #[cfg(feature = "rate-limit")]
661 rate_limit_key: self.rate_limit_key.clone(),
662 }
663 }
664}
665
666impl Pool {
667 pub fn new(config: PoolConfig, factory: Arc<dyn ConnectionFactory>) -> Result<Self, PoolError> {
691 config.validate()?;
692 let max_size = config.max_size as usize;
695 let dynamic_max = config.max_size;
696 Ok(Self {
697 config,
698 factory,
699 idle: Arc::new(ArrayQueue::new(max_size)),
700 total_count: Arc::new(AtomicU32::new(0)),
701 closed: Arc::new(AtomicBool::new(false)),
702 notify: Arc::new(Notify::new()),
703 waiters_count: Arc::new(AtomicU32::new(0)),
704 dynamic_max_size: Arc::new(AtomicU32::new(dynamic_max)),
705 #[cfg(feature = "circuit-breaker")]
708 circuit_breaker: Arc::new(PlMutex::new(DefaultCircuitBreaker::new(
709 5,
710 std::time::Duration::from_secs(30),
711 ))),
712 #[cfg(feature = "rate-limit")]
715 rate_limiter: Arc::new(PlRwLock::new(None)),
716 #[cfg(feature = "rate-limit")]
717 rate_limit_key: "pool".to_string(),
718 })
719 }
720
721 pub fn config(&self) -> &PoolConfig {
723 &self.config
724 }
725
726 #[cfg(feature = "circuit-breaker")]
740 pub fn configure_circuit_breaker(
741 &self,
742 failure_threshold: usize,
743 reset_timeout: std::time::Duration,
744 ) {
745 let new_cb = DefaultCircuitBreaker::new(failure_threshold, reset_timeout);
746 let mut guard = self.circuit_breaker.lock();
748 *guard = new_cb;
749 }
750
751 #[cfg(feature = "circuit-breaker")]
756 pub fn reset_circuit_breaker(&self) -> bool {
757 let mut guard = self.circuit_breaker.lock();
759 guard.reset()
760 }
761
762 #[cfg(feature = "circuit-breaker")]
764 pub fn circuit_state(&self) -> CircuitState {
765 let guard = self.circuit_breaker.lock();
767 guard.state()
768 }
769
770 #[cfg(feature = "rate-limit")]
779 pub fn set_rate_limiter(&self, limiter: Option<Arc<dyn RateLimiter>>) {
780 let mut guard = self.rate_limiter.write();
782 *guard = limiter;
783 }
784
785 #[cfg(feature = "rate-limit")]
787 pub fn with_rate_limit_key(mut self, key: impl Into<String>) -> Self {
788 self.rate_limit_key = key.into();
789 self
790 }
791
792 fn emit_event(&self, event: PoolEvent) {
794 if let Some(ref callback) = self.config.on_event {
795 callback(event);
796 }
797 }
798
799 #[tracing::instrument(skip(self), fields(max_size = self.config.max_size, acquire_timeout = ?self.config.acquire_timeout))]
819 pub async fn acquire(&self) -> Result<PooledConnection, PoolError> {
820 if self.closed.load(Ordering::Acquire) {
822 return Err(PoolError::Closed);
823 }
824
825 #[cfg(feature = "circuit-breaker")]
829 {
830 let mut guard = self.circuit_breaker.lock();
831 if !guard.can_execute() {
832 return Err(PoolError::CircuitOpen);
833 }
834 }
835
836 #[cfg(feature = "rate-limit")]
840 {
841 let guard = self.rate_limiter.read();
842 if let Some(ref limiter) = *guard {
843 match limiter.try_acquire(&self.rate_limit_key) {
844 Ok(result) if !result.allowed => {
845 return Err(PoolError::RateLimited {
846 remaining: result.remaining,
847 reset_at: result.reset_at,
848 });
849 }
850 Ok(_) => {} Err(_) => {
852 }
854 }
855 }
856 }
857
858 let deadline = Instant::now() + self.config.acquire_timeout;
859 let mut backoff = Duration::from_millis(1);
861 const MAX_BACKOFF: Duration = Duration::from_millis(100);
863
864 loop {
865 let mut to_close: Vec<PooledConnection> = Vec::new();
871 let acquired: Option<PooledConnection> = {
872 let mut found: Option<PooledConnection> = None;
873 while let Some(pooled) = self.idle.pop() {
874 if pooled.is_expired(self.config.max_lifetime) {
876 to_close.push(pooled);
877 continue;
878 }
879 if pooled.is_idle_too_long(self.config.idle_timeout) {
881 to_close.push(pooled);
882 continue;
883 }
884 if !pooled.conn.is_connected() {
887 to_close.push(pooled);
888 continue;
889 }
890 found = Some(pooled);
891 break;
892 }
893 found
894 };
895
896 for mut pooled in to_close {
898 let _ = pooled.conn.close().await;
899 self.total_count.fetch_sub(1, Ordering::SeqCst);
901 }
902
903 if let Some(mut pooled) = acquired {
904 pooled.pool = Some(self.clone());
907 return Ok(pooled);
908 }
909
910 let current_max = self.dynamic_max_size.load(Ordering::Acquire);
915 let created = loop {
916 let current = self.total_count.load(Ordering::Acquire);
917 if current >= current_max {
918 break None; }
920 match self.total_count.compare_exchange(
921 current,
922 current + 1,
923 Ordering::SeqCst,
924 Ordering::Acquire,
925 ) {
926 Ok(_) => break Some(()), Err(_) => continue, }
929 };
930
931 if created.is_some() {
932 match tokio::time::timeout(self.config.connection_timeout, self.factory.create())
933 .await
934 {
935 Ok(Ok(conn)) => {
936 #[cfg(feature = "circuit-breaker")]
939 {
940 self.circuit_breaker.lock().record_success();
941 }
942 self.emit_event(PoolEvent::ConnectionCreated);
943 self.emit_event(PoolEvent::ConnectionAcquired);
944 return Ok(PooledConnection::new(conn, self.clone()));
945 }
946 Ok(Err(e)) => {
947 self.total_count.fetch_sub(1, Ordering::SeqCst);
949 #[cfg(feature = "circuit-breaker")]
952 {
953 self.circuit_breaker.lock().record_failure();
954 }
955 return Err(PoolError::ConnectionFailed(e.to_string()));
956 }
957 Err(_) => {
958 self.total_count.fetch_sub(1, Ordering::SeqCst);
960 #[cfg(feature = "circuit-breaker")]
963 {
964 self.circuit_breaker.lock().record_failure();
965 }
966 return Err(PoolError::Timeout);
967 }
968 }
969 }
970
971 let now = Instant::now();
973 if now >= deadline {
974 self.emit_event(PoolEvent::AcquireTimeout);
975 return Err(PoolError::Timeout);
976 }
977 self.waiters_count.fetch_add(1, Ordering::SeqCst);
979 let wait = std::cmp::min(backoff, deadline - now);
980 match tokio::time::timeout(wait, self.notify.notified()).await {
981 Ok(()) => {
982 backoff = Duration::from_millis(1);
984 }
985 Err(_) => {
986 backoff = std::cmp::min(backoff * 2, MAX_BACKOFF);
988 }
989 }
990 self.waiters_count.fetch_sub(1, Ordering::SeqCst);
992 }
993 }
994
995 #[tracing::instrument(skip(self, pooled))]
1003 pub async fn release(&self, mut pooled: PooledConnection) {
1004 pooled.pool = None;
1006
1007 if self.closed.load(Ordering::Acquire) {
1009 let _ = pooled.conn.close().await;
1010 self.total_count.fetch_sub(1, Ordering::SeqCst);
1012 self.emit_event(PoolEvent::ConnectionClosed);
1013 return;
1014 }
1015
1016 if !pooled.conn.is_connected() {
1018 let _ = pooled.conn.close().await;
1019 self.total_count.fetch_sub(1, Ordering::SeqCst);
1020 self.emit_event(PoolEvent::ConnectionClosed);
1021 return;
1022 }
1023
1024 if pooled.conn.in_transaction() {
1029 let _ = pooled.conn.rollback().await;
1030 }
1031
1032 pooled.last_used_at = Instant::now();
1034
1035 if let Err(mut rejected) = self.idle.push(pooled) {
1041 let _ = rejected.conn.close().await;
1043 self.total_count.fetch_sub(1, Ordering::SeqCst);
1044 self.emit_event(PoolEvent::ConnectionClosed);
1045 } else {
1046 self.emit_event(PoolEvent::ConnectionReleased);
1047 }
1048 self.notify.notify_one();
1049 }
1050
1051 pub async fn status(&self) -> PoolStatus {
1056 let idle_count = self.idle.len() as u32;
1057 let active = self.total_count.load(Ordering::Acquire);
1059 let waiters = self.waiters_count.load(Ordering::Acquire);
1060 PoolStatus {
1061 idle: idle_count,
1062 active,
1063 max: self.dynamic_max_size.load(Ordering::Acquire),
1064 min: self.config.min_idle,
1065 waiters,
1066 }
1067 }
1068
1069 #[tracing::instrument(skip(self))]
1071 pub async fn reap_idle(&self) {
1072 let mut all: Vec<PooledConnection> = Vec::new();
1076 while let Some(pooled) = self.idle.pop() {
1077 all.push(pooled);
1078 }
1079
1080 let mut to_close = Vec::new();
1082 for pooled in all {
1083 if pooled.is_idle_too_long(self.config.idle_timeout)
1084 || pooled.is_expired(self.config.max_lifetime)
1085 {
1086 to_close.push(pooled);
1087 } else {
1088 if let Err(mut rejected) = self.idle.push(pooled) {
1090 let _ = rejected.conn.close().await;
1091 self.total_count.fetch_sub(1, Ordering::SeqCst);
1092 }
1093 }
1094 }
1095
1096 for mut pooled in to_close {
1098 let _ = pooled.conn.close().await;
1099 self.total_count.fetch_sub(1, Ordering::SeqCst);
1101 }
1102 }
1103
1104 pub async fn close_all(&self) {
1108 self.closed.store(true, Ordering::Release);
1110 let mut to_close: Vec<PooledConnection> = Vec::new();
1113 while let Some(pooled) = self.idle.pop() {
1114 to_close.push(pooled);
1115 }
1116 let closed_count: u32 = to_close.len() as u32;
1118 for mut pooled in to_close {
1119 let _ = pooled.conn.close().await;
1120 }
1121 self.total_count.fetch_sub(closed_count, Ordering::SeqCst);
1124 }
1125
1126 pub async fn health_check(&self) -> u32 {
1142 let mut to_check: Vec<PooledConnection> = Vec::new();
1144 while let Some(pooled) = self.idle.pop() {
1145 to_check.push(pooled);
1146 }
1147
1148 let mut removed: u32 = 0;
1149 let mut alive: Vec<PooledConnection> = Vec::with_capacity(to_check.len());
1150 for mut pooled in to_check.drain(..) {
1151 if !pooled.conn.is_connected() {
1153 let _ = pooled.conn.close().await;
1154 removed += 1;
1155 continue;
1156 }
1157 let ping_timeout = self.config.connection_timeout / 2;
1159 match tokio::time::timeout(ping_timeout, pooled.conn.ping()).await {
1160 Ok(true) => alive.push(pooled),
1161 Ok(false) => {
1162 let _ = pooled.conn.close().await;
1164 removed += 1;
1165 }
1166 Err(_) => {
1167 let _ = pooled.conn.close().await;
1169 removed += 1;
1170 }
1171 }
1172 }
1173
1174 let alive_count: u32 = alive.len() as u32;
1176 for pooled in alive {
1177 if let Err(mut rejected) = self.idle.push(pooled) {
1179 let _ = rejected.conn.close().await;
1180 removed += 1;
1181 }
1182 }
1183
1184 if removed > 0 {
1186 self.total_count.fetch_sub(removed, Ordering::SeqCst);
1187 }
1188
1189 if alive_count > 0 {
1191 self.notify.notify_one();
1192 }
1193
1194 removed
1195 }
1196
1197 pub async fn shutdown(&self) {
1204 self.closed.store(true, Ordering::SeqCst);
1206 self.notify.notify_waiters();
1208 self.close_all().await;
1210 let deadline = Instant::now() + Duration::from_secs(30);
1212 while self.total_count.load(Ordering::SeqCst) > 0 {
1213 if Instant::now() >= deadline {
1214 break;
1215 }
1216 tokio::time::sleep(Duration::from_millis(100)).await;
1217 }
1218 }
1219
1220 pub fn resize(&self, new_max: usize) {
1228 self.set_max_size(new_max as u32);
1229 }
1230
1231 pub fn set_max_size(&self, new_max: u32) {
1233 self.dynamic_max_size.store(new_max, Ordering::SeqCst);
1234 }
1235
1236 pub fn max_size(&self) -> u32 {
1238 self.dynamic_max_size.load(Ordering::Acquire)
1239 }
1240
1241 pub async fn warmup(&self, min_idle: usize) -> Result<(), PoolError> {
1245 for _ in 0..min_idle {
1246 let current_max = self.dynamic_max_size.load(Ordering::Acquire);
1247 let current = self.total_count.load(Ordering::Acquire);
1248 if current >= current_max {
1249 break;
1250 }
1251 match self.total_count.compare_exchange(
1253 current,
1254 current + 1,
1255 Ordering::SeqCst,
1256 Ordering::Acquire,
1257 ) {
1258 Ok(_) => {}
1259 Err(_) => continue, }
1261 match self.factory.create().await {
1262 Ok(conn) => {
1263 let now = Instant::now();
1264 let pooled = PooledConnection {
1265 conn,
1266 created_at: now,
1267 last_used_at: now,
1268 pool: None,
1269 };
1270 if let Err(mut rejected) = self.idle.push(pooled) {
1271 let _ = rejected.conn.close().await;
1273 self.total_count.fetch_sub(1, Ordering::SeqCst);
1274 }
1275 self.emit_event(PoolEvent::ConnectionCreated);
1276 }
1277 Err(_) => {
1278 self.total_count.fetch_sub(1, Ordering::SeqCst);
1280 break;
1281 }
1282 }
1283 }
1284 Ok(())
1285 }
1286
1287 pub async fn query_with_timeout(&self, sql: &str) -> Result<QueryRows, sz_orm_model::DbError> {
1292 let timeout = self.config.query_timeout.unwrap_or(Duration::from_secs(30));
1293 let mut conn = self
1294 .acquire()
1295 .await
1296 .map_err(sz_orm_model::DbError::PoolError)?;
1297 tokio::time::timeout(timeout, conn.query(sql))
1298 .await
1299 .map_err(|_| {
1300 sz_orm_model::DbError::QueryError(format!("Query timeout after {:?}", timeout))
1301 })?
1302 }
1303}
1304
1305#[cfg(test)]
1306mod tests {
1307 use super::*;
1308
1309 struct MockConnection {
1311 connected: bool,
1312 }
1313
1314 impl MockConnection {
1315 fn new() -> Self {
1316 Self { connected: true }
1317 }
1318 }
1319
1320 impl Connection for MockConnection {
1321 fn execute<'a>(
1322 &'a mut self,
1323 _sql: &'a str,
1324 ) -> Pin<Box<dyn Future<Output = Result<u64, sz_orm_model::DbError>> + Send + 'a>> {
1325 Box::pin(async move { Ok(1) })
1326 }
1327
1328 fn query<'a>(
1329 &'a mut self,
1330 _sql: &'a str,
1331 ) -> Pin<
1332 Box<
1333 dyn Future<
1334 Output = Result<
1335 Vec<std::collections::HashMap<String, sz_orm_model::Value>>,
1336 sz_orm_model::DbError,
1337 >,
1338 > + Send
1339 + 'a,
1340 >,
1341 > {
1342 Box::pin(async move { Ok(vec![]) })
1343 }
1344
1345 fn begin_transaction<'a>(
1346 &'a mut self,
1347 ) -> Pin<Box<dyn Future<Output = Result<(), sz_orm_model::DbError>> + Send + 'a>> {
1348 Box::pin(async move { Ok(()) })
1349 }
1350
1351 fn commit<'a>(
1352 &'a mut self,
1353 ) -> Pin<Box<dyn Future<Output = Result<(), sz_orm_model::DbError>> + Send + 'a>> {
1354 Box::pin(async move { Ok(()) })
1355 }
1356
1357 fn rollback<'a>(
1358 &'a mut self,
1359 ) -> Pin<Box<dyn Future<Output = Result<(), sz_orm_model::DbError>> + Send + 'a>> {
1360 Box::pin(async move { Ok(()) })
1361 }
1362
1363 fn is_connected(&self) -> bool {
1364 self.connected
1365 }
1366
1367 fn ping<'a>(&'a mut self) -> Pin<Box<dyn Future<Output = bool> + Send + 'a>> {
1368 Box::pin(async move { true })
1369 }
1370
1371 fn close<'a>(
1372 &'a mut self,
1373 ) -> Pin<Box<dyn Future<Output = Result<(), sz_orm_model::DbError>> + Send + 'a>> {
1374 Box::pin(async move {
1375 self.connected = false;
1376 Ok(())
1377 })
1378 }
1379 }
1380
1381 struct MockConnectionFactory;
1382
1383 #[async_trait]
1384 impl ConnectionFactory for MockConnectionFactory {
1385 async fn create(&self) -> Result<Box<dyn Connection>, sz_orm_model::DbError> {
1386 Ok(Box::new(MockConnection::new()))
1387 }
1388 }
1389
1390 #[tokio::test]
1391 async fn test_pool_config_builder() -> Result<(), Box<dyn std::error::Error>> {
1392 let config = PoolConfigBuilder::new().max_size(50).min_idle(10).build()?;
1393
1394 assert_eq!(config.max_size, 50);
1395 assert_eq!(config.min_idle, 10);
1396 Ok(())
1397 }
1398
1399 #[test]
1400 fn test_pool_status_display() {
1401 let status = PoolStatus {
1402 idle: 5,
1403 active: 10,
1404 max: 100,
1405 min: 5,
1406 waiters: 0,
1407 };
1408
1409 let display = format!("{:?}", status);
1410 assert!(display.contains("idle"));
1411 assert!(display.contains("active"));
1412 }
1413
1414 #[test]
1415 fn test_default_pool_config() {
1416 let config = PoolConfig::default();
1417 assert_eq!(config.max_size, 100);
1418 assert_eq!(config.min_idle, 0);
1419 assert_eq!(config.acquire_timeout.as_secs(), 30);
1420 assert_eq!(config.idle_timeout.as_secs(), 600);
1421 assert_eq!(config.max_lifetime.as_secs(), 1800);
1422 }
1423
1424 #[tokio::test]
1425 async fn test_pool_config_clone() {
1426 let config = PoolConfig::default();
1427 let cloned = config.clone();
1428 assert_eq!(cloned.max_size, config.max_size);
1429 assert_eq!(cloned.min_idle, config.min_idle);
1430 }
1431
1432 #[test]
1433 fn test_pool_config_builder_default() -> Result<(), Box<dyn std::error::Error>> {
1434 let builder = PoolConfigBuilder::new();
1435 let config = builder.build()?;
1436 assert_eq!(config.max_size, 100);
1437 Ok(())
1438 }
1439
1440 #[test]
1441 fn test_pool_config_validate() {
1442 let result = PoolConfigBuilder::new().max_size(0).build();
1443 assert!(result.is_err());
1444
1445 let result = PoolConfigBuilder::new().max_size(10).min_idle(20).build();
1446 assert!(result.is_err());
1447 }
1448
1449 #[tokio::test]
1450 async fn test_pool_acquire_and_release() -> Result<(), Box<dyn std::error::Error>> {
1451 let config = PoolConfigBuilder::new().max_size(5).min_idle(1).build()?;
1452 let factory = Arc::new(MockConnectionFactory);
1453 let pool = Pool::new(config, factory)?;
1454
1455 let conn = pool.acquire().await?;
1456 let status = pool.status().await;
1457 assert_eq!(status.active, 1);
1458 assert_eq!(status.idle, 0);
1459
1460 pool.release(conn).await;
1461 let status = pool.status().await;
1462 assert_eq!(status.idle, 1);
1463
1464 let _conn2 = pool.acquire().await?;
1466 let status = pool.status().await;
1467 assert_eq!(status.idle, 0);
1468 Ok(())
1469 }
1470
1471 #[tokio::test]
1472 async fn test_pool_status() -> Result<(), Box<dyn std::error::Error>> {
1473 let config = PoolConfigBuilder::new().max_size(10).min_idle(2).build()?;
1474 let factory = Arc::new(MockConnectionFactory);
1475 let pool = Pool::new(config, factory)?;
1476
1477 let status = pool.status().await;
1478 assert_eq!(status.max, 10);
1479 assert_eq!(status.min, 2);
1480 assert_eq!(status.active, 0);
1481 Ok(())
1482 }
1483
1484 #[tokio::test]
1485 async fn test_pool_close_all() -> Result<(), Box<dyn std::error::Error>> {
1486 let config = PoolConfigBuilder::new().max_size(5).build()?;
1487 let factory = Arc::new(MockConnectionFactory);
1488 let pool = Pool::new(config, factory)?;
1489
1490 let conn1 = pool.acquire().await?;
1492 let conn2 = pool.acquire().await?;
1493 pool.release(conn1).await;
1494 pool.release(conn2).await;
1495
1496 pool.close_all().await;
1497 let status = pool.status().await;
1498 assert_eq!(status.idle, 0);
1499 assert_eq!(status.active, 0);
1500 Ok(())
1501 }
1502
1503 #[tokio::test]
1504 async fn test_pool_reap_idle() -> Result<(), Box<dyn std::error::Error>> {
1505 let config = PoolConfigBuilder::new()
1506 .max_size(5)
1507 .idle_timeout(0) .build()?;
1509 let factory = Arc::new(MockConnectionFactory);
1510 let pool = Pool::new(config, factory)?;
1511
1512 let conn = pool.acquire().await?;
1513 pool.release(conn).await;
1514
1515 tokio::time::sleep(tokio::time::Duration::from_millis(10)).await;
1517
1518 pool.reap_idle().await;
1519 let status = pool.status().await;
1520 assert_eq!(status.idle, 0);
1521 Ok(())
1522 }
1523
1524 #[tokio::test]
1530 async fn test_h7_acquire_timeout_default_30s() {
1531 let config = PoolConfig::default();
1532 assert_eq!(
1533 config.acquire_timeout,
1534 Duration::from_secs(30),
1535 "H-7: acquire_timeout 默认应为 30s"
1536 );
1537 }
1538
1539 #[tokio::test]
1541 async fn test_h7_acquire_timeout_configurable() -> Result<(), Box<dyn std::error::Error>> {
1542 let config = PoolConfigBuilder::new()
1543 .max_size(1)
1544 .acquire_timeout(5) .build()?;
1546 assert_eq!(config.acquire_timeout, Duration::from_secs(5));
1547
1548 let factory = Arc::new(MockConnectionFactory);
1550 let pool = Pool::new(config, factory)?;
1551 let _conn1 = pool.acquire().await?;
1552
1553 let fast_config = PoolConfigBuilder::new()
1555 .max_size(1)
1556 .acquire_timeout(0) .build()?;
1558 let fast_pool = Pool::new(fast_config, Arc::new(MockConnectionFactory))?;
1561 let _fast_conn = fast_pool.acquire().await?; let result = fast_pool.acquire().await;
1563 assert!(
1564 matches!(result, Err(PoolError::Timeout)),
1565 "H-7: 应返回 Timeout"
1566 );
1567 Ok(())
1568 }
1569
1570 #[tokio::test]
1573 async fn test_m7_health_check_removes_nothing_when_all_healthy(
1574 ) -> Result<(), Box<dyn std::error::Error>> {
1575 let config = PoolConfigBuilder::new().max_size(5).build()?;
1577 let factory = Arc::new(MockConnectionFactory);
1578 let pool = Pool::new(config, factory)?;
1579
1580 let conn1 = pool.acquire().await?;
1582 let conn2 = pool.acquire().await?;
1583 let conn3 = pool.acquire().await?;
1584 pool.release(conn1).await;
1585 pool.release(conn2).await;
1586 pool.release(conn3).await;
1587
1588 let removed = pool.health_check().await;
1589 assert_eq!(removed, 0, "Healthy connections should not be removed");
1590
1591 let status = pool.status().await;
1592 assert_eq!(status.idle, 3);
1593 assert_eq!(status.active, 3);
1594 Ok(())
1595 }
1596
1597 #[tokio::test]
1598 async fn test_m7_health_check_returns_zero_for_empty_pool(
1599 ) -> Result<(), Box<dyn std::error::Error>> {
1600 let config = PoolConfigBuilder::new().max_size(5).build()?;
1601 let factory = Arc::new(MockConnectionFactory);
1602 let pool = Pool::new(config, factory)?;
1603
1604 let removed = pool.health_check().await;
1605 assert_eq!(removed, 0);
1606 Ok(())
1607 }
1608
1609 struct CountingFactory {
1613 count: AtomicU32,
1614 }
1615
1616 impl CountingFactory {
1617 fn new() -> Self {
1618 Self {
1619 count: AtomicU32::new(0),
1620 }
1621 }
1622 fn created_count(&self) -> u32 {
1623 self.count.load(Ordering::SeqCst)
1624 }
1625 }
1626
1627 #[async_trait]
1628 impl ConnectionFactory for CountingFactory {
1629 async fn create(&self) -> Result<Box<dyn Connection>, sz_orm_model::DbError> {
1630 self.count.fetch_add(1, Ordering::SeqCst);
1631 Ok(Box::new(MockConnection::new()))
1632 }
1633 }
1634
1635 #[tokio::test]
1641 async fn test_production_bug_max_lifetime_never_expires(
1642 ) -> Result<(), Box<dyn std::error::Error>> {
1643 let config = PoolConfig {
1646 max_size: 5,
1647 min_idle: 0,
1648 acquire_timeout: Duration::from_secs(30),
1649 idle_timeout: Duration::from_secs(600),
1650 max_lifetime: Duration::from_millis(100), connection_timeout: Duration::from_secs(10),
1652 tls: None,
1653 query_timeout: None,
1654 max_rows: None,
1655 memory_limit: None,
1656 on_event: None,
1657 };
1658 let factory = Arc::new(CountingFactory::new());
1659 let pool = Pool::new(config, factory.clone())?;
1660
1661 let conn = pool.acquire().await?;
1663 assert_eq!(factory.created_count(), 1, "应创建 1 个连接");
1664
1665 pool.release(conn).await;
1667
1668 tokio::time::sleep(Duration::from_millis(150)).await;
1670
1671 let conn2 = pool.acquire().await?;
1673
1674 assert_eq!(
1677 factory.created_count(),
1678 2,
1679 "超过 max_lifetime 后应创建新连接(旧连接应被回收)"
1680 );
1681
1682 pool.release(conn2).await;
1683 Ok(())
1684 }
1685
1686 #[tokio::test]
1693 async fn test_drop_auto_release_connection() -> Result<(), Box<dyn std::error::Error>> {
1694 let config = PoolConfigBuilder::new().max_size(2).build()?;
1695 let factory = Arc::new(CountingFactory::new());
1696 let pool = Pool::new(config, factory.clone())?;
1697
1698 {
1700 let _conn = pool.acquire().await?;
1701 assert_eq!(factory.created_count(), 1, "应创建 1 个连接");
1702 let status = pool.status().await;
1703 assert_eq!(status.active, 1, "active 应为 1");
1704 assert_eq!(status.idle, 0, "idle 应为 0");
1705 }
1707
1708 tokio::time::sleep(Duration::from_millis(50)).await;
1710
1711 let status = pool.status().await;
1713 assert_eq!(status.idle, 1, "Drop 后连接应自动归还,idle 应为 1");
1714 assert_eq!(status.active, 1, "total_count 应为 1");
1715 assert_eq!(factory.created_count(), 1, "应复用归还的连接,不创建新连接");
1716 Ok(())
1717 }
1718
1719 #[tokio::test]
1721 async fn test_drop_auto_release_then_reuse() -> Result<(), Box<dyn std::error::Error>> {
1722 let config = PoolConfigBuilder::new().max_size(1).build()?;
1723 let factory = Arc::new(CountingFactory::new());
1724 let pool = Pool::new(config, factory.clone())?;
1725
1726 {
1728 let _conn = pool.acquire().await?;
1729 }
1730
1731 tokio::time::sleep(Duration::from_millis(50)).await;
1733
1734 let conn = pool.acquire().await?;
1736 assert_eq!(factory.created_count(), 1, "应复用归还的连接,不创建新连接");
1737
1738 pool.release(conn).await;
1739 Ok(())
1740 }
1741
1742 #[tokio::test]
1744 async fn test_into_inner_does_not_return_to_pool() -> Result<(), Box<dyn std::error::Error>> {
1745 let config = PoolConfigBuilder::new().max_size(2).build()?;
1746 let factory = Arc::new(CountingFactory::new());
1747 let pool = Pool::new(config, factory.clone())?;
1748
1749 let conn = pool.acquire().await?;
1750 assert_eq!(factory.created_count(), 1);
1751
1752 let _raw_conn = conn.into_inner();
1754
1755 tokio::time::sleep(Duration::from_millis(50)).await;
1757
1758 let status = pool.status().await;
1759 assert_eq!(status.idle, 0, "into_inner 后连接不应归还");
1760 assert_eq!(status.active, 1, "total_count 仍为 1(连接被外部持有)");
1761 Ok(())
1762 }
1763
1764 #[tokio::test]
1766 async fn test_explicit_release_no_double_return() -> Result<(), Box<dyn std::error::Error>> {
1767 let config = PoolConfigBuilder::new().max_size(2).build()?;
1768 let factory = Arc::new(CountingFactory::new());
1769 let pool = Pool::new(config, factory.clone())?;
1770
1771 let conn = pool.acquire().await?;
1772 pool.release(conn).await;
1773
1774 let status = pool.status().await;
1775 assert_eq!(status.idle, 1, "release 后 idle 应为 1");
1776
1777 let conn = pool.acquire().await?;
1779 pool.release(conn).await;
1780
1781 let status = pool.status().await;
1782 assert_eq!(status.idle, 1, "再次 release 后 idle 仍应为 1(不重复)");
1783 assert_eq!(status.active, 1, "total_count 应为 1");
1784 Ok(())
1785 }
1786}