Merge pull request #319 from OffchainLabs/stylus_cache_tag

Stylus cache tag
OffchainLabs · May 21, 2024 · 3ad1488 · 3ad1488
2 parents 247b930 + b8d4ced
commit 3ad1488
Show file tree

Hide file tree

Showing 8 changed files with 43 additions and 21 deletions.
diff --git a/arbitrum/apibackend.go b/arbitrum/apibackend.go
@@ -45,6 +45,8 @@ var (
 type APIBackend struct {
 	b *Backend
 
+	dbForAPICalls ethdb.Database
+
 	fallbackClient types.FallbackClient
 	sync           SyncProgressBackend
 }
@@ -101,8 +103,15 @@ func createRegisterAPIBackend(backend *Backend, filterConfig filters.Config, fal
 	if err != nil {
 		return nil, err
 	}
+	// discard stylus-tag on any call made from api database
+	dbForAPICalls := backend.chainDb
+	wasmStore, tag := backend.chainDb.WasmDataBase()
+	if tag != 0 {
+		dbForAPICalls = rawdb.WrapDatabaseWithWasm(backend.chainDb, wasmStore, 0)
+	}
 	backend.apiBackend = &APIBackend{
 		b:              backend,
+		dbForAPICalls:  dbForAPICalls,
 		fallbackClient: fallbackClient,
 	}
 	filterSystem := filters.NewFilterSystem(backend.apiBackend, filterConfig)
@@ -314,7 +323,7 @@ func (a *APIBackend) FeeHistory(
 }
 
 func (a *APIBackend) ChainDb() ethdb.Database {
-	return a.b.chainDb
+	return a.dbForAPICalls
 }
 
 func (a *APIBackend) AccountManager() *accounts.Manager {

diff --git a/arbitrum/recordingdb.go b/arbitrum/recordingdb.go
@@ -269,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.WrapDatabaseWithWasm(rawdb.NewDatabase(recordingKeyValue), r.db.WasmStore()))
+	recordingStateDatabase := state.NewDatabase(rawdb.WrapDatabaseWithWasm(rawdb.NewDatabase(recordingKeyValue), r.db.WasmStore(), 0))
 	var prevRoot common.Hash
 	if lastBlockHeader != nil {
 		prevRoot = lastBlockHeader.Root

diff --git a/core/rawdb/database.go b/core/rawdb/database.go
@@ -43,8 +43,8 @@ type freezerdb struct {
 }
 
 // AncientDatadir returns the path of root ancient directory.
-func (frdb *freezerdb) WasmDataBase() ethdb.KeyValueStore {
-	return frdb
+func (frdb *freezerdb) WasmDataBase() (ethdb.KeyValueStore, uint32) {
+	return frdb, 0
 }
 
 // AncientDatadir returns the path of root ancient directory.
@@ -171,8 +171,8 @@ func (db *nofreezedb) AncientDatadir() (string, error) {
 }
 
 // AncientDatadir returns the path of root ancient directory.
-func (db *nofreezedb) WasmDataBase() ethdb.KeyValueStore {
-	return db
+func (db *nofreezedb) WasmDataBase() (ethdb.KeyValueStore, uint32) {
+	return db, 0
 }
 
 // NewDatabase creates a high level database on top of a given key-value data
@@ -183,11 +183,12 @@ func NewDatabase(db ethdb.KeyValueStore) ethdb.Database {
 
 type dbWithWasmEntry struct {
 	ethdb.Database
-	wasmDb ethdb.KeyValueStore
+	wasmDb       ethdb.KeyValueStore
+	wasmCacheTag uint32
 }
 
-func (db *dbWithWasmEntry) WasmDataBase() ethdb.KeyValueStore {
-	return db.wasmDb
+func (db *dbWithWasmEntry) WasmDataBase() (ethdb.KeyValueStore, uint32) {
+	return db.wasmDb, db.wasmCacheTag
 }
 
 func (db *dbWithWasmEntry) Close() error {
@@ -199,8 +200,8 @@ func (db *dbWithWasmEntry) Close() error {
 	return wasmErr
 }
 
-func WrapDatabaseWithWasm(db ethdb.Database, wasm ethdb.KeyValueStore) ethdb.Database {
-	return &dbWithWasmEntry{db, wasm}
+func WrapDatabaseWithWasm(db ethdb.Database, wasm ethdb.KeyValueStore, cacheTag uint32) ethdb.Database {
+	return &dbWithWasmEntry{db, wasm, cacheTag}
 }
 
 // resolveChainFreezerDir is a helper function which resolves the absolute path

diff --git a/core/rawdb/table.go b/core/rawdb/table.go
@@ -40,7 +40,7 @@ func (t *table) Close() error {
 	return nil
 }
 
-func (t *table) WasmDataBase() ethdb.KeyValueStore {
+func (t *table) WasmDataBase() (ethdb.KeyValueStore, uint32) {
 	return t.db.WasmDataBase()
 }
 

diff --git a/core/state/database.go b/core/state/database.go
@@ -55,6 +55,7 @@ type Database interface {
 	ActivatedAsm(moduleHash common.Hash) (asm []byte, err error)
 	ActivatedModule(moduleHash common.Hash) (module []byte, err error)
 	WasmStore() ethdb.KeyValueStore
+	WasmCacheTag() uint32
 
 	// OpenTrie opens the main account trie.
 	OpenTrie(root common.Hash) (Trie, error)
@@ -159,13 +160,15 @@ func NewDatabase(db ethdb.Database) Database {
 // is safe for concurrent use and retains a lot of collapsed RLP trie nodes in a
 // large memory cache.
 func NewDatabaseWithConfig(db ethdb.Database, config *trie.Config) Database {
+	wasmdb, wasmTag := db.WasmDataBase()
 	cdb := &cachingDB{
 		// Arbitrum only
 		activatedAsmCache:    lru.NewSizeConstrainedCache[common.Hash, []byte](activatedWasmCacheSize),
 		activatedModuleCache: lru.NewSizeConstrainedCache[common.Hash, []byte](activatedWasmCacheSize),
+		wasmTag:              wasmTag,
 
 		disk:          db,
-		wasmdb:        db.WasmDataBase(),
+		wasmdb:        wasmdb,
 		codeSizeCache: lru.NewCache[common.Hash, int](codeSizeCacheSize),
 		codeCache:     lru.NewSizeConstrainedCache[common.Hash, []byte](codeCacheSize),
 		triedb:        trie.NewDatabase(db, config),
@@ -175,13 +178,15 @@ func NewDatabaseWithConfig(db ethdb.Database, config *trie.Config) Database {
 
 // NewDatabaseWithNodeDB creates a state database with an already initialized node database.
 func NewDatabaseWithNodeDB(db ethdb.Database, triedb *trie.Database) Database {
+	wasmdb, wasmTag := db.WasmDataBase()
 	cdb := &cachingDB{
 		// Arbitrum only
 		activatedAsmCache:    lru.NewSizeConstrainedCache[common.Hash, []byte](activatedWasmCacheSize),
 		activatedModuleCache: lru.NewSizeConstrainedCache[common.Hash, []byte](activatedWasmCacheSize),
+		wasmTag:              wasmTag,
 
 		disk:          db,
-		wasmdb:        db.WasmDataBase(),
+		wasmdb:        wasmdb,
 		codeSizeCache: lru.NewCache[common.Hash, int](codeSizeCacheSize),
 		codeCache:     lru.NewSizeConstrainedCache[common.Hash, []byte](codeCacheSize),
 		triedb:        triedb,
@@ -193,6 +198,7 @@ type cachingDB struct {
 	// Arbitrum
 	activatedAsmCache    *lru.SizeConstrainedCache[common.Hash, []byte]
 	activatedModuleCache *lru.SizeConstrainedCache[common.Hash, []byte]
+	wasmTag              uint32
 
 	disk          ethdb.KeyValueStore
 	wasmdb        ethdb.KeyValueStore
@@ -205,6 +211,10 @@ func (db *cachingDB) WasmStore() ethdb.KeyValueStore {
 	return db.wasmdb
 }
 
+func (db *cachingDB) WasmCacheTag() uint32 {
+	return db.wasmTag
+}
+
 // OpenTrie opens the main account trie at a specific root hash.
 func (db *cachingDB) OpenTrie(root common.Hash) (Trie, error) {
 	if db.triedb.IsVerkle() {

diff --git a/core/state/journal_arbitrum.go b/core/state/journal_arbitrum.go
@@ -17,17 +17,18 @@ func (ch wasmActivation) dirtied() *common.Address {
 }
 
 // Updates the Rust-side recent program cache
-var CacheWasmRust func(asm []byte, moduleHash common.Hash, version uint16, debug bool) = func([]byte, common.Hash, uint16, bool) {}
-var EvictWasmRust func(moduleHash common.Hash, version uint16, debug bool) = func(common.Hash, uint16, bool) {}
+var CacheWasmRust func(asm []byte, moduleHash common.Hash, version uint16, tag uint32, debug bool) = func([]byte, common.Hash, uint16, uint32, bool) {}
+var EvictWasmRust func(moduleHash common.Hash, version uint16, tag uint32, debug bool) = func(common.Hash, uint16, uint32, bool) {}
 
 type CacheWasm struct {
 	ModuleHash common.Hash
 	Version    uint16
+	Tag        uint32
 	Debug      bool
 }
 
 func (ch CacheWasm) revert(s *StateDB) {
-	EvictWasmRust(ch.ModuleHash, ch.Version, ch.Debug)
+	EvictWasmRust(ch.ModuleHash, ch.Version, ch.Tag, ch.Debug)
 }
 
 func (ch CacheWasm) dirtied() *common.Address {
@@ -37,14 +38,15 @@ func (ch CacheWasm) dirtied() *common.Address {
 type EvictWasm struct {
 	ModuleHash common.Hash
 	Version    uint16
+	Tag        uint32
 	Debug      bool
 }
 
 func (ch EvictWasm) revert(s *StateDB) {
 	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)
+		CacheWasmRust(asm, ch.ModuleHash, ch.Version, ch.Tag, ch.Debug)
 	}
 }
 

diff --git a/ethdb/database.go b/ethdb/database.go
@@ -179,7 +179,7 @@ type AncientStore interface {
 }
 
 type WasmDataBaseRetriever interface {
-	WasmDataBase() KeyValueStore
+	WasmDataBase() (KeyValueStore, uint32)
 }
 
 // Database contains all the methods required by the high level database to not

diff --git a/ethdb/remotedb/remotedb.go b/ethdb/remotedb/remotedb.go
@@ -39,8 +39,8 @@ func (db *Database) Has(key []byte) (bool, error) {
 	return true, nil
 }
 
-func (t *Database) WasmDataBase() ethdb.KeyValueStore {
-	return t
+func (t *Database) WasmDataBase() (ethdb.KeyValueStore, uint32) {
+	return t, 0
 }
 
 func (db *Database) Get(key []byte) ([]byte, error) {