Adds SHUTDOWN command as alternate option to sending SIGINT (#331)

* Adds SHUTDOWN command to PgCat as alternate option to sending SIGINT

* Check if we're already in SHUTDOWN sequence

* Send signal directly from shutdown instead of using channel

* Add tests

* trigger build

* Lowercase response and boolean change

* Update tests

* Fix tests

* typo

src/admin.rs

@@ -1,6 +1,8 @@
 /// Admin database.
 use bytes::{Buf, BufMut, BytesMut};
-use log::{info, trace};
+use log::{error, info, trace};
+use nix::sys::signal::{self, Signal};
+use nix::unistd::Pid;
 use std::collections::HashMap;
 use tokio::time::Instant;
 
@@ -67,6 +69,10 @@ where
             trace!("RESUME");
             resume(stream, query_parts[1]).await
         }
+        "SHUTDOWN" => {
+            trace!("SHUTDOWN");
+            shutdown(stream).await
+        }
         "SHOW" => match query_parts[1].to_ascii_uppercase().as_str() {
             "CONFIG" => {
                 trace!("SHOW CONFIG");
@@ -671,6 +677,34 @@ where
     }
 }
 
+/// Send response packets for shutdown.
+async fn shutdown<T>(stream: &mut T) -> Result<(), Error>
+where
+    T: tokio::io::AsyncWrite + std::marker::Unpin,
+{
+    let mut res = BytesMut::new();
+
+    res.put(row_description(&vec![("success", DataType::Text)]));
+
+    let mut shutdown_success = "t";
+
+    let pid = std::process::id();
+    if signal::kill(Pid::from_raw(pid.try_into().unwrap()), Signal::SIGINT).is_err() {
+        error!("Unable to send SIGINT to PID: {}", pid);
+        shutdown_success = "f";
+    }
+
+    res.put(data_row(&vec![shutdown_success.to_string()]));
+
+    res.put(command_complete("SHUTDOWN"));
+
+    res.put_u8(b'Z');
+    res.put_i32(5);
+    res.put_u8(b'I');
+
+    write_all_half(stream, &res).await
+}
+
 /// Show Users.
 async fn show_users<T>(stream: &mut T) -> Result<(), Error>
 where

src/main.rs

@@ -232,7 +232,13 @@ fn main() -> Result<(), Box<dyn std::error::Error>> {
 
                 // Initiate graceful shutdown sequence on sig int
                 _ = interrupt_signal.recv() => {
-                    info!("Got SIGINT, waiting for client connection drain now");
+                    info!("Got SIGINT");
+
+                    // Don't want this to happen more than once
+                    if admin_only {
+                        continue;
+                    }
+
                     admin_only = true;
 
                     // Broadcast that client tasks need to finish
@@ -241,98 +247,98 @@ fn main() -> Result<(), Box<dyn std::error::Error>> {
                     let _ = drain_tx.send(0).await;
 
                     tokio::task::spawn(async move {
-                      let mut interval = tokio::time::interval(tokio::time::Duration::from_millis(config.general.shutdown_timeout));
+                        let mut interval = tokio::time::interval(tokio::time::Duration::from_millis(config.general.shutdown_timeout));
 
-                      // First tick fires immediately.
-                      interval.tick().await;
+                        // First tick fires immediately.
+                        interval.tick().await;
 
-                      // Second one in the interval time.
-                      interval.tick().await;
+                        // Second one in the interval time.
+                        interval.tick().await;
 
-                      // We're done waiting.
-                      error!("Graceful shutdown timed out. {} active clients being closed", total_clients);
+                        // We're done waiting.
+                        error!("Graceful shutdown timed out. {} active clients being closed", total_clients);
 
-                      let _ = exit_tx.send(()).await;
+                        let _ = exit_tx.send(()).await;
                     });
                 },
 
                 _ = term_signal.recv() => {
                     info!("Got SIGTERM, closing with {} clients active", total_clients);
                     break;
-            },
-
-            new_client = listener.accept() => {
-                let (socket, addr) = match new_client {
-                    Ok((socket, addr)) => (socket, addr),
-                    Err(err) => {
-                        error!("{:?}", err);
-                        continue;
-                    }
-                };
-
-                let shutdown_rx = shutdown_tx.subscribe();
-                let drain_tx = drain_tx.clone();
-                let client_server_map = client_server_map.clone();
-
-                let tls_certificate = config.general.tls_certificate.clone();
-
-                tokio::task::spawn(async move {
-                    let start = chrono::offset::Utc::now().naive_utc();
-
-                    match client::client_entrypoint(
-                        socket,
-                        client_server_map,
-                        shutdown_rx,
-                        drain_tx,
-                        admin_only,
-                        tls_certificate.clone(),
-                        config.general.log_client_connections,
-                    )
-                    .await
-                    {
-                        Ok(()) => {
-                            let duration = chrono::offset::Utc::now().naive_utc() - start;
-
-                            if config.general.log_client_disconnections {
-                                info!(
-                                    "Client {:?} disconnected, session duration: {}",
-                                    addr,
-                                    format_duration(&duration)
-                                );
-                            } else {
-                                debug!(
-                                    "Client {:?} disconnected, session duration: {}",
-                                    addr,
-                                    format_duration(&duration)
-                                );
-                            }
-                        }
+                },
 
+                new_client = listener.accept() => {
+                    let (socket, addr) = match new_client {
+                        Ok((socket, addr)) => (socket, addr),
                         Err(err) => {
-                            match err {
-                                errors::Error::ClientBadStartup => debug!("Client disconnected with error {:?}", err),
-                                _ => warn!("Client disconnected with error {:?}", err),
-                            }
-
+                            error!("{:?}", err);
+                            continue;
                         }
                     };
-                });
-            }
 
-            _ = exit_rx.recv() => {
-                break;
-            }
+                    let shutdown_rx = shutdown_tx.subscribe();
+                    let drain_tx = drain_tx.clone();
+                    let client_server_map = client_server_map.clone();
 
-            client_ping = drain_rx.recv() => {
-                let client_ping = client_ping.unwrap();
-                total_clients += client_ping;
+                    let tls_certificate = config.general.tls_certificate.clone();
 
-                if total_clients == 0 && admin_only {
-                    let _ = exit_tx.send(()).await;
+                    tokio::task::spawn(async move {
+                        let start = chrono::offset::Utc::now().naive_utc();
+
+                        match client::client_entrypoint(
+                            socket,
+                            client_server_map,
+                            shutdown_rx,
+                            drain_tx,
+                            admin_only,
+                            tls_certificate.clone(),
+                            config.general.log_client_connections,
+                        )
+                        .await
+                        {
+                            Ok(()) => {
+                                let duration = chrono::offset::Utc::now().naive_utc() - start;
+
+                                if config.general.log_client_disconnections {
+                                    info!(
+                                        "Client {:?} disconnected, session duration: {}",
+                                        addr,
+                                        format_duration(&duration)
+                                    );
+                                } else {
+                                    debug!(
+                                        "Client {:?} disconnected, session duration: {}",
+                                        addr,
+                                        format_duration(&duration)
+                                    );
+                                }
+                            }
+
+                            Err(err) => {
+                                match err {
+                                    errors::Error::ClientBadStartup => debug!("Client disconnected with error {:?}", err),
+                                    _ => warn!("Client disconnected with error {:?}", err),
+                                }
+
+                            }
+                        };
+                    });
+                }
+
+                _ = exit_rx.recv() => {
+                    break;
+                }
+
+                client_ping = drain_rx.recv() => {
+                    let client_ping = client_ping.unwrap();
+                    total_clients += client_ping;
+
+                    if total_clients == 0 && admin_only {
+                        let _ = exit_tx.send(()).await;
+                    }
                 }
             }
         }
-    }
 
     info!("Shutting down...");
     });