From fb1a8532f064d2e5d926e2f524cca7655c363e45 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Karol=20Bary=C5=82a?= <karol.baryla@scylladb.com>
Date: Wed, 11 Oct 2023 16:37:23 +0200
Subject: [PATCH] Add test for TimestampedAverage panic

When Instant::now() returns the same value during 2 calls of
TimestampedAverage::compute_next, division by 0 occurs,
and Duration::from_secs_f64 panics because of NaN input.

This commit adds regression test for this problem.
---
 scylla/Cargo.toml                             |  2 +-
 .../src/transport/load_balancing/default.rs   | 25 ++++++++++++++++---
 2 files changed, 23 insertions(+), 4 deletions(-)

diff --git a/scylla/Cargo.toml b/scylla/Cargo.toml
index 2460e020b9..6b017c9d10 100644
--- a/scylla/Cargo.toml
+++ b/scylla/Cargo.toml
@@ -27,7 +27,7 @@ bytes = "1.0.1"
 futures = "0.3.6"
 histogram = "0.6.9"
 num_enum = "0.6"
-tokio = { version = "1.27", features = ["net", "time", "io-util", "sync", "rt", "macros"] }
+tokio = { version = "1.27", features = ["net", "time", "io-util", "sync", "rt", "macros", "test-util"] }
 snap = "1.0"
 uuid = { version = "1.0", features = ["v4"] }
 rand = "0.8.3"
diff --git a/scylla/src/transport/load_balancing/default.rs b/scylla/src/transport/load_balancing/default.rs
index 3c2097c9dc..3667ee1f77 100644
--- a/scylla/src/transport/load_balancing/default.rs
+++ b/scylla/src/transport/load_balancing/default.rs
@@ -2132,6 +2132,7 @@ mod latency_awareness {
     use futures::{future::RemoteHandle, FutureExt};
     use itertools::Either;
     use scylla_cql::errors::{DbError, QueryError};
+    use tokio::time::{Duration, Instant};
     use tracing::{instrument::WithSubscriber, trace};
     use uuid::Uuid;
 
@@ -2143,7 +2144,6 @@ mod latency_awareness {
             atomic::{AtomicU64, Ordering},
             Arc, RwLock,
         },
-        time::{Duration, Instant},
     };
 
     #[derive(Debug)]
@@ -2168,7 +2168,7 @@ mod latency_awareness {
         }
     }
 
-    #[derive(Debug, Clone, Copy)]
+    #[derive(Debug, Clone, Copy, PartialEq, Eq)]
     pub(super) struct TimestampedAverage {
         pub(super) timestamp: Instant,
         pub(super) average: Duration,
@@ -2733,7 +2733,7 @@ mod latency_awareness {
             },
             ExecutionProfile,
         };
-        use std::time::Instant;
+        use tokio::time::Instant;
 
         trait DefaultPolicyTestExt {
             fn set_nodes_latency_stats(
@@ -3453,5 +3453,24 @@ mod latency_awareness {
 
             session.query("whatever", ()).await.unwrap_err();
         }
+
+        #[tokio::test]
+        async fn timestamped_average_works_when_clock_stops() {
+            tokio::time::pause();
+            let avg = Some(TimestampedAverage {
+                timestamp: Instant::now(),
+                average: Duration::from_secs(123),
+                num_measures: 1,
+            });
+            let new_avg = TimestampedAverage::compute_next(avg, Duration::from_secs(456), 10.0);
+            assert_eq!(
+                new_avg,
+                Some(TimestampedAverage {
+                    timestamp: Instant::now(),
+                    average: Duration::from_secs(123),
+                    num_measures: 2,
+                }),
+            );
+        }
     }
 }