From db97fea13ee7600f2ff771a07fb20bfcc6ab7477 Mon Sep 17 00:00:00 2001
From: antoniolocascio <antonio.locascio1@gmail.com>
Date: Fri, 5 Jul 2024 13:34:25 +0200
Subject: [PATCH] feat(prover): store time taken by prover with ms precision

---
 core/lib/db_connection/src/utils.rs     | 26 ++++++++++++++++++-------
 prover/prover_dal/src/fri_prover_dal.rs |  4 ++--
 prover/prover_dal/src/lib.rs            |  2 +-
 prover/prover_fri/src/utils.rs          |  5 +++--
 4 files changed, 25 insertions(+), 12 deletions(-)

diff --git a/core/lib/db_connection/src/utils.rs b/core/lib/db_connection/src/utils.rs
index 80cf0a5cbb35..532a355d916c 100644
--- a/core/lib/db_connection/src/utils.rs
+++ b/core/lib/db_connection/src/utils.rs
@@ -13,14 +13,26 @@ const MICROSECONDS_IN_A_SECOND: i64 = 1_000_000;
 const MICROSECONDS_IN_A_MINUTE: i64 = MICROSECONDS_IN_A_SECOND * 60;
 const MICROSECONDS_IN_AN_HOUR: i64 = MICROSECONDS_IN_A_MINUTE * 60;
 
+fn duration_to_naive_time_common(duration: Duration, use_ms: bool) -> NaiveTime {
+    let total_ms = duration.as_millis();
+    let total_seconds = (total_ms / 1_000) as u32;
+    let hours = total_seconds / 3_600;
+    let minutes = (total_seconds / 60) % 60;
+    let seconds = total_seconds % 60;
+    let ms = (total_ms % 1_000) as u32;
+
+    if use_ms {
+        NaiveTime::from_hms_milli_opt(hours, minutes, seconds, ms).unwrap()
+    } else {
+        NaiveTime::from_hms_opt(hours, minutes, seconds).unwrap()
+    }
+}
 pub fn duration_to_naive_time(duration: Duration) -> NaiveTime {
-    let total_seconds = duration.as_secs() as u32;
-    NaiveTime::from_hms_opt(
-        total_seconds / 3600,
-        (total_seconds / 60) % 60,
-        total_seconds % 60,
-    )
-    .unwrap()
+    duration_to_naive_time_common(duration, false)
+}
+
+pub fn duration_to_naive_time_ms(duration: Duration) -> NaiveTime {
+    duration_to_naive_time_common(duration, true)
 }
 
 pub const fn pg_interval_from_duration(processing_timeout: Duration) -> PgInterval {
diff --git a/prover/prover_dal/src/fri_prover_dal.rs b/prover/prover_dal/src/fri_prover_dal.rs
index 419cb635ac53..d94e830e668a 100644
--- a/prover/prover_dal/src/fri_prover_dal.rs
+++ b/prover/prover_dal/src/fri_prover_dal.rs
@@ -13,7 +13,7 @@ use zksync_db_connection::{
     connection::Connection, instrument::InstrumentExt, metrics::MethodLatency,
 };
 
-use crate::{duration_to_naive_time, pg_interval_from_duration, Prover};
+use crate::{duration_to_naive_time_ms, pg_interval_from_duration, Prover};
 
 #[derive(Debug)]
 pub struct FriProverDal<'a, 'c> {
@@ -265,7 +265,7 @@ impl FriProverDal<'_, '_> {
                 prover_jobs_fri.depth,
                 prover_jobs_fri.is_node_final_proof
             "#,
-            duration_to_naive_time(time_taken),
+            duration_to_naive_time_ms(time_taken),
             blob_url,
             i64::from(id)
         )
diff --git a/prover/prover_dal/src/lib.rs b/prover/prover_dal/src/lib.rs
index bb552b899e90..3df1e4098e9b 100644
--- a/prover/prover_dal/src/lib.rs
+++ b/prover/prover_dal/src/lib.rs
@@ -2,7 +2,7 @@ use zksync_db_connection::connection::DbMarker;
 pub use zksync_db_connection::{
     connection::Connection,
     connection_pool::ConnectionPool,
-    utils::{duration_to_naive_time, pg_interval_from_duration},
+    utils::{duration_to_naive_time, duration_to_naive_time_ms, pg_interval_from_duration},
 };
 
 use crate::{
diff --git a/prover/prover_fri/src/utils.rs b/prover/prover_fri/src/utils.rs
index 15a2a6c18bb2..e53e0810d693 100644
--- a/prover/prover_fri/src/utils.rs
+++ b/prover/prover_fri/src/utils.rs
@@ -67,10 +67,11 @@ pub async fn save_proof(
     connection: &mut Connection<'_, Prover>,
     protocol_version: ProtocolSemanticVersion,
 ) {
+    let time_taken = started_at.elapsed();
     tracing::info!(
         "Successfully proven job: {}, total time taken: {:?}",
         job_id,
-        started_at.elapsed()
+        time_taken,
     );
     let proof = artifacts.proof_wrapper;
 
@@ -101,7 +102,7 @@ pub async fn save_proof(
     let mut transaction = connection.start_transaction().await.unwrap();
     transaction
         .fri_prover_jobs_dal()
-        .save_proof(job_id, started_at.elapsed(), &blob_url)
+        .save_proof(job_id, time_taken, &blob_url)
         .await;
     if is_scheduler_proof {
         transaction