br_pgsql/

pools.rs

use crate::config::Config;
use crate::connect::Connect;
use crate::error::PgsqlError;
use log::{error, info, warn};
use std::collections::VecDeque;
use std::sync::{Arc, Mutex, MutexGuard, PoisonError};
use std::thread;
use std::time::Duration;

pub static DB_POOL: std::sync::LazyLock<Mutex<VecDeque<Connect>>> =
    std::sync::LazyLock::new(|| Mutex::new(VecDeque::new()));

fn lock_pool<'a>() -> MutexGuard<'a, VecDeque<Connect>> {
    DB_POOL.lock().unwrap_or_else(PoisonError::into_inner)
}

#[derive(Clone)]
pub struct Pools {
    pub config: Config,
    max_pools: usize,
    total_connections: Arc<Mutex<usize>>,
}

fn lock_counter(counter: &Mutex<usize>) -> MutexGuard<'_, usize> {
    counter.lock().unwrap_or_else(PoisonError::into_inner)
}

pub struct ConnectionGuard {
    pool: Pools,
    conn: Option<Connect>,
}

impl ConnectionGuard {
    pub fn new(pool: Pools) -> Result<Self, PgsqlError> {
        let conn = pool.clone().get_connect()?;
        Ok(Self {
            pool,
            conn: Some(conn),
        })
    }

    pub fn conn(&mut self) -> &mut Connect {
        self.conn.as_mut().expect("connection already released")
    }
}

impl Drop for ConnectionGuard {
    fn drop(&mut self) {
        if let Some(conn) = self.conn.take() {
            self.pool.release_conn(conn);
        }
    }
}

impl Pools {
    pub fn get_guard(&mut self) -> Result<ConnectionGuard, PgsqlError> {
        ConnectionGuard::new(self.clone())
    }

    pub fn new(config: Config, size: usize) -> Result<Self, PgsqlError> {
        let mut pool_guard = lock_pool();
        let init_size = 2.min(size);
        let mut created = 0;
        for _ in 0..init_size {
            match Connect::new(config.clone()) {
                Ok(conn) => {
                    pool_guard.push_back(conn);
                    created += 1;
                }
                Err(e) => warn!("初始化连接失败: {e}"),
            }
        }

        let pools = Self {
            config,
            max_pools: size,
            total_connections: Arc::new(Mutex::new(created)),
        };

        Ok(pools)
    }

    pub fn get_connect(&mut self) -> Result<Connect, PgsqlError> {
        let mut attempts = 0;

        loop {
            if attempts >= 20 {
                return Err(PgsqlError::Pool("无法连接数据库，重试超时".into()));
            }
            let maybe_conn = {
                let mut pool = lock_pool();
                pool.pop_front()
            };

            if let Some(mut conn) = maybe_conn {
                if conn.is_valid() {
                    return Ok(conn);
                } else {
                    *lock_counter(&self.total_connections) -= 1;
                    warn!(
                        "连接失效，尝试重建，当前总连接数量: {}",
                        self.total_connections()
                    );
                    match Connect::new(self.config.clone()) {
                        Ok(new_conn) => {
                            *lock_counter(&self.total_connections) += 1;
                            return Ok(new_conn);
                        }
                        Err(e) => {
                            error!("重建连接失败: {}", e);
                            attempts += 1;
                            thread::sleep(Duration::from_secs(1));
                            continue;
                        }
                    }
                }
            } else if self.total_connections() < self.max_pools {
                match Connect::new(self.config.clone()) {
                    Ok(new_conn) => {
                        *lock_counter(&self.total_connections) += 1;
                        return Ok(new_conn);
                    }
                    Err(e) => {
                        error!("创建新连接失败: {}", e);
                        attempts += 1;
                        thread::sleep(Duration::from_secs(1));
                        continue;
                    }
                }
            } else {
                thread::sleep(Duration::from_millis(50));
            }
        }
    }

    /// 事务专用连接，不归还到池
    pub fn get_connect_for_transaction(&mut self) -> Result<Connect, PgsqlError> {
        let mut attempts = 0;

        loop {
            if attempts >= 20 {
                return Err(PgsqlError::Pool("无法获取事务连接，重试超时".into()));
            }

            let maybe_conn = {
                let mut pool = lock_pool();
                pool.pop_front()
            };

            if let Some(mut conn) = maybe_conn {
                if conn.is_valid() {
                    return Ok(conn);
                } else {
                    *lock_counter(&self.total_connections) -= 1;
                    warn!(
                        "事务连接失效，尝试重建，当前总连接数量: {}",
                        self.total_connections()
                    );
                }
            }

            match Connect::new(self.config.clone()) {
                Ok(new_conn) => {
                    *lock_counter(&self.total_connections) += 1;
                    return Ok(new_conn);
                }
                Err(e) => {
                    error!("创建事务连接失败: {}", e);
                    attempts += 1;
                    thread::sleep(Duration::from_secs(1));
                    continue;
                }
            }
        }
    }

    pub fn release_transaction_conn(&self) {
        *lock_counter(&self.total_connections) -= 1;
    }

    pub fn release_conn(&self, mut conn: Connect) {
        if conn.is_valid() {
            let mut pool = lock_pool();
            if pool.len() < self.max_pools {
                pool.push_back(conn);
            } else {
                *lock_counter(&self.total_connections) -= 1;
                warn!("连接池已满，丢弃连接");
            }
        } else {
            *lock_counter(&self.total_connections) -= 1;
            warn!("释放时检测到坏连接，已丢弃");
        }
    }

    #[allow(dead_code)]
    fn start_maintainer(&self) {
        let config = self.config.clone();
        let max_pools = self.max_pools;

        thread::spawn(move || loop {
            {
                let mut pool_guard = lock_pool();
                let current = pool_guard.len();
                if current < max_pools {
                    let need = max_pools - current;
                    for _ in 0..need {
                        match Connect::new(config.clone()) {
                            Ok(conn) => {
                                pool_guard.push_back(conn);
                                info!("补齐一个连接，当前连接数 {}", pool_guard.len());
                            }
                            Err(e) => {
                                warn!("补齐连接失败: {e}");
                            }
                        }
                    }
                }
            }
            thread::sleep(Duration::from_secs(5));
        });
    }

    pub fn idle_pool_size(&self) -> usize {
        let pool = lock_pool();
        pool.len()
    }

    pub fn total_connections(&self) -> usize {
        *lock_counter(&self.total_connections)
    }

    pub fn borrowed_connections(&self) -> usize {
        self.total_connections() - self.idle_pool_size()
    }

    #[allow(dead_code)]
    pub fn _cleanup_idle_connections(&self) {
        let mut pool = lock_pool();
        println!("当前连接池中的连接数量（清理前）: {}", pool.len());

        pool.retain(|conn| {
            let is_ok = conn.stream.peer_addr().is_ok();
            if !is_ok {
                println!("检测到无效连接，已移除");
            }
            is_ok
        });

        println!("当前连接池中的连接数量（清理后）: {}", pool.len());
    }
}