Merge pull request #2059 from CosmWasm/aw/pinned-cache-metrics

Metrics for the pinned cache

Author: aumetra

CHANGELOG.md

@@ -10,6 +10,7 @@ and this project adheres to
 
 - cosmwasm-vm: Add `secp256r1_verify` and `secp256r1_recover_pubkey` imports for
   ECDSA signature verification over secp256r1. ([#1983], [#2057], [#2058])
+- cosmwasm-vm: Add metrics for the pinned memory cache ([#2059])
 
 [#1983]: https://github.com/CosmWasm/cosmwasm/pull/1983
 [#2057]: https://github.com/CosmWasm/cosmwasm/pull/2057
@@ -23,6 +24,7 @@ and this project adheres to
 
 [#2044]: https://github.com/CosmWasm/cosmwasm/pull/2044
 [#2051]: https://github.com/CosmWasm/cosmwasm/pull/2051
+[#2059]: https://github.com/CosmWasm/cosmwasm/pull/2059
 
 ## [2.0.0] - 2024-03-12
 

packages/vm/src/cache.rs

@@ -53,6 +53,22 @@ pub struct Metrics {
     pub size_memory_cache: usize,
 }
 
+#[derive(Debug, Clone)]
+pub struct PerModuleMetrics {
+    /// Hits (i.e. loads) of the module from the cache
+    pub hits: u32,
+    /// Size the module takes up in memory
+    pub size: usize,
+}
+
+#[derive(Debug, Clone)]
+pub struct PinnedMetrics {
+    // It is *intentional* that this is only a vector
+    // We don't need a potentially expensive hashing algorithm here
+    // The checksums are sourced from a hashmap already, ensuring uniqueness of the checksums
+    pub per_module: Vec<(Checksum, PerModuleMetrics)>,
+}
+
 #[derive(Clone, Debug)]
 #[non_exhaustive]
 pub struct CacheOptions {
@@ -182,6 +198,24 @@ where
         self.inner.lock().unwrap().stats
     }
 
+    pub fn pinned_metrics(&self) -> PinnedMetrics {
+        let cache = self.inner.lock().unwrap();
+        let per_module = cache
+            .pinned_memory_cache
+            .iter()
+            .map(|(checksum, module)| {
+                let metrics = PerModuleMetrics {
+                    hits: module.hits,
+                    size: module.module.size_estimate,
+                };
+
+                (*checksum, metrics)
+            })
+            .collect();
+
+        PinnedMetrics { per_module }
+    }
+
     pub fn metrics(&self) -> Metrics {
         let cache = self.inner.lock().unwrap();
         Metrics {
@@ -1409,6 +1443,48 @@ mod tests {
         );
     }
 
+    #[test]
+    fn pinned_metrics_works() {
+        let cache = unsafe { Cache::new(make_testing_options()).unwrap() };
+        let checksum = cache.save_wasm(CONTRACT).unwrap();
+
+        cache.pin(&checksum).unwrap();
+
+        let pinned_metrics = cache.pinned_metrics();
+        assert_eq!(pinned_metrics.per_module.len(), 1);
+        assert_eq!(pinned_metrics.per_module[0].0, checksum);
+        assert_eq!(pinned_metrics.per_module[0].1.hits, 0);
+
+        let backend = mock_backend(&[]);
+        let _ = cache
+            .get_instance(&checksum, backend, TESTING_OPTIONS)
+            .unwrap();
+
+        let pinned_metrics = cache.pinned_metrics();
+        assert_eq!(pinned_metrics.per_module.len(), 1);
+        assert_eq!(pinned_metrics.per_module[0].0, checksum);
+        assert_eq!(pinned_metrics.per_module[0].1.hits, 1);
+
+        let empty_checksum = cache.save_wasm(EMPTY_CONTRACT).unwrap();
+        cache.pin(&empty_checksum).unwrap();
+
+        let pinned_metrics = cache.pinned_metrics();
+        assert_eq!(pinned_metrics.per_module.len(), 2);
+
+        let get_module_hits = |checksum| {
+            pinned_metrics
+                .per_module
+                .iter()
+                .find(|(iter_checksum, _module)| *iter_checksum == checksum)
+                .map(|(_checksum, module)| module)
+                .cloned()
+                .unwrap()
+        };
+
+        assert_eq!(get_module_hits(checksum).hits, 1);
+        assert_eq!(get_module_hits(empty_checksum).hits, 0);
+    }
+
     #[test]
     fn pin_unpin_works() {
         let cache = unsafe { Cache::new(make_testing_options()).unwrap() };

packages/vm/src/lib.rs

@@ -23,7 +23,9 @@ mod wasm_backend;
 pub use crate::backend::{
     Backend, BackendApi, BackendError, BackendResult, GasInfo, Querier, Storage,
 };
-pub use crate::cache::{AnalysisReport, Cache, CacheOptions, Metrics, Stats};
+pub use crate::cache::{
+    AnalysisReport, Cache, CacheOptions, Metrics, PerModuleMetrics, PinnedMetrics, Stats,
+};
 pub use crate::calls::{
     call_execute, call_execute_raw, call_instantiate, call_instantiate_raw, call_migrate,
     call_migrate_raw, call_query, call_query_raw, call_reply, call_reply_raw, call_sudo,

packages/vm/src/modules/pinned_memory_cache.rs

@@ -4,9 +4,18 @@ use std::collections::HashMap;
 use super::cached_module::CachedModule;
 use crate::VmResult;
 
+/// Struct storing some additional metadata, which is only of interest for the pinned cache,
+/// alongside the cached module.
+pub struct InstrumentedModule {
+    /// Number of loads from memory this module received
+    pub hits: u32,
+    /// The actual cached module
+    pub module: CachedModule,
+}
+
 /// An pinned in memory module cache
 pub struct PinnedMemoryCache {
-    modules: HashMap<Checksum, CachedModule>,
+    modules: HashMap<Checksum, InstrumentedModule>,
 }
 
 impl PinnedMemoryCache {
@@ -17,8 +26,19 @@ impl PinnedMemoryCache {
         }
     }
 
+    pub fn iter(&self) -> impl Iterator<Item = (&Checksum, &InstrumentedModule)> {
+        self.modules.iter()
+    }
+
     pub fn store(&mut self, checksum: &Checksum, cached_module: CachedModule) -> VmResult<()> {
-        self.modules.insert(*checksum, cached_module);
+        self.modules.insert(
+            *checksum,
+            InstrumentedModule {
+                hits: 0,
+                module: cached_module,
+            },
+        );
+
         Ok(())
     }
 
@@ -31,8 +51,11 @@ impl PinnedMemoryCache {
 
     /// Looks up a module in the cache and creates a new module
     pub fn load(&mut self, checksum: &Checksum) -> VmResult<Option<CachedModule>> {
-        match self.modules.get(checksum) {
-            Some(cached) => Ok(Some(cached.clone())),
+        match self.modules.get_mut(checksum) {
+            Some(cached) => {
+                cached.hits = cached.hits.saturating_add(1);
+                Ok(Some(cached.module.clone()))
+            }
             None => Ok(None),
         }
     }
@@ -54,7 +77,7 @@ impl PinnedMemoryCache {
     pub fn size(&self) -> usize {
         self.modules
             .iter()
-            .map(|(key, module)| std::mem::size_of_val(key) + module.size_estimate)
+            .map(|(key, module)| std::mem::size_of_val(key) + module.module.size_estimate)
             .sum()
     }
 }
@@ -166,6 +189,51 @@ mod tests {
         assert!(!cache.has(&checksum));
     }
 
+    #[test]
+    fn hit_metric_works() {
+        let mut cache = PinnedMemoryCache::new();
+
+        // Create module
+        let wasm = wat::parse_str(
+            r#"(module
+            (type $t0 (func (param i32) (result i32)))
+            (func $add_one (export "add_one") (type $t0) (param $p0 i32) (result i32)
+                local.get $p0
+                i32.const 1
+                i32.add)
+            )"#,
+        )
+        .unwrap();
+        let checksum = Checksum::generate(&wasm);
+
+        assert!(!cache.has(&checksum));
+
+        // Add
+        let engine = make_compiling_engine(TESTING_MEMORY_LIMIT);
+        let original = compile(&engine, &wasm).unwrap();
+        let module = CachedModule {
+            module: original,
+            engine: make_runtime_engine(TESTING_MEMORY_LIMIT),
+            size_estimate: 0,
+        };
+        cache.store(&checksum, module).unwrap();
+
+        let (_checksum, module) = cache
+            .iter()
+            .find(|(iter_checksum, _module)| **iter_checksum == checksum)
+            .unwrap();
+
+        assert_eq!(module.hits, 0);
+
+        let _ = cache.load(&checksum).unwrap();
+        let (_checksum, module) = cache
+            .iter()
+            .find(|(iter_checksum, _module)| **iter_checksum == checksum)
+            .unwrap();
+
+        assert_eq!(module.hits, 1);
+    }
+
     #[test]
     fn len_works() {
         let mut cache = PinnedMemoryCache::new();