Per-shard statistics (#57)

* per shard stats

* aight

* cleaner

* fix show lists

* comments

* more friendly

* case-insensitive

* test all shards

* ok

* HUH?

Author: levkk

.circleci/run_tests.sh

@@ -13,6 +13,9 @@ 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
@@ -28,9 +31,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 -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 -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 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

Cargo.lock

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

Cargo.toml

@@ -31,7 +31,7 @@ pub async fn handle_admin(
 
     if query.starts_with("SHOW STATS") {
         trace!("SHOW STATS");
-        show_stats(stream).await
+        show_stats(stream, &pool).await
     } else if query.starts_with("RELOAD") {
         trace!("RELOAD");
         reload(stream).await
@@ -77,23 +77,39 @@ async fn show_lists(stream: &mut OwnedWriteHalf, pool: &ConnectionPool) -> Resul
     ])); // but admin tools that work with pgbouncer want this
     res.put(data_row(&vec![
         "free_clients".to_string(),
-        stats["cl_idle"].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["cl_active"].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["sv_idle"].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["sv_active"].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()]));
@@ -125,7 +141,7 @@ async fn show_version(stream: &mut OwnedWriteHalf) -> Result<(), Error> {
 }
 
 /// SHOW POOLS
-async fn show_pools(stream: &mut OwnedWriteHalf, _pool: &ConnectionPool) -> Result<(), Error> {
+async fn show_pools(stream: &mut OwnedWriteHalf, pool: &ConnectionPool) -> Result<(), Error> {
     let stats = get_stats();
     let config = {
         let guard = get_config();
@@ -151,16 +167,26 @@ async fn show_pools(stream: &mut OwnedWriteHalf, _pool: &ConnectionPool) -> Resu
     let mut res = BytesMut::new();
     res.put(row_description(&columns));
 
-    let mut row = vec![String::from("all"), config.user.name.clone()];
+    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(),
+            };
 
-    for column in &columns[2..columns.len() - 1] {
-        let value = stats.get(column.0).unwrap_or(&0).to_string();
-        row.push(value);
-    }
+            let mut row = vec![address.name(), config.user.name.clone()];
 
-    row.push(config.general.pool_mode.to_string());
+            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(data_row(&row));
     res.put(command_complete("SHOW"));
 
     res.put_u8(b'Z');
@@ -309,7 +335,7 @@ async fn show_config(stream: &mut OwnedWriteHalf) -> Result<(), Error> {
 }
 
 /// SHOW STATS
-async fn show_stats(stream: &mut OwnedWriteHalf) -> Result<(), Error> {
+async fn show_stats(stream: &mut OwnedWriteHalf, pool: &ConnectionPool) -> Result<(), Error> {
     let columns = vec![
         ("database", DataType::Text),
         ("total_xact_count", DataType::Numeric),
@@ -332,15 +358,24 @@ async fn show_stats(stream: &mut OwnedWriteHalf) -> Result<(), Error> {
     let mut res = BytesMut::new();
     res.put(row_description(&columns));
 
-    let mut row = vec![
-        String::from("all"), // TODO: per-database stats,
-    ];
+    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(),
+            };
 
-    for column in &columns[1..] {
-        row.push(stats.get(column.0).unwrap_or(&0).to_string());
+            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(data_row(&row));
     res.put(command_complete("SHOW"));
 
     res.put_u8(b'Z');

src/admin.rs

@@ -58,6 +58,12 @@ pub struct Client {
 
     // 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 {
@@ -147,6 +153,8 @@ impl Client {
                         parameters: parameters,
                         stats: stats,
                         admin: admin,
+                        last_address_id: None,
+                        last_server_id: None,
                     });
                 }
 
@@ -169,6 +177,8 @@ impl Client {
                         parameters: HashMap::new(),
                         stats: stats,
                         admin: false,
+                        last_address_id: None,
+                        last_server_id: None,
                     });
                 }
 
@@ -219,9 +229,6 @@ 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,
@@ -292,13 +299,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.get(query_router.shard(), query_router.role()).await {
+            let connection = match pool
+                .get(query_router.shard(), query_router.role(), self.process_id)
+                .await
+            {
                 Ok(conn) => {
                     debug!("Got connection from pool");
                     conn
@@ -312,15 +319,23 @@ 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);
-            self.stats.server_active(server.process_id());
+            self.stats.client_active(self.process_id, address.id);
+            self.stats.server_active(server.process_id(), address.id);
+            self.last_address_id = Some(address.id);
+            self.last_server_id = Some(server.process_id());
 
             debug!(
                 "Client {:?} talking to server {:?}",
@@ -392,17 +407,17 @@ impl Client {
                         }
 
                         // Report query executed statistics.
-                        self.stats.query();
+                        self.stats.query(address.id);
 
                         // 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.stats.transaction(address.id);
 
                             // 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());
+                                self.stats.server_idle(server.process_id(), address.id);
                                 break;
                             }
                         }
@@ -478,15 +493,15 @@ impl Client {
                         }
 
                         // Report query executed statistics.
-                        self.stats.query();
+                        self.stats.query(address.id);
 
                         // 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.stats.transaction(address.id);
 
                             if self.transaction_mode {
-                                self.stats.server_idle(server.process_id());
+                                self.stats.server_idle(server.process_id(), address.id);
                                 break;
                             }
                         }
@@ -517,10 +532,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.stats.transaction(address.id);
 
                             if self.transaction_mode {
-                                self.stats.server_idle(server.process_id());
+                                self.stats.server_idle(server.process_id(), address.id);
                                 break;
                             }
                         }
@@ -537,6 +552,7 @@ 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);
         }
     }
 
@@ -549,6 +565,14 @@ impl Client {
 
 impl Drop for Client {
     fn drop(&mut self) {
-        self.stats.client_disconnecting(self.process_id);
+        // Disconnect the client
+        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/client.rs

@@ -48,6 +48,7 @@ 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,
@@ -58,6 +59,7 @@ pub struct Address {
 impl Default for Address {
     fn default() -> Address {
         Address {
+            id: 0,
             host: String::from("127.0.0.1"),
             port: String::from("5432"),
             shard: 0,