br_pgsql/

pools.rs

use std::collections::VecDeque;
use std::sync::{Arc, Mutex};
use std::thread;
use std::time::Duration;
use log::{error, warn};
use crate::config::{Config, PoolConstraints, PoolOpts};
use crate::connect::Connect as PgsqlConnect;

/// 每个 Pools 实例使用独立的连接池，避免不同实例之间的锁竞争
/// 参考 mysql crate 的 Pool 设计
#[derive(Clone)]
pub struct Pools {
    /// 配置
    pub config: Config,
    /// 连接池选项
    pub pool_opts: PoolOpts,
    /// 最大连接数
    max_pools: usize,
    /// 最小连接数
    min_pools: usize,
    /// 连接池（每个实例独立）
    pool: Arc<Mutex<VecDeque<PgsqlConnect>>>,
    /// 总连接数（包括已借出的）
    total_connections: Arc<Mutex<usize>>,
}

pub struct ConnectionGuard<'a> {
    pool: &'a Pools,
    conn: Option<PgsqlConnect>,
}

impl<'a> ConnectionGuard<'a> {
    pub fn new(pool: &'a mut Pools) -> Result<Self, String> {
        let conn = pool.get_connect()?;
        Ok(Self { pool, conn: Some(conn) })
    }

    pub fn conn(&mut self) -> &mut PgsqlConnect {
        self.conn.as_mut().unwrap()
    }
}

impl<'a> Drop for ConnectionGuard<'a> {
    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<'_>, String> {
        ConnectionGuard::new(self)
    }
    
    /// 创建新的连接池（参考 mysql 的 Pool::new 方式）
    /// 
    /// # 参数
    /// * `config` - 数据库配置
    /// * `pool_opts` - 连接池选项（可选，默认使用 PoolOpts::default()）
    /// 
    /// # 示例
    /// ```no_run
    /// use br_pgsql::pools::Pools;
    /// use br_pgsql::config::{Config, PoolConstraints, PoolOpts};
    /// 
    /// let config = Config::new(&json::object! {
    ///     hostname: "localhost",
    ///     hostport: 5432,
    ///     // ...
    /// })?;
    /// 
    /// let constraints = PoolConstraints::new(0, 400)?;
    /// let pool_opts = PoolOpts::default()
    ///     .with_constraints(constraints);
    /// 
    /// let pools = Pools::new(config, pool_opts)?;
    /// # Ok::<(), String>(())
    /// ```
    pub fn new(config: Config, pool_opts: PoolOpts) -> Result<Self, String> {
        let constraints = pool_opts.constraints.clone();
        let max_pools = constraints.max;
        let min_pools = constraints.min;
        
        if max_pools == 0 {
            return Err("连接池最大连接数必须大于 0".into());
        }
        
        // 创建独立的连接池
        let pool = Arc::new(Mutex::new(VecDeque::with_capacity(max_pools)));
        let total_connections = Arc::new(Mutex::new(0));
        
        // 尝试初始化连接（在锁外创建，减少锁持有时间）
        // 优化：至少创建 min_pools 个连接，最多创建 max_pools 个
        let mut connections = Vec::with_capacity(max_pools);
        let mut success_count = 0;
        let mut failed_count = 0;
        
        // 优化：分批创建连接，避免一次性创建过多导致数据库压力
        // 至少尝试创建 min_pools 个连接，或直到成功创建至少 1 个
        let min_attempts = min_pools.max(1);
        let batch_size = max_pools.min(5); // 每批最多创建 5 个
        let init_size = min_pools.max(1).min(max_pools); // 初始化时至少创建 min_pools 个
        
        for i in 0..init_size {
            match PgsqlConnect::new(config.clone()) {
                Ok(conn) => {
                    connections.push(conn);
                    success_count += 1;
                    // 如果已经创建了足够的连接，可以提前退出
                    if success_count >= min_attempts && i >= batch_size - 1 {
                        break;
                    }
                }
                Err(e) => {
                    failed_count += 1;
                    warn!("初始化连接失败 ({}/{}): {e}", i + 1, init_size);
                    // 如果失败次数过多且已有连接，提前退出
                    if failed_count > init_size / 2 && success_count > 0 {
                        warn!("已创建 {} 个连接，但失败次数过多，停止初始化", success_count);
                        break;
                    }
                    // 如果前几个连接都失败，提前退出
                    if i < min_attempts && failed_count >= min_attempts && success_count == 0 {
                        break;
                    }
                }
            }
        }

        if success_count == 0 {
            return Err(format!("无法创建任何连接，请检查配置（尝试了 {} 次）", init_size));
        }
        
        if success_count < min_pools {
            warn!("连接池初始化：期望至少 {} 个连接，实际创建 {} 个", min_pools, success_count);
        }

        // 快速将连接放入池中
        {
            let mut pool_guard = pool.lock()
                .map_err(|e| format!("获取连接池锁失败: {}", e))?;
            pool_guard.extend(connections);
        }
        
        // 更新总连接数
        if let Ok(mut total) = total_connections.lock() {
            *total = success_count;
        }

        let pools = Self {
            config,
            pool_opts: pool_opts.clone(),
            max_pools,
            min_pools,
            pool,
            total_connections,
        };

        // 启动后台补齐和维护线程
        pools.start_maintainer();

        Ok(pools)
    }


    /// 获取连接 —— 优化：减少锁持有时间，优先复用有效连接，改进等待策略
    pub fn get_connect(&mut self) -> Result<PgsqlConnect, String> {
        const MAX_ATTEMPTS: usize = 10;
        const RETRY_DELAY: Duration = Duration::from_millis(500);
        const WAIT_DELAY: Duration = Duration::from_millis(50); // 增加等待延迟，减少 CPU 占用
        
        let mut attempts = 0;
        let mut consecutive_failures = 0;
        let mut last_wait_time = Duration::from_millis(10);

        loop {
            if attempts >= MAX_ATTEMPTS {
                return Err(format!("无法连接数据库，重试超时（已尝试 {} 次）", attempts));
            }
            
            // 快速获取连接（最小化锁持有时间）
            let maybe_conn = {
                let mut pool = self.pool.lock()
                    .map_err(|e| format!("获取连接池锁失败: {}", e))?;
                pool.pop_front()
            };

            if let Some(conn) = maybe_conn {
                // 快速检查连接有效性（只检查 TCP 状态，不执行查询）
                if conn.is_quick_valid() {
                    return Ok(conn);
                } else {
                    // 连接无效，减少总计数
                    if let Ok(mut total) = self.total_connections.lock() {
                        *total = total.saturating_sub(1);
                    }
                    consecutive_failures += 1;
                    if consecutive_failures > 3 {
                        warn!("连续 {} 次连接失效，可能数据库不可用", consecutive_failures);
                    }
                }
            }
            
            // 尝试创建新连接（如果还有容量）
            let current_total = self.total_connections();
            if current_total < self.max_pools {
                // 在锁外创建连接，避免阻塞
                match PgsqlConnect::new(self.config.clone()) {
                    Ok(new_conn) => {
                        // 新建成功，增加总计数
                        if let Ok(mut total) = self.total_connections.lock() {
                            *total += 1;
                        }
                        return Ok(new_conn);
                    }
                    Err(e) => {
                        error!("创建新连接失败: {}", e);
                        attempts += 1;
                        // 使用指数退避，但限制最大延迟
                        thread::sleep(RETRY_DELAY.min(Duration::from_secs(2)));
                        continue;
                    }
                }
            } else {
                // 连接池已满，使用指数退避等待
                thread::sleep(last_wait_time);
                last_wait_time = (last_wait_time * 2).min(WAIT_DELAY * 8); // 最大等待 400ms
                attempts += 1;
            }
        }
    }

    /// 把连接放回池子（显式归还），如果连接无效或池子已满则丢弃
    /// 优化：使用快速检查，减少锁持有时间，优化池满处理
    pub fn release_conn(&self, conn: PgsqlConnect) {
        // 使用快速检查（只检查 TCP 状态），避免执行查询的开销
        if conn.is_quick_valid() {
            let pool_len = {
                let pool_guard = self.pool.lock();
                match pool_guard {
                    Ok(p) => p.len(),
                    Err(e) => {
                        warn!("获取连接池锁失败，丢弃连接: {}", e);
                        if let Ok(mut total) = self.total_connections.lock() {
                            *total = total.saturating_sub(1);
                        }
                        return;
                    }
                }
            };
            
            // 检查池是否已满（在锁外检查，减少锁持有时间）
            if pool_len < self.max_pools {
                // 快速放入池中
                if let Ok(mut pool) = self.pool.lock() {
                    pool.push_back(conn);
                } else {
                    // 锁失败，减少总连接数
                    if let Ok(mut total) = self.total_connections.lock() {
                        *total = total.saturating_sub(1);
                    }
                }
            } else {
                // 池子已满，减少总连接数（不持有锁）
                if let Ok(mut total) = self.total_connections.lock() {
                    *total = total.saturating_sub(1);
                }
                // 不记录警告，因为这是正常情况（连接池已满）
            }
        } else {
            // 连接无效，减少总连接数
            if let Ok(mut total) = self.total_connections.lock() {
                *total = total.saturating_sub(1);
            }
            warn!("释放时检测到坏连接，已丢弃");
        }
    }

    // 连接池保持连接数量和维护
    // 优化：减少维护频率，批量处理，避免频繁操作
    fn start_maintainer(&self) {
        let config = Arc::new(self.config.clone());
        let pool = Arc::clone(&self.pool);
        let total_connections = Arc::clone(&self.total_connections);
        let max_pools = self.max_pools;
        let min_pools = self.min_pools;
        const MAINTAIN_INTERVAL: Duration = Duration::from_secs(10); // 增加维护间隔，减少开销
        const CLEANUP_INTERVAL: Duration = Duration::from_secs(60); // 增加清理间隔
        let mut last_cleanup = std::time::Instant::now();
        let mut consecutive_failures = 0;

        thread::spawn(move || loop {
            // 补齐连接
            {
                let current = {
                    let pool_guard = pool.lock();
                    match pool_guard {
                        Ok(p) => p.len(),
                        Err(e) => {
                            warn!("连接池维护线程获取锁失败: {}", e);
                            thread::sleep(MAINTAIN_INTERVAL);
                            continue;
                        }
                    }
                };
                
                // 确保至少保持 min_pools 个连接
                let target_size = if current < min_pools {
                    min_pools
                } else {
                    max_pools
                };
                
                if current < target_size {
                    let need = target_size - current;
                    // 限制每次补齐的连接数，避免一次性创建过多连接
                    let batch_size = need.min(2); // 减少批量大小，更温和
                    let mut connections = Vec::with_capacity(batch_size);
                    let mut success_count = 0;
                    
                    // 在锁外创建连接，减少锁持有时间
                    for _ in 0..batch_size {
                        match PgsqlConnect::new((*config).clone()) {
                            Ok(conn) => {
                                connections.push(conn);
                                success_count += 1;
                            }
                            Err(e) => {
                                warn!("补齐连接失败: {e}");
                                consecutive_failures += 1;
                                // 如果连续失败次数过多，停止尝试
                                if consecutive_failures > 3 {
                                    break;
                                }
                            }
                        }
                    }
                    
                    // 重置失败计数
                    if success_count > 0 {
                        consecutive_failures = 0;
                    }
                    
                    // 快速将连接放入池中
                    if !connections.is_empty() {
                        let mut pool_guard = match pool.lock() {
                            Ok(p) => p,
                            Err(e) => {
                                warn!("连接池维护线程获取锁失败: {}", e);
                                continue;
                            }
                        };
                        
                        let before = pool_guard.len();
                        pool_guard.extend(connections);
                        let after = pool_guard.len();
                        let added = after.saturating_sub(before);
                        
                        // 更新总连接数
                        if let Ok(mut total) = total_connections.lock() {
                            *total = total.saturating_add(added);
                        }
                        
                        // 连接补齐完成（已移除正常日志）
                    }
                } else {
                    // 连接池已满，重置失败计数
                    consecutive_failures = 0;
                }
            }
            
            // 定期清理无效连接
            if last_cleanup.elapsed() >= CLEANUP_INTERVAL {
                {
                    let mut pool_guard = match pool.lock() {
                        Ok(p) => p,
                        Err(_) => {
                            thread::sleep(MAINTAIN_INTERVAL);
                            continue;
                        }
                    };
                    
                    let before = pool_guard.len();
                    pool_guard.retain(|conn| {
                        conn.is_quick_valid()
                    });
                    let after = pool_guard.len();
                    let removed = before.saturating_sub(after);
                    
                    // 更新总连接数
                    if removed > 0 {
                        if let Ok(mut total) = total_connections.lock() {
                            *total = total.saturating_sub(removed);
                        }
                        // 清理无效连接完成（已移除正常日志）
                    }
                }
                
                last_cleanup = std::time::Instant::now();
            }
            
            thread::sleep(MAINTAIN_INTERVAL);
        });
    }

    /// 获取当前连接池中的空闲连接数量
    pub fn idle_pool_size(&self) -> usize {
        self.pool.lock()
            .map(|pool| pool.len())
            .unwrap_or(0)
    }

    /// 获取总连接数量（包括已借出的）
    pub fn total_connections(&self) -> usize {
        self.total_connections.lock()
            .map(|total| *total)
            .unwrap_or(0)
    }

    /// 获取已借出的连接数量
    pub fn borrowed_connections(&self) -> usize {
        let total = self.total_connections();
        let idle = self.idle_pool_size();
        total.saturating_sub(idle)
    }
    
    /// 获取连接池使用率（0.0 - 1.0）
    pub fn usage_rate(&self) -> f64 {
        let total = self.total_connections();
        if total == 0 {
            return 0.0;
        }
        let idle = self.idle_pool_size();
        let borrowed = total.saturating_sub(idle);
        borrowed as f64 / total as f64
    }
    
    /// 获取连接池健康状态
    pub fn health_status(&self) -> String {
        let total = self.total_connections();
        let idle = self.idle_pool_size();
        let borrowed = total.saturating_sub(idle);
        let usage = self.usage_rate();
        
        format!(
            "连接池状态: 总数={}, 空闲={}, 已借出={}, 使用率={:.1}%, 最小={}, 最大={}",
            total, idle, borrowed, usage * 100.0, self.min_pools, self.max_pools
        )
    }
    
    /// 兼容旧版本的 new 方法（使用 size 参数）
    /// 推荐使用 new(config, pool_opts) 方法
    pub fn new_with_size(config: Config, size: usize) -> Result<Self, String> {
        let constraints = PoolConstraints::new(0, size)
            .map_err(|e| format!("连接池约束配置失败: {}", e))?;
        let pool_opts = PoolOpts::default()
            .with_constraints(constraints);
        Self::new(config, pool_opts)
    }

    /// 清除无效连接的方法
    /// 返回清理的连接数量
    /// 优化：使用快速检查，减少开销
    pub fn cleanup_idle_connections(&self) -> usize {
        match self.pool.lock() {
            Ok(mut pool) => {
                let before = pool.len();
                pool.retain(|conn| {
                    conn.is_quick_valid()
                });
                let after = pool.len();
                let removed = before.saturating_sub(after);
                
                // 更新总连接数
                if removed > 0 {
                    if let Ok(mut total) = self.total_connections.lock() {
                        *total = total.saturating_sub(removed);
                    }
                    // 清理无效连接完成（已移除正常日志）
                }
                removed
            }
            Err(e) => {
                warn!("获取连接池锁失败，无法清理: {}", e);
                0
            }
        }
    }
}