From e572061428df29042db0f3763f095c4f031d54a7 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:45 +0200
Subject: [PATCH] load_balancing: Prevent 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 fixes the issue and adds additional logging
in case another conversion problem occurs in the future.
---
 .../src/transport/load_balancing/default.rs   | 32 ++++++++++++++++---
 1 file changed, 28 insertions(+), 4 deletions(-)

diff --git a/scylla/src/transport/load_balancing/default.rs b/scylla/src/transport/load_balancing/default.rs
index 3667ee1f77..d58d26febe 100644
--- a/scylla/src/transport/load_balancing/default.rs
+++ b/scylla/src/transport/load_balancing/default.rs
@@ -2133,7 +2133,7 @@ mod latency_awareness {
     use itertools::Either;
     use scylla_cql::errors::{DbError, QueryError};
     use tokio::time::{Duration, Instant};
-    use tracing::{instrument::WithSubscriber, trace};
+    use tracing::{instrument::WithSubscriber, trace, warn};
     use uuid::Uuid;
 
     use crate::{load_balancing::NodeRef, transport::node::Node};
@@ -2192,13 +2192,37 @@ mod latency_awareness {
                 Some(prev_avg) => Some({
                     let delay = (now - prev_avg.timestamp).as_secs_f64();
                     let scaled_delay = delay / scale_secs;
-                    let prev_weight = (scaled_delay + 1.).ln() / scaled_delay;
+                    let prev_weight = if scaled_delay <= 0. {
+                        1.
+                    } else {
+                        (scaled_delay + 1.).ln() / scaled_delay
+                    };
 
                     let last_latency_secs = last_latency.as_secs_f64();
                     let prev_avg_secs = prev_avg.average.as_secs_f64();
-                    let average = Duration::from_secs_f64(
+                    let average = match Duration::try_from_secs_f64(
                         (1. - prev_weight) * last_latency_secs + prev_weight * prev_avg_secs,
-                    );
+                    ) {
+                        Ok(ts) => ts,
+                        Err(e) => {
+                            warn!(
+                                "Error while calculating average: {e}. \
+                            prev_avg_secs: {prev_avg_secs}, \
+                            last_latency_secs: {last_latency_secs}, \
+                            prev_weight: {prev_weight}, \
+                            scaled_delay: {scaled_delay}, \
+                            delay: {delay}, \
+                            prev_avg.timestamp: {:?}, \
+                            now: {now:?}",
+                                prev_avg.timestamp
+                            );
+
+                            // Not sure when we could enter this branch,
+                            // so I have no idea what would be a sensible value to return here,
+                            // this does not seem like a very bad choice.
+                            prev_avg.average
+                        }
+                    };
                     Self {
                         num_measures: prev_avg.num_measures + 1,
                         timestamp: now,