add LRU cache & journaling with Rust stub

arbitrator/stylus/src/lib.rs

@@ -160,7 +160,7 @@ pub unsafe extern "C" fn stylus_activate(
     UserOutcomeKind::Success
 }
 
-/// Calls a compiled user program.
+/// Calls an activated user program.
 ///
 /// # Safety
 ///
@@ -204,6 +204,31 @@ pub unsafe extern "C" fn stylus_call(
     status
 }
 
+/// Caches an activated user program.
+///
+/// # Safety
+///
+/// `module` must represent a valid module produced from `stylus_activate`.
+#[no_mangle]
+pub unsafe extern "C" fn stylus_cache_module(
+    module: GoSliceData,
+    module_hash: Bytes32,
+    version: u16,
+    debug: bool,
+) {
+    println!("caching module {}", module_hash);
+}
+
+/// Evicts an activated user program from the init cache.
+///
+/// # Safety
+///
+/// `module` must represent a valid module produced from `stylus_activate`.
+#[no_mangle]
+pub unsafe extern "C" fn stylus_evict_module(module_hash: Bytes32) {
+    println!("evicting module {}", module_hash);
+}
+
 /// Frees the vector. Does nothing when the vector is null.
 ///
 /// # Safety

arbos/programs/native.go

@@ -23,6 +23,7 @@ import (
 	"fmt"
 
 	"github.com/ethereum/go-ethereum/common"
+	"github.com/ethereum/go-ethereum/core"
 	"github.com/ethereum/go-ethereum/core/state"
 	"github.com/ethereum/go-ethereum/core/vm"
 	"github.com/ethereum/go-ethereum/log"
@@ -99,18 +100,19 @@ func callProgram(
 	address common.Address,
 	moduleHash common.Hash,
 	scope *vm.ScopeContext,
-	db vm.StateDB,
 	interpreter *vm.EVMInterpreter,
 	tracingInfo *util.TracingInfo,
 	calldata []byte,
 	evmData *evmData,
 	stylusParams *goParams,
 	memoryModel *MemoryModel,
 ) ([]byte, error) {
+	db := interpreter.Evm().StateDB
+	asm := db.GetActivatedAsm(moduleHash)
+
 	if db, ok := db.(*state.StateDB); ok {
 		db.RecordProgram(moduleHash)
 	}
-	asm := db.GetActivatedAsm(moduleHash)
 
 	evmApi := newApi(interpreter, tracingInfo, scope, memoryModel)
 	defer evmApi.drop()
@@ -147,6 +149,40 @@ func handleReqImpl(apiId usize, req_type u32, data *rustSlice, costPtr *u64, out
 	api.pinAndRef(raw_data, out_raw_data)
 }
 
+// Caches a program in Rust. We write a record so that we can undo on revert.
+// For gas estimation and eth_call, we ignore permanent updates and rely on Rust's LRU.
+func cacheProgram(db vm.StateDB, module common.Hash, version uint16, debug bool, runMode core.MessageRunMode) {
+	if runMode == core.MessageCommitMode {
+		asm := db.GetActivatedAsm(module)
+		state.CacheWasmRust(asm, module, version, debug)
+		db.RecordCacheWasm(state.CacheWasm{ModuleHash: module})
+	}
+}
+
+// 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)
+		if !forever {
+			db.RecordEvictWasm(state.EvictWasm{
+				ModuleHash: module,
+				Version:    version,
+				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.EvictWasmRust = func(moduleHash common.Hash) {
+		C.stylus_evict_module(hashToBytes32(moduleHash))
+	}
+}
+
 func (value bytes32) toHash() common.Hash {
 	hash := common.Hash{}
 	for index, b := range value.bytes {

arbos/programs/params.go

@@ -25,7 +25,7 @@ const initialMinInitGas = 0         // assume pricer is correct (update in case
 const initialMinCachedInitGas = 0   // assume pricer is correct (update in case of emergency)
 const initialExpiryDays = 365       // deactivate after 1 year.
 const initialKeepaliveDays = 31     // wait a month before allowing reactivation
-const initialTxCacheSize = 32       // cache the 32 most recent programs
+const initialRecentCacheSize = 32   // cache the 32 most recent programs
 
 // This struct exists to collect the many Stylus configuration parameters into a single word.
 // The items here must only be modified in ArbOwner precompile methods (or in ArbOS upgrades).
@@ -42,7 +42,7 @@ type StylusParams struct {
 	MinCachedInitGas uint8 // measured in 64-gas increments
 	ExpiryDays       uint16
 	KeepaliveDays    uint16
-	TxCacheSize      uint8
+	BlockCacheSize   uint16
 }
 
 // Provides a view of the Stylus parameters. Call Save() to persist.
@@ -83,7 +83,7 @@ func (p Programs) Params() (*StylusParams, error) {
 		MinCachedInitGas: am.BytesToUint8(take(1)),
 		ExpiryDays:       am.BytesToUint16(take(2)),
 		KeepaliveDays:    am.BytesToUint16(take(2)),
-		TxCacheSize:      am.BytesToUint8(take(1)),
+		BlockCacheSize:   am.BytesToUint16(take(2)),
 	}, nil
 }
 
@@ -107,7 +107,7 @@ func (p *StylusParams) Save() error {
 		am.Uint8ToBytes(p.MinCachedInitGas),
 		am.Uint16ToBytes(p.ExpiryDays),
 		am.Uint16ToBytes(p.KeepaliveDays),
-		am.Uint8ToBytes(p.TxCacheSize),
+		am.Uint16ToBytes(p.BlockCacheSize),
 	)
 
 	slot := uint64(0)
@@ -140,7 +140,7 @@ func initStylusParams(sto *storage.Storage) {
 		MinCachedInitGas: initialMinCachedInitGas,
 		ExpiryDays:       initialExpiryDays,
 		KeepaliveDays:    initialKeepaliveDays,
-		TxCacheSize:      initialTxCacheSize,
+		BlockCacheSize:   initialRecentCacheSize,
 	}
 	_ = params.Save()
 }

arbos/programs/programs.go

@@ -9,6 +9,7 @@ import (
 	"math/big"
 
 	"github.com/ethereum/go-ethereum/common"
+	"github.com/ethereum/go-ethereum/core"
 	"github.com/ethereum/go-ethereum/core/state"
 	"github.com/ethereum/go-ethereum/core/vm"
 	"github.com/ethereum/go-ethereum/log"
@@ -80,7 +81,7 @@ func (p Programs) CacheManagers() *addressSet.AddressSet {
 	return p.cacheManagers
 }
 
-func (p Programs) ActivateProgram(evm *vm.EVM, address common.Address, debugMode bool) (
+func (p Programs) ActivateProgram(evm *vm.EVM, address common.Address, runMode core.MessageRunMode, debugMode bool) (
 	uint16, common.Hash, common.Hash, *big.Int, bool, error,
 ) {
 	statedb := evm.StateDB
@@ -118,6 +119,16 @@ func (p Programs) ActivateProgram(evm *vm.EVM, address common.Address, debugMode
 	if err != nil {
 		return 0, codeHash, common.Hash{}, nil, true, err
 	}
+
+	// replace the cached asm
+	if cached {
+		oldModuleHash, err := p.moduleHashes.Get(codeHash)
+		if err != nil {
+			return 0, codeHash, common.Hash{}, nil, true, err
+		}
+		evictProgram(statedb, oldModuleHash, currentVersion, debugMode, runMode, expired)
+		cacheProgram(statedb, info.moduleHash, stylusVersion, debugMode, runMode)
+	}
 	if err := p.moduleHashes.Set(codeHash, info.moduleHash); err != nil {
 		return 0, codeHash, common.Hash{}, nil, true, err
 	}
@@ -139,7 +150,7 @@ func (p Programs) ActivateProgram(evm *vm.EVM, address common.Address, debugMode
 		footprint:     info.footprint,
 		asmEstimateKb: estimateKb,
 		activatedAt:   hoursSinceArbitrum(time),
-		cached:        cached, // TODO: propagate to Rust
+		cached:        cached,
 	}
 	return stylusVersion, codeHash, info.moduleHash, dataFee, false, p.setProgram(codeHash, programData)
 }
@@ -150,7 +161,6 @@ func (p Programs) CallProgram(
 	interpreter *vm.EVMInterpreter,
 	tracingInfo *util.TracingInfo,
 	calldata []byte,
-	recentPrograms *RecentPrograms,
 	reentrant bool,
 ) ([]byte, error) {
 	evm := interpreter.Evm()
@@ -183,7 +193,7 @@ func (p Programs) CallProgram(
 	callCost := model.GasCost(program.footprint, open, ever)
 
 	// pay for program init
-	cached := program.cached || recentPrograms.Insert(codeHash, params)
+	cached := program.cached || statedb.GetRecentWasms().Insert(codeHash, params.BlockCacheSize)
 	if cached {
 		callCost = arbmath.SaturatingUAdd(callCost, 64*uint64(params.MinCachedInitGas))
 		callCost = arbmath.SaturatingUAdd(callCost, uint64(program.cachedInitGas))
@@ -217,10 +227,7 @@ func (p Programs) CallProgram(
 	if contract.CodeAddr != nil {
 		address = *contract.CodeAddr
 	}
-	return callProgram(
-		address, moduleHash, scope, statedb, interpreter,
-		tracingInfo, calldata, evmData, goParams, model,
-	)
+	return callProgram(address, moduleHash, scope, interpreter, tracingInfo, calldata, evmData, goParams, model)
 }
 
 func getWasm(statedb vm.StateDB, program common.Address) ([]byte, error) {
@@ -335,7 +342,16 @@ func (p Programs) ProgramCached(codeHash common.Hash) (bool, error) {
 }
 
 // Sets whether a program is cached. Errors if trying to cache an expired program.
-func (p Programs) SetProgramCached(emitEvent func() error, codeHash common.Hash, cache bool, time uint64, params *StylusParams) error {
+func (p Programs) SetProgramCached(
+	emitEvent func() error,
+	db vm.StateDB,
+	codeHash common.Hash,
+	cache bool,
+	time uint64,
+	params *StylusParams,
+	runMode core.MessageRunMode,
+	debug bool,
+) error {
 	program, err := p.getProgram(codeHash, time)
 	if err != nil {
 		return err
@@ -354,13 +370,21 @@ func (p Programs) SetProgramCached(emitEvent func() error, codeHash common.Hash,
 	if err := emitEvent(); err != nil {
 		return err
 	}
+
+	// pay to cache the program, or to re-cache in case of upcoming revert
+	if err := p.programs.Burner().Burn(uint64(program.initGas)); err != nil {
+		return err
+	}
+	moduleHash, err := p.moduleHashes.Get(codeHash)
+	if err != nil {
+		return err
+	}
 	if cache {
-		// pay to cache the program
-		if err := p.programs.Burner().Burn(uint64(program.initGas)); err != nil {
-			return err
-		}
+		cacheProgram(db, moduleHash, program.version, debug, runMode)
+	} else {
+		evictProgram(db, moduleHash, program.version, debug, runMode, expired)
 	}
-	program.cached = cache // TODO: propagate to Rust
+	program.cached = cache
 	return p.setProgram(codeHash, program)
 }
 

arbos/programs/wasm.go

@@ -11,6 +11,7 @@ import (
 	"unsafe"
 
 	"github.com/ethereum/go-ethereum/common"
+	"github.com/ethereum/go-ethereum/core"
 	"github.com/ethereum/go-ethereum/core/vm"
 	"github.com/ethereum/go-ethereum/log"
 	"github.com/offchainlabs/nitro/arbos/burn"
@@ -90,6 +91,12 @@ func activateProgram(
 	return &activationInfo{moduleHash, initGas, cachedInitGas, asmEstimate, footprint}, nil
 }
 
+// stub any non-consensus, Rust-side caching updates
+func cacheProgram(db vm.StateDB, module common.Hash, version uint16, debug bool, mode core.MessageRunMode) {
+}
+func evictProgram(db vm.StateDB, module common.Hash, version uint16, debug bool, mode core.MessageRunMode, forever bool) {
+}
+
 //go:wasmimport programs new_program
 func newProgram(
 	hashPtr unsafe.Pointer,
@@ -122,7 +129,6 @@ func callProgram(
 	address common.Address,
 	moduleHash common.Hash,
 	scope *vm.ScopeContext,
-	db vm.StateDB,
 	interpreter *vm.EVMInterpreter,
 	tracingInfo *util.TracingInfo,
 	calldata []byte,

arbos/tx_processor.go

@@ -9,7 +9,6 @@ import (
 	"math/big"
 
 	"github.com/offchainlabs/nitro/arbos/l1pricing"
-	"github.com/offchainlabs/nitro/arbos/programs"
 
 	"github.com/offchainlabs/nitro/arbos/util"
 	"github.com/offchainlabs/nitro/util/arbmath"
@@ -42,9 +41,8 @@ type TxProcessor struct {
 	computeHoldGas   uint64 // amount of gas temporarily held to prevent compute from exceeding the gas limit
 	delayedInbox     bool   // whether this tx was submitted through the delayed inbox
 	Contracts        []*vm.Contract
-	Programs         map[common.Address]uint  // # of distinct context spans for each program
-	RecentPrograms   *programs.RecentPrograms // programs recently accessed this tx
-	TopTxType        *byte                    // set once in StartTxHook
+	Programs         map[common.Address]uint // # of distinct context spans for each program
+	TopTxType        *byte                   // set once in StartTxHook
 	evm              *vm.EVM
 	CurrentRetryable *common.Hash
 	CurrentRefundTo  *common.Address
@@ -66,7 +64,6 @@ func NewTxProcessor(evm *vm.EVM, msg *core.Message) *TxProcessor {
 		delayedInbox:        evm.Context.Coinbase != l1pricing.BatchPosterAddress,
 		Contracts:           []*vm.Contract{},
 		Programs:            make(map[common.Address]uint),
-		RecentPrograms:      programs.NewRecentProgramsTracker(),
 		TopTxType:           nil,
 		evm:                 evm,
 		CurrentRetryable:    nil,
@@ -128,7 +125,6 @@ func (p *TxProcessor) ExecuteWASM(scope *vm.ScopeContext, input []byte, interpre
 		interpreter,
 		tracingInfo,
 		input,
-		p.RecentPrograms,
 		reentrant,
 	)
 }
@@ -480,6 +476,10 @@ func (p *TxProcessor) GasChargingHook(gasRemaining *uint64) (common.Address, err
 	return tipReceipient, nil
 }
 
+func (p *TxProcessor) RunMode() core.MessageRunMode {
+	return p.msg.TxRunMode
+}
+
 func (p *TxProcessor) NonrefundableGas() uint64 {
 	// EVM-incentivized activity like freeing storage should only refund amounts paid to the network address,
 	// which represents the overall burden to node operators. A poster's costs, then, should not be eligible