From 571b02e178ab2fc98dd8cd89f65284f153fe6288 Mon Sep 17 00:00:00 2001
From: Lev Kokotov <levkk@users.noreply.github.com>
Date: Mon, 8 May 2023 10:06:16 -0700
Subject: [PATCH] Calculate averages correctly and preserve totals like before
 (#429)

* Reset totals after avg calculation

* like it used to be
---
 Cargo.lock           | 16 ++++++++++++++
 Cargo.toml           |  1 +
 src/stats/address.rs | 50 ++++++++++++++++++++++++++++++++------------
 3 files changed, 54 insertions(+), 13 deletions(-)

diff --git a/Cargo.lock b/Cargo.lock
index 6b2d76f..9c553d4 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -250,6 +250,12 @@ dependencies = [
  "subtle",
 ]
 
+[[package]]
+name = "either"
+version = "1.8.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "7fcaabb2fef8c910e7f4c7ce9f67a1283a1715879a7c230ca9d6d1ae31f16d91"
+
 [[package]]
 name = "enum-as-inner"
 version = "0.5.1"
@@ -658,6 +664,15 @@ dependencies = [
  "windows-sys",
 ]
 
+[[package]]
+name = "itertools"
+version = "0.10.5"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "b0fd2260e829bddf4cb6ea802289de2f86d6a7a690192fbe91b3f46e0f2c8473"
+dependencies = [
+ "either",
+]
+
 [[package]]
 name = "itoa"
 version = "1.0.5"
@@ -897,6 +912,7 @@ dependencies = [
  "futures",
  "hmac",
  "hyper",
+ "itertools",
  "jemallocator",
  "log",
  "md-5",
diff --git a/Cargo.toml b/Cargo.toml
index 1b3c8b1..5428ab3 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -45,6 +45,7 @@ rustls = { version = "0.21", features = ["dangerous_configuration"] }
 trust-dns-resolver = "0.22.0"
 tokio-test = "0.4.2"
 serde_json = "1"
+itertools = "0.10"
 
 [target.'cfg(not(target_env = "msvc"))'.dependencies]
 jemallocator = "0.5.0"
diff --git a/src/stats/address.rs b/src/stats/address.rs
index a5759e1..51d6a68 100644
--- a/src/stats/address.rs
+++ b/src/stats/address.rs
@@ -1,4 +1,3 @@
-use log::warn;
 use std::sync::atomic::*;
 use std::sync::Arc;
 
@@ -13,6 +12,16 @@ pub struct AddressStats {
     pub total_query_time: Arc<AtomicU64>,
     pub total_wait_time: Arc<AtomicU64>,
     pub total_errors: Arc<AtomicU64>,
+
+    pub old_total_xact_count: Arc<AtomicU64>,
+    pub old_total_query_count: Arc<AtomicU64>,
+    pub old_total_received: Arc<AtomicU64>,
+    pub old_total_sent: Arc<AtomicU64>,
+    pub old_total_xact_time: Arc<AtomicU64>,
+    pub old_total_query_time: Arc<AtomicU64>,
+    pub old_total_wait_time: Arc<AtomicU64>,
+    pub old_total_errors: Arc<AtomicU64>,
+
     pub avg_query_count: Arc<AtomicU64>,
     pub avg_query_time: Arc<AtomicU64>,
     pub avg_recv: Arc<AtomicU64>,
@@ -104,16 +113,16 @@ impl AddressStats {
     }
 
     pub fn update_averages(&self) {
-        let (totals, averages) = self.fields_iterators();
-        for data in totals.iter().zip(averages.iter()) {
-            let (total, average) = data;
-            if let Err(err) = average.fetch_update(Ordering::Relaxed, Ordering::Relaxed, |avg| {
-                let total = total.load(Ordering::Relaxed);
-                let avg = (total - avg) / (crate::stats::STAT_PERIOD / 1_000); // Avg / second
-                Some(avg)
-            }) {
-                warn!("Could not update averages for addresses stats, {:?}", err);
-            }
+        let (totals, averages, old_totals) = self.fields_iterators();
+        for data in itertools::izip!(totals, averages, old_totals) {
+            let (total, average, old_total) = data;
+            let total = total.load(Ordering::Relaxed);
+            let old = old_total.load(Ordering::Relaxed);
+            average.store(
+                (total - old) / (crate::stats::STAT_PERIOD / 1_000),
+                Ordering::Relaxed,
+            ); // Avg / second
+            old_total.store(total, Ordering::Relaxed);
         }
     }
 
@@ -123,27 +132,42 @@ impl AddressStats {
         }
     }
 
-    fn fields_iterators(&self) -> (Vec<Arc<AtomicU64>>, Vec<Arc<AtomicU64>>) {
+    fn fields_iterators(
+        &self,
+    ) -> (
+        Vec<Arc<AtomicU64>>,
+        Vec<Arc<AtomicU64>>,
+        Vec<Arc<AtomicU64>>,
+    ) {
         let mut totals: Vec<Arc<AtomicU64>> = Vec::new();
         let mut averages: Vec<Arc<AtomicU64>> = Vec::new();
+        let mut old_totals: Vec<Arc<AtomicU64>> = Vec::new();
 
         totals.push(self.total_xact_count.clone());
+        old_totals.push(self.old_total_xact_count.clone());
         averages.push(self.avg_xact_count.clone());
         totals.push(self.total_query_count.clone());
+        old_totals.push(self.old_total_query_count.clone());
         averages.push(self.avg_query_count.clone());
         totals.push(self.total_received.clone());
+        old_totals.push(self.old_total_received.clone());
         averages.push(self.avg_recv.clone());
         totals.push(self.total_sent.clone());
+        old_totals.push(self.old_total_sent.clone());
         averages.push(self.avg_sent.clone());
         totals.push(self.total_xact_time.clone());
+        old_totals.push(self.old_total_xact_time.clone());
         averages.push(self.avg_xact_time.clone());
         totals.push(self.total_query_time.clone());
+        old_totals.push(self.old_total_query_time.clone());
         averages.push(self.avg_query_time.clone());
         totals.push(self.total_wait_time.clone());
+        old_totals.push(self.old_total_wait_time.clone());
         averages.push(self.avg_wait_time.clone());
         totals.push(self.total_errors.clone());
+        old_totals.push(self.old_total_errors.clone());
         averages.push(self.avg_errors.clone());
 
-        (totals, averages)
+        (totals, averages, old_totals)
     }
 }