Stylus cache improvements

arbitrator/stylus/src/cache.rs

@@ -29,8 +29,16 @@ pub struct LruCounters {
     pub does_not_fit: u32,
 }
 
+pub struct LongTermCounters {
+    pub hits: u32,
+    pub misses: u32,
+}
+
 pub struct InitCache {
     long_term: HashMap<CacheKey, CacheItem>,
+    long_term_size_bytes: usize,
+    long_term_counters: LongTermCounters,
+
     lru: CLruCache<CacheKey, CacheItem, RandomState, CustomWeightScale>,
     lru_counters: LruCounters,
 }
@@ -91,6 +99,20 @@ pub struct LruCacheMetrics {
     pub does_not_fit: u32,
 }
 
+#[repr(C)]
+pub struct LongTermCacheMetrics {
+    pub size_bytes: u64,
+    pub count: u32,
+    pub hits: u32,
+    pub misses: u32,
+}
+
+#[repr(C)]
+pub struct CacheMetrics {
+    pub lru: LruCacheMetrics,
+    pub long_term: LongTermCacheMetrics,
+}
+
 pub fn deserialize_module(
     module: &[u8],
     version: u16,
@@ -117,6 +139,9 @@ impl InitCache {
     fn new(size_bytes: usize) -> Self {
         Self {
             long_term: HashMap::new(),
+            long_term_size_bytes: 0,
+            long_term_counters: LongTermCounters { hits: 0, misses: 0 },
+
             lru: CLruCache::with_config(
                 CLruCacheConfig::new(NonZeroUsize::new(size_bytes).unwrap())
                     .with_scale(CustomWeightScale),
@@ -142,8 +167,11 @@ impl InitCache {
 
         // See if the item is in the long term cache
         if let Some(item) = cache.long_term.get(&key) {
-            return Some(item.data());
+            let data = item.data();
+            cache.long_term_counters.hits += 1;
+            return Some(data);
         }
+        cache.long_term_counters.misses += 1;
 
         // See if the item is in the LRU cache, promoting if so
         if let Some(item) = cache.lru.get(&key) {
@@ -174,6 +202,7 @@ impl InitCache {
         if let Some(item) = cache.lru.peek(&key).cloned() {
             if long_term_tag == Self::ARBOS_TAG {
                 cache.long_term.insert(key, item.clone());
+                cache.long_term_size_bytes += item.entry_size_estimate_bytes;
             } else {
                 // only calls get to move the key to the head of the LRU list
                 cache.lru.get(&key);
@@ -195,6 +224,7 @@ impl InitCache {
             };
         } else {
             cache.long_term.insert(key, item);
+            cache.long_term_size_bytes += entry_size_estimate_bytes;
         }
         Ok(data)
     }
@@ -207,6 +237,7 @@ impl InitCache {
         let key = CacheKey::new(module_hash, version, debug);
         let mut cache = cache!();
         if let Some(item) = cache.long_term.remove(&key) {
+            cache.long_term_size_bytes -= item.entry_size_estimate_bytes;
             if cache.lru.put_with_weight(key, item).is_err() {
                 eprintln!("{}", Self::DOES_NOT_FIT_MSG);
             }
@@ -225,23 +256,32 @@ impl InitCache {
                 eprintln!("{}", Self::DOES_NOT_FIT_MSG);
             }
         }
+        cache.long_term_size_bytes = 0;
     }
 
-    pub fn get_lru_metrics() -> LruCacheMetrics {
+    pub fn get_metrics() -> CacheMetrics {
         let mut cache = cache!();
 
-        let count = cache.lru.len();
-        let metrics = LruCacheMetrics {
-            // add 1 to each entry to account that we subtracted 1 in the weight calculation
-            size_bytes: (cache.lru.weight() + count).try_into().unwrap(),
+        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: 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,
+        };
+
         // Empty counters.
         // go side, which is the only consumer of this function besides tests,
         // will read those counters and increment its own prometheus counters with them.
@@ -250,8 +290,12 @@ impl InitCache {
             misses: 0,
             does_not_fit: 0,
         };
+        cache.long_term_counters = LongTermCounters { hits: 0, misses: 0 };
 
-        metrics
+        CacheMetrics {
+            lru: lru_metrics,
+            long_term: long_term_metrics,
+        }
     }
 
     // only used for testing

arbitrator/stylus/src/lib.rs

@@ -11,7 +11,7 @@ use arbutil::{
     format::DebugBytes,
     Bytes32,
 };
-use cache::{deserialize_module, InitCache, LruCacheMetrics};
+use cache::{deserialize_module, CacheMetrics, InitCache};
 use evm_api::NativeRequestHandler;
 use eyre::ErrReport;
 use native::NativeInstance;
@@ -364,10 +364,10 @@ pub unsafe extern "C" fn stylus_drop_vec(vec: RustBytes) {
     }
 }
 
-/// Gets lru cache metrics.
+/// Gets cache metrics.
 #[no_mangle]
-pub extern "C" fn stylus_get_lru_cache_metrics() -> LruCacheMetrics {
-    InitCache::get_lru_metrics()
+pub extern "C" fn stylus_get_cache_metrics() -> CacheMetrics {
+    InitCache::get_metrics()
 }
 
 /// Clears lru cache.
@@ -377,18 +377,16 @@ pub extern "C" fn stylus_clear_lru_cache() {
     InitCache::clear_lru_cache()
 }
 
-/// Gets lru entry size in bytes.
+/// Gets entry size in bytes.
 /// Only used for testing purposes.
 #[no_mangle]
-pub extern "C" fn stylus_get_lru_entry_size_estimate_bytes(
+pub extern "C" fn stylus_get_entry_size_estimate_bytes(
     module: GoSliceData,
     version: u16,
     debug: bool,
 ) -> u64 {
     match deserialize_module(module.slice(), version, debug) {
         Err(error) => panic!("tried to get invalid asm!: {error}"),
-        Ok((_, _, lru_entry_size_estimate_bytes)) => {
-            lru_entry_size_estimate_bytes.try_into().unwrap()
-        }
+        Ok((_, _, entry_size_estimate_bytes)) => entry_size_estimate_bytes.try_into().unwrap(),
     }
 }

arbos/programs/native.go

@@ -52,6 +52,11 @@ var (
 	stylusLRUCacheSizeHitsCounter       = metrics.NewRegisteredCounter("arb/arbos/stylus/cache/lru/hits", nil)
 	stylusLRUCacheSizeMissesCounter     = metrics.NewRegisteredCounter("arb/arbos/stylus/cache/lru/misses", nil)
 	stylusLRUCacheSizeDoesNotFitCounter = metrics.NewRegisteredCounter("arb/arbos/stylus/cache/lru/does_not_fit", nil)
+
+	stylusLongTermCacheSizeBytesGauge    = metrics.NewRegisteredGauge("arb/arbos/stylus/cache/long_term/size_bytes", nil)
+	stylusLongTermCacheSizeCountGauge    = metrics.NewRegisteredGauge("arb/arbos/stylus/cache/long_term/count", nil)
+	stylusLongTermCacheSizeHitsCounter   = metrics.NewRegisteredCounter("arb/arbos/stylus/cache/long_term/hits", nil)
+	stylusLongTermCacheSizeMissesCounter = metrics.NewRegisteredCounter("arb/arbos/stylus/cache/long_term/misses", nil)
 )
 
 func activateProgram(
@@ -333,24 +338,52 @@ func SetWasmLruCacheCapacity(capacityBytes uint64) {
 	C.stylus_set_cache_lru_capacity(u64(capacityBytes))
 }
 
-// exported for testing
+func UpdateWasmCacheMetrics() {
+	metrics := C.stylus_get_cache_metrics()
+
+	stylusLRUCacheSizeBytesGauge.Update(int64(metrics.lru.size_bytes))
+	stylusLRUCacheSizeCountGauge.Update(int64(metrics.lru.count))
+	stylusLRUCacheSizeHitsCounter.Inc(int64(metrics.lru.hits))
+	stylusLRUCacheSizeMissesCounter.Inc(int64(metrics.lru.misses))
+	stylusLRUCacheSizeDoesNotFitCounter.Inc(int64(metrics.lru.does_not_fit))
+
+	stylusLongTermCacheSizeBytesGauge.Update(int64(metrics.long_term.size_bytes))
+	stylusLongTermCacheSizeCountGauge.Update(int64(metrics.long_term.count))
+	stylusLongTermCacheSizeHitsCounter.Inc(int64(metrics.long_term.hits))
+	stylusLongTermCacheSizeMissesCounter.Inc(int64(metrics.long_term.misses))
+}
+
+// Used for testing
 type WasmLruCacheMetrics struct {
 	SizeBytes uint64
 	Count     uint32
 }
 
-func GetWasmLruCacheMetrics() *WasmLruCacheMetrics {
-	metrics := C.stylus_get_lru_cache_metrics()
+// Used for testing
+type WasmLongTermCacheMetrics struct {
+	SizeBytes uint64
+	Count     uint32
+}
 
-	stylusLRUCacheSizeBytesGauge.Update(int64(metrics.size_bytes))
-	stylusLRUCacheSizeCountGauge.Update(int64(metrics.count))
-	stylusLRUCacheSizeHitsCounter.Inc(int64(metrics.hits))
-	stylusLRUCacheSizeMissesCounter.Inc(int64(metrics.misses))
-	stylusLRUCacheSizeDoesNotFitCounter.Inc(int64(metrics.does_not_fit))
+// Used for testing
+type WasmCacheMetrics struct {
+	Lru      WasmLruCacheMetrics
+	LongTerm WasmLongTermCacheMetrics
+}
 
-	return &WasmLruCacheMetrics{
-		SizeBytes: uint64(metrics.size_bytes),
-		Count:     uint32(metrics.count),
+// Used for testing
+func GetWasmCacheMetrics() *WasmCacheMetrics {
+	metrics := C.stylus_get_cache_metrics()
+
+	return &WasmCacheMetrics{
+		Lru: WasmLruCacheMetrics{
+			SizeBytes: uint64(metrics.lru.size_bytes),
+			Count:     uint32(metrics.lru.count),
+		},
+		LongTerm: WasmLongTermCacheMetrics{
+			SizeBytes: uint64(metrics.long_term.size_bytes),
+			Count:     uint32(metrics.long_term.count),
+		},
 	}
 }
 
@@ -360,8 +393,8 @@ func ClearWasmLruCache() {
 }
 
 // Used for testing
-func GetLruEntrySizeEstimateBytes(module []byte, version uint16, debug bool) uint64 {
-	return uint64(C.stylus_get_lru_entry_size_estimate_bytes(goSlice(module), u16(version), cbool(debug)))
+func GetEntrySizeEstimateBytes(module []byte, version uint16, debug bool) uint64 {
+	return uint64(C.stylus_get_entry_size_estimate_bytes(goSlice(module), u16(version), cbool(debug)))
 }
 
 const DefaultTargetDescriptionArm = "arm64-linux-unknown+neon"

execution/gethexec/blockchain.go

@@ -26,20 +26,21 @@ import (
 )
 
 type CachingConfig struct {
-	Archive                            bool          `koanf:"archive"`
-	BlockCount                         uint64        `koanf:"block-count"`
-	BlockAge                           time.Duration `koanf:"block-age"`
-	TrieTimeLimit                      time.Duration `koanf:"trie-time-limit"`
-	TrieDirtyCache                     int           `koanf:"trie-dirty-cache"`
-	TrieCleanCache                     int           `koanf:"trie-clean-cache"`
-	SnapshotCache                      int           `koanf:"snapshot-cache"`
-	DatabaseCache                      int           `koanf:"database-cache"`
-	SnapshotRestoreGasLimit            uint64        `koanf:"snapshot-restore-gas-limit"`
-	MaxNumberOfBlocksToSkipStateSaving uint32        `koanf:"max-number-of-blocks-to-skip-state-saving"`
-	MaxAmountOfGasToSkipStateSaving    uint64        `koanf:"max-amount-of-gas-to-skip-state-saving"`
-	StylusLRUCacheCapacity             uint32        `koanf:"stylus-lru-cache-capacity"`
-	StateScheme                        string        `koanf:"state-scheme"`
-	StateHistory                       uint64        `koanf:"state-history"`
+	Archive                             bool          `koanf:"archive"`
+	BlockCount                          uint64        `koanf:"block-count"`
+	BlockAge                            time.Duration `koanf:"block-age"`
+	TrieTimeLimit                       time.Duration `koanf:"trie-time-limit"`
+	TrieDirtyCache                      int           `koanf:"trie-dirty-cache"`
+	TrieCleanCache                      int           `koanf:"trie-clean-cache"`
+	SnapshotCache                       int           `koanf:"snapshot-cache"`
+	DatabaseCache                       int           `koanf:"database-cache"`
+	SnapshotRestoreGasLimit             uint64        `koanf:"snapshot-restore-gas-limit"`
+	MaxNumberOfBlocksToSkipStateSaving  uint32        `koanf:"max-number-of-blocks-to-skip-state-saving"`
+	MaxAmountOfGasToSkipStateSaving     uint64        `koanf:"max-amount-of-gas-to-skip-state-saving"`
+	StylusLRUCacheCapacity              uint32        `koanf:"stylus-lru-cache-capacity"`
+	DisableStylusCacheMetricsCollection bool          `koanf:"disable-stylus-cache-metrics-collection"`
+	StateScheme                         string        `koanf:"state-scheme"`
+	StateHistory                        uint64        `koanf:"state-history"`
 }
 
 func CachingConfigAddOptions(prefix string, f *flag.FlagSet) {
@@ -55,6 +56,7 @@ func CachingConfigAddOptions(prefix string, f *flag.FlagSet) {
 	f.Uint32(prefix+".max-number-of-blocks-to-skip-state-saving", DefaultCachingConfig.MaxNumberOfBlocksToSkipStateSaving, "maximum number of blocks to skip state saving to persistent storage (archive node only) -- warning: this option seems to cause issues")
 	f.Uint64(prefix+".max-amount-of-gas-to-skip-state-saving", DefaultCachingConfig.MaxAmountOfGasToSkipStateSaving, "maximum amount of gas in blocks to skip saving state to Persistent storage (archive node only) -- warning: this option seems to cause issues")
 	f.Uint32(prefix+".stylus-lru-cache-capacity", DefaultCachingConfig.StylusLRUCacheCapacity, "capacity, in megabytes, of the LRU cache that keeps initialized stylus programs")
+	f.Bool(prefix+".disable-stylus-cache-metrics-collection", DefaultCachingConfig.DisableStylusCacheMetricsCollection, "disable metrics collection for the stylus cache")
 	f.String(prefix+".state-scheme", DefaultCachingConfig.StateScheme, "scheme to use for state trie storage (hash, path)")
 	f.Uint64(prefix+".state-history", DefaultCachingConfig.StateHistory, "number of recent blocks to retain state history for (path state-scheme only)")
 }

execution/gethexec/executionengine.go

@@ -87,6 +87,8 @@ type ExecutionEngine struct {
 
 	reorgSequencing bool
 
+	disableStylusCacheMetricsCollection bool
+
 	prefetchBlock bool
 
 	cachedL1PriceData *L1PriceData
@@ -212,6 +214,16 @@ func (s *ExecutionEngine) EnableReorgSequencing() {
 	s.reorgSequencing = true
 }
 
+func (s *ExecutionEngine) DisableStylusCacheMetricsCollection() {
+	if s.Started() {
+		panic("trying to disable stylus cache metrics collection after start")
+	}
+	if s.disableStylusCacheMetricsCollection {
+		panic("trying to disable stylus cache metrics collection when already set")
+	}
+	s.disableStylusCacheMetricsCollection = true
+}
+
 func (s *ExecutionEngine) EnablePrefetchBlock() {
 	if s.Started() {
 		panic("trying to enable prefetch block after start")
@@ -963,15 +975,17 @@ func (s *ExecutionEngine) Start(ctx_in context.Context) {
 			}
 		}
 	})
-	// periodically update stylus lru cache metrics
-	s.LaunchThread(func(ctx context.Context) {
-		for {
-			select {
-			case <-ctx.Done():
-				return
-			case <-time.After(time.Minute):
-				programs.GetWasmLruCacheMetrics()
+	if !s.disableStylusCacheMetricsCollection {
+		// periodically update stylus cache metrics
+		s.LaunchThread(func(ctx context.Context) {
+			for {
+				select {
+				case <-ctx.Done():
+					return
+				case <-time.After(time.Minute):
+					programs.UpdateWasmCacheMetrics()
+				}
 			}
-		}
-	})
+		})
+	}
 }

execution/gethexec/node.go

@@ -187,6 +187,9 @@ func CreateExecutionNode(
 	if config.EnablePrefetchBlock {
 		execEngine.EnablePrefetchBlock()
 	}
+	if config.Caching.DisableStylusCacheMetricsCollection {
+		execEngine.DisableStylusCacheMetricsCollection()
+	}
 	if err != nil {
 		return nil, err
 	}