Test solution for the feature request #473

pgdog/Cargo.toml

@@ -1,6 +1,6 @@
 [package]
 name = "pgdog"
-version = "0.1.7"
+version = "0.1.700"
 edition = "2021"
 description = "Modern PostgreSQL proxy, pooler and load balancer."
 authors = ["Lev Kokotov <[email protected]>"]
@@ -23,6 +23,8 @@ tracing = "0.1"
 tracing-subscriber = { version = "0.3", features = ["env-filter", "std"] }
 parking_lot = "0.12"
 thiserror = "2"
+libc = "0.2"
+nix = {version = "0.30.1", features = ["process"]}
 bytes = "1"
 clap = { version = "4", features = ["derive"] }
 serde = { version = "1", features = ["derive"] }

pgdog/src/frontend/listener.rs

@@ -1,9 +1,4 @@
 //! Connection listener. Handles all client connections.
-
-use std::io::ErrorKind;
-use std::net::SocketAddr;
-use std::sync::Arc;
-
 use crate::backend::databases::{databases, reload, shutdown};
 use crate::config::config;
 use crate::frontend::client::query_engine::two_pc::Manager;
@@ -12,12 +7,18 @@ use crate::net::messages::{hello::SslReply, Startup};
 use crate::net::{self, tls::acceptor};
 use crate::net::{tweak, Stream};
 use crate::sighup::Sighup;
+use crate::sigusr::Sigusr;
+use nix::unistd::execv;
+use std::env::current_exe;
+use std::ffi::CString;
+use std::io::ErrorKind;
+use std::net::SocketAddr;
+use std::sync::Arc;
 use tokio::net::{TcpListener, TcpStream};
 use tokio::signal::ctrl_c;
 use tokio::sync::Notify;
 use tokio::time::timeout;
 use tokio::{select, spawn};
-
 use tracing::{error, info, warn};
 
 use super::{
@@ -42,16 +43,17 @@ impl Listener {
     }
 
     /// Listen for client connections and handle them.
+    #[allow(unreachable_code)]
     pub async fn listen(&mut self) -> Result<(), Error> {
         info!("🐕 PgDog listening on {}", self.addr);
         let listener = TcpListener::bind(&self.addr).await?;
         let comms = comms();
         let shutdown_signal = comms.shutting_down();
         let mut sighup = Sighup::new()?;
+        let mut sigusr = Sigusr::new()?;
 
         loop {
             let comms = comms.clone();
-
             select! {
                 connection = listener.accept() => {
                    let (stream, addr) = connection?;
@@ -82,6 +84,23 @@ impl Listener {
                     self.start_shutdown();
                 }
 
+                _ = sigusr.listen() =>
+                {
+                    info!("Hot-patching signal received");
+                    match current_exe()
+                    {
+                        Ok(path) => {
+                            info!("Hot-loading from {}",path.display());
+                            let program = CString::new(path.into_os_string().into_encoded_bytes()).unwrap();
+                            let args: Vec<CString> = vec![];
+                            execv(&program, &args).expect("execv failed");
+                        }
+                        Err(e) => {
+                            error!("Failed to get executable path: {}", e);
+                        }
+                    }
+                }
+
                 _ = sighup.listen() => {
                     if let Err(err) = reload() {
                         error!("configuration reload error: {}", err);
@@ -96,7 +115,6 @@ impl Listener {
 
         Ok(())
     }
-
     /// Shutdown this listener.
     pub fn shutdown(&self) {
         self.shutdown.notify_one();

pgdog/src/lib.rs

@@ -12,6 +12,7 @@ pub mod healthcheck;
 pub mod net;
 pub mod plugin;
 pub mod sighup;
+pub mod sigusr;
 pub mod state;
 pub mod stats;
 #[cfg(feature = "tui")]

pgdog/src/sigusr.rs

@@ -0,0 +1,33 @@
+#[cfg(target_family = "unix")]
+use tokio::signal::unix::*;
+
+pub struct Sigusr {
+    #[cfg(target_family = "unix")]
+    sig: Signal,
+}
+
+impl Sigusr {
+    #[cfg(target_family = "unix")]
+    pub(crate) fn new() -> std::io::Result<Self> {
+        let sig = signal(SignalKind::user_defined1())?;
+        Ok(Self { sig })
+    }
+
+    #[cfg(not(target_family = "unix"))]
+    pub(crate) fn new() -> std::io::Result<Self> {
+        Self {}
+    }
+
+    pub(crate) async fn listen(&mut self) {
+        #[cfg(target_family = "unix")]
+        self.sig.recv().await;
+
+        #[cfg(not(target_family = "unix"))]
+        loop {
+            use std::time::Duration;
+            use tokio::time::sleep;
+
+            sleep(Duration::MAX).await;
+        }
+    }
+}