From 3c0da906d9d6b5e7e65b7052aaf3178ecdcc5131 Mon Sep 17 00:00:00 2001
From: smoczy123 <michal.hudobski@poczta.onet.pl>
Date: Wed, 20 Nov 2024 21:57:24 +0100
Subject: [PATCH] transport/timestamp_generator: Added
 MonotonicTimestampGenerator

Added TimestampGenerator trait and MonotonicTimestampGenerator based on c++ driver's implementation
---
 scylla/src/transport/mod.rs                 |  1 +
 scylla/src/transport/timestamp_generator.rs | 93 +++++++++++++++++++++
 2 files changed, 94 insertions(+)
 create mode 100644 scylla/src/transport/timestamp_generator.rs

diff --git a/scylla/src/transport/mod.rs b/scylla/src/transport/mod.rs
index 55184aadc..75b8266d8 100644
--- a/scylla/src/transport/mod.rs
+++ b/scylla/src/transport/mod.rs
@@ -18,6 +18,7 @@ pub mod retry_policy;
 pub mod session;
 pub mod session_builder;
 pub mod speculative_execution;
+pub mod timestamp_generator;
 pub mod topology;
 
 pub use crate::frame::{Authenticator, Compression};
diff --git a/scylla/src/transport/timestamp_generator.rs b/scylla/src/transport/timestamp_generator.rs
new file mode 100644
index 000000000..f7c6178cb
--- /dev/null
+++ b/scylla/src/transport/timestamp_generator.rs
@@ -0,0 +1,93 @@
+use std::{
+    ops::DerefMut,
+    sync::atomic::AtomicI64,
+    time::{SystemTime, UNIX_EPOCH},
+};
+
+use futures::lock::Mutex;
+use std::sync::atomic::Ordering;
+use tokio::time::{Duration, Instant};
+use tracing::warn;
+
+pub(crate) trait TimestampGenerator {
+    async fn next(&self) -> i64;
+}
+
+pub struct MonotonicTimestampGenerator {
+    last: AtomicI64,
+    last_warning: Mutex<Instant>,
+    warning_threshold_us: i64,
+    warning_interval_ms: i64,
+}
+
+impl MonotonicTimestampGenerator {
+    pub fn new_with_settings(warning_threshold_us: i64, warning_interval_ms: i64) -> Self {
+        MonotonicTimestampGenerator {
+            last: AtomicI64::new(0),
+            last_warning: Mutex::new(Instant::now()),
+            warning_threshold_us,
+            warning_interval_ms,
+        }
+    }
+    pub fn new() -> Self {
+        MonotonicTimestampGenerator::new_with_settings(1000000, 1000)
+    }
+
+    // This is guaranteed to return a monotonic timestamp. If clock skew is detected
+    // then this method will increment the last timestamp.
+    async fn compute_next(&self, last: i64) -> i64 {
+        let current = SystemTime::now().duration_since(UNIX_EPOCH);
+        if let Ok(cur_time) = current {
+            let u_cur = cur_time.as_micros() as i64;
+            if u_cur > last {
+                return u_cur;
+            } else if self.warning_threshold_us >= 0 && last - u_cur > self.warning_threshold_us {
+                let mut last_warn = self.last_warning.lock().await;
+                let now = Instant::now();
+                if now
+                    >= last_warn
+                        .checked_add(Duration::from_millis(self.warning_interval_ms as u64))
+                        .unwrap()
+                {
+                    *last_warn.deref_mut() = now;
+                    drop(last_warn);
+                    warn!(
+                        "Clock skew detected. The current time ({}) was {} \
+                    microseconds behind the last generated timestamp ({}). \
+                    The next generated timestamp will be artificially incremented \
+                    to guarantee monotonicity.",
+                        u_cur,
+                        last - u_cur,
+                        last
+                    )
+                }
+            }
+        } else {
+            warn!("Clock skew detected. The current time was behind UNIX epoch.");
+        }
+
+        last + 1
+    }
+}
+
+impl Default for MonotonicTimestampGenerator {
+    fn default() -> Self {
+        Self::new()
+    }
+}
+
+impl TimestampGenerator for MonotonicTimestampGenerator {
+    async fn next(&self) -> i64 {
+        loop {
+            let last = self.last.load(Ordering::SeqCst);
+            let cur = self.compute_next(last).await;
+            if self
+                .last
+                .compare_exchange(last, cur, Ordering::SeqCst, Ordering::SeqCst)
+                .is_ok()
+            {
+                return cur;
+            }
+        }
+    }
+}