br_pgsql/

connect.rs

use crate::config::Config;
use crate::packet::{AuthStatus, Packet, SuccessMessage};
use log::warn;
use std::io::{Read, Write};
use std::net::TcpStream;
use std::sync::Arc;
use std::time::Duration;

#[derive(Clone, Debug)]
pub struct Connect {
    /// 基础配置
    pub(crate) stream: Arc<TcpStream>,
    packet: Packet,
    /// 认证状态
    auth_status: AuthStatus,
}

impl Connect {
    // 判断连接是否正常（完整检查，包括查询）
    pub fn is_valid(&mut self) -> bool {
        // 先检查 TCP 连接状态
        if self.stream.peer_addr().is_err() {
            return false;
        }
        // 发送一个简单心跳，使用更轻量的查询
        self.query("SELECT 1").is_ok()
    }

    /// 快速检查连接是否可用（只检查 TCP 状态，不执行查询）
    /// 用于连接池，避免执行查询的开销
    pub fn is_quick_valid(&self) -> bool {
        // 只检查 TCP 连接状态，避免执行查询的开销
        self.stream.peer_addr().is_ok()
    }

    // 关闭程序和pgsql的连接
    pub fn _close(&mut self) {
        let _ = self.stream.as_ref().write_all(&Packet::pack_terminate());
        let _ = self.stream.shutdown(std::net::Shutdown::Both);
    }

    pub fn new(config: Config) -> Result<Connect, String> {
        let stream = match TcpStream::connect(config.url()) {
            Ok(stream) => stream,
            Err(e) => return Err(format!("数据库连接失败: {}", e)),
        };
        stream
            .set_read_timeout(Some(Duration::from_secs(5)))
            .map_err(|e| format!("设置读取超时失败: {}", e))?;
        stream
            .set_write_timeout(Some(Duration::from_secs(5)))
            .map_err(|e| format!("设置写入超时失败: {}", e))?;

        match stream.peer_addr() {
            Ok(_) => {}
            Err(e) => warn!("无法获取对端地址: {}", e),
        }

        let mut connect = Self {
            stream: Arc::new(stream),
            packet: Packet::new(config),
            auth_status: AuthStatus::None,
        };

        connect.startup_message()?;
        connect.sasl_initial_response_message()?;

        Ok(connect)
    }

    fn read(&mut self) -> Result<Vec<u8>, String> {
        const BUFFER_SIZE: usize = 4096; // 增大缓冲区以提高性能
        const SLEEP_DURATION: Duration = Duration::from_millis(5);
        
        let mut msg = Vec::new();
        let mut buf = vec![0u8; BUFFER_SIZE];

        loop {
            match self.stream.as_ref().read(&mut buf) {
                Ok(0) => return Err("连接已关闭或服务端断开".into()),
                Ok(n) => {
                    msg.extend_from_slice(&buf[..n]);
                    // 如果读取的数据少于缓冲区大小，可能已经读取完毕
                    if n < BUFFER_SIZE {
                        // 检查是否已经读取完整消息
                        if msg.len() >= 5 {
                            let len_bytes = &msg[1..=4];
                            if let Ok(len) = len_bytes.try_into().map(u32::from_be_bytes) {
                                if msg.len() >= len as usize {
                                    break;
                                }
                            }
                        }
                    }
                }
                Err(ref e)
                    if e.kind() == std::io::ErrorKind::WouldBlock
                        || e.kind() == std::io::ErrorKind::TimedOut =>
                {
                    // 避免空循环占用 CPU
                    std::thread::sleep(SLEEP_DURATION);
                    continue;
                }
                Err(e) => return Err(format!("读取失败: {}", e)),
            };

            if let AuthStatus::AuthenticationOk = self.auth_status {
                if msg.ends_with(&[90, 0, 0, 0, 5, 73])
                    || msg.ends_with(&[90, 0, 0, 0, 5, 84])
                    || msg.ends_with(&[90, 0, 0, 0, 5, 69])
                {
                    break;
                }
            } else if msg.len() >= 5 {
                // 取长度字节
                let len_bytes = &msg[1..=4];
                if let Ok(len) = len_bytes.try_into().map(u32::from_be_bytes) {
                    if msg.len() >= len as usize {
                        break;
                    }
                }
            }
        }

        Ok(msg)
    }
    /// Startup Message
    fn startup_message(&mut self) -> Result<(), String> {
        // try_clone -> as_ref Arc<TcpStream> 是安全的
        self.stream
            .as_ref()
            .write_all(&self.packet.pack_first())
            .map_err(|e| format!("发送 startup message 失败: {}", e))?;
        let data = self.read()?;
        self.packet.unpack(data, 0)?;
        Ok(())
    }
    /// `SASLInitialResponse` message
    fn sasl_initial_response_message(&mut self) -> Result<(), String> {
        self.stream
            .as_ref()
            .write_all(&self.packet.pack_auth())
            .map_err(|e| format!("发送 SASL Initial Response 失败: {}", e))?;

        let data = self.read()?;
        self.packet.unpack(data, 0)?;

        self.stream
            .as_ref()
            .write_all(&self.packet.pack_auth_verify())
            .map_err(|e| format!("发送 SASL Verify 失败: {}", e))?;

        let data = self.read()?;
        self.packet.unpack(data, 0)?;
        self.auth_status = AuthStatus::AuthenticationOk;
        Ok(())
    }
    /// 查询
    pub fn query(&mut self, sql: &str) -> Result<SuccessMessage, String> {
        // self.stream.try_clone().unwrap().write_all(&self.packet.pack_query(sql)).unwrap();
        // let data = self.read()?;
        // let mut packet = self.packet.clone();

        // 直接使用原始stream，避免try_clone
        self.stream
            .as_ref()
            .write_all(&self.packet.pack_query(sql))
            .map_err(|e| format!("query error: {}", e))?;

        let data = self
            .read()
            .map_err(|e| format!("query read error: {}", e))?;

        self.packet.unpack(data, 0)

        // 唯一的区别是 packet.unpack 运行在一个新栈上（你显式设置了 8 MB stack size）。可能是为了防止深递归或大数据解析时栈溢出。
        // 如果你的 packet.unpack 是字节流解析（状态机 + 循环），那 2 MB 绰绰有余，不用担心。
        // 如果它内部藏了大数组或递归，那才需要考虑开大栈。
        // std::thread::Builder::new().stack_size(8 * 1024 * 1024).spawn(move || -> Result<SuccessMessage, String> {
        //     packet.unpack(data, 0)
        // }).unwrap().join().unwrap()
    }
    /// 执行
    pub fn execute(&mut self, sql: &str) -> Result<SuccessMessage, String> {
        // self.stream.try_clone().unwrap().write_all(&self.packet.pack_execute(sql)).unwrap();
        // let data = self.read()?;
        // let mut packet = self.packet.clone();
        //
        // std::thread::Builder::new().stack_size(8 * 1024 * 1024).spawn(move || -> Result<SuccessMessage, String> {
        //     packet.unpack(data, 0)
        // }).unwrap().join().unwrap()

        self.stream
            .as_ref()
            .write_all(&self.packet.pack_execute(sql))
            .map_err(|e| format!("execute error: {}", e))?;
        let data = self.read()?;
        self.packet.unpack(data, 0)
    }
}
impl Drop for Connect {
    fn drop(&mut self) {
        // 如果链接关闭的情况下，unwrap会失败，所以此处用match处理
        // 只做资源清理，不处理连接池归还
        match self.stream.peer_addr() {
            Ok(_addr) => {
                // 连接关闭（已移除正常日志）
            },
            Err(e) => warn!("[DROP] 获取 peer_addr 失败: {e}"),
        }

        // 尝试发送 Terminate 包，不强制 unwrap
        if let Err(e) = self.stream.as_ref().write_all(&Packet::pack_terminate()) {
            warn!("[DROP] 发送 Terminate 包失败: {e}");
        }

        if let Err(e) = self.stream.shutdown(std::net::Shutdown::Both) {
            warn!("[DROP] shutdown 失败: {e}");
        }
    }

    // fn drop(&mut self) {
    //     println!("[DROP] 丢弃连接: {}", self.id); // 调试日志
    //
    //     let _ = self.stream.as_ref().write_all(&Packet::pack_terminate());
    //     let _ = self.stream.shutdown(std::net::Shutdown::Both);
    //
    //     // // 显式发送 PostgreSQL 协议 Terminate 包，告诉服务端“我要断开了”
    //     // if let Err(e) = self.stream.as_ref().write_all(&Packet::pack_terminate()) {
    //     //     println!("[DROP] 发送 Terminate 出错: {}", e);
    //     // }
    //     //
    //     // if self.is_valid() {
    //     //     match crate::pools::DB_POOL.lock() {
    //     //         Ok(mut pool) => {
    //     //             if pool.len() <= 20 {
    //     //                 pool.push(self.clone());
    //     //                 println!("连接成功归还");
    //     //             } else {
    //     //                 println!("连接池已满，丢弃连接");
    //     //             }
    //     //         }
    //     //         Err(e) => println!("获取连接池锁失败: {}", e),
    //     //     }
    //     // } else {
    //     //     println!("丢弃无效连接");
    //     // }
    // }
}