query routing docs

.circleci/run_tests.sh

@@ -13,9 +13,6 @@ function start_pgcat() {
 
 # Setup the database with shards and user
 psql -e -h 127.0.0.1 -p 5432 -U postgres -f tests/sharding/query_routing_setup.sql
-PGPASSWORD=sharding_user pgbench -h 127.0.0.1 -U sharding_user shard0 -i
-PGPASSWORD=sharding_user pgbench -h 127.0.0.1 -U sharding_user shard1 -i
-PGPASSWORD=sharding_user pgbench -h 127.0.0.1 -U sharding_user shard2 -i
 
 # Install Toxiproxy to simulate a downed/slow database
 wget -O toxiproxy-2.1.4.deb https://github.com/Shopify/toxiproxy/releases/download/v2.1.4/toxiproxy_2.1.4_amd64.deb
@@ -31,9 +28,9 @@ toxiproxy-cli create -l 127.0.0.1:5433 -u 127.0.0.1:5432 postgres_replica
 start_pgcat "info"
 
 # pgbench test
-pgbench -i -h 127.0.0.1 -p 6432
+pgbench -i -h 127.0.0.1 -p 6432 && \
-pgbench -h 127.0.0.1 -p 6432 -t 500 -c 2 --protocol simple -f tests/pgbench/simple.sql
+    pgbench -h 127.0.0.1 -p 6432 -t 500 -c 2 --protocol simple && \
-pgbench -h 127.0.0.1 -p 6432 -t 500 -c 2 --protocol extended
+    pgbench -h 127.0.0.1 -p 6432 -t 500 -c 2 --protocol extended
 
 # COPY TO STDOUT test
 psql -h 127.0.0.1 -p 6432 -c 'COPY (SELECT * FROM pgbench_accounts LIMIT 15) TO STDOUT;' > /dev/null
@@ -60,17 +57,6 @@ cd tests/ruby && \
     ruby tests.rb && \
 cd ../..
 
-# Admin tests
-psql -e -h 127.0.0.1 -p 6432 -d pgbouncer -c 'SHOW STATS' > /dev/null
-psql -h 127.0.0.1 -p 6432 -d pgbouncer -c 'RELOAD' > /dev/null
-psql -h 127.0.0.1 -p 6432 -d pgbouncer -c 'SHOW CONFIG' > /dev/null
-psql -h 127.0.0.1 -p 6432 -d pgbouncer -c 'SHOW DATABASES' > /dev/null
-psql -h 127.0.0.1 -p 6432 -d pgbouncer -c 'SHOW LISTS' > /dev/null
-psql -h 127.0.0.1 -p 6432 -d pgbouncer -c 'SHOW POOLS' > /dev/null
-psql -h 127.0.0.1 -p 6432 -d pgbouncer -c 'SHOW VERSION' > /dev/null
-psql -h 127.0.0.1 -p 6432 -d pgbouncer -c "SET client_encoding TO 'utf8'" > /dev/null # will ignore
-(! psql -e -h 127.0.0.1 -p 6432 -d random_db -c 'SHOW STATS' > /dev/null)
-
 # Start PgCat in debug to demonstrate failover better
 start_pgcat "debug"
 

Cargo.lock

@@ -371,6 +371,7 @@ dependencies = [
  "serde_derive",
  "sha-1",
  "sqlparser",
+ "statsd",
  "tokio",
  "toml",
 ]
@@ -541,6 +542,15 @@ dependencies = [
  "log",
 ]
 
+[[package]]
+name = "statsd"
+version = "0.15.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "df1efceb4bf2c0b5ebec94354285a43bbbed1375605bdf2ebe4132299434a330"
+dependencies = [
+ "rand",
+]
+
 [[package]]
 name = "syn"
 version = "1.0.86"

Cargo.toml

@@ -20,6 +20,7 @@ serde_derive = "1"
 regex = "1"
 num_cpus = "1"
 once_cell = "1"
+statsd = "0.15"
 sqlparser = "0.14"
 log = "0.4"
 arc-swap = "1"

src/admin.rs

@@ -1,386 +0,0 @@
-use bytes::{Buf, BufMut, BytesMut};
-use log::{info, trace};
-use tokio::net::tcp::OwnedWriteHalf;
-
-use std::collections::HashMap;
-
-use crate::config::{get_config, parse};
-use crate::errors::Error;
-use crate::messages::*;
-use crate::pool::ConnectionPool;
-use crate::stats::get_stats;
-
-/// Handle admin client
-pub async fn handle_admin(
-    stream: &mut OwnedWriteHalf,
-    mut query: BytesMut,
-    pool: ConnectionPool,
-) -> Result<(), Error> {
-    let code = query.get_u8() as char;
-
-    if code != 'Q' {
-        return Err(Error::ProtocolSyncError);
-    }
-
-    let len = query.get_i32() as usize;
-    let query = String::from_utf8_lossy(&query[..len - 5])
-        .to_string()
-        .to_ascii_uppercase();
-
-    trace!("Admin query: {}", query);
-
-    if query.starts_with("SHOW STATS") {
-        trace!("SHOW STATS");
-        show_stats(stream, &pool).await
-    } else if query.starts_with("RELOAD") {
-        trace!("RELOAD");
-        reload(stream).await
-    } else if query.starts_with("SHOW CONFIG") {
-        trace!("SHOW CONFIG");
-        show_config(stream).await
-    } else if query.starts_with("SHOW DATABASES") {
-        trace!("SHOW DATABASES");
-        show_databases(stream, &pool).await
-    } else if query.starts_with("SHOW POOLS") {
-        trace!("SHOW POOLS");
-        show_pools(stream, &pool).await
-    } else if query.starts_with("SHOW LISTS") {
-        trace!("SHOW LISTS");
-        show_lists(stream, &pool).await
-    } else if query.starts_with("SHOW VERSION") {
-        trace!("SHOW VERSION");
-        show_version(stream).await
-    } else if query.starts_with("SET ") {
-        trace!("SET");
-        ignore_set(stream).await
-    } else {
-        error_response(stream, "Unsupported query against the admin database").await
-    }
-}
-
-/// SHOW LISTS
-async fn show_lists(stream: &mut OwnedWriteHalf, pool: &ConnectionPool) -> Result<(), Error> {
-    let stats = get_stats();
-
-    let columns = vec![("list", DataType::Text), ("items", DataType::Int4)];
-
-    let mut res = BytesMut::new();
-    res.put(row_description(&columns));
-    res.put(data_row(&vec![
-        "databases".to_string(),
-        (pool.databases() + 1).to_string(), // see comment below
-    ]));
-    res.put(data_row(&vec!["users".to_string(), "1".to_string()]));
-    res.put(data_row(&vec![
-        "pools".to_string(),
-        (pool.databases() + 1).to_string(), // +1 for the pgbouncer admin db pool which isn't real
-    ])); // but admin tools that work with pgbouncer want this
-    res.put(data_row(&vec![
-        "free_clients".to_string(),
-        stats
-            .keys()
-            .map(|address_id| stats[&address_id]["cl_idle"])
-            .sum::<i64>()
-            .to_string(),
-    ]));
-    res.put(data_row(&vec![
-        "used_clients".to_string(),
-        stats
-            .keys()
-            .map(|address_id| stats[&address_id]["cl_active"])
-            .sum::<i64>()
-            .to_string(),
-    ]));
-    res.put(data_row(&vec![
-        "login_clients".to_string(),
-        "0".to_string(),
-    ]));
-    res.put(data_row(&vec![
-        "free_servers".to_string(),
-        stats
-            .keys()
-            .map(|address_id| stats[&address_id]["sv_idle"])
-            .sum::<i64>()
-            .to_string(),
-    ]));
-    res.put(data_row(&vec![
-        "used_servers".to_string(),
-        stats
-            .keys()
-            .map(|address_id| stats[&address_id]["sv_active"])
-            .sum::<i64>()
-            .to_string(),
-    ]));
-    res.put(data_row(&vec!["dns_names".to_string(), "0".to_string()]));
-    res.put(data_row(&vec!["dns_zones".to_string(), "0".to_string()]));
-    res.put(data_row(&vec!["dns_queries".to_string(), "0".to_string()]));
-    res.put(data_row(&vec!["dns_pending".to_string(), "0".to_string()]));
-
-    res.put(command_complete("SHOW"));
-
-    res.put_u8(b'Z');
-    res.put_i32(5);
-    res.put_u8(b'I');
-
-    write_all_half(stream, res).await
-}
-
-/// SHOW VERSION
-async fn show_version(stream: &mut OwnedWriteHalf) -> Result<(), Error> {
-    let mut res = BytesMut::new();
-
-    res.put(row_description(&vec![("version", DataType::Text)]));
-    res.put(data_row(&vec!["PgCat 0.1.0".to_string()]));
-    res.put(command_complete("SHOW"));
-
-    res.put_u8(b'Z');
-    res.put_i32(5);
-    res.put_u8(b'I');
-
-    write_all_half(stream, res).await
-}
-
-/// SHOW POOLS
-async fn show_pools(stream: &mut OwnedWriteHalf, pool: &ConnectionPool) -> Result<(), Error> {
-    let stats = get_stats();
-    let config = {
-        let guard = get_config();
-        &*guard.clone()
-    };
-
-    let columns = vec![
-        ("database", DataType::Text),
-        ("user", DataType::Text),
-        ("cl_active", DataType::Numeric),
-        ("cl_waiting", DataType::Numeric),
-        ("cl_cancel_req", DataType::Numeric),
-        ("sv_active", DataType::Numeric),
-        ("sv_idle", DataType::Numeric),
-        ("sv_used", DataType::Numeric),
-        ("sv_tested", DataType::Numeric),
-        ("sv_login", DataType::Numeric),
-        ("maxwait", DataType::Numeric),
-        ("maxwait_us", DataType::Numeric),
-        ("pool_mode", DataType::Text),
-    ];
-
-    let mut res = BytesMut::new();
-    res.put(row_description(&columns));
-
-    for shard in 0..pool.shards() {
-        for server in 0..pool.servers(shard) {
-            let address = pool.address(shard, server);
-            let stats = match stats.get(&address.id) {
-                Some(stats) => stats.clone(),
-                None => HashMap::new(),
-            };
-
-            let mut row = vec![address.name(), config.user.name.clone()];
-
-            for column in &columns[2..columns.len() - 1] {
-                let value = stats.get(column.0).unwrap_or(&0).to_string();
-                row.push(value);
-            }
-
-            row.push(config.general.pool_mode.to_string());
-            res.put(data_row(&row));
-        }
-    }
-
-    res.put(command_complete("SHOW"));
-
-    res.put_u8(b'Z');
-    res.put_i32(5);
-    res.put_u8(b'I');
-
-    write_all_half(stream, res).await
-}
-
-/// SHOW DATABASES
-async fn show_databases(stream: &mut OwnedWriteHalf, pool: &ConnectionPool) -> Result<(), Error> {
-    let guard = get_config();
-    let config = &*guard.clone();
-    drop(guard);
-
-    // Columns
-    let columns = vec![
-        ("name", DataType::Text),
-        ("host", DataType::Text),
-        ("port", DataType::Text),
-        ("database", DataType::Text),
-        ("force_user", DataType::Text),
-        ("pool_size", DataType::Int4),
-        ("min_pool_size", DataType::Int4),
-        ("reserve_pool", DataType::Int4),
-        ("pool_mode", DataType::Text),
-        ("max_connections", DataType::Int4),
-        ("current_connections", DataType::Int4),
-        ("paused", DataType::Int4),
-        ("disabled", DataType::Int4),
-    ];
-
-    let mut res = BytesMut::new();
-
-    // RowDescription
-    res.put(row_description(&columns));
-
-    for shard in 0..pool.shards() {
-        let database_name = &config.shards[&shard.to_string()].database;
-
-        for server in 0..pool.servers(shard) {
-            let address = pool.address(shard, server);
-            let pool_state = pool.pool_state(shard, server);
-
-            res.put(data_row(&vec![
-                address.name(),                       // name
-                address.host.to_string(),             // host
-                address.port.to_string(),             // port
-                database_name.to_string(),            // database
-                config.user.name.to_string(),         // force_user
-                config.general.pool_size.to_string(), // pool_size
-                "0".to_string(),                      // min_pool_size
-                "0".to_string(),                      // reserve_pool
-                config.general.pool_mode.to_string(), // pool_mode
-                config.general.pool_size.to_string(), // max_connections
-                pool_state.connections.to_string(),   // current_connections
-                "0".to_string(),                      // paused
-                "0".to_string(),                      // disabled
-            ]));
-        }
-    }
-
-    res.put(command_complete("SHOW"));
-
-    // ReadyForQuery
-    res.put_u8(b'Z');
-    res.put_i32(5);
-    res.put_u8(b'I');
-
-    write_all_half(stream, res).await
-}
-
-/// Ignore any SET commands the client sends.
-/// This is common initialization done by ORMs.
-async fn ignore_set(stream: &mut OwnedWriteHalf) -> Result<(), Error> {
-    custom_protocol_response_ok(stream, "SET").await
-}
-
-/// RELOAD
-async fn reload(stream: &mut OwnedWriteHalf) -> Result<(), Error> {
-    info!("Reloading config");
-
-    let config = get_config();
-    let path = config.path.clone().unwrap();
-
-    parse(&path).await?;
-
-    let config = get_config();
-
-    config.show();
-
-    let mut res = BytesMut::new();
-
-    // CommandComplete
-    res.put(command_complete("RELOAD"));
-
-    // ReadyForQuery
-    res.put_u8(b'Z');
-    res.put_i32(5);
-    res.put_u8(b'I');
-
-    write_all_half(stream, res).await
-}
-
-async fn show_config(stream: &mut OwnedWriteHalf) -> Result<(), Error> {
-    let guard = get_config();
-    let config = &*guard.clone();
-    let config: HashMap<String, String> = config.into();
-    drop(guard);
-
-    // Configs that cannot be changed dynamically.
-    let immutables = ["host", "port", "connect_timeout"];
-
-    // Columns
-    let columns = vec![
-        ("key", DataType::Text),
-        ("value", DataType::Text),
-        ("default", DataType::Text),
-        ("changeable", DataType::Text),
-    ];
-
-    // Response data
-    let mut res = BytesMut::new();
-    res.put(row_description(&columns));
-
-    // DataRow rows
-    for (key, value) in config {
-        let changeable = if immutables.iter().filter(|col| *col == &key).count() == 1 {
-            "no".to_string()
-        } else {
-            "yes".to_string()
-        };
-
-        let row = vec![key, value, "-".to_string(), changeable];
-
-        res.put(data_row(&row));
-    }
-
-    res.put(command_complete("SHOW"));
-
-    res.put_u8(b'Z');
-    res.put_i32(5);
-    res.put_u8(b'I');
-
-    write_all_half(stream, res).await
-}
-
-/// SHOW STATS
-async fn show_stats(stream: &mut OwnedWriteHalf, pool: &ConnectionPool) -> Result<(), Error> {
-    let columns = vec![
-        ("database", DataType::Text),
-        ("total_xact_count", DataType::Numeric),
-        ("total_query_count", DataType::Numeric),
-        ("total_received", DataType::Numeric),
-        ("total_sent", DataType::Numeric),
-        ("total_xact_time", DataType::Numeric),
-        ("total_query_time", DataType::Numeric),
-        ("total_wait_time", DataType::Numeric),
-        ("avg_xact_count", DataType::Numeric),
-        ("avg_query_count", DataType::Numeric),
-        ("avg_recv", DataType::Numeric),
-        ("avg_sent", DataType::Numeric),
-        ("avg_xact_time", DataType::Numeric),
-        ("avg_query_time", DataType::Numeric),
-        ("avg_wait_time", DataType::Numeric),
-    ];
-
-    let stats = get_stats();
-    let mut res = BytesMut::new();
-    res.put(row_description(&columns));
-
-    for shard in 0..pool.shards() {
-        for server in 0..pool.servers(shard) {
-            let address = pool.address(shard, server);
-            let stats = match stats.get(&address.id) {
-                Some(stats) => stats.clone(),
-                None => HashMap::new(),
-            };
-
-            let mut row = vec![address.name()];
-
-            for column in &columns[1..] {
-                row.push(stats.get(column.0).unwrap_or(&0).to_string());
-            }
-
-            res.put(data_row(&row));
-        }
-    }
-
-    res.put(command_complete("SHOW"));
-
-    res.put_u8(b'Z');
-    res.put_i32(5);
-    res.put_u8(b'I');
-
-    write_all_half(stream, res).await
-}

src/client.rs

@@ -11,7 +11,6 @@ use tokio::net::{
 
 use std::collections::HashMap;
 
-use crate::admin::handle_admin;
 use crate::config::get_config;
 use crate::constants::*;
 use crate::errors::Error;
@@ -55,15 +54,6 @@ pub struct Client {
 
     // Statistics
     stats: Reporter,
-
-    // Clients want to talk to admin
-    admin: bool,
-
-    // Last address the client talked to
-    last_address_id: Option<usize>,
-
-    // Last server process id we talked to
-    last_server_id: Option<i32>,
 }
 
 impl Client {
@@ -128,15 +118,6 @@ impl Client {
                     ready_for_query(&mut stream).await?;
                     trace!("Startup OK");
 
-                    let database = parameters
-                        .get("database")
-                        .unwrap_or(parameters.get("user").unwrap());
-                    let admin = ["pgcat", "pgbouncer"]
-                        .iter()
-                        .filter(|db| *db == &database)
-                        .count()
-                        == 1;
-
                     // Split the read and write streams
                     // so we can control buffering.
                     let (read, write) = stream.into_split();
@@ -152,9 +133,6 @@ impl Client {
                         client_server_map: client_server_map,
                         parameters: parameters,
                         stats: stats,
-                        admin: admin,
-                        last_address_id: None,
-                        last_server_id: None,
                     });
                 }
 
@@ -176,9 +154,6 @@ impl Client {
                         client_server_map: client_server_map,
                         parameters: HashMap::new(),
                         stats: stats,
-                        admin: false,
-                        last_address_id: None,
-                        last_server_id: None,
                     });
                 }
 
@@ -229,6 +204,9 @@ impl Client {
         loop {
             trace!("Client idle, waiting for message");
 
+            // Client idle, waiting for messages.
+            self.stats.client_idle(self.process_id);
+
             // Read a complete message from the client, which normally would be
             // either a `Q` (query) or `P` (prepare, extended protocol).
             // We can parse it here before grabbing a server from the pool,
@@ -242,13 +220,6 @@ impl Client {
                 return Ok(());
             }
 
-            // Handle admin database real quick
-            if self.admin {
-                trace!("Handling admin command");
-                handle_admin(&mut self.write, message, pool.clone()).await?;
-                continue;
-            }
-
             // Handle all custom protocol commands here.
             match query_router.try_execute_command(message.clone()) {
                 // Normal query
@@ -299,13 +270,13 @@ impl Client {
                 continue;
             }
 
+            // Waiting for server connection.
+            self.stats.client_waiting(self.process_id);
+
             debug!("Waiting for connection from pool");
 
             // Grab a server from the pool: the client issued a regular query.
-            let connection = match pool
+            let connection = match pool.get(query_router.shard(), query_router.role()).await {
-                .get(query_router.shard(), query_router.role(), self.process_id)
-                .await
-            {
                 Ok(conn) => {
                     debug!("Got connection from pool");
                     conn
@@ -319,23 +290,15 @@ impl Client {
             };
 
             let mut reference = connection.0;
-            let address = connection.1;
+            let _address = connection.1;
             let server = &mut *reference;
 
             // Claim this server as mine for query cancellation.
             server.claim(self.process_id, self.secret_key);
 
-            // "disconnect" from the previous server stats-wise
-            if let Some(last_address_id) = self.last_address_id {
-                self.stats
-                    .client_disconnecting(self.process_id, last_address_id);
-            }
-
             // Client active & server active
-            self.stats.client_active(self.process_id, address.id);
+            self.stats.client_active(self.process_id);
-            self.stats.server_active(server.process_id(), address.id);
+            self.stats.server_active(server.process_id());
-            self.last_address_id = Some(address.id);
-            self.last_server_id = Some(server.process_id());
 
             debug!(
                 "Client {:?} talking to server {:?}",
@@ -407,17 +370,17 @@ impl Client {
                         }
 
                         // Report query executed statistics.
-                        self.stats.query(self.process_id, address.id);
+                        self.stats.query();
 
                         // The transaction is over, we can release the connection back to the pool.
                         if !server.in_transaction() {
                             // Report transaction executed statistics.
-                            self.stats.transaction(self.process_id, address.id);
+                            self.stats.transaction();
 
                             // Release server back to the pool if we are in transaction mode.
                             // If we are in session mode, we keep the server until the client disconnects.
                             if self.transaction_mode {
-                                self.stats.server_idle(server.process_id(), address.id);
+                                self.stats.server_idle(server.process_id());
                                 break;
                             }
                         }
@@ -493,15 +456,15 @@ impl Client {
                         }
 
                         // Report query executed statistics.
-                        self.stats.query(self.process_id, address.id);
+                        self.stats.query();
 
                         // Release server back to the pool if we are in transaction mode.
                         // If we are in session mode, we keep the server until the client disconnects.
                         if !server.in_transaction() {
-                            self.stats.transaction(self.process_id, address.id);
+                            self.stats.transaction();
 
                             if self.transaction_mode {
-                                self.stats.server_idle(server.process_id(), address.id);
+                                self.stats.server_idle(server.process_id());
                                 break;
                             }
                         }
@@ -532,10 +495,10 @@ impl Client {
                         // Release server back to the pool if we are in transaction mode.
                         // If we are in session mode, we keep the server until the client disconnects.
                         if !server.in_transaction() {
-                            self.stats.transaction(self.process_id, address.id);
+                            self.stats.transaction();
 
                             if self.transaction_mode {
-                                self.stats.server_idle(server.process_id(), address.id);
+                                self.stats.server_idle(server.process_id());
                                 break;
                             }
                         }
@@ -552,7 +515,6 @@ impl Client {
             // The server is no longer bound to us, we can't cancel it's queries anymore.
             debug!("Releasing server back into the pool");
             self.release();
-            self.stats.client_idle(self.process_id, address.id);
         }
     }
 
@@ -565,14 +527,6 @@ impl Client {
 
 impl Drop for Client {
     fn drop(&mut self) {
-        // Disconnect the client
+        self.stats.client_disconnecting(self.process_id);
-        if let Some(address_id) = self.last_address_id {
-            self.stats.client_disconnecting(self.process_id, address_id);
-
-            // The server is now idle
-            if let Some(process_id) = self.last_server_id {
-                self.stats.server_idle(process_id, address_id);
-            }
-        }
     }
 }

src/config.rs

@@ -19,15 +19,6 @@ pub enum Role {
     Replica,
 }
 
-impl ToString for Role {
-    fn to_string(&self) -> String {
-        match *self {
-            Role::Primary => "primary".to_string(),
-            Role::Replica => "replica".to_string(),
-        }
-    }
-}
-
 impl PartialEq<Option<Role>> for Role {
     fn eq(&self, other: &Option<Role>) -> bool {
         match other {
@@ -48,37 +39,23 @@ impl PartialEq<Role> for Option<Role> {
 
 #[derive(Clone, PartialEq, Hash, std::cmp::Eq, Debug)]
 pub struct Address {
-    pub id: usize,
     pub host: String,
     pub port: String,
     pub shard: usize,
     pub role: Role,
-    pub replica_number: usize,
 }
 
 impl Default for Address {
     fn default() -> Address {
         Address {
-            id: 0,
             host: String::from("127.0.0.1"),
             port: String::from("5432"),
             shard: 0,
-            replica_number: 0,
             role: Role::Replica,
         }
     }
 }
 
-impl Address {
-    pub fn name(&self) -> String {
-        match self.role {
-            Role::Primary => format!("shard_{}_primary", self.shard),
-
-            Role::Replica => format!("shard_{}_replica_{}", self.shard, self.replica_number),
-        }
-    }
-}
-
 #[derive(Clone, PartialEq, Hash, std::cmp::Eq, Deserialize, Debug)]
 pub struct User {
     pub name: String,
@@ -157,7 +134,6 @@ impl Default for QueryRouter {
 
 #[derive(Deserialize, Debug, Clone)]
 pub struct Config {
-    pub path: Option<String>,
     pub general: General,
     pub user: User,
     pub shards: HashMap<String, Shard>,
@@ -167,7 +143,6 @@ pub struct Config {
 impl Default for Config {
     fn default() -> Config {
         Config {
-            path: Some(String::from("pgcat.toml")),
             general: General::default(),
             user: User::default(),
             shards: HashMap::from([(String::from("1"), Shard::default())]),
@@ -176,52 +151,6 @@ impl Default for Config {
     }
 }
 
-impl From<&Config> for std::collections::HashMap<String, String> {
-    fn from(config: &Config) -> HashMap<String, String> {
-        HashMap::from([
-            ("host".to_string(), config.general.host.to_string()),
-            ("port".to_string(), config.general.port.to_string()),
-            (
-                "pool_size".to_string(),
-                config.general.pool_size.to_string(),
-            ),
-            (
-                "pool_mode".to_string(),
-                config.general.pool_mode.to_string(),
-            ),
-            (
-                "connect_timeout".to_string(),
-                config.general.connect_timeout.to_string(),
-            ),
-            (
-                "healthcheck_timeout".to_string(),
-                config.general.healthcheck_timeout.to_string(),
-            ),
-            ("ban_time".to_string(), config.general.ban_time.to_string()),
-            (
-                "statsd_address".to_string(),
-                config.general.statsd_address.to_string(),
-            ),
-            (
-                "default_role".to_string(),
-                config.query_router.default_role.to_string(),
-            ),
-            (
-                "query_parser_enabled".to_string(),
-                config.query_router.query_parser_enabled.to_string(),
-            ),
-            (
-                "primary_reads_enabled".to_string(),
-                config.query_router.primary_reads_enabled.to_string(),
-            ),
-            (
-                "sharding_function".to_string(),
-                config.query_router.sharding_function.to_string(),
-            ),
-        ])
-    }
-}
-
 impl Config {
     pub fn show(&self) {
         info!("Pool size: {}", self.general.pool_size);
@@ -260,7 +189,7 @@ pub async fn parse(path: &str) -> Result<(), Error> {
         }
     };
 
-    let mut config: Config = match toml::from_str(&contents) {
+    let config: Config = match toml::from_str(&contents) {
         Ok(config) => config,
         Err(err) => {
             error!("Could not parse config file: {}", err.to_string());
@@ -350,8 +279,6 @@ pub async fn parse(path: &str) -> Result<(), Error> {
         }
     };
 
-    config.path = Some(path.to_string());
-
     CONFIG.store(Arc::new(config.clone()));
 
     Ok(())
@@ -369,6 +296,5 @@ mod test {
         assert_eq!(get_config().shards["1"].servers[0].0, "127.0.0.1");
         assert_eq!(get_config().shards["0"].servers[0].2, "primary");
         assert_eq!(get_config().query_router.default_role, "any");
-        assert_eq!(get_config().path, Some("pgcat.toml".to_string()));
     }
 }

src/constants.rs

@@ -20,8 +20,3 @@ pub const AUTHENTICATION_SUCCESSFUL: i32 = 0;
 
 // ErrorResponse: A code identifying the field type; if zero, this is the message terminator and no string follows.
 pub const MESSAGE_TERMINATOR: u8 = 0;
-
-//
-// Data types
-//
-pub const _OID_INT8: i32 = 20; // bigint

src/main.rs

@@ -31,6 +31,7 @@ extern crate once_cell;
 extern crate serde;
 extern crate serde_derive;
 extern crate sqlparser;
+extern crate statsd;
 extern crate tokio;
 extern crate toml;
 
@@ -46,7 +47,6 @@ use tokio::{
 use std::collections::HashMap;
 use std::sync::Arc;
 
-mod admin;
 mod client;
 mod config;
 mod constants;
@@ -112,19 +112,15 @@ async fn main() {
 
     // Collect statistics and send them to StatsD
     let (tx, rx) = mpsc::channel(100);
+    let collector_tx = tx.clone();
+    tokio::task::spawn(async move {
+        let mut stats_collector = Collector::new(rx, collector_tx);
+        stats_collector.collect().await;
+    });
 
-    // Connection pool for all shards and replicas
     let mut pool =
         ConnectionPool::from_config(client_server_map.clone(), Reporter::new(tx.clone())).await;
 
-    let collector_tx = tx.clone();
-    let addresses = pool.databases();
-    tokio::task::spawn(async move {
-        let mut stats_collector = Collector::new(rx, collector_tx);
-        stats_collector.collect(addresses).await;
-    });
-
-    // Connect to all servers and validate their versions.
     let server_info = match pool.validate().await {
         Ok(info) => info,
         Err(err) => {

src/messages.rs

@@ -8,26 +8,9 @@ use tokio::net::{
     TcpStream,
 };
 
-use crate::errors::Error;
 use std::collections::HashMap;
 
-/// Postgres data type mappings
+use crate::errors::Error;
-/// used in RowDescription ('T') message.
-pub enum DataType {
-    Text,
-    Int4,
-    Numeric,
-}
-
-impl From<&DataType> for i32 {
-    fn from(data_type: &DataType) -> i32 {
-        match data_type {
-            DataType::Text => 25,
-            DataType::Int4 => 23,
-            DataType::Numeric => 1700,
-        }
-    }
-}
 
 /// Tell the client that authentication handshake completed successfully.
 pub async fn auth_ok(stream: &mut TcpStream) -> Result<(), Error> {
@@ -108,8 +91,9 @@ pub async fn startup(stream: &mut TcpStream, user: &str, database: &str) -> Resu
     }
 }
 
-/// Parse the params the server sends as a key/value format.
+/// Parse StartupMessage parameters.
-pub fn parse_params(mut bytes: BytesMut) -> Result<HashMap<String, String>, Error> {
+/// e.g. user, database, application_name, etc.
+pub fn parse_startup(mut bytes: BytesMut) -> Result<HashMap<String, String>, Error> {
     let mut result = HashMap::new();
     let mut buf = Vec::new();
     let mut tmp = String::new();
@@ -131,7 +115,7 @@ pub fn parse_params(mut bytes: BytesMut) -> Result<HashMap<String, String>, Erro
 
     // Expect pairs of name and value
     // and at least one pair to be present.
-    if buf.len() % 2 != 0 || buf.len() < 2 {
+    if buf.len() % 2 != 0 && buf.len() >= 2 {
         return Err(Error::ClientBadStartup);
     }
 
@@ -143,14 +127,6 @@ pub fn parse_params(mut bytes: BytesMut) -> Result<HashMap<String, String>, Erro
         i += 2;
     }
 
-    Ok(result)
-}
-
-/// Parse StartupMessage parameters.
-/// e.g. user, database, application_name, etc.
-pub fn parse_startup(bytes: BytesMut) -> Result<HashMap<String, String>, Error> {
-    let result = parse_params(bytes)?;
-
     // Minimum required parameters
     // I want to have the user at the very minimum, according to the protocol spec.
     if !result.contains_key("user") {
@@ -276,17 +252,68 @@ pub async fn show_response(
     // 3. CommandComplete
     // 4. ReadyForQuery
 
+    // RowDescription
+    let mut row_desc = BytesMut::new();
+
+    // Number of columns: 1
+    row_desc.put_i16(1);
+
+    // Column name
+    row_desc.put_slice(&format!("{}\0", name).as_bytes());
+
+    // Doesn't belong to any table
+    row_desc.put_i32(0);
+
+    // Doesn't belong to any table
+    row_desc.put_i16(0);
+
+    // Text
+    row_desc.put_i32(25);
+
+    // Text size = variable (-1)
+    row_desc.put_i16(-1);
+
+    // Type modifier: none that I know
+    row_desc.put_i32(0);
+
+    // Format being used: text (0), binary (1)
+    row_desc.put_i16(0);
+
+    // DataRow
+    let mut data_row = BytesMut::new();
+
+    // Number of columns
+    data_row.put_i16(1);
+
+    // Size of the column content (length of the string really)
+    data_row.put_i32(value.len() as i32);
+
+    // The content
+    data_row.put_slice(value.as_bytes());
+
+    // CommandComplete
+    let mut command_complete = BytesMut::new();
+
+    // Number of rows returned (just one)
+    command_complete.put_slice(&b"SELECT 1\0"[..]);
+
     // The final messages sent to the client
     let mut res = BytesMut::new();
 
     // RowDescription
-    res.put(row_description(&vec![(name, DataType::Text)]));
+    res.put_u8(b'T');
+    res.put_i32(row_desc.len() as i32 + 4);
+    res.put(row_desc);
 
     // DataRow
-    res.put(data_row(&vec![value.to_string()]));
+    res.put_u8(b'D');
+    res.put_i32(data_row.len() as i32 + 4);
+    res.put(data_row);
 
     // CommandComplete
-    res.put(command_complete("SELECT 1"));
+    res.put_u8(b'C');
+    res.put_i32(command_complete.len() as i32 + 4);
+    res.put(command_complete);
 
     // ReadyForQuery
     res.put_u8(b'Z');
@@ -296,77 +323,6 @@ pub async fn show_response(
     write_all_half(stream, res).await
 }
 
-pub fn row_description(columns: &Vec<(&str, DataType)>) -> BytesMut {
-    let mut res = BytesMut::new();
-    let mut row_desc = BytesMut::new();
-
-    // how many colums we are storing
-    row_desc.put_i16(columns.len() as i16);
-
-    for (name, data_type) in columns {
-        // Column name
-        row_desc.put_slice(&format!("{}\0", name).as_bytes());
-
-        // Doesn't belong to any table
-        row_desc.put_i32(0);
-
-        // Doesn't belong to any table
-        row_desc.put_i16(0);
-
-        // Text
-        row_desc.put_i32(data_type.into());
-
-        // Text size = variable (-1)
-        let type_size = match data_type {
-            DataType::Text => -1,
-            DataType::Int4 => 4,
-            DataType::Numeric => -1,
-        };
-
-        row_desc.put_i16(type_size);
-
-        // Type modifier: none that I know
-        row_desc.put_i32(-1);
-
-        // Format being used: text (0), binary (1)
-        row_desc.put_i16(0);
-    }
-
-    res.put_u8(b'T');
-    res.put_i32(row_desc.len() as i32 + 4);
-    res.put(row_desc);
-
-    res
-}
-
-pub fn data_row(row: &Vec<String>) -> BytesMut {
-    let mut res = BytesMut::new();
-    let mut data_row = BytesMut::new();
-
-    data_row.put_i16(row.len() as i16);
-
-    for column in row {
-        let column = column.as_bytes();
-        data_row.put_i32(column.len() as i32);
-        data_row.put_slice(&column);
-    }
-
-    res.put_u8(b'D');
-    res.put_i32(data_row.len() as i32 + 4);
-    res.put(data_row);
-
-    res
-}
-
-pub fn command_complete(command: &str) -> BytesMut {
-    let cmd = BytesMut::from(format!("{}\0", command).as_bytes());
-    let mut res = BytesMut::new();
-    res.put_u8(b'C');
-    res.put_i32(cmd.len() as i32 + 4);
-    res.put(cmd);
-    res
-}
-
 /// Write all data in the buffer to the TcpStream.
 pub async fn write_all(stream: &mut TcpStream, buf: BytesMut) -> Result<(), Error> {
     match stream.write_all(&buf).await {

src/pool.rs

@@ -38,7 +38,6 @@ impl ConnectionPool {
         let mut shards = Vec::new();
         let mut addresses = Vec::new();
         let mut banlist = Vec::new();
-        let mut address_id = 0;
         let mut shard_ids = config
             .shards
             .clone()
@@ -50,8 +49,7 @@ impl ConnectionPool {
         for shard_idx in shard_ids {
             let shard = &config.shards[&shard_idx];
             let mut pools = Vec::new();
-            let mut servers = Vec::new();
+            let mut replica_addresses = Vec::new();
-            let mut replica_number = 0;
 
             for server in shard.servers.iter() {
                 let role = match server.2.as_ref() {
@@ -64,20 +62,12 @@ impl ConnectionPool {
                 };
 
                 let address = Address {
-                    id: address_id,
                     host: server.0.clone(),
                     port: server.1.to_string(),
                     role: role,
-                    replica_number,
                     shard: shard_idx.parse::<usize>().unwrap(),
                 };
 
-                address_id += 1;
-
-                if role == Role::Replica {
-                    replica_number += 1;
-                }
-
                 let manager = ServerPool::new(
                     address.clone(),
                     config.user.clone(),
@@ -97,11 +87,11 @@ impl ConnectionPool {
                     .unwrap();
 
                 pools.push(pool);
-                servers.push(address);
+                replica_addresses.push(address);
             }
 
             shards.push(pools);
-            addresses.push(servers);
+            addresses.push(replica_addresses);
             banlist.push(HashMap::new());
         }
 
@@ -125,13 +115,9 @@ impl ConnectionPool {
     pub async fn validate(&mut self) -> Result<BytesMut, Error> {
         let mut server_infos = Vec::new();
 
-        let stats = self.stats.clone();
         for shard in 0..self.shards() {
             for _ in 0..self.servers(shard) {
-                // To keep stats consistent.
+                let connection = match self.get(shard, None).await {
-                let fake_process_id = 0;
-
-                let connection = match self.get(shard, None, fake_process_id).await {
                     Ok(conn) => conn,
                     Err(err) => {
                         error!("Shard {} down or misconfigured: {:?}", shard, err);
@@ -140,24 +126,10 @@ impl ConnectionPool {
                 };
 
                 let mut proxy = connection.0;
-                let address = connection.1;
+                let _address = connection.1;
                 let server = &mut *proxy;
 
-                let server_info = server.server_info();
+                server_infos.push(server.server_info());
-
-                stats.client_disconnecting(fake_process_id, address.id);
-
-                if server_infos.len() > 0 {
-                    // Compare against the last server checked.
-                    if server_info != server_infos[server_infos.len() - 1] {
-                        warn!(
-                            "{:?} has different server configuration than the last server",
-                            address
-                        );
-                    }
-                }
-
-                server_infos.push(server_info);
             }
         }
 
@@ -175,7 +147,6 @@ impl ConnectionPool {
         &mut self,
         shard: usize,
         role: Option<Role>,
-        process_id: i32,
     ) -> Result<(PooledConnection<'_, ServerPool>, Address), Error> {
         let now = Instant::now();
         let addresses = &self.addresses[shard];
@@ -211,8 +182,6 @@ impl ConnectionPool {
             let index = self.round_robin % addresses.len();
             let address = &addresses[index];
 
-            self.stats.client_waiting(process_id, address.id);
-
             // Make sure you're getting a primary or a replica
             // as per request. If no specific role is requested, the first
             // available will be chosen.
@@ -232,9 +201,6 @@ impl ConnectionPool {
                 Err(err) => {
                     error!("Banning replica {}, error: {:?}", index, err);
                     self.ban(address, shard);
-                    self.stats.client_disconnecting(process_id, address.id);
-                    self.stats
-                        .checkout_time(now.elapsed().as_micros(), process_id, address.id);
                     continue;
                 }
             };
@@ -243,7 +209,7 @@ impl ConnectionPool {
             let server = &mut *conn;
             let healthcheck_timeout = get_config().general.healthcheck_timeout;
 
-            self.stats.server_tested(server.process_id(), address.id);
+            self.stats.server_tested(server.process_id());
 
             match tokio::time::timeout(
                 tokio::time::Duration::from_millis(healthcheck_timeout),
@@ -254,9 +220,8 @@ impl ConnectionPool {
                 // Check if health check succeeded
                 Ok(res) => match res {
                     Ok(_) => {
-                        self.stats
+                        self.stats.checkout_time(now.elapsed().as_micros());
-                            .checkout_time(now.elapsed().as_micros(), process_id, address.id);
+                        self.stats.server_idle(conn.process_id());
-                        self.stats.server_idle(conn.process_id(), address.id);
                         return Ok((conn, address.clone()));
                     }
                     Err(_) => {
@@ -265,9 +230,6 @@ impl ConnectionPool {
                         server.mark_bad();
 
                         self.ban(address, shard);
-                        self.stats.client_disconnecting(process_id, address.id);
-                        self.stats
-                            .checkout_time(now.elapsed().as_micros(), process_id, address.id);
                         continue;
                     }
                 },
@@ -278,9 +240,6 @@ impl ConnectionPool {
                     server.mark_bad();
 
                     self.ban(address, shard);
-                    self.stats.client_disconnecting(process_id, address.id);
-                    self.stats
-                        .checkout_time(now.elapsed().as_micros(), process_id, address.id);
                     continue;
                 }
             }
@@ -365,22 +324,6 @@ impl ConnectionPool {
     pub fn servers(&self, shard: usize) -> usize {
         self.addresses[shard].len()
     }
-
-    pub fn databases(&self) -> usize {
-        let mut databases = 0;
-        for shard in 0..self.shards() {
-            databases += self.servers(shard);
-        }
-        databases
-    }
-
-    pub fn pool_state(&self, shard: usize, server: usize) -> bb8::State {
-        self.databases[shard][server].state()
-    }
-
-    pub fn address(&self, shard: usize, server: usize) -> &Address {
-        &self.addresses[shard][server]
-    }
 }
 
 pub struct ServerPool {
@@ -418,14 +361,13 @@ impl ManageConnection for ServerPool {
     async fn connect(&self) -> Result<Self::Connection, Self::Error> {
         info!(
             "Creating a new connection to {:?} using user {:?}",
-            self.address.name(),
+            self.address, self.user.name
-            self.user.name
         );
 
         // Put a temporary process_id into the stats
         // for server login.
         let process_id = rand::random::<i32>();
-        self.stats.server_login(process_id, self.address.id);
+        self.stats.server_login(process_id);
 
         match Server::startup(
             &self.address,
@@ -438,12 +380,12 @@ impl ManageConnection for ServerPool {
         {
             Ok(conn) => {
                 // Remove the temporary process_id from the stats.
-                self.stats.server_disconnecting(process_id, self.address.id);
+                self.stats.server_disconnecting(process_id);
                 Ok(conn)
             }
             Err(err) => {
                 // Remove the temporary process_id from the stats.
-                self.stats.server_disconnecting(process_id, self.address.id);
+                self.stats.server_disconnecting(process_id);
                 Err(err)
             }
         }

src/query_router.rs

@@ -12,7 +12,7 @@ use sqlparser::parser::Parser;
 
 const CUSTOM_SQL_REGEXES: [&str; 5] = [
     r"(?i)^ *SET SHARDING KEY TO '?([0-9]+)'? *;? *$",
-    r"(?i)^ *SET SHARD TO '?([0-9]+|ANY)'? *;? *$",
+    r"(?i)^ *SET SHARD TO '?([0-9]+)'? *;? *$",
     r"(?i)^ *SHOW SHARD *;? *$",
     r"(?i)^ *SET SERVER ROLE TO '(PRIMARY|REPLICA|ANY|AUTO|DEFAULT)' *;? *$",
     r"(?i)^ *SHOW SERVER ROLE *;? *$",
@@ -192,10 +192,7 @@ impl QueryRouter {
             }
 
             Command::SetShard => {
-                self.active_shard = match value.to_ascii_uppercase().as_ref() {
+                self.active_shard = Some(value.parse::<usize>().unwrap());
-                    "ANY" => Some(rand::random::<usize>() % self.shards),
-                    _ => Some(value.parse::<usize>().unwrap()),
-                };
             }
 
             Command::SetServerRole => {

src/server.rs

@@ -268,8 +268,7 @@ impl Server {
 
     /// Send messages to the server from the client.
     pub async fn send(&mut self, messages: BytesMut) -> Result<(), Error> {
-        self.stats
+        self.stats.data_sent(messages.len());
-            .data_sent(messages.len(), self.process_id, self.address.id);
 
         match write_all_half(&mut self.write, messages).await {
             Ok(_) => Ok(()),
@@ -375,8 +374,7 @@ impl Server {
         let bytes = self.buffer.clone();
 
         // Keep track of how much data we got from the server for stats.
-        self.stats
+        self.stats.data_received(bytes.len());
-            .data_received(bytes.len(), self.process_id, self.address.id);
 
         // Clear the buffer for next query.
         self.buffer.clear();
@@ -472,8 +470,7 @@ impl Drop for Server {
     /// the socket is in non-blocking mode, so it may not be ready
     /// for a write.
     fn drop(&mut self) {
-        self.stats
+        self.stats.server_disconnecting(self.process_id());
-            .server_disconnecting(self.process_id(), self.address.id);
 
         let mut bytes = BytesMut::with_capacity(4);
         bytes.put_u8(b'X');

src/stats.rs

@@ -1,20 +1,12 @@
-/// Statistics and reporting.
 use log::info;
-use once_cell::sync::Lazy;
+use statsd::Client;
-use parking_lot::Mutex;
+/// Events collector and publisher.
-use std::collections::HashMap;
 use tokio::sync::mpsc::{Receiver, Sender};
 
-// Latest stats updated every second; used in SHOW STATS and other admin commands.
+use std::collections::HashMap;
-static LATEST_STATS: Lazy<Mutex<HashMap<usize, HashMap<String, i64>>>> =
-    Lazy::new(|| Mutex::new(HashMap::new()));
 
-// Statistics period used for average calculations.
+use crate::config::get_config;
-// 15 seconds.
-static STAT_PERIOD: u64 = 15000;
 
-/// The names for the events reported
-/// to the statistics collector.
 #[derive(Debug, Clone, Copy)]
 enum EventName {
     CheckoutTime,
@@ -31,266 +23,201 @@ enum EventName {
     ServerTested,
     ServerLogin,
     ServerDisconnecting,
-    UpdateStats,
+    FlushStatsToStatsD,
-    UpdateAverages,
 }
 
-/// Event data sent to the collector
-/// from clients and servers.
 #[derive(Debug)]
 pub struct Event {
-    /// The name of the event being reported.
     name: EventName,
-
-    /// The value being reported. Meaning differs based on event name.
     value: i64,
-
+    process_id: Option<i32>,
-    /// The client or server connection reporting the event.
-    process_id: i32,
-
-    /// The server the client is connected to.
-    address_id: usize,
 }
 
-/// The statistics reporter. An instance is given
-/// to each possible source of statistics,
-/// e.g. clients, servers, connection pool.
 #[derive(Clone, Debug)]
 pub struct Reporter {
     tx: Sender<Event>,
 }
 
 impl Reporter {
-    /// Create a new Reporter instance.
     pub fn new(tx: Sender<Event>) -> Reporter {
         Reporter { tx: tx }
     }
 
-    /// Report a query executed by a client against
+    pub fn query(&self) {
-    /// a server identified by the `address_id`.
-    pub fn query(&self, process_id: i32, address_id: usize) {
         let event = Event {
             name: EventName::Query,
             value: 1,
-            process_id: process_id,
+            process_id: None,
-            address_id: address_id,
         };
 
         let _ = self.tx.try_send(event);
     }
 
-    /// Report a transaction executed by a client against
+    pub fn transaction(&self) {
-    /// a server identified by the `address_id`.
-    pub fn transaction(&self, process_id: i32, address_id: usize) {
         let event = Event {
             name: EventName::Transaction,
             value: 1,
-            process_id: process_id,
+            process_id: None,
-            address_id: address_id,
         };
 
         let _ = self.tx.try_send(event);
     }
 
-    /// Report data sent to a server identified by `address_id`.
+    pub fn data_sent(&self, amount: usize) {
-    /// The `amount` is measured in bytes.
-    pub fn data_sent(&self, amount: usize, process_id: i32, address_id: usize) {
         let event = Event {
             name: EventName::DataSent,
             value: amount as i64,
-            process_id: process_id,
+            process_id: None,
-            address_id: address_id,
         };
 
         let _ = self.tx.try_send(event);
     }
 
-    /// Report data received from a server identified by `address_id`.
+    pub fn data_received(&self, amount: usize) {
-    /// The `amount` is measured in bytes.
-    pub fn data_received(&self, amount: usize, process_id: i32, address_id: usize) {
         let event = Event {
             name: EventName::DataReceived,
             value: amount as i64,
-            process_id: process_id,
+            process_id: None,
-            address_id: address_id,
         };
 
         let _ = self.tx.try_send(event);
     }
 
-    /// Time spent waiting to get a healthy connection from the pool
+    pub fn checkout_time(&self, ms: u128) {
-    /// for a server identified by `address_id`.
-    /// Measured in milliseconds.
-    pub fn checkout_time(&self, ms: u128, process_id: i32, address_id: usize) {
         let event = Event {
             name: EventName::CheckoutTime,
             value: ms as i64,
-            process_id: process_id,
+            process_id: None,
-            address_id: address_id,
         };
 
         let _ = self.tx.try_send(event);
     }
 
-    /// Reports a client identified by `process_id` waiting for a connection
+    pub fn client_waiting(&self, process_id: i32) {
-    /// to a server identified by `address_id`.
-    pub fn client_waiting(&self, process_id: i32, address_id: usize) {
         let event = Event {
             name: EventName::ClientWaiting,
             value: 1,
-            process_id: process_id,
+            process_id: Some(process_id),
-            address_id: address_id,
         };
 
         let _ = self.tx.try_send(event);
     }
 
-    /// Reports a client identified by `process_id` is done waiting for a connection
+    pub fn client_active(&self, process_id: i32) {
-    /// to a server identified by `address_id` and is about to query the server.
-    pub fn client_active(&self, process_id: i32, address_id: usize) {
         let event = Event {
             name: EventName::ClientActive,
             value: 1,
-            process_id: process_id,
+            process_id: Some(process_id),
-            address_id: address_id,
         };
 
         let _ = self.tx.try_send(event);
     }
 
-    /// Reports a client identified by `process_id` is done querying the server
+    pub fn client_idle(&self, process_id: i32) {
-    /// identified by `address_id` and is no longer active.
-    pub fn client_idle(&self, process_id: i32, address_id: usize) {
         let event = Event {
             name: EventName::ClientIdle,
             value: 1,
-            process_id: process_id,
+            process_id: Some(process_id),
-            address_id: address_id,
         };
 
         let _ = self.tx.try_send(event);
     }
 
-    /// Reports a client identified by `process_id` is disconecting from the pooler.
+    pub fn client_disconnecting(&self, process_id: i32) {
-    /// The last server it was connected to is identified by `address_id`.
-    pub fn client_disconnecting(&self, process_id: i32, address_id: usize) {
         let event = Event {
             name: EventName::ClientDisconnecting,
             value: 1,
-            process_id: process_id,
+            process_id: Some(process_id),
-            address_id: address_id,
         };
 
         let _ = self.tx.try_send(event);
     }
 
-    /// Reports a server connection identified by `process_id` for
+    pub fn server_active(&self, process_id: i32) {
-    /// a configured server identified by `address_id` is actively used
-    /// by a client.
-    pub fn server_active(&self, process_id: i32, address_id: usize) {
         let event = Event {
             name: EventName::ServerActive,
             value: 1,
-            process_id: process_id,
+            process_id: Some(process_id),
-            address_id: address_id,
         };
 
         let _ = self.tx.try_send(event);
     }
 
-    /// Reports a server connection identified by `process_id` for
+    pub fn server_idle(&self, process_id: i32) {
-    /// a configured server identified by `address_id` is no longer
-    /// actively used by a client and is now idle.
-    pub fn server_idle(&self, process_id: i32, address_id: usize) {
         let event = Event {
             name: EventName::ServerIdle,
             value: 1,
-            process_id: process_id,
+            process_id: Some(process_id),
-            address_id: address_id,
         };
 
         let _ = self.tx.try_send(event);
     }
 
-    /// Reports a server connection identified by `process_id` for
+    pub fn server_login(&self, process_id: i32) {
-    /// a configured server identified by `address_id` is attempting
-    /// to login.
-    pub fn server_login(&self, process_id: i32, address_id: usize) {
         let event = Event {
             name: EventName::ServerLogin,
             value: 1,
-            process_id: process_id,
+            process_id: Some(process_id),
-            address_id: address_id,
         };
 
         let _ = self.tx.try_send(event);
     }
 
-    /// Reports a server connection identified by `process_id` for
+    pub fn server_tested(&self, process_id: i32) {
-    /// a configured server identified by `address_id` is being
-    /// tested before being given to a client.
-    pub fn server_tested(&self, process_id: i32, address_id: usize) {
         let event = Event {
             name: EventName::ServerTested,
             value: 1,
-            process_id: process_id,
+            process_id: Some(process_id),
-            address_id: address_id,
         };
 
         let _ = self.tx.try_send(event);
     }
 
-    /// Reports a server connection identified by `process_id` is disconecting from the pooler.
+    pub fn server_disconnecting(&self, process_id: i32) {
-    /// The configured server it was connected to is identified by `address_id`.
-    pub fn server_disconnecting(&self, process_id: i32, address_id: usize) {
         let event = Event {
             name: EventName::ServerDisconnecting,
             value: 1,
-            process_id: process_id,
+            process_id: Some(process_id),
-            address_id: address_id,
         };
 
         let _ = self.tx.try_send(event);
     }
+
+    // pub fn flush_to_statsd(&self) {
+    //     let event = Event {
+    //         name: EventName::FlushStatsToStatsD,
+    //         value: 0,
+    //         process_id: None,
+    //     };
+
+    //     let _ = self.tx.try_send(event);
+    // }
 }
 
-/// The statistics collector which is receiving statistics
-/// from clients, servers, and the connection pool. There is
-/// only one collector (kind of like a singleton).
-/// The collector can trigger events on its own, e.g.
-/// it updates aggregates every second and averages every
-/// 15 seconds.
 pub struct Collector {
     rx: Receiver<Event>,
     tx: Sender<Event>,
+    client: Client,
 }
 
 impl Collector {
-    /// Create a new collector instance. There should only be one instance
-    /// at a time. This is ensured by mpsc which allows only one receiver.
     pub fn new(rx: Receiver<Event>, tx: Sender<Event>) -> Collector {
-        Collector { rx, tx }
+        Collector {
+            rx,
+            tx,
+            client: Client::new(&get_config().general.statsd_address, "pgcat").unwrap(),
+        }
     }
 
-    /// The statistics collection handler. It will collect statistics
+    pub async fn collect(&mut self) {
-    /// for `address_id`s starting at 0 up to `addresses`.
-    pub async fn collect(&mut self, addresses: usize) {
         info!("Events reporter started");
 
-        let stats_template = HashMap::from([
+        let mut stats = HashMap::from([
             ("total_query_count", 0),
             ("total_xact_count", 0),
             ("total_sent", 0),
             ("total_received", 0),
-            ("total_xact_time", 0),
-            ("total_query_time", 0),
             ("total_wait_time", 0),
-            ("avg_xact_time", 0),
-            ("avg_query_time", 0),
-            ("avg_xact_count", 0),
-            ("avg_sent", 0),
-            ("avg_received", 0),
-            ("avg_wait_time", 0),
             ("maxwait_us", 0),
             ("maxwait", 0),
             ("cl_waiting", 0),
@@ -302,51 +229,21 @@ impl Collector {
             ("sv_tested", 0),
         ]);
 
-        let mut stats = HashMap::new();
+        let mut client_server_states: HashMap<i32, EventName> = HashMap::new();
-
-        // Stats saved after each iteration of the flush event. Used in calculation
-        // of averages in the last flush period.
-        let mut old_stats: HashMap<usize, HashMap<String, i64>> = HashMap::new();
-
-        // Track which state the client and server are at any given time.
-        let mut client_server_states: HashMap<usize, HashMap<i32, EventName>> = HashMap::new();
-
-        // Flush stats to StatsD and calculate averages every 15 seconds.
         let tx = self.tx.clone();
+
         tokio::task::spawn(async move {
-            let mut interval =
+            let mut interval = tokio::time::interval(tokio::time::Duration::from_millis(15000));
-                tokio::time::interval(tokio::time::Duration::from_millis(STAT_PERIOD / 15));
             loop {
                 interval.tick().await;
-                for address_id in 0..addresses {
+                let _ = tx.try_send(Event {
-                    let _ = tx.try_send(Event {
+                    name: EventName::FlushStatsToStatsD,
-                        name: EventName::UpdateStats,
+                    value: 0,
-                        value: 0,
+                    process_id: None,
-                        process_id: -1,
+                });
-                        address_id: address_id,
-                    });
-                }
             }
         });
 
-        let tx = self.tx.clone();
-        tokio::task::spawn(async move {
-            let mut interval =
-                tokio::time::interval(tokio::time::Duration::from_millis(STAT_PERIOD));
-            loop {
-                interval.tick().await;
-                for address_id in 0..addresses {
-                    let _ = tx.try_send(Event {
-                        name: EventName::UpdateAverages,
-                        value: 0,
-                        process_id: -1,
-                        address_id: address_id,
-                    });
-                }
-            }
-        });
-
-        // The collector loop
         loop {
             let stat = match self.rx.recv().await {
                 Some(stat) => stat,
@@ -356,14 +253,6 @@ impl Collector {
                 }
             };
 
-            let stats = stats
-                .entry(stat.address_id)
-                .or_insert(stats_template.clone());
-            let client_server_states = client_server_states
-                .entry(stat.address_id)
-                .or_insert(HashMap::new());
-            let old_stats = old_stats.entry(stat.address_id).or_insert(HashMap::new());
-
             // Some are counters, some are gauges...
             match stat.name {
                 EventName::Query => {
@@ -391,11 +280,10 @@ impl Collector {
                     *counter += stat.value;
 
                     let counter = stats.entry("maxwait_us").or_insert(0);
-                    let mic_part = stat.value % 1_000_000;
 
                     // Report max time here
-                    if mic_part > *counter {
+                    if stat.value > *counter {
-                        *counter = mic_part;
+                        *counter = stat.value;
                     }
 
                     let counter = stats.entry("maxwait").or_insert(0);
@@ -413,16 +301,15 @@ impl Collector {
                 | EventName::ServerIdle
                 | EventName::ServerTested
                 | EventName::ServerLogin => {
-                    client_server_states.insert(stat.process_id, stat.name);
+                    client_server_states.insert(stat.process_id.unwrap(), stat.name);
                 }
 
                 EventName::ClientDisconnecting | EventName::ServerDisconnecting => {
-                    client_server_states.remove(&stat.process_id);
+                    client_server_states.remove(&stat.process_id.unwrap());
                 }
 
-                EventName::UpdateStats => {
+                EventName::FlushStatsToStatsD => {
-                    // Calculate connection states
+                    for (_, state) in &client_server_states {
-                    for (_, state) in client_server_states.iter() {
                         match state {
                             EventName::ClientActive => {
                                 let counter = stats.entry("cl_active").or_insert(0);
@@ -434,6 +321,11 @@ impl Collector {
                                 *counter += 1;
                             }
 
+                            EventName::ClientIdle => {
+                                let counter = stats.entry("cl_idle").or_insert(0);
+                                *counter += 1;
+                            }
+
                             EventName::ServerIdle => {
                                 let counter = stats.entry("sv_idle").or_insert(0);
                                 *counter += 1;
@@ -454,64 +346,22 @@ impl Collector {
                                 *counter += 1;
                             }
 
-                            EventName::ClientIdle => {
-                                let counter = stats.entry("cl_idle").or_insert(0);
-                                *counter += 1;
-                            }
-
                             _ => unreachable!(),
                         };
                     }
 
-                    // Update latest stats used in SHOW STATS
+                    info!("{:?}", stats);
-                    let mut guard = LATEST_STATS.lock();
+
-                    for (key, value) in stats.iter() {
+                    let mut pipeline = self.client.pipeline();
-                        let entry = guard.entry(stat.address_id).or_insert(HashMap::new());
+
-                        entry.insert(key.to_string(), value.clone());
+                    for (key, value) in stats.iter_mut() {
+                        pipeline.gauge(key, *value as f64);
+                        *value = 0;
                     }
 
-                    // These are re-calculated every iteration of the loop, so we don't want to add values
+                    pipeline.send(&self.client);
-                    // from the last iteration.
-                    for stat in &[
-                        "cl_active",
-                        "cl_waiting",
-                        "cl_idle",
-                        "sv_idle",
-                        "sv_active",
-                        "sv_tested",
-                        "sv_login",
-                        "maxwait",
-                        "maxwait_us",
-                    ] {
-                        stats.insert(stat, 0);
-                    }
-                }
-
-                EventName::UpdateAverages => {
-                    // Calculate averages
-                    for stat in &[
-                        "avg_query_count",
-                        "avgxact_count",
-                        "avg_sent",
-                        "avg_received",
-                        "avg_wait_time",
-                    ] {
-                        let total_name = stat.replace("avg_", "total_");
-                        let old_value = old_stats.entry(total_name.clone()).or_insert(0);
-                        let new_value = stats.get(total_name.as_str()).unwrap_or(&0).to_owned();
-                        let avg = (new_value - *old_value) / (STAT_PERIOD as i64 / 1_000); // Avg / second
-
-                        stats.insert(stat, avg);
-                        *old_value = new_value;
-                    }
                 }
             };
         }
     }
 }
-
-/// Get a snapshot of statistics. Updated once a second
-/// by the `Collector`.
-pub fn get_stats() -> HashMap<usize, HashMap<String, i64>> {
-    LATEST_STATS.lock().clone()
-}

tests/pgbench/simple.sql

@@ -1,28 +0,0 @@
-
--- \setrandom aid 1 :naccounts
-\set aid random(1, 100000)
--- \setrandom bid 1 :nbranches
-\set bid random(1, 100000)
--- \setrandom tid 1 :ntellers
-\set tid random(1, 100000)
--- \setrandom delta -5000 5000
-\set delta random(-5000,5000)
-
-\set shard random(0, 2)
-
-SET SHARD TO :shard;
-
-BEGIN;
-
-UPDATE pgbench_accounts SET abalance = abalance + :delta WHERE aid = :aid;
-
-SELECT abalance FROM pgbench_accounts WHERE aid = :aid;
-
-UPDATE pgbench_tellers SET tbalance = tbalance + :delta WHERE tid = :tid;
-
-UPDATE pgbench_branches SET bbalance = bbalance + :delta WHERE bid = :bid;
-
-INSERT INTO pgbench_history (tid, bid, aid, delta, mtime) VALUES (:tid, :bid, :aid, :delta, CURRENT_TIMESTAMP);
-
-END;
-