Separate Wasmstore from normal db: initial

arbitrum/recordingdb.go

@@ -55,12 +55,6 @@ func (db *RecordingKV) Get(key []byte) ([]byte, error) {
 		// Retrieving code
 		copy(hash[:], key[len(rawdb.CodePrefix):])
 		res, err = db.diskDb.Get(key)
-	} else if ok, _ := rawdb.IsActivatedAsmKey(key); ok {
-		// Arbitrum: the asm is non-consensus
-		return db.diskDb.Get(key)
-	} else if ok, _ := rawdb.IsActivatedModuleKey(key); ok {
-		// Arbitrum: the module is non-consensus (only its hash is)
-		return db.diskDb.Get(key)
 	} else {
 		err = fmt.Errorf("recording KV attempted to access non-hash key %v", hex.EncodeToString(key))
 	}
@@ -275,7 +269,7 @@ func (r *RecordingDatabase) PrepareRecording(ctx context.Context, lastBlockHeade
 	defer func() { r.Dereference(finalDereference) }()
 	recordingKeyValue := newRecordingKV(r.db.TrieDB(), r.db.DiskDB())
 
-	recordingStateDatabase := state.NewDatabase(rawdb.NewDatabase(recordingKeyValue))
+	recordingStateDatabase := state.NewDatabase(rawdb.WrapDatabaseWithWasm(rawdb.NewDatabase(recordingKeyValue), r.db.WasmStore()))
 	var prevRoot common.Hash
 	if lastBlockHeader != nil {
 		prevRoot = lastBlockHeader.Root

core/rawdb/database.go

@@ -42,6 +42,11 @@ type freezerdb struct {
 	ethdb.AncientStore
 }
 
+// AncientDatadir returns the path of root ancient directory.
+func (frdb *freezerdb) WasmDataBase() ethdb.KeyValueStore {
+	return frdb
+}
+
 // AncientDatadir returns the path of root ancient directory.
 func (frdb *freezerdb) AncientDatadir() (string, error) {
 	return frdb.ancientRoot, nil
@@ -165,12 +170,39 @@ func (db *nofreezedb) AncientDatadir() (string, error) {
 	return "", errNotSupported
 }
 
+// AncientDatadir returns the path of root ancient directory.
+func (db *nofreezedb) WasmDataBase() ethdb.KeyValueStore {
+	return db
+}
+
 // NewDatabase creates a high level database on top of a given key-value data
 // store without a freezer moving immutable chain segments into cold storage.
 func NewDatabase(db ethdb.KeyValueStore) ethdb.Database {
 	return &nofreezedb{KeyValueStore: db}
 }
 
+type dbWithWasmEntry struct {
+	ethdb.Database
+	wasmDb ethdb.KeyValueStore
+}
+
+func (db *dbWithWasmEntry) WasmDataBase() ethdb.KeyValueStore {
+	return db.wasmDb
+}
+
+func (db *dbWithWasmEntry) Close() error {
+	dbErr := db.Database.Close()
+	wasmErr := db.wasmDb.Close()
+	if dbErr != nil {
+		return dbErr
+	}
+	return wasmErr
+}
+
+func WrapDatabaseWithWasm(db ethdb.Database, wasm ethdb.KeyValueStore) ethdb.Database {
+	return &dbWithWasmEntry{db, wasm}
+}
+
 // resolveChainFreezerDir is a helper function which resolves the absolute path
 // of chain freezer by considering backward compatibility.
 func resolveChainFreezerDir(ancient string) string {

core/rawdb/table.go

@@ -40,6 +40,10 @@ func (t *table) Close() error {
 	return nil
 }
 
+func (t *table) WasmDataBase() ethdb.KeyValueStore {
+	return t.db.WasmDataBase()
+}
+
 // Has retrieves if a prefixed version of a key is present in the database.
 func (t *table) Has(key []byte) (bool, error) {
 	return t.db.Has(append([]byte(t.prefix), key...))

core/state/database.go

@@ -54,6 +54,7 @@ type Database interface {
 	// Arbitrum: Read activated Stylus contracts
 	ActivatedAsm(moduleHash common.Hash) (asm []byte, err error)
 	ActivatedModule(moduleHash common.Hash) (module []byte, err error)
+	WasmStore() ethdb.KeyValueStore
 
 	// OpenTrie opens the main account trie.
 	OpenTrie(root common.Hash) (Trie, error)
@@ -164,6 +165,7 @@ func NewDatabaseWithConfig(db ethdb.Database, config *trie.Config) Database {
 		activatedModuleCache: lru.NewSizeConstrainedCache[common.Hash, []byte](activatedWasmCacheSize),
 
 		disk:          db,
+		wasmdb:        db.WasmDataBase(),
 		codeSizeCache: lru.NewCache[common.Hash, int](codeSizeCacheSize),
 		codeCache:     lru.NewSizeConstrainedCache[common.Hash, []byte](codeCacheSize),
 		triedb:        trie.NewDatabase(db, config),
@@ -179,6 +181,7 @@ func NewDatabaseWithNodeDB(db ethdb.Database, triedb *trie.Database) Database {
 		activatedModuleCache: lru.NewSizeConstrainedCache[common.Hash, []byte](activatedWasmCacheSize),
 
 		disk:          db,
+		wasmdb:        db.WasmDataBase(),
 		codeSizeCache: lru.NewCache[common.Hash, int](codeSizeCacheSize),
 		codeCache:     lru.NewSizeConstrainedCache[common.Hash, []byte](codeCacheSize),
 		triedb:        triedb,
@@ -192,11 +195,16 @@ type cachingDB struct {
 	activatedModuleCache *lru.SizeConstrainedCache[common.Hash, []byte]
 
 	disk          ethdb.KeyValueStore
+	wasmdb        ethdb.KeyValueStore
 	codeSizeCache *lru.Cache[common.Hash, int]
 	codeCache     *lru.SizeConstrainedCache[common.Hash, []byte]
 	triedb        *trie.Database
 }
 
+func (db *cachingDB) WasmStore() ethdb.KeyValueStore {
+	return db.wasmdb
+}
+
 // OpenTrie opens the main account trie at a specific root hash.
 func (db *cachingDB) OpenTrie(root common.Hash) (Trie, error) {
 	if db.triedb.IsVerkle() {

core/state/database_arbitrum.go

@@ -12,7 +12,7 @@ func (db *cachingDB) ActivatedAsm(moduleHash common.Hash) ([]byte, error) {
 		return asm, nil
 	}
 	wasmKey := rawdb.ActivatedAsmKey(moduleHash)
-	asm, err := db.disk.Get(wasmKey[:])
+	asm, err := db.wasmdb.Get(wasmKey[:])
 	if err != nil {
 		return nil, err
 	}
@@ -28,7 +28,7 @@ func (db *cachingDB) ActivatedModule(moduleHash common.Hash) ([]byte, error) {
 		return module, nil
 	}
 	wasmKey := rawdb.ActivatedModuleKey(moduleHash)
-	module, err := db.disk.Get(wasmKey[:])
+	module, err := db.wasmdb.Get(wasmKey[:])
 	if err != nil {
 		return nil, err
 	}

core/state/journal_arbitrum.go

@@ -41,8 +41,11 @@ type EvictWasm struct {
 }
 
 func (ch EvictWasm) revert(s *StateDB) {
-	asm := s.GetActivatedAsm(ch.ModuleHash) // only happens in native mode
-	CacheWasmRust(asm, ch.ModuleHash, ch.Version, ch.Debug)
+	asm, err := s.TryGetActivatedAsm(ch.ModuleHash) // only happens in native mode
+	if err == nil && len(asm) != 0 {
+		//if we failed to get it - it's not in the current rust cache
+		CacheWasmRust(asm, ch.ModuleHash, ch.Version, ch.Debug)
+	}
 }
 
 func (ch EvictWasm) dirtied() *common.Address {

core/state/statedb.go

@@ -1249,6 +1249,7 @@ func (s *StateDB) Commit(block uint64, deleteEmptyObjects bool) (common.Hash, er
 		storageTrieNodesDeleted int
 		nodes                   = trienode.NewMergedNodeSet()
 		codeWriter              = s.db.DiskDB().NewBatch()
+		wasmCodeWriter          = s.db.WasmStore().NewBatch()
 	)
 	// Handle all state deletions first
 	incomplete, err := s.handleDestruction(nodes)
@@ -1286,7 +1287,7 @@ func (s *StateDB) Commit(block uint64, deleteEmptyObjects bool) (common.Hash, er
 
 	// Arbitrum: write Stylus programs to disk
 	for moduleHash, info := range s.arbExtraData.activatedWasms {
-		rawdb.WriteActivation(codeWriter, moduleHash, info.Asm, info.Module)
+		rawdb.WriteActivation(wasmCodeWriter, moduleHash, info.Asm, info.Module)
 	}
 	if len(s.arbExtraData.activatedWasms) > 0 {
 		s.arbExtraData.activatedWasms = make(map[common.Hash]*ActivatedWasm)
@@ -1297,6 +1298,11 @@ func (s *StateDB) Commit(block uint64, deleteEmptyObjects bool) (common.Hash, er
 			log.Crit("Failed to commit dirty codes", "error", err)
 		}
 	}
+	if wasmCodeWriter.ValueSize() > 0 {
+		if err := wasmCodeWriter.Write(); err != nil {
+			log.Crit("Failed to commit dirty stylus codes", "error", err)
+		}
+	}
 	// Write the account trie changes, measuring the amount of wasted time
 	var start time.Time
 	if metrics.EnabledExpensive {

core/state/statedb_arbitrum.go

@@ -28,6 +28,7 @@ import (
 	"github.com/ethereum/go-ethereum/common"
 	"github.com/ethereum/go-ethereum/common/lru"
 	"github.com/ethereum/go-ethereum/core/types"
+	"github.com/ethereum/go-ethereum/log"
 	"github.com/ethereum/go-ethereum/rlp"
 	"github.com/ethereum/go-ethereum/trie"
 )
@@ -89,16 +90,12 @@ func (s *StateDB) ActivateWasm(moduleHash common.Hash, asm, module []byte) {
 	})
 }
 
-func (s *StateDB) GetActivatedAsm(moduleHash common.Hash) []byte {
+func (s *StateDB) TryGetActivatedAsm(moduleHash common.Hash) ([]byte, error) {
 	info, exists := s.arbExtraData.activatedWasms[moduleHash]
 	if exists {
-		return info.Asm
+		return info.Asm, nil
 	}
-	asm, err := s.db.ActivatedAsm(moduleHash)
-	if err != nil {
-		s.setError(fmt.Errorf("failed to load asm for %x: %v", moduleHash, err))
-	}
-	return asm
+	return s.db.ActivatedAsm(moduleHash)
 }
 
 func (s *StateDB) GetActivatedModule(moduleHash common.Hash) []byte {
@@ -237,9 +234,13 @@ func (s *StateDB) StartRecording() {
 }
 
 func (s *StateDB) RecordProgram(moduleHash common.Hash) {
+	asm, err := s.TryGetActivatedAsm(moduleHash)
+	if err != nil {
+		log.Crit("can't find activated wasm while recording", "modulehash", moduleHash)
+	}
 	if s.arbExtraData.userWasms != nil {
 		s.arbExtraData.userWasms[moduleHash] = ActivatedWasm{
-			Asm:    s.GetActivatedAsm(moduleHash),
+			Asm:    asm,
 			Module: s.GetActivatedModule(moduleHash),
 		}
 	}

core/vm/interface.go

@@ -30,7 +30,7 @@ import (
 type StateDB interface {
 	// Arbitrum: manage Stylus wasms
 	ActivateWasm(moduleHash common.Hash, asm, module []byte)
-	GetActivatedAsm(moduleHash common.Hash) (asm []byte)
+	TryGetActivatedAsm(moduleHash common.Hash) (asm []byte, err error)
 	GetActivatedModule(moduleHash common.Hash) (module []byte)
 	RecordCacheWasm(wasm state.CacheWasm)
 	RecordEvictWasm(wasm state.EvictWasm)

ethdb/database.go

@@ -178,6 +178,10 @@ type AncientStore interface {
 	io.Closer
 }
 
+type WasmDataBaseRetriever interface {
+	WasmDataBase() KeyValueStore
+}
+
 // Database contains all the methods required by the high level database to not
 // only access the key-value data store but also the chain freezer.
 type Database interface {
@@ -189,4 +193,5 @@ type Database interface {
 	Compacter
 	Snapshotter
 	io.Closer
+	WasmDataBaseRetriever
 }

ethdb/remotedb/remotedb.go

@@ -39,6 +39,10 @@ func (db *Database) Has(key []byte) (bool, error) {
 	return true, nil
 }
 
+func (t *Database) WasmDataBase() ethdb.KeyValueStore {
+	return t
+}
+
 func (db *Database) Get(key []byte) ([]byte, error) {
 	var resp hexutil.Bytes
 	err := db.remote.Call(&resp, "debug_dbGet", hexutil.Bytes(key))

internal/ethapi/api.go

@@ -968,6 +968,9 @@ func (diff *StateOverride) Apply(state *state.StateDB) error {
 		return nil
 	}
 	for addr, account := range *diff {
+		if addr == types.ArbosStateAddress {
+			return fmt.Errorf("overriding address %v  not allowed", types.ArbosStateAddress)
+		}
 		// Override account nonce.
 		if account.Nonce != nil {
 			state.SetNonce(addr, uint64(*account.Nonce))

node/node.go

@@ -799,7 +799,6 @@ func (n *Node) OpenDatabaseWithFreezer(name string, cache, handles int, ancient
 			ReadOnly:          readonly,
 		})
 	}
-
 	if err == nil {
 		db = n.wrapDatabase(db)
 	}