Merge pull request #2734 from OffchainLabs/fix_mem_leak_stylus_cache_…

…metrics Fix memory leak when getting stylus cache metrics
OffchainLabs · Oct 10, 2024 · a2ae1c5 · a2ae1c5
2 parents 65196bb + c5625d6
commit a2ae1c5
Show file tree

Hide file tree

Showing 3 changed files with 23 additions and 28 deletions.
diff --git a/arbitrator/stylus/src/cache.rs b/arbitrator/stylus/src/cache.rs
@@ -275,28 +275,21 @@ impl InitCache {
         cache.long_term_size_bytes = 0;
     }
 
-    pub fn get_metrics() -> CacheMetrics {
+    pub fn get_metrics(output: &mut CacheMetrics) {
         let mut cache = cache!();
 
         let lru_count = cache.lru.len();
-        let lru_metrics = LruCacheMetrics {
-            // adds 1 to each entry to account that we subtracted 1 in the weight calculation
-            size_bytes: (cache.lru.weight() + lru_count).try_into().unwrap(),
-
-            count: lru_count.try_into().unwrap(),
-
-            hits: cache.lru_counters.hits,
-            misses: cache.lru_counters.misses,
-            does_not_fit: cache.lru_counters.does_not_fit,
-        };
-
-        let long_term_metrics = LongTermCacheMetrics {
-            size_bytes: cache.long_term_size_bytes.try_into().unwrap(),
-            count: cache.long_term.len().try_into().unwrap(),
-
-            hits: cache.long_term_counters.hits,
-            misses: cache.long_term_counters.misses,
-        };
+        // adds 1 to each entry to account that we subtracted 1 in the weight calculation
+        output.lru.size_bytes = (cache.lru.weight() + lru_count).try_into().unwrap();
+        output.lru.count = lru_count.try_into().unwrap();
+        output.lru.hits = cache.lru_counters.hits;
+        output.lru.misses = cache.lru_counters.misses;
+        output.lru.does_not_fit = cache.lru_counters.does_not_fit;
+
+        output.long_term.size_bytes = cache.long_term_size_bytes.try_into().unwrap();
+        output.long_term.count = cache.long_term.len().try_into().unwrap();
+        output.long_term.hits = cache.long_term_counters.hits;
+        output.long_term.misses = cache.long_term_counters.misses;
 
         // Empty counters.
         // go side, which is the only consumer of this function besides tests,
@@ -307,11 +300,6 @@ impl InitCache {
             does_not_fit: 0,
         };
         cache.long_term_counters = LongTermCounters { hits: 0, misses: 0 };
-
-        CacheMetrics {
-            lru: lru_metrics,
-            long_term: long_term_metrics,
-        }
     }
 
     // only used for testing

diff --git a/arbitrator/stylus/src/lib.rs b/arbitrator/stylus/src/lib.rs
@@ -365,9 +365,14 @@ pub unsafe extern "C" fn stylus_drop_vec(vec: RustBytes) {
 }
 
 /// Gets cache metrics.
+///
+/// # Safety
+///
+/// `output` must not be null.
 #[no_mangle]
-pub extern "C" fn stylus_get_cache_metrics() -> CacheMetrics {
-    InitCache::get_metrics()
+pub unsafe extern "C" fn stylus_get_cache_metrics(output: *mut CacheMetrics) {
+    let output = &mut *output;
+    InitCache::get_metrics(output);
 }
 
 /// Clears lru cache.

diff --git a/arbos/programs/native.go b/arbos/programs/native.go
@@ -339,7 +339,8 @@ func SetWasmLruCacheCapacity(capacityBytes uint64) {
 }
 
 func UpdateWasmCacheMetrics() {
-	metrics := C.stylus_get_cache_metrics()
+	metrics := &C.CacheMetrics{}
+	C.stylus_get_cache_metrics(metrics)
 
 	stylusLRUCacheSizeBytesGauge.Update(int64(metrics.lru.size_bytes))
 	stylusLRUCacheCountGauge.Update(int64(metrics.lru.count))
@@ -373,7 +374,8 @@ type WasmCacheMetrics struct {
 
 // Used for testing
 func GetWasmCacheMetrics() *WasmCacheMetrics {
-	metrics := C.stylus_get_cache_metrics()
+	metrics := &C.CacheMetrics{}
+	C.stylus_get_cache_metrics(metrics)
 
 	return &WasmCacheMetrics{
 		Lru: WasmLruCacheMetrics{