Implementing graceful shutdown (#105)

* Initial commit for graceful shutdown

* fmt

* Add .vscode to gitignore

* Updates shutdown logic to use channels

* fmt

* fmt

* Adds shutdown timeout

* Fmt and updates tomls

* Updates readme

* fmt and updates log levels

* Update python tests to test shutdown

* merge changes

* Rename listener rx and update bash to be in line with master

* Update python test bash script ordering

* Adds error response message before shutdown

* Add details on shutdown event loop

* Fixes response length for error

* Adds handler for sigterm

* Uses ready for query function and fixes number of bytes

* fmt

src/client.rs

@@ -4,6 +4,7 @@ use log::{debug, error, info, trace};
 use std::collections::HashMap;
 use tokio::io::{split, AsyncReadExt, BufReader, ReadHalf, WriteHalf};
 use tokio::net::TcpStream;
+use tokio::sync::broadcast::Receiver;
 
 use crate::admin::{generate_server_info_for_admin, handle_admin};
 use crate::config::get_config;
@@ -73,12 +74,15 @@ pub struct Client<S, T> {
     last_server_id: Option<i32>,
 
     target_pool: ConnectionPool,
+
+    shutdown_event_receiver: Receiver<()>,
 }
 
 /// Client entrypoint.
 pub async fn client_entrypoint(
     mut stream: TcpStream,
     client_server_map: ClientServerMap,
+    shutdown_event_receiver: Receiver<()>,
 ) -> Result<(), Error> {
     // Figure out if the client wants TLS or not.
     let addr = stream.peer_addr().unwrap();
@@ -97,7 +101,7 @@ pub async fn client_entrypoint(
                 write_all(&mut stream, yes).await?;
 
                 // Negotiate TLS.
-                match startup_tls(stream, client_server_map).await {
+                match startup_tls(stream, client_server_map, shutdown_event_receiver).await {
                     Ok(mut client) => {
                         info!("Client {:?} connected (TLS)", addr);
 
@@ -121,7 +125,16 @@ pub async fn client_entrypoint(
                         let (read, write) = split(stream);
 
                         // Continue with regular startup.
-                        match Client::startup(read, write, addr, bytes, client_server_map).await {
+                        match Client::startup(
+                            read,
+                            write,
+                            addr,
+                            bytes,
+                            client_server_map,
+                            shutdown_event_receiver,
+                        )
+                        .await
+                        {
                             Ok(mut client) => {
                                 info!("Client {:?} connected (plain)", addr);
 
@@ -142,7 +155,16 @@ pub async fn client_entrypoint(
             let (read, write) = split(stream);
 
             // Continue with regular startup.
-            match Client::startup(read, write, addr, bytes, client_server_map).await {
+            match Client::startup(
+                read,
+                write,
+                addr,
+                bytes,
+                client_server_map,
+                shutdown_event_receiver,
+            )
+            .await
+            {
                 Ok(mut client) => {
                     info!("Client {:?} connected (plain)", addr);
 
@@ -157,7 +179,16 @@ pub async fn client_entrypoint(
             let (read, write) = split(stream);
 
             // Continue with cancel query request.
-            match Client::cancel(read, write, addr, bytes, client_server_map).await {
+            match Client::cancel(
+                read,
+                write,
+                addr,
+                bytes,
+                client_server_map,
+                shutdown_event_receiver,
+            )
+            .await
+            {
                 Ok(mut client) => {
                     info!("Client {:?} issued a cancel query request", addr);
 
@@ -214,6 +245,7 @@ where
 pub async fn startup_tls(
     stream: TcpStream,
     client_server_map: ClientServerMap,
+    shutdown_event_receiver: Receiver<()>,
 ) -> Result<Client<ReadHalf<TlsStream<TcpStream>>, WriteHalf<TlsStream<TcpStream>>>, Error> {
     // Negotiate TLS.
     let tls = Tls::new()?;
@@ -237,7 +269,15 @@ pub async fn startup_tls(
         Ok((ClientConnectionType::Startup, bytes)) => {
             let (read, write) = split(stream);
 
-            Client::startup(read, write, addr, bytes, client_server_map).await
+            Client::startup(
+                read,
+                write,
+                addr,
+                bytes,
+                client_server_map,
+                shutdown_event_receiver,
+            )
+            .await
         }
 
         // Bad Postgres client.
@@ -258,6 +298,7 @@ where
         addr: std::net::SocketAddr,
         bytes: BytesMut, // The rest of the startup message.
         client_server_map: ClientServerMap,
+        shutdown_event_receiver: Receiver<()>,
     ) -> Result<Client<S, T>, Error> {
         let config = get_config();
         let stats = get_reporter();
@@ -384,6 +425,7 @@ where
             last_address_id: None,
             last_server_id: None,
             target_pool: target_pool,
+            shutdown_event_receiver: shutdown_event_receiver,
         });
     }
 
@@ -394,6 +436,7 @@ where
         addr: std::net::SocketAddr,
         mut bytes: BytesMut, // The rest of the startup message.
         client_server_map: ClientServerMap,
+        shutdown_event_receiver: Receiver<()>,
     ) -> Result<Client<S, T>, Error> {
         let process_id = bytes.get_i32();
         let secret_key = bytes.get_i32();
@@ -413,6 +456,7 @@ where
             last_address_id: None,
             last_server_id: None,
             target_pool: ConnectionPool::default(),
+            shutdown_event_receiver: shutdown_event_receiver,
         });
     }
 
@@ -467,7 +511,14 @@ where
             // We can parse it here before grabbing a server from the pool,
             // in case the client is sending some custom protocol messages, e.g.
             // SET SHARDING KEY TO 'bigint';
-            let mut message = read_message(&mut self.read).await?;
+
+            let mut message = tokio::select! {
+                _ = self.shutdown_event_receiver.recv() => {
+                    error_response_terminal(&mut self.write, &format!("terminating connection due to administrator command")).await?;
+                    return Ok(())
+                },
+                message_result = read_message(&mut self.read) => message_result?
+            };
 
             // Get a pool instance referenced by the most up-to-date
             // pointer. This ensures we always read the latest config