From 2f96176f6e1f46c3e3ac1629035941c01e43f07b Mon Sep 17 00:00:00 2001
From: Dvir Yosef <dvir@starkware.co>
Date: Sun, 16 Jun 2024 15:18:12 +0300
Subject: [PATCH] feat(storage): add latency metrics, minor debug message
 change and minor fix

---
 crates/papyrus_storage/src/lib.rs           | 5 +++++
 crates/papyrus_storage/src/mmap_file/mod.rs | 4 ++--
 crates/papyrus_storage/src/state/mod.rs     | 2 +-
 3 files changed, 8 insertions(+), 3 deletions(-)

diff --git a/crates/papyrus_storage/src/lib.rs b/crates/papyrus_storage/src/lib.rs
index 7b0a8097e5..2a5bfcf073 100644
--- a/crates/papyrus_storage/src/lib.rs
+++ b/crates/papyrus_storage/src/lib.rs
@@ -480,6 +480,7 @@ pub struct StorageTxn<'env, Mode: TransactionKind> {
 
 impl<'env> StorageTxn<'env, RW> {
     /// Commits the changes made in the transaction to the storage.
+    #[latency_histogram("storage_commit_latency_seconds", false)]
     pub fn commit(self) -> StorageResult<()> {
         self.file_handlers.flush();
         Ok(self.txn.commit()?)
@@ -701,11 +702,15 @@ impl FileHandlers<RW> {
     }
 
     // TODO(dan): Consider 1. flushing only the relevant files, 2. flushing concurrently.
+    #[latency_histogram("storage_file_handler_flush_latency_seconds", false)]
     fn flush(&self) {
+        debug!("Flushing the mmap files.");
         self.thin_state_diff.flush();
         self.contract_class.flush();
         self.casm.flush();
         self.deprecated_contract_class.flush();
+        self.transaction_output.flush();
+        self.transaction.flush();
     }
 
     // Appends a thin transaction output to the corresponding file and returns its location.
diff --git a/crates/papyrus_storage/src/mmap_file/mod.rs b/crates/papyrus_storage/src/mmap_file/mod.rs
index 578d7ed976..0c61e6e25c 100644
--- a/crates/papyrus_storage/src/mmap_file/mod.rs
+++ b/crates/papyrus_storage/src/mmap_file/mod.rs
@@ -157,7 +157,7 @@ impl<V: ValueSerde> MMapFile<V> {
 
     /// Flushes the mmap to the file.
     fn flush(&mut self) {
-        debug!("Flushing mmap to file");
+        trace!("Flushing mmap to file");
         self.mmap.flush().expect("Failed to flush the mmap");
         self.should_flush = false;
     }
@@ -231,7 +231,7 @@ impl<V: ValueSerde + Debug> Writer<V> for FileHandler<V, RW> {
         {
             let mut mmap_file = self.mmap_file.lock().expect("Lock should not be poisoned");
             offset = mmap_file.offset;
-            debug!("Inserting object at offset: {}", offset);
+            trace!("Inserting object at offset: {}", offset);
             let mmap_slice = &mut mmap_file.mmap[offset..];
             mmap_slice[..len].copy_from_slice(&serialized);
             mmap_file
diff --git a/crates/papyrus_storage/src/state/mod.rs b/crates/papyrus_storage/src/state/mod.rs
index fe009c59a1..212bd83259 100644
--- a/crates/papyrus_storage/src/state/mod.rs
+++ b/crates/papyrus_storage/src/state/mod.rs
@@ -682,7 +682,7 @@ fn write_replaced_classes<'env>(
     Ok(())
 }
 
-#[latency_histogram("storage_write_storage_diffs_latency_seconds", true)]
+#[latency_histogram("storage_write_storage_diffs_latency_seconds", false)]
 fn write_storage_diffs<'env>(
     storage_diffs: &IndexMap<ContractAddress, IndexMap<StorageKey, Felt>>,
     txn: &DbTransaction<'env, RW>,