Flush stats (#38)

* flush stats

* stats

* refactor

src/main.rs

@@ -111,9 +111,9 @@ 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);
+        let mut stats_collector = Collector::new(rx, collector_tx);
         stats_collector.collect().await;
     });
 

src/stats.rs

@@ -4,7 +4,6 @@ use statsd::Client;
 use tokio::sync::mpsc::{Receiver, Sender};
 
 use std::collections::HashMap;
-use std::time::Instant;
 
 use crate::config::get_config;
 
@@ -24,6 +23,7 @@ enum EventName {
     ServerTested,
     ServerLogin,
     ServerDisconnecting,
+    FlushStatsToStatsD,
 }
 
 #[derive(Debug)]
@@ -44,155 +44,167 @@ impl Reporter {
     }
 
     pub fn query(&self) {
-        let statistic = Event {
+        let event = Event {
             name: EventName::Query,
             value: 1,
             process_id: None,
         };
 
-        let _ = self.tx.try_send(statistic);
+        let _ = self.tx.try_send(event);
     }
 
     pub fn transaction(&self) {
-        let statistic = Event {
+        let event = Event {
             name: EventName::Transaction,
             value: 1,
             process_id: None,
         };
 
-        let _ = self.tx.try_send(statistic);
+        let _ = self.tx.try_send(event);
     }
 
     pub fn data_sent(&self, amount: usize) {
-        let statistic = Event {
+        let event = Event {
             name: EventName::DataSent,
             value: amount as i64,
             process_id: None,
         };
 
-        let _ = self.tx.try_send(statistic);
+        let _ = self.tx.try_send(event);
     }
 
     pub fn data_received(&self, amount: usize) {
-        let statistic = Event {
+        let event = Event {
             name: EventName::DataReceived,
             value: amount as i64,
             process_id: None,
         };
 
-        let _ = self.tx.try_send(statistic);
+        let _ = self.tx.try_send(event);
     }
 
     pub fn checkout_time(&self, ms: u128) {
-        let statistic = Event {
+        let event = Event {
             name: EventName::CheckoutTime,
             value: ms as i64,
             process_id: None,
         };
 
-        let _ = self.tx.try_send(statistic);
+        let _ = self.tx.try_send(event);
     }
 
     pub fn client_waiting(&self, process_id: i32) {
-        let statistic = Event {
+        let event = Event {
             name: EventName::ClientWaiting,
             value: 1,
             process_id: Some(process_id),
         };
 
-        let _ = self.tx.try_send(statistic);
+        let _ = self.tx.try_send(event);
     }
 
     pub fn client_active(&self, process_id: i32) {
-        let statistic = Event {
+        let event = Event {
             name: EventName::ClientActive,
             value: 1,
             process_id: Some(process_id),
         };
 
-        let _ = self.tx.try_send(statistic);
+        let _ = self.tx.try_send(event);
     }
 
     pub fn client_idle(&self, process_id: i32) {
-        let statistic = Event {
+        let event = Event {
             name: EventName::ClientIdle,
             value: 1,
             process_id: Some(process_id),
         };
 
-        let _ = self.tx.try_send(statistic);
+        let _ = self.tx.try_send(event);
     }
 
     pub fn client_disconnecting(&self, process_id: i32) {
-        let statistic = Event {
+        let event = Event {
             name: EventName::ClientDisconnecting,
             value: 1,
             process_id: Some(process_id),
         };
 
-        let _ = self.tx.try_send(statistic);
+        let _ = self.tx.try_send(event);
     }
 
     pub fn server_active(&self, process_id: i32) {
-        let statistic = Event {
+        let event = Event {
             name: EventName::ServerActive,
             value: 1,
             process_id: Some(process_id),
         };
 
-        let _ = self.tx.try_send(statistic);
+        let _ = self.tx.try_send(event);
     }
 
     pub fn server_idle(&self, process_id: i32) {
-        let statistic = Event {
+        let event = Event {
             name: EventName::ServerIdle,
             value: 1,
             process_id: Some(process_id),
         };
 
-        let _ = self.tx.try_send(statistic);
+        let _ = self.tx.try_send(event);
     }
 
     pub fn server_login(&self, process_id: i32) {
-        let statistic = Event {
+        let event = Event {
             name: EventName::ServerLogin,
             value: 1,
             process_id: Some(process_id),
         };
 
-        let _ = self.tx.try_send(statistic);
+        let _ = self.tx.try_send(event);
     }
 
     pub fn server_tested(&self, process_id: i32) {
-        let statistic = Event {
+        let event = Event {
             name: EventName::ServerTested,
             value: 1,
             process_id: Some(process_id),
         };
 
-        let _ = self.tx.try_send(statistic);
+        let _ = self.tx.try_send(event);
     }
 
     pub fn server_disconnecting(&self, process_id: i32) {
-        let statistic = Event {
+        let event = Event {
             name: EventName::ServerDisconnecting,
             value: 1,
             process_id: Some(process_id),
         };
 
-        let _ = self.tx.try_send(statistic);
+        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);
+    // }
 }
 
 pub struct Collector {
     rx: Receiver<Event>,
+    tx: Sender<Event>,
     client: Client,
 }
 
 impl Collector {
-    pub fn new(rx: Receiver<Event>) -> Collector {
+    pub fn new(rx: Receiver<Event>, tx: Sender<Event>) -> Collector {
         Collector {
-            rx: rx,
+            rx,
+            tx,
             client: Client::new(&get_config().general.statsd_address, "pgcat").unwrap(),
         }
     }
@@ -218,8 +230,19 @@ impl Collector {
         ]);
 
         let mut client_server_states: HashMap<i32, EventName> = HashMap::new();
+        let tx = self.tx.clone();
 
-        let mut now = Instant::now();
+        tokio::task::spawn(async move {
+            let mut interval = tokio::time::interval(tokio::time::Duration::from_millis(15000));
+            loop {
+                interval.tick().await;
+                let _ = tx.try_send(Event {
+                    name: EventName::FlushStatsToStatsD,
+                    value: 0,
+                    process_id: None,
+                });
+            }
+        });
 
         loop {
             let stat = match self.rx.recv().await {
@@ -284,65 +307,61 @@ impl Collector {
                 EventName::ClientDisconnecting | EventName::ServerDisconnecting => {
                     client_server_states.remove(&stat.process_id.unwrap());
                 }
+
+                EventName::FlushStatsToStatsD => {
+                    for (_, state) in &client_server_states {
+                        match state {
+                            EventName::ClientActive => {
+                                let counter = stats.entry("cl_active").or_insert(0);
+                                *counter += 1;
+                            }
+
+                            EventName::ClientWaiting => {
+                                let counter = stats.entry("cl_waiting").or_insert(0);
+                                *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;
+                            }
+
+                            EventName::ServerActive => {
+                                let counter = stats.entry("sv_active").or_insert(0);
+                                *counter += 1;
+                            }
+
+                            EventName::ServerTested => {
+                                let counter = stats.entry("sv_tested").or_insert(0);
+                                *counter += 1;
+                            }
+
+                            EventName::ServerLogin => {
+                                let counter = stats.entry("sv_login").or_insert(0);
+                                *counter += 1;
+                            }
+
+                            _ => unreachable!(),
+                        };
+                    }
+
+                    info!("{:?}", stats);
+
+                    let mut pipeline = self.client.pipeline();
+
+                    for (key, value) in stats.iter_mut() {
+                        pipeline.gauge(key, *value as f64);
+                        *value = 0;
+                    }
+
+                    pipeline.send(&self.client);
+                }
             };
-
-            // It's been 15 seconds. If there is no traffic, it won't publish anything,
-            // but it also doesn't matter then.
-            if now.elapsed().as_secs() > 15 {
-                for (_, state) in &client_server_states {
-                    match state {
-                        EventName::ClientActive => {
-                            let counter = stats.entry("cl_active").or_insert(0);
-                            *counter += 1;
-                        }
-
-                        EventName::ClientWaiting => {
-                            let counter = stats.entry("cl_waiting").or_insert(0);
-                            *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;
-                        }
-
-                        EventName::ServerActive => {
-                            let counter = stats.entry("sv_active").or_insert(0);
-                            *counter += 1;
-                        }
-
-                        EventName::ServerTested => {
-                            let counter = stats.entry("sv_tested").or_insert(0);
-                            *counter += 1;
-                        }
-
-                        EventName::ServerLogin => {
-                            let counter = stats.entry("sv_login").or_insert(0);
-                            *counter += 1;
-                        }
-
-                        _ => unreachable!(),
-                    };
-                }
-
-                info!("{:?}", stats);
-
-                let mut pipeline = self.client.pipeline();
-
-                for (key, value) in stats.iter_mut() {
-                    pipeline.gauge(key, *value as f64);
-                    *value = 0;
-                }
-
-                pipeline.send(&self.client);
-
-                now = Instant::now();
-            }
         }
     }
 }