use std::process::{Command, Output, Stdio}; use std::time::{Duration, Instant, SystemTime, UNIX_EPOCH}; const SQL_TIMEOUT: Duration = Duration::from_secs(10); const PROBE_TIMEOUT: Duration = Duration::from_millis(1200); fn output_with_timeout(mut command: Command, timeout: Duration) -> Option { command .stdin(Stdio::null()) .stdout(Stdio::piped()) .stderr(Stdio::piped()); let mut child = command.spawn().ok()?; let deadline = Instant::now() + timeout; loop { match child.try_wait() { Ok(Some(_)) => return child.wait_with_output().ok(), Ok(None) if Instant::now() >= deadline => { let _ = child.kill(); let _ = child.wait(); return None; } Ok(None) => std::thread::sleep(Duration::from_millis(20)), Err(_) => { let _ = child.kill(); let _ = child.wait(); return None; } } } } pub fn split_host_port(addr: &str) -> (String, String) { match addr.rsplit_once(':') { Some((host, port)) => (host.to_string(), port.to_string()), None => (addr.to_string(), String::from("5432")), } } pub fn parent_dir(path: &str) -> String { match path.rsplit_once('/') { Some((dir, _)) => dir.to_string(), None => String::from("."), } } pub fn spawn_rejoin( script: &str, pgbin: &str, datadir: &str, leader_host: &str, leader_port: &str, node_id: u64, passfile: &str, standby_conninfo: &str, ) -> bool { Command::new("setsid") .arg("bash") .arg(script) .arg(pgbin) .arg(datadir) .arg(leader_host) .arg(leader_port) .arg(node_id.to_string()) .arg(passfile) .arg(standby_conninfo) .stdin(Stdio::null()) .stdout(Stdio::null()) .stderr(Stdio::null()) .spawn() .is_ok() } pub fn write_heartbeat(path: &str) { let millis = SystemTime::now() .duration_since(UNIX_EPOCH) .map(|elapsed| elapsed.as_millis()) .unwrap_or(0); let tmp = format!("{}.tmp", path); if std::fs::write(&tmp, millis.to_string()).is_ok() { let _ = std::fs::rename(&tmp, path); } } pub fn spawn_watchdog( script: &str, psql: &str, host: &str, port: &str, heartbeat: &str, node_id: u64, user: &str, ) -> bool { Command::new("setsid") .arg("bash") .arg(script) .arg(psql) .arg(host) .arg(port) .arg(heartbeat) .arg(node_id.to_string()) .arg(user) .stdin(Stdio::null()) .stdout(Stdio::null()) .stderr(Stdio::null()) .spawn() .is_ok() } pub fn peer_wal_lsn(psql: &str, host: &str, port: &str) -> Option { let mut command = Command::new(psql); command .args([ "-h", host, "-p", port, "-U", "replicator", "-d", "postgres", "-w", "-tAc", "SELECT pg_wal_lsn_diff(COALESCE(pg_last_wal_receive_lsn(), '0/0'), '0/0')::bigint", ]) .env("PGCONNECT_TIMEOUT", "1"); let output = output_with_timeout(command, PROBE_TIMEOUT)?; if output.status.success() { String::from_utf8_lossy(&output.stdout) .trim() .parse::() .ok() .map(|value| value.max(0) as u64) } else { None } } pub fn wal_lsn(psql: &str, host: &str, port: &str, in_recovery: bool) -> u64 { let sql = if in_recovery { "SELECT pg_wal_lsn_diff(COALESCE(pg_last_wal_receive_lsn(), '0/0'), '0/0')::bigint" } else { "SELECT pg_wal_lsn_diff(pg_current_wal_lsn(), '0/0')::bigint" }; run_sql_with_timeout(psql, host, port, sql, PROBE_TIMEOUT) .ok() .and_then(|value| value.trim().parse::().ok()) .map(|value| value.max(0) as u64) .unwrap_or(0) } pub fn run_sql(psql: &str, host: &str, port: &str, sql: &str) -> Result { run_sql_with_timeout(psql, host, port, sql, SQL_TIMEOUT) } fn run_sql_with_timeout( psql: &str, host: &str, port: &str, sql: &str, timeout: Duration, ) -> Result { let user = crate::config::apply_user(); let db = crate::config::apply_db(); let passfile = crate::config::passfile(); let mut command = Command::new(psql); command.args([ "-h", host, "-p", port, "-U", user.as_str(), "-d", db.as_str(), "-w", "-v", "ON_ERROR_STOP=1", "-tAc", sql, ]); command.env("PGCONNECT_TIMEOUT", "2"); command.env("PGOPTIONS", "-c statement_timeout=5000"); if !passfile.is_empty() { command.env("PGPASSFILE", &passfile); } let output = output_with_timeout(command, timeout) .ok_or_else(|| String::from("psql failed to spawn or timed out"))?; if output.status.success() { Ok(String::from_utf8_lossy(&output.stdout).trim().to_string()) } else { Err(String::from_utf8_lossy(&output.stderr).trim().to_string()) } }