Cache storage keys

arbos/addressSet/addressSet.go

@@ -26,9 +26,9 @@ func Initialize(sto *storage.Storage) error {
 
 func OpenAddressSet(sto *storage.Storage) *AddressSet {
 	return &AddressSet{
-		sto,
+		sto.NoCacheCopy(),
 		sto.OpenStorageBackedUint64(0),
-		sto.OpenSubStorage([]byte{0}),
+		sto.OpenSubStorage([]byte{0}, false),
 	}
 }
 

arbos/addressTable/addressTable.go

@@ -25,7 +25,7 @@ func Initialize(sto *storage.Storage) {
 
 func Open(sto *storage.Storage) *AddressTable {
 	numItems := sto.OpenStorageBackedUint64(0)
-	return &AddressTable{sto, sto.OpenSubStorage([]byte{}), numItems}
+	return &AddressTable{sto.NoCacheCopy(), sto.OpenSubStorage([]byte{}, false), numItems}
 }
 
 func (atab *AddressTable) Register(addr common.Address) (uint64, error) {

arbos/arbosState/arbosstate.go

@@ -71,13 +71,13 @@ func OpenArbosState(stateDB vm.StateDB, burner burn.Burner) (*ArbosState, error)
 		backingStorage.OpenStorageBackedUint64(uint64(upgradeVersionOffset)),
 		backingStorage.OpenStorageBackedUint64(uint64(upgradeTimestampOffset)),
 		backingStorage.OpenStorageBackedAddress(uint64(networkFeeAccountOffset)),
-		l1pricing.OpenL1PricingState(backingStorage.OpenSubStorage(l1PricingSubspace)),
-		l2pricing.OpenL2PricingState(backingStorage.OpenSubStorage(l2PricingSubspace)),
-		retryables.OpenRetryableState(backingStorage.OpenSubStorage(retryablesSubspace), stateDB),
-		addressTable.Open(backingStorage.OpenSubStorage(addressTableSubspace)),
-		addressSet.OpenAddressSet(backingStorage.OpenSubStorage(chainOwnerSubspace)),
-		merkleAccumulator.OpenMerkleAccumulator(backingStorage.OpenSubStorage(sendMerkleSubspace)),
-		blockhash.OpenBlockhashes(backingStorage.OpenSubStorage(blockhashesSubspace)),
+		l1pricing.OpenL1PricingState(backingStorage.OpenSubStorage(l1PricingSubspace, true)),
+		l2pricing.OpenL2PricingState(backingStorage.OpenSubStorage(l2PricingSubspace, true)),
+		retryables.OpenRetryableState(backingStorage.OpenSubStorage(retryablesSubspace, true), stateDB),
+		addressTable.Open(backingStorage.OpenSubStorage(addressTableSubspace, true)),
+		addressSet.OpenAddressSet(backingStorage.OpenSubStorage(chainOwnerSubspace, true)),
+		merkleAccumulator.OpenMerkleAccumulator(backingStorage.OpenSubStorage(sendMerkleSubspace, true)),
+		blockhash.OpenBlockhashes(backingStorage.OpenSubStorage(blockhashesSubspace, true)),
 		backingStorage.OpenStorageBackedBigInt(uint64(chainIdOffset)),
 		backingStorage.OpenStorageBackedBytes(chainConfigSubspace),
 		backingStorage.OpenStorageBackedUint64(uint64(genesisBlockNumOffset)),
@@ -220,14 +220,14 @@ func InitializeArbosState(stateDB vm.StateDB, burner burn.Burner, chainConfig *p
 	if desiredArbosVersion >= 2 {
 		initialRewardsRecipient = initialChainOwner
 	}
-	_ = l1pricing.InitializeL1PricingState(sto.OpenSubStorage(l1PricingSubspace), initialRewardsRecipient, initMessage.InitialL1BaseFee)
-	_ = l2pricing.InitializeL2PricingState(sto.OpenSubStorage(l2PricingSubspace))
-	_ = retryables.InitializeRetryableState(sto.OpenSubStorage(retryablesSubspace))
-	addressTable.Initialize(sto.OpenSubStorage(addressTableSubspace))
-	merkleAccumulator.InitializeMerkleAccumulator(sto.OpenSubStorage(sendMerkleSubspace))
-	blockhash.InitializeBlockhashes(sto.OpenSubStorage(blockhashesSubspace))
-
-	ownersStorage := sto.OpenSubStorage(chainOwnerSubspace)
+	_ = l1pricing.InitializeL1PricingState(sto.OpenSubStorage(l1PricingSubspace, true), initialRewardsRecipient, initMessage.InitialL1BaseFee)
+	_ = l2pricing.InitializeL2PricingState(sto.OpenSubStorage(l2PricingSubspace, true))
+	_ = retryables.InitializeRetryableState(sto.OpenSubStorage(retryablesSubspace, true))
+	addressTable.Initialize(sto.OpenSubStorage(addressTableSubspace, true))
+	merkleAccumulator.InitializeMerkleAccumulator(sto.OpenSubStorage(sendMerkleSubspace, true))
+	blockhash.InitializeBlockhashes(sto.OpenSubStorage(blockhashesSubspace, true))
+
+	ownersStorage := sto.OpenSubStorage(chainOwnerSubspace, true)
 	_ = addressSet.Initialize(ownersStorage)
 	_ = addressSet.OpenAddressSet(ownersStorage).Add(initialChainOwner)
 
@@ -401,7 +401,7 @@ func (state *ArbosState) ChainOwners() *addressSet.AddressSet {
 
 func (state *ArbosState) SendMerkleAccumulator() *merkleAccumulator.MerkleAccumulator {
 	if state.sendMerkle == nil {
-		state.sendMerkle = merkleAccumulator.OpenMerkleAccumulator(state.backingStorage.OpenSubStorage(sendMerkleSubspace))
+		state.sendMerkle = merkleAccumulator.OpenMerkleAccumulator(state.backingStorage.OpenSubStorage(sendMerkleSubspace, true))
 	}
 	return state.sendMerkle
 }

arbos/arbosState/arbosstate_test.go

@@ -64,7 +64,7 @@ func TestStorageBackedInt64(t *testing.T) {
 
 func TestStorageSlots(t *testing.T) {
 	state, _ := NewArbosMemoryBackedArbOSState()
-	sto := state.BackingStorage().OpenSubStorage([]byte{})
+	sto := state.BackingStorage().OpenSubStorage([]byte{}, true)
 
 	println("nil address", colors.Blue, storage.NilAddressRepresentation.String(), colors.Clear)
 

arbos/blockhash/blockhash.go

@@ -21,7 +21,7 @@ func InitializeBlockhashes(backingStorage *storage.Storage) {
 }
 
 func OpenBlockhashes(backingStorage *storage.Storage) *Blockhashes {
-	return &Blockhashes{backingStorage, backingStorage.OpenStorageBackedUint64(0)}
+	return &Blockhashes{backingStorage.NoCacheCopy(), backingStorage.OpenStorageBackedUint64(0)}
 }
 
 func (bh *Blockhashes) L1BlockNumber() (uint64, error) {

arbos/l1pricing/batchPoster.go

@@ -42,13 +42,13 @@ func InitializeBatchPostersTable(storage *storage.Storage) error {
 	if err := totalFundsDue.SetChecked(common.Big0); err != nil {
 		return err
 	}
-	return addressSet.Initialize(storage.OpenSubStorage(PosterAddrsKey))
+	return addressSet.Initialize(storage.OpenSubStorage(PosterAddrsKey, true))
 }
 
 func OpenBatchPostersTable(storage *storage.Storage) *BatchPostersTable {
 	return &BatchPostersTable{
-		posterAddrs:   addressSet.OpenAddressSet(storage.OpenSubStorage(PosterAddrsKey)),
-		posterInfo:    storage.OpenSubStorage(PosterInfoKey),
+		posterAddrs:   addressSet.OpenAddressSet(storage.OpenSubStorage(PosterAddrsKey, true)),
+		posterInfo:    storage.OpenSubStorage(PosterInfoKey, false),
 		totalFundsDue: storage.OpenStorageBackedBigInt(totalFundsDueOffset),
 	}
 }
@@ -68,7 +68,7 @@ func (bpt *BatchPostersTable) OpenPoster(poster common.Address, createIfNotExist
 }
 
 func (bpt *BatchPostersTable) internalOpen(poster common.Address) *BatchPosterState {
-	bpStorage := bpt.posterInfo.OpenSubStorage(poster.Bytes())
+	bpStorage := bpt.posterInfo.OpenSubStorage(poster.Bytes(), false)
 	return &BatchPosterState{
 		fundsDue:     bpStorage.OpenStorageBackedBigInt(0),
 		payTo:        bpStorage.OpenStorageBackedAddress(1),

arbos/l1pricing/l1pricing.go

@@ -83,7 +83,7 @@ var InitialEquilibrationUnitsV0 = arbmath.UintToBig(60 * params.TxDataNonZeroGas
 var InitialEquilibrationUnitsV6 = arbmath.UintToBig(params.TxDataNonZeroGasEIP2028 * 10000000)
 
 func InitializeL1PricingState(sto *storage.Storage, initialRewardsRecipient common.Address, initialL1BaseFee *big.Int) error {
-	bptStorage := sto.OpenSubStorage(BatchPosterTableKey)
+	bptStorage := sto.OpenSubStorage(BatchPosterTableKey, true)
 	if err := InitializeBatchPostersTable(bptStorage); err != nil {
 		return err
 	}
@@ -118,7 +118,7 @@ func InitializeL1PricingState(sto *storage.Storage, initialRewardsRecipient comm
 func OpenL1PricingState(sto *storage.Storage) *L1PricingState {
 	return &L1PricingState{
 		sto,
-		OpenBatchPostersTable(sto.OpenSubStorage(BatchPosterTableKey)),
+		OpenBatchPostersTable(sto.OpenSubStorage(BatchPosterTableKey, true)),
 		sto.OpenStorageBackedAddress(payRewardsToOffset),
 		sto.OpenStorageBackedBigUint(equilibrationUnitsOffset),
 		sto.OpenStorageBackedUint64(inertiaOffset),

arbos/queue_test.go

@@ -14,7 +14,7 @@ import (
 
 func TestQueue(t *testing.T) {
 	state, statedb := arbosState.NewArbosMemoryBackedArbOSState()
-	sto := state.BackingStorage().OpenSubStorage([]byte{})
+	sto := state.BackingStorage().OpenSubStorage([]byte{}, true)
 	Require(t, storage.InitializeQueue(sto))
 	q := storage.OpenQueue(sto)
 

arbos/retryables/retryable.go

@@ -31,13 +31,13 @@ var (
 )
 
 func InitializeRetryableState(sto *storage.Storage) error {
-	return storage.InitializeQueue(sto.OpenSubStorage(timeoutQueueKey))
+	return storage.InitializeQueue(sto.OpenSubStorage(timeoutQueueKey, true))
 }
 
 func OpenRetryableState(sto *storage.Storage, statedb vm.StateDB) *RetryableState {
 	return &RetryableState{
 		sto,
-		storage.OpenQueue(sto.OpenSubStorage(timeoutQueueKey)),
+		storage.OpenQueue(sto.OpenSubStorage(timeoutQueueKey, true)),
 	}
 }
 
@@ -73,7 +73,7 @@ func (rs *RetryableState) CreateRetryable(
 	beneficiary common.Address,
 	calldata []byte,
 ) (*Retryable, error) {
-	sto := rs.retryables.OpenSubStorage(id.Bytes())
+	sto := rs.retryables.OpenSubStorage(id.Bytes(), false)
 	ret := &Retryable{
 		id,
 		sto,
@@ -100,7 +100,7 @@ func (rs *RetryableState) CreateRetryable(
 }
 
 func (rs *RetryableState) OpenRetryable(id common.Hash, currentTimestamp uint64) (*Retryable, error) {
-	sto := rs.retryables.OpenSubStorage(id.Bytes())
+	sto := rs.retryables.OpenSubStorage(id.Bytes(), false)
 	timeoutStorage := sto.OpenStorageBackedUint64(timeoutOffset)
 	timeout, err := timeoutStorage.Get()
 	if timeout == 0 || timeout < currentTimestamp || err != nil {
@@ -134,7 +134,7 @@ func (rs *RetryableState) RetryableSizeBytes(id common.Hash, currentTime uint64)
 }
 
 func (rs *RetryableState) DeleteRetryable(id common.Hash, evm *vm.EVM, scenario util.TracingScenario) (bool, error) {
-	retStorage := rs.retryables.OpenSubStorage(id.Bytes())
+	retStorage := rs.retryables.OpenSubStorage(id.Bytes(), false)
 	timeout, err := retStorage.GetByUint64(timeoutOffset)
 	if timeout == (common.Hash{}) || err != nil {
 		return false, err
@@ -157,7 +157,7 @@ func (rs *RetryableState) DeleteRetryable(id common.Hash, evm *vm.EVM, scenario
 	_ = retStorage.ClearByUint64(beneficiaryOffset)
 	_ = retStorage.ClearByUint64(timeoutOffset)
 	_ = retStorage.ClearByUint64(timeoutWindowsLeftOffset)
-	err = retStorage.OpenSubStorage(calldataKey).ClearBytes()
+	err = retStorage.OpenSubStorage(calldataKey, false).ClearBytes()
 	return true, err
 }
 
@@ -291,7 +291,7 @@ func (rs *RetryableState) TryToReapOneRetryable(currentTimestamp uint64, evm *vm
 	if err != nil || id == nil {
 		return err
 	}
-	retryableStorage := rs.retryables.OpenSubStorage(id.Bytes())
+	retryableStorage := rs.retryables.OpenSubStorage(id.Bytes(), false)
 	timeoutStorage := retryableStorage.OpenStorageBackedUint64(timeoutOffset)
 	timeout, err := timeoutStorage.Get()
 	if err != nil {

arbos/storage/queue.go

@@ -25,7 +25,7 @@ func InitializeQueue(sto *Storage) error {
 
 func OpenQueue(sto *Storage) *Queue {
 	return &Queue{
-		sto,
+		sto.NoCacheCopy(),
 		sto.OpenStorageBackedUint64(0),
 		sto.OpenStorageBackedUint64(1),
 	}

arbos/storage/storage.go

@@ -4,8 +4,10 @@
 package storage
 
 import (
+	"bytes"
 	"fmt"
 	"math/big"
+	"sync/atomic"
 
 	"github.com/ethereum/go-ethereum/common"
 	"github.com/ethereum/go-ethereum/core/rawdb"
@@ -17,6 +19,7 @@ import (
 	"github.com/offchainlabs/nitro/arbos/burn"
 	"github.com/offchainlabs/nitro/arbos/util"
 	"github.com/offchainlabs/nitro/util/arbmath"
+	"github.com/offchainlabs/nitro/util/containers"
 )
 
 // Storage allows ArbOS to store data persistently in the Ethereum-compatible stateDB. This is represented in
@@ -43,12 +46,18 @@ type Storage struct {
 	db         vm.StateDB
 	storageKey []byte
 	burner     burn.Burner
+	hashCache  *containers.SafeLruCache[string, []byte]
 }
 
 const StorageReadCost = params.SloadGasEIP2200
 const StorageWriteCost = params.SstoreSetGasEIP2200
 const StorageWriteZeroCost = params.SstoreResetGasEIP2200
 
+const storageKeyCacheSize = 1024
+
+var storageHashCache = containers.NewSafeLruCache[string, []byte](storageKeyCacheSize)
+var cacheFullLogged atomic.Bool
+
 // NewGeth uses a Geth database to create an evm key-value store
 func NewGeth(statedb vm.StateDB, burner burn.Burner) *Storage {
 	account := common.HexToAddress("0xA4B05FFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFF")
@@ -58,6 +67,7 @@ func NewGeth(statedb vm.StateDB, burner burn.Burner) *Storage {
 		db:         statedb,
 		storageKey: []byte{},
 		burner:     burner,
+		hashCache:  storageHashCache,
 	}
 }
 
@@ -81,15 +91,13 @@ func NewMemoryBackedStateDB() vm.StateDB {
 // a page, to preserve contiguity within a page. This will reduce cost if/when Ethereum switches to storage
 // representations that reward contiguity.
 // Because page numbers are 248 bits, this gives us 124-bit security against collision attacks, which is good enough.
-func mapAddress(storageKey []byte, key common.Hash) common.Hash {
+func (store *Storage) mapAddress(key common.Hash) common.Hash {
 	keyBytes := key.Bytes()
 	boundary := common.HashLength - 1
-	return common.BytesToHash(
-		append(
-			crypto.Keccak256(storageKey, keyBytes[:boundary])[:boundary],
-			keyBytes[boundary],
-		),
-	)
+	mapped := make([]byte, 0, common.HashLength)
+	mapped = append(mapped, store.cachedKeccak(store.storageKey, keyBytes[:boundary])[:boundary]...)
+	mapped = append(mapped, keyBytes[boundary])
+	return common.BytesToHash(mapped)
 }
 
 func writeCost(value common.Hash) uint64 {
@@ -111,11 +119,11 @@ func (store *Storage) Get(key common.Hash) (common.Hash, error) {
 	if info := store.burner.TracingInfo(); info != nil {
 		info.RecordStorageGet(key)
 	}
-	return store.db.GetState(store.account, mapAddress(store.storageKey, key)), nil
+	return store.db.GetState(store.account, store.mapAddress(key)), nil
 }
 
 func (store *Storage) GetStorageSlot(key common.Hash) common.Hash {
-	return mapAddress(store.storageKey, key)
+	return store.mapAddress(key)
 }
 
 func (store *Storage) GetUint64(key common.Hash) (uint64, error) {
@@ -143,7 +151,7 @@ func (store *Storage) Set(key common.Hash, value common.Hash) error {
 	if info := store.burner.TracingInfo(); info != nil {
 		info.RecordStorageSet(key, value)
 	}
-	store.db.SetState(store.account, mapAddress(store.storageKey, key), value)
+	store.db.SetState(store.account, store.mapAddress(key), value)
 	return nil
 }
 
@@ -171,12 +179,27 @@ func (store *Storage) Swap(key common.Hash, newValue common.Hash) (common.Hash,
 	return oldValue, store.Set(key, newValue)
 }
 
-func (store *Storage) OpenSubStorage(id []byte) *Storage {
+func (store *Storage) OpenSubStorage(id []byte, cacheKeys bool) *Storage {
+	var hashCache *containers.SafeLruCache[string, []byte]
+	if cacheKeys {
+		hashCache = storageHashCache
+	}
 	return &Storage{
 		store.account,
 		store.db,
-		crypto.Keccak256(store.storageKey, id),
+		store.cachedKeccak(store.storageKey, id),
 		store.burner,
+		hashCache,
+	}
+}
+
+func (store *Storage) NoCacheCopy() *Storage {
+	return &Storage{
+		store.account,
+		store.db,
+		store.storageKey,
+		store.burner,
+		nil,
 	}
 }
 
@@ -266,6 +289,24 @@ func (store *Storage) Keccak(data ...[]byte) ([]byte, error) {
 	return crypto.Keccak256(data...), nil
 }
 
+// note: returned slice is not thread-safe
+func (store *Storage) cachedKeccak(data ...[]byte) []byte {
+	if store.hashCache == nil {
+		return crypto.Keccak256(data...)
+	}
+	keyString := string(bytes.Join(data, []byte{}))
+	hash, wasCached := store.hashCache.Get(keyString)
+	if wasCached {
+		return hash
+	}
+	hash = crypto.Keccak256(data...)
+	evicted := store.hashCache.Add(keyString, hash)
+	if evicted && cacheFullLogged.CompareAndSwap(false, true) {
+		log.Warn("Hash cache full, we didn't expect that. Some non-static storage keys may fill up the cache.")
+	}
+	return hash
+}
+
 func (store *Storage) KeccakHash(data ...[]byte) (common.Hash, error) {
 	bytes, err := store.Keccak(data...)
 	return common.BytesToHash(bytes), err
@@ -279,7 +320,7 @@ type StorageSlot struct {
 }
 
 func (store *Storage) NewSlot(offset uint64) StorageSlot {
-	return StorageSlot{store.account, store.db, mapAddress(store.storageKey, util.UintToHash(offset)), store.burner}
+	return StorageSlot{store.account, store.db, store.mapAddress(util.UintToHash(offset)), store.burner}
 }
 
 func (ss *StorageSlot) Get() (common.Hash, error) {
@@ -590,7 +631,7 @@ type StorageBackedBytes struct {
 
 func (store *Storage) OpenStorageBackedBytes(id []byte) StorageBackedBytes {
 	return StorageBackedBytes{
-		*store.OpenSubStorage(id),
+		*store.OpenSubStorage(id, false),
 	}
 }