From f7a951745ca002cc571d7f064c5705f3fa9e9cf3 Mon Sep 17 00:00:00 2001
From: Mostafa Abdelraouf <mostafa.abdelraouf@instacart.com>
Date: Thu, 15 Sep 2022 01:21:45 -0500
Subject: [PATCH] Report Query times (#166)

* Report avg and total query timing

* Report query times

* fmt
---
 src/client.rs            |  8 +++++++-
 src/stats.rs             | 11 +++++++++--
 tests/ruby/admin_spec.rb | 23 +++++++++++++++++++++++
 3 files changed, 39 insertions(+), 3 deletions(-)

diff --git a/src/client.rs b/src/client.rs
index c07a21f..3b0b0ea 100644
--- a/src/client.rs
+++ b/src/client.rs
@@ -2,6 +2,7 @@
 use bytes::{Buf, BufMut, BytesMut};
 use log::{debug, error, info, trace};
 use std::collections::HashMap;
+use std::time::Instant;
 use tokio::io::{split, AsyncReadExt, BufReader, ReadHalf, WriteHalf};
 use tokio::net::TcpStream;
 use tokio::sync::broadcast::Receiver;
@@ -994,6 +995,7 @@ where
         self.send_server_message(server, message, &address, &pool)
             .await?;
 
+        let query_start = Instant::now();
         // Read all data the server has to offer, which can be multiple messages
         // buffered in 8196 bytes chunks.
         loop {
@@ -1013,7 +1015,11 @@ where
         }
 
         // Report query executed statistics.
-        self.stats.query(self.process_id, server.server_id());
+        self.stats.query(
+            self.process_id,
+            server.server_id(),
+            Instant::now().duration_since(query_start).as_millis(),
+        );
 
         Ok(())
     }
diff --git a/src/stats.rs b/src/stats.rs
index 9ab7167..90b55dd 100644
--- a/src/stats.rs
+++ b/src/stats.rs
@@ -9,7 +9,6 @@ use tokio::sync::mpsc::{channel, Receiver, Sender};
 use tokio::time::Instant;
 
 use crate::pool::{get_all_pools, get_number_of_addresses};
-use crate::server;
 
 /// Convenience types for various stats
 type ClientStatesLookup = HashMap<i32, ClientInformation>;
@@ -138,6 +137,7 @@ enum EventName {
     Query {
         client_id: i32,
         server_id: i32,
+        duration_ms: u128,
     },
     Transaction {
         client_id: i32,
@@ -269,11 +269,12 @@ impl Reporter {
     }
 
     /// Report a query executed by a client against a server
-    pub fn query(&self, client_id: i32, server_id: i32) {
+    pub fn query(&self, client_id: i32, server_id: i32, duration_ms: u128) {
         let event = Event {
             name: EventName::Query {
                 client_id,
                 server_id,
+                duration_ms,
             },
             value: 1,
         };
@@ -562,6 +563,7 @@ impl Collector {
                 EventName::Query {
                     client_id,
                     server_id,
+                    duration_ms,
                 } => {
                     // Update client stats
                     let app_name = match client_states.get_mut(&client_id) {
@@ -585,6 +587,11 @@ impl Collector {
                                 .entry("total_query_count".to_string())
                                 .or_insert(0);
                             *counter += stat.value;
+
+                            let duration = pool_stats
+                                .entry("total_query_time".to_string())
+                                .or_insert(0);
+                            *duration += duration_ms as i64;
                         }
                         None => (),
                     }
diff --git a/tests/ruby/admin_spec.rb b/tests/ruby/admin_spec.rb
index b570984..a348146 100644
--- a/tests/ruby/admin_spec.rb
+++ b/tests/ruby/admin_spec.rb
@@ -11,6 +11,29 @@ describe "Admin" do
     processes.pgcat.shutdown
   end
 
+  describe "SHOW STATS" do
+    context "clients connect and make one query" do
+      it "updates *_query_time and *_wait_time" do
+        connection = PG::connect("#{pgcat_conn_str}?application_name=one_query")
+        connection.async_exec("SELECT pg_sleep(0.25)")
+        connection.async_exec("SELECT pg_sleep(0.25)")
+        connection.async_exec("SELECT pg_sleep(0.25)")
+        connection.close
+
+        # wait for averages to be calculated, we shouldn't do this too often
+        sleep(15.5)
+        admin_conn = PG::connect(processes.pgcat.admin_connection_string)
+        results = admin_conn.async_exec("SHOW STATS")[0]
+        admin_conn.close
+        expect(results["total_query_time"].to_i).to be_within(200).of(750)
+        expect(results["avg_query_time"].to_i).to_not eq(0)
+
+        expect(results["total_wait_time"].to_i).to_not eq(0)
+        expect(results["avg_wait_time"].to_i).to_not eq(0)
+      end
+    end
+  end
+
   describe "SHOW POOLS" do
     context "bad credentials" do
       it "does not change any stats" do