Some stats (#19)

Author: levkk

Cargo.lock

@@ -20,3 +20,4 @@ serde_derive = "1"
 regex = "1"
 num_cpus = "1"
 once_cell = "1"
+statsd = "0.14"

Cargo.toml

@@ -16,6 +16,7 @@ use crate::messages::*;
 use crate::pool::{ClientServerMap, ConnectionPool};
 use crate::server::Server;
 use crate::sharding::Sharder;
+use crate::stats::Reporter;
 
 pub const SHARDING_REGEX: &str = r"SET SHARDING KEY TO '[0-9]+';";
 pub const ROLE_REGEX: &str = r"SET SERVER ROLE TO '(PRIMARY|REPLICA)';";
@@ -57,6 +58,9 @@ pub struct Client {
 
     // Client parameters, e.g. user, client_encoding, etc.
     parameters: HashMap<String, String>,
+
+    // Statistics
+    stats: Reporter,
 }
 
 impl Client {
@@ -69,6 +73,7 @@ impl Client {
         transaction_mode: bool,
         default_server_role: Option<Role>,
         server_info: BytesMut,
+        stats: Reporter,
     ) -> Result<Client, Error> {
         loop {
             // Could be StartupMessage or SSLRequest
@@ -127,6 +132,7 @@ impl Client {
                         client_server_map: client_server_map,
                         default_server_role: default_server_role,
                         parameters: parameters,
+                        stats: stats,
                     });
                 }
 
@@ -148,6 +154,7 @@ impl Client {
                         client_server_map: client_server_map,
                         default_server_role: default_server_role,
                         parameters: HashMap::new(),
+                        stats: stats,
                     });
                 }
 
@@ -220,7 +227,6 @@ impl Client {
             };
 
             // Grab a server from the pool.
-            // None = any shard
             let connection = match pool.get(shard, role).await {
                 Ok(conn) => conn,
                 Err(err) => {
@@ -287,11 +293,19 @@ impl Client {
                             }
                         }
 
-                        // Release server
-                        if !server.in_transaction() && self.transaction_mode {
-                            shard = None;
-                            role = self.default_server_role;
-                            break;
+                        // Send statistic
+                        self.stats.query();
+
+                        // Transaction over
+                        if !server.in_transaction() {
+                            self.stats.transaction();
+
+                            // Release server
+                            if self.transaction_mode {
+                                shard = None;
+                                role = self.default_server_role;
+                                break;
+                            }
                         }
                     }
 
@@ -350,11 +364,17 @@ impl Client {
                             }
                         }
 
+                        self.stats.query();
+
                         // Release server
-                        if !server.in_transaction() && self.transaction_mode {
-                            shard = None;
-                            role = self.default_server_role;
-                            break;
+                        if !server.in_transaction() {
+                            self.stats.transaction();
+
+                            if self.transaction_mode {
+                                shard = None;
+                                role = self.default_server_role;
+                                break;
+                            }
                         }
                     }
 
@@ -378,11 +398,14 @@ impl Client {
                         };
 
                         // Release the server
-                        if !server.in_transaction() && self.transaction_mode {
-                            println!("Releasing after copy done");
-                            shard = None;
-                            role = self.default_server_role;
-                            break;
+                        if !server.in_transaction() {
+                            self.stats.transaction();
+
+                            if self.transaction_mode {
+                                shard = None;
+                                role = self.default_server_role;
+                                break;
+                            }
                         }
                     }
 

src/client.rs

@@ -21,6 +21,7 @@ extern crate num_cpus;
 extern crate once_cell;
 extern crate serde;
 extern crate serde_derive;
+extern crate statsd;
 extern crate tokio;
 extern crate toml;
 
@@ -30,6 +31,7 @@ use tokio::signal;
 
 use std::collections::HashMap;
 use std::sync::{Arc, Mutex};
+use tokio::sync::mpsc;
 
 mod client;
 mod config;
@@ -38,11 +40,13 @@ mod messages;
 mod pool;
 mod server;
 mod sharding;
+mod stats;
 
 // Support for query cancellation: this maps our process_ids and
 // secret keys to the backend's.
 use config::Role;
 use pool::{ClientServerMap, ConnectionPool};
+use stats::{Collector, Reporter};
 
 /// Main!
 #[tokio::main(worker_threads = 4)]
@@ -87,7 +91,23 @@ async fn main() {
     );
     println!("> Connection timeout: {}ms", config.general.connect_timeout);
 
-    let mut pool = ConnectionPool::from_config(config.clone(), client_server_map.clone()).await;
+    // Collect statistics and send them to StatsD
+    let (tx, rx) = mpsc::channel(100);
+
+    tokio::task::spawn(async move {
+        println!("> Statistics reporter started");
+
+        let mut stats_collector = Collector::new(rx);
+        stats_collector.collect().await;
+    });
+
+    let mut pool = ConnectionPool::from_config(
+        config.clone(),
+        client_server_map.clone(),
+        Reporter::new(tx.clone()),
+    )
+    .await;
+
     let transaction_mode = config.general.pool_mode == "transaction";
     let default_server_role = match config.query_router.default_role.as_ref() {
         "any" => None,
@@ -115,6 +135,7 @@ async fn main() {
             let pool = pool.clone();
             let client_server_map = client_server_map.clone();
             let server_info = server_info.clone();
+            let reporter = Reporter::new(tx.clone());
 
             let (socket, addr) = match listener.accept().await {
                 Ok((socket, addr)) => (socket, addr),
@@ -136,6 +157,7 @@ async fn main() {
                     transaction_mode,
                     default_server_role,
                     server_info,
+                    reporter,
                 )
                 .await
                 {
@@ -170,7 +192,10 @@ async fn main() {
 
     // Setup shut down sequence
     match signal::ctrl_c().await {
-        Ok(()) => {}
+        Ok(()) => {
+            println!("> Shutting down...");
+        }
+
         Err(err) => {
             eprintln!("Unable to listen for shutdown signal: {}", err);
         }