Cleaner shutdown (#12)

* Cleaner shutdown

* mark as bad just in case although im pretty sure we dont need it

* server session duration

* test clean shutdown

* ah

.circleci/config.yml

@@ -25,7 +25,7 @@ jobs:
           key: cargo-lock-2-{{ checksum "Cargo.lock" }}
       - run:
           name: "Install dependencies"
-          command: "sudo apt-get update && sudo apt-get install -y postgresql-contrib-12 postgresql-client-12"
+          command: "sudo apt-get update && sudo apt-get install -y psmisc postgresql-contrib-12 postgresql-client-12"
       - run:
           name: "Build" 
           command: "cargo build"

.circleci/run_tests.sh

@@ -26,3 +26,6 @@ psql -e -h 127.0.0.1 -p 6432 -f tests/sharding/query_routing_test_select.sql > /
 
 # Replica/primary selection & more sharding tests
 psql -e -h 127.0.0.1 -p 6432 -f tests/sharding/query_routing_test_primary_replica.sql > /dev/null
+
+# Attempt clean shut down
+killall pgcat -s SIGINT

src/main.rs

@@ -26,6 +26,7 @@ extern crate toml;
 
 use regex::Regex;
 use tokio::net::TcpListener;
+use tokio::signal;
 
 use std::collections::HashMap;
 use std::sync::{Arc, Mutex};
@@ -108,6 +109,8 @@ async fn main() {
 
     println!("> Waiting for clients...");
 
+    // Main app runs here.
+    tokio::task::spawn(async move {
         loop {
             let pool = pool.clone();
             let client_server_map = client_server_map.clone();
@@ -163,6 +166,15 @@ async fn main() {
                 };
             });
         }
+    });
+
+    // Setup shut down sequence
+    match signal::ctrl_c().await {
+        Ok(()) => {}
+        Err(err) => {
+            eprintln!("Unable to listen for shutdown signal: {}", err);
+        }
+    };
 }
 
 /// Format chrono::Duration to be more human-friendly.

src/server.rs

@@ -49,7 +49,11 @@ pub struct Server {
     // Mapping of clients and servers used for query cancellation.
     client_server_map: ClientServerMap,
 
+    // Server role, e.g. primary or replica.
     role: Role,
+
+    // Server connected at
+    connected_at: chrono::naive::NaiveDateTime,
 }
 
 impl Server {
@@ -193,6 +197,7 @@ impl Server {
                         bad: false,
                         client_server_map: client_server_map,
                         role: role,
+                        connected_at: chrono::offset::Utc::now().naive_utc(),
                     });
                 }
 
@@ -417,3 +422,27 @@ impl Server {
         }
     }
 }
+
+impl Drop for Server {
+    // Try to do a clean shut down.
+    fn drop(&mut self) {
+        let mut bytes = BytesMut::with_capacity(4);
+        bytes.put_u8(b'X');
+        bytes.put_i32(4);
+
+        match self.write.try_write(&bytes) {
+            Ok(n) => (),
+            Err(_) => (),
+        };
+
+        self.bad = true;
+
+        let now = chrono::offset::Utc::now().naive_utc();
+        let duration = now - self.connected_at;
+
+        println!(
+            ">> Server connection closed, session duration: {}",
+            crate::format_duration(&duration)
+        );
+    }
+}