starknet_patricia_storage: atomic counters for read stats

crates/starknet_patricia_storage/src/map_storage.rs

@@ -1,6 +1,7 @@
 use std::collections::BTreeMap;
 use std::fmt::Display;
 use std::num::NonZeroUsize;
+use std::sync::atomic::{AtomicU64, Ordering};
 
 use apollo_config::dumping::{prepend_sub_config_name, ser_param, SerializeConfig};
 use apollo_config::{ParamPath, ParamPrivacyInput, SerializedParam};
@@ -76,8 +77,8 @@ pub struct CachedStorage<S: Storage> {
     pub storage: S,
     pub cache: LruCache<DbKey, Option<DbValue>>,
     pub cache_on_write: bool,
-    reads: u128,
-    cached_reads: u128,
+    reads: AtomicU64,
+    cached_reads: AtomicU64,
     writes: u128,
     include_inner_stats: bool,
 }
@@ -204,8 +205,8 @@ impl<S: Storage> CachedStorage<S> {
             storage,
             cache: LruCache::new(config.cache_size),
             cache_on_write: config.cache_on_write,
-            reads: 0,
-            cached_reads: 0,
+            reads: AtomicU64::new(0),
+            cached_reads: AtomicU64::new(0),
             writes: 0,
             include_inner_stats: config.include_inner_stats,
         }
@@ -217,14 +218,6 @@ impl<S: Storage> CachedStorage<S> {
         }
     }
 
-    pub fn total_reads(&self) -> u128 {
-        self.reads
-    }
-
-    pub fn total_cached_reads(&self) -> u128 {
-        self.cached_reads
-    }
-
     pub fn total_writes(&self) -> u128 {
         self.writes
     }
@@ -235,9 +228,9 @@ impl<S: Storage> Storage for CachedStorage<S> {
     type Config = CachedStorageConfig<S::Config>;
 
     async fn get(&mut self, key: &DbKey) -> PatriciaStorageResult<Option<DbValue>> {
-        self.reads += 1;
+        self.reads.fetch_add(1, Ordering::Relaxed);
         if let Some(cached_value) = self.cache.get(key) {
-            self.cached_reads += 1;
+            self.cached_reads.fetch_add(1, Ordering::Relaxed);
             return Ok(cached_value.clone());
         }
 
@@ -267,9 +260,11 @@ impl<S: Storage> Storage for CachedStorage<S> {
             }
         }
 
-        self.reads += u128::try_from(keys.len()).expect("usize should fit in u128");
-        self.cached_reads +=
-            u128::try_from(keys.len() - keys_to_fetch.len()).expect("usize should fit in u128");
+        let n_keys = u64::try_from(keys.len()).expect("keys length should fit in u64");
+        let cached_reads =
+            u64::try_from(keys.len() - keys_to_fetch.len()).expect("usize should fit in u64");
+        self.reads.fetch_add(n_keys, Ordering::Relaxed);
+        self.cached_reads.fetch_add(cached_reads, Ordering::Relaxed);
 
         let fetched_values = self.storage.mget(keys_to_fetch.as_slice()).await?;
         indices_to_fetch.iter().zip(keys_to_fetch).zip(fetched_values).for_each(
@@ -297,9 +292,11 @@ impl<S: Storage> Storage for CachedStorage<S> {
     }
 
     fn get_stats(&self) -> PatriciaStorageResult<Self::Stats> {
+        let reads = u128::from(self.reads.load(Ordering::Acquire));
+        let cached_reads = u128::from(self.cached_reads.load(Ordering::Acquire));
         Ok(CachedStorageStats {
-            reads: self.reads,
-            cached_reads: self.cached_reads,
+            reads,
+            cached_reads,
             writes: self.writes,
             inner_stats: if self.include_inner_stats {
                 Some(self.storage.get_stats()?)
@@ -310,8 +307,8 @@ impl<S: Storage> Storage for CachedStorage<S> {
     }
 
     fn reset_stats(&mut self) -> PatriciaStorageResult<()> {
-        self.reads = 0;
-        self.cached_reads = 0;
+        self.reads.store(0, Ordering::Release);
+        self.cached_reads.store(0, Ordering::Release);
         self.writes = 0;
         self.storage.reset_stats()
     }