Merge pull request #2315 from OffchainLabs/stylus_cache_tag

Stylus cache tag

arbitrator/stylus/src/cache.rs

@@ -21,7 +21,7 @@ macro_rules! cache {
 }
 
 pub struct InitCache {
-    arbos: HashMap<CacheKey, CacheItem>,
+    long_term: HashMap<CacheKey, CacheItem>,
     lru: LruCache<CacheKey, CacheItem>,
 }
 
@@ -59,20 +59,31 @@ impl CacheItem {
 }
 
 impl InitCache {
+    // current implementation only has one tag that stores to the long_term
+    // future implementations might have more, but 0 is a reserved tag
+    // that will never modify long_term state
+    const ARBOS_TAG: u32 = 1;
+
     fn new(size: usize) -> Self {
         Self {
-            arbos: HashMap::new(),
+            long_term: HashMap::new(),
             lru: LruCache::new(NonZeroUsize::new(size).unwrap()),
         }
     }
 
+    pub fn set_lru_size(size: u32) {
+        cache!()
+            .lru
+            .resize(NonZeroUsize::new(size.try_into().unwrap()).unwrap())
+    }
+
     /// Retrieves a cached value, updating items as necessary.
     pub fn get(module_hash: Bytes32, version: u16, debug: bool) -> Option<(Module, Store)> {
         let mut cache = cache!();
         let key = CacheKey::new(module_hash, version, debug);
 
         // See if the item is in the long term cache
-        if let Some(item) = cache.arbos.get(&key) {
+        if let Some(item) = cache.long_term.get(&key) {
             return Some(item.data());
         }
 
@@ -84,18 +95,27 @@ impl InitCache {
     }
 
     /// Inserts an item into the long term cache, cloning from the LRU cache if able.
+    /// If long_term_tag is 0 will only insert to LRU
     pub fn insert(
         module_hash: Bytes32,
         module: &[u8],
         version: u16,
+        long_term_tag: u32,
         debug: bool,
     ) -> Result<(Module, Store)> {
         let key = CacheKey::new(module_hash, version, debug);
 
         // if in LRU, add to ArbOS
         let mut cache = cache!();
+        if let Some(item) = cache.long_term.get(&key) {
+            return Ok(item.data());
+        }
         if let Some(item) = cache.lru.peek(&key).cloned() {
-            cache.arbos.insert(key, item.clone());
+            if long_term_tag == Self::ARBOS_TAG {
+                cache.long_term.insert(key, item.clone());
+            } else {
+                cache.lru.promote(&key)
+            }
             return Ok(item.data());
         }
         drop(cache);
@@ -105,37 +125,34 @@ impl InitCache {
 
         let item = CacheItem::new(module, engine);
         let data = item.data();
-        cache!().arbos.insert(key, item);
+        let mut cache = cache!();
+        if long_term_tag != Self::ARBOS_TAG {
+            cache.lru.put(key, item);
+        } else {
+            cache.long_term.insert(key, item);
+        }
         Ok(data)
     }
 
-    /// Inserts an item into the short-lived LRU cache.
-    pub fn insert_lru(
-        module_hash: Bytes32,
-        module: &[u8],
-        version: u16,
-        debug: bool,
-    ) -> Result<(Module, Store)> {
-        let engine = CompileConfig::version(version, debug).engine();
-        let module = unsafe { Module::deserialize_unchecked(&engine, module)? };
-
-        let key = CacheKey::new(module_hash, version, debug);
-        let item = CacheItem::new(module, engine);
-        cache!().lru.put(key, item.clone());
-        Ok(item.data())
-    }
-
     /// Evicts an item in the long-term cache.
-    pub fn evict(module_hash: Bytes32, version: u16, debug: bool) {
+    pub fn evict(module_hash: Bytes32, version: u16, long_term_tag: u32, debug: bool) {
+        if long_term_tag != Self::ARBOS_TAG {
+            return;
+        }
         let key = CacheKey::new(module_hash, version, debug);
-        cache!().arbos.remove(&key);
+        let mut cache = cache!();
+        if let Some(item) = cache.long_term.remove(&key) {
+            cache.lru.put(key, item);
+        }
     }
 
-    /// Modifies the cache for reorg, dropping the long-term cache.
-    pub fn reorg(_block: u64) {
+    pub fn clear_long_term(long_term_tag: u32) {
+        if long_term_tag != Self::ARBOS_TAG {
+            return;
+        }
         let mut cache = cache!();
         let cache = &mut *cache;
-        for (key, item) in cache.arbos.drain() {
+        for (key, item) in cache.long_term.drain() {
             cache.lru.put(key, item); // not all will fit, just a heuristic
         }
     }

arbitrator/stylus/src/lib.rs

@@ -183,6 +183,7 @@ pub unsafe extern "C" fn stylus_call(
     debug_chain: bool,
     output: *mut RustBytes,
     gas: *mut u64,
+    long_term_tag: u32,
 ) -> UserOutcomeKind {
     let module = module.slice();
     let calldata = calldata.slice().to_vec();
@@ -193,7 +194,14 @@ pub unsafe extern "C" fn stylus_call(
 
     // Safety: module came from compile_user_wasm and we've paid for memory expansion
     let instance = unsafe {
-        NativeInstance::deserialize_cached(module, config.version, evm_api, evm_data, debug_chain)
+        NativeInstance::deserialize_cached(
+            module,
+            config.version,
+            evm_api,
+            evm_data,
+            long_term_tag,
+            debug_chain,
+        )
     };
     let mut instance = match instance {
         Ok(instance) => instance,
@@ -212,33 +220,47 @@ pub unsafe extern "C" fn stylus_call(
     status
 }
 
+/// resize lru
+#[no_mangle]
+pub extern "C" fn stylus_cache_lru_resize(size: u32) {
+    InitCache::set_lru_size(size);
+}
+
 /// Caches an activated user program.
 ///
 /// # Safety
 ///
 /// `module` must represent a valid module produced from `stylus_activate`.
+/// arbos_tag: a tag for arbos cache. 0 won't affect real caching
+/// currently only if tag==1 caching will be affected
 #[no_mangle]
 pub unsafe extern "C" fn stylus_cache_module(
     module: GoSliceData,
     module_hash: Bytes32,
     version: u16,
+    arbos_tag: u32,
     debug: bool,
 ) {
-    if let Err(error) = InitCache::insert(module_hash, module.slice(), version, debug) {
+    if let Err(error) = InitCache::insert(module_hash, module.slice(), version, arbos_tag, debug) {
         panic!("tried to cache invalid asm!: {error}");
     }
 }
 
 /// Evicts an activated user program from the init cache.
 #[no_mangle]
-pub extern "C" fn stylus_evict_module(module_hash: Bytes32, version: u16, debug: bool) {
-    InitCache::evict(module_hash, version, debug);
+pub extern "C" fn stylus_evict_module(
+    module_hash: Bytes32,
+    version: u16,
+    arbos_tag: u32,
+    debug: bool,
+) {
+    InitCache::evict(module_hash, version, arbos_tag, debug);
 }
 
 /// Reorgs the init cache. This will likely never happen.
 #[no_mangle]
-pub extern "C" fn stylus_reorg_vm(block: u64) {
-    InitCache::reorg(block);
+pub extern "C" fn stylus_reorg_vm(_block: u64, arbos_tag: u32) {
+    InitCache::clear_long_term(arbos_tag);
 }
 
 /// Frees the vector. Does nothing when the vector is null.

arbitrator/stylus/src/native.rs

@@ -113,6 +113,7 @@ impl<D: DataReader, E: EvmApi<D>> NativeInstance<D, E> {
         version: u16,
         evm: E,
         evm_data: EvmData,
+        mut long_term_tag: u32,
         debug: bool,
     ) -> Result<Self> {
         let compile = CompileConfig::version(version, debug);
@@ -122,10 +123,11 @@ impl<D: DataReader, E: EvmApi<D>> NativeInstance<D, E> {
         if let Some((module, store)) = InitCache::get(module_hash, version, debug) {
             return Self::from_module(module, store, env);
         }
-        let (module, store) = match env.evm_data.cached {
-            true => InitCache::insert(module_hash, module, version, debug)?,
-            false => InitCache::insert_lru(module_hash, module, version, debug)?,
-        };
+        if !env.evm_data.cached {
+            long_term_tag = 0;
+        }
+        let (module, store) =
+            InitCache::insert(module_hash, module, version, long_term_tag, debug)?;
         Self::from_module(module, store, env)
     }
 

arbnode/inbox_test.go

@@ -65,6 +65,7 @@ func NewTransactionStreamerForTest(t *testing.T, ownerAddress common.Address) (*
 	if err != nil {
 		Fail(t, err)
 	}
+	execEngine.Initialize(gethexec.DefaultCachingConfig.StylusLRUCache)
 	execSeq := &execClientWrapper{execEngine, t}
 	inbox, err := NewTransactionStreamer(arbDb, bc.Config(), execSeq, nil, make(chan error, 1), transactionStreamerConfigFetcher)
 	if err != nil {

arbos/programs/native.go

@@ -172,6 +172,7 @@ func callProgram(
 	evmData *EvmData,
 	stylusParams *ProgParams,
 	memoryModel *MemoryModel,
+	arbos_tag uint32,
 ) ([]byte, error) {
 	db := interpreter.Evm().StateDB
 	debug := stylusParams.DebugMode
@@ -198,6 +199,7 @@ func callProgram(
 		cbool(debug),
 		output,
 		(*u64)(&scope.Contract.Gas),
+		u32(arbos_tag),
 	))
 
 	depth := interpreter.Depth()
@@ -228,31 +230,37 @@ func cacheProgram(db vm.StateDB, module common.Hash, program Program, params *St
 		if err != nil {
 			panic("unable to recreate wasm")
 		}
-		state.CacheWasmRust(asm, module, program.version, debug)
-		db.RecordCacheWasm(state.CacheWasm{ModuleHash: module, Version: program.version, Debug: debug})
+		tag := db.Database().WasmCacheTag()
+		state.CacheWasmRust(asm, module, program.version, tag, debug)
+		db.RecordCacheWasm(state.CacheWasm{ModuleHash: module, Version: program.version, Tag: tag, Debug: debug})
 	}
 }
 
 // Evicts a program in Rust. We write a record so that we can undo on revert, unless we don't need to (e.g. expired)
 // For gas estimation and eth_call, we ignore permanent updates and rely on Rust's LRU.
 func evictProgram(db vm.StateDB, module common.Hash, version uint16, debug bool, runMode core.MessageRunMode, forever bool) {
 	if runMode == core.MessageCommitMode {
-		state.EvictWasmRust(module, version, debug)
+		tag := db.Database().WasmCacheTag()
+		state.EvictWasmRust(module, version, tag, debug)
 		if !forever {
-			db.RecordEvictWasm(state.EvictWasm{ModuleHash: module, Version: version, Debug: debug})
+			db.RecordEvictWasm(state.EvictWasm{ModuleHash: module, Version: version, Tag: tag, Debug: debug})
 		}
 	}
 }
 
 func init() {
-	state.CacheWasmRust = func(asm []byte, moduleHash common.Hash, version uint16, debug bool) {
-		C.stylus_cache_module(goSlice(asm), hashToBytes32(moduleHash), u16(version), cbool(debug))
+	state.CacheWasmRust = func(asm []byte, moduleHash common.Hash, version uint16, tag uint32, debug bool) {
+		C.stylus_cache_module(goSlice(asm), hashToBytes32(moduleHash), u16(version), u32(tag), cbool(debug))
 	}
-	state.EvictWasmRust = func(moduleHash common.Hash, version uint16, debug bool) {
-		C.stylus_evict_module(hashToBytes32(moduleHash), u16(version), cbool(debug))
+	state.EvictWasmRust = func(moduleHash common.Hash, version uint16, tag uint32, debug bool) {
+		C.stylus_evict_module(hashToBytes32(moduleHash), u16(version), u32(tag), cbool(debug))
 	}
 }
 
+func ResizeWasmLruCache(size uint32) {
+	C.stylus_cache_lru_resize(u32(size))
+}
+
 func (value bytes32) toHash() common.Hash {
 	hash := common.Hash{}
 	for index, b := range value.bytes {

arbos/programs/programs.go

@@ -166,6 +166,7 @@ func (p Programs) CallProgram(
 	tracingInfo *util.TracingInfo,
 	calldata []byte,
 	reentrant bool,
+	runmode core.MessageRunMode,
 ) ([]byte, error) {
 	evm := interpreter.Evm()
 	contract := scope.Contract
@@ -237,7 +238,11 @@ func (p Programs) CallProgram(
 	if contract.CodeAddr != nil {
 		address = *contract.CodeAddr
 	}
-	return callProgram(address, moduleHash, localAsm, scope, interpreter, tracingInfo, calldata, evmData, goParams, model)
+	var arbos_tag uint32
+	if runmode == core.MessageCommitMode {
+		arbos_tag = statedb.Database().WasmCacheTag()
+	}
+	return callProgram(address, moduleHash, localAsm, scope, interpreter, tracingInfo, calldata, evmData, goParams, model, arbos_tag)
 }
 
 func getWasm(statedb vm.StateDB, program common.Address) ([]byte, error) {

arbos/programs/wasm.go

@@ -143,6 +143,7 @@ func callProgram(
 	evmData *EvmData,
 	params *ProgParams,
 	memoryModel *MemoryModel,
+	_arbos_tag uint32,
 ) ([]byte, error) {
 	reqHandler := newApiClosures(interpreter, tracingInfo, scope, memoryModel)
 	gasLeft, retData, err := CallProgramLoop(moduleHash, calldata, scope.Contract.Gas, evmData, params, reqHandler)

arbos/tx_processor.go

@@ -127,6 +127,7 @@ func (p *TxProcessor) ExecuteWASM(scope *vm.ScopeContext, input []byte, interpre
 		tracingInfo,
 		input,
 		reentrant,
+		p.RunMode(),
 	)
 }
 

cmd/nitro/init.go

@@ -186,7 +186,7 @@ func openInitializeChainDb(ctx context.Context, stack *node.Node, config *NodeCo
 				if err != nil {
 					return nil, nil, err
 				}
-				chainDb := rawdb.WrapDatabaseWithWasm(chainData, wasmDb)
+				chainDb := rawdb.WrapDatabaseWithWasm(chainData, wasmDb, 1)
 				err = pruning.PruneChainDb(ctx, chainDb, stack, &config.Init, cacheConfig, l1Client, rollupAddrs, config.Node.ValidatorRequired())
 				if err != nil {
 					return chainDb, nil, fmt.Errorf("error pruning: %w", err)
@@ -243,7 +243,7 @@ func openInitializeChainDb(ctx context.Context, stack *node.Node, config *NodeCo
 	if err != nil {
 		return nil, nil, err
 	}
-	chainDb := rawdb.WrapDatabaseWithWasm(chainData, wasmDb)
+	chainDb := rawdb.WrapDatabaseWithWasm(chainData, wasmDb, 1)
 
 	if config.Init.ImportFile != "" {
 		initDataReader, err = statetransfer.NewJsonInitDataReader(config.Init.ImportFile)

execution/gethexec/blockchain.go

@@ -37,6 +37,7 @@ type CachingConfig struct {
 	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"`
+	StylusLRUCache                     uint32        `koanf:"stylus-lru-cache"`
 }
 
 func CachingConfigAddOptions(prefix string, f *flag.FlagSet) {
@@ -51,6 +52,7 @@ func CachingConfigAddOptions(prefix string, f *flag.FlagSet) {
 	f.Uint64(prefix+".snapshot-restore-gas-limit", DefaultCachingConfig.SnapshotRestoreGasLimit, "maximum gas rolled back to recover snapshot")
 	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", DefaultCachingConfig.StylusLRUCache, "initialized stylus programs to keep in LRU cache")
 }
 
 var DefaultCachingConfig = CachingConfig{
@@ -65,6 +67,22 @@ var DefaultCachingConfig = CachingConfig{
 	SnapshotRestoreGasLimit:            300_000_000_000,
 	MaxNumberOfBlocksToSkipStateSaving: 0,
 	MaxAmountOfGasToSkipStateSaving:    0,
+	StylusLRUCache:                     256,
+}
+
+var TestCachingConfig = CachingConfig{
+	Archive:                            false,
+	BlockCount:                         128,
+	BlockAge:                           30 * time.Minute,
+	TrieTimeLimit:                      time.Hour,
+	TrieDirtyCache:                     1024,
+	TrieCleanCache:                     600,
+	SnapshotCache:                      400,
+	DatabaseCache:                      2048,
+	SnapshotRestoreGasLimit:            300_000_000_000,
+	MaxNumberOfBlocksToSkipStateSaving: 0,
+	MaxAmountOfGasToSkipStateSaving:    0,
+	StylusLRUCache:                     0,
 }
 
 // TODO remove stack from parameters as it is no longer needed here