From 2e6d2ef08281416074b1ab44bb8bd89cd551f4e2 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Piotr=20Ko=C5=82aczkowski?= <pkolaczk@gmail.com>
Date: Sat, 3 Aug 2024 19:28:05 +0200
Subject: [PATCH] Add ctx.elapsed_secs() function

---
 README.md       |  4 ++++
 src/context.rs  | 14 ++++++++++----
 src/workload.rs |  5 ++++-
 3 files changed, 18 insertions(+), 5 deletions(-)

diff --git a/README.md b/README.md
index cde656b..3a36560 100644
--- a/README.md
+++ b/README.md
@@ -296,3 +296,7 @@ Errors during execution of a workload script are divided into three classes:
   errors terminate  
   the benchmark immediately. Overload errors (e.g. timeouts) that happen during the main run phase are counted and
   reported in the benchmark report.
+
+### Other functions
+
+- `ctx.elapsed_secs()` – returns the number of seconds elapsed since starting the workload, as float
diff --git a/src/context.rs b/src/context.rs
index b5835bc..fdbc637 100644
--- a/src/context.rs
+++ b/src/context.rs
@@ -389,6 +389,7 @@ pub fn get_exponential_retry_interval(
 /// It also tracks query execution metrics such as number of requests, rows, response times etc.
 #[derive(Any)]
 pub struct Context {
+    start_time: TryLock<Instant>,
     session: Arc<scylla::Session>,
     statements: HashMap<String, Arc<PreparedStatement>>,
     stats: TryLock<SessionStats>,
@@ -412,6 +413,7 @@ unsafe impl Sync for Context {}
 impl Context {
     pub fn new(session: scylla::Session, retry_number: u64, retry_interval: RetryDelay) -> Context {
         Context {
+            start_time: TryLock::new(Instant::now()),
             session: Arc::new(session),
             statements: HashMap::new(),
             stats: TryLock::new(SessionStats::new()),
@@ -433,10 +435,9 @@ impl Context {
             session: self.session.clone(),
             statements: self.statements.clone(),
             stats: TryLock::new(SessionStats::default()),
-            retry_number: self.retry_number,
-            retry_interval: self.retry_interval,
-            load_cycle_count: self.load_cycle_count,
             data: deserialized,
+            start_time: TryLock::new(*self.start_time.try_lock().unwrap()),
+            ..*self
         })
     }
 
@@ -532,6 +533,10 @@ impl Context {
         rs
     }
 
+    pub fn elapsed_secs(&self) -> f64 {
+        self.start_time.try_lock().unwrap().elapsed().as_secs_f64()
+    }
+
     fn should_retry<R>(result: &Result<R, QueryError>) -> bool {
         matches!(
             result,
@@ -555,8 +560,9 @@ impl Context {
     }
 
     /// Resets query and request counters
-    pub fn reset_session_stats(&self) {
+    pub fn reset(&self) {
         self.stats.try_lock().unwrap().reset();
+        *self.start_time.try_lock().unwrap() = Instant::now();
     }
 }
 
diff --git a/src/workload.rs b/src/workload.rs
index 444aade..f5cbd42 100644
--- a/src/workload.rs
+++ b/src/workload.rs
@@ -110,6 +110,9 @@ impl Program {
         context_module
             .async_inst_fn("execute_prepared", Context::execute_prepared)
             .unwrap();
+        context_module
+            .inst_fn("elapsed_secs", Context::elapsed_secs)
+            .unwrap();
 
         let mut err_module = Module::default();
         err_module.ty::<CassError>().unwrap();
@@ -495,7 +498,7 @@ impl Workload {
         let mut state = self.state.try_lock().unwrap();
         state.fn_stats = FnStats::default();
         state.start_time = start_time;
-        self.context.reset_session_stats();
+        self.context.reset();
     }
 
     /// Returns statistics of the operations invoked by this workload so far.