core, trie, eth, cmd: rework preimage store (#205)

* core, trie, eth, cmd: rework preimage store

* fix ci

* fix ci

* Update zk_trie_database.go

---------

Co-authored-by: rjl493456442 <garyrong0905@gmail.com>
Co-authored-by: maskpp <maskpp266@gmail.com>

Author: 3 people

cmd/evm/internal/t8ntool/execution.go

@@ -263,7 +263,7 @@ func (pre *Prestate) Apply(vmConfig vm.Config, chainConfig *params.ChainConfig,
 }
 
 func MakePreState(db ethdb.Database, accounts core.GenesisAlloc) *state.StateDB {
-	sdb := state.NewDatabase(db)
+	sdb := state.NewDatabaseWithConfig(db, &trie.Config{Preimages: true})
 	statedb, _ := state.New(common.Hash{}, sdb, nil)
 	for addr, a := range accounts {
 		statedb.SetCode(addr, a.Code)

core/state/state_test.go

@@ -24,6 +24,7 @@ import (
 	"github.com/scroll-tech/go-ethereum/common"
 	"github.com/scroll-tech/go-ethereum/core/rawdb"
 	"github.com/scroll-tech/go-ethereum/ethdb"
+	"github.com/scroll-tech/go-ethereum/trie"
 )
 
 type stateTest struct {
@@ -39,7 +40,7 @@ func newStateTest() *stateTest {
 
 func TestDump(t *testing.T) {
 	db := rawdb.NewMemoryDatabase()
-	sdb, _ := New(common.Hash{}, NewDatabaseWithConfig(db, nil), nil)
+	sdb, _ := New(common.Hash{}, NewDatabaseWithConfig(db, &trie.Config{Preimages: true}), nil)
 	s := &stateTest{db: db, state: sdb}
 
 	// generate a few entries

eth/api_test.go

@@ -30,6 +30,7 @@ import (
 	"github.com/scroll-tech/go-ethereum/core/rawdb"
 	"github.com/scroll-tech/go-ethereum/core/state"
 	"github.com/scroll-tech/go-ethereum/crypto"
+	"github.com/scroll-tech/go-ethereum/trie"
 )
 
 var dumper = spew.ConfigState{Indent: "    "}
@@ -67,7 +68,7 @@ func TestAccountRange(t *testing.T) {
 	t.Parallel()
 
 	var (
-		statedb  = state.NewDatabaseWithConfig(rawdb.NewMemoryDatabase(), nil)
+		statedb  = state.NewDatabaseWithConfig(rawdb.NewMemoryDatabase(), &trie.Config{Preimages: true})
 		state, _ = state.New(common.Hash{}, statedb, nil)
 		addrs    = [AccountRangeMaxResults * 2]common.Address{}
 		m        = map[common.Address]bool{}

trie/database.go

@@ -80,8 +80,6 @@ type Database struct {
 	oldest  common.Hash                 // Oldest tracked node, flush-list head
 	newest  common.Hash                 // Newest tracked node, flush-list tail
 
-	preimages map[common.Hash][]byte // Preimages of nodes from the secure trie
-
 	gctime  time.Duration      // Time spent on garbage collection since last commit
 	gcnodes uint64             // Nodes garbage collected since last commit
 	gcsize  common.StorageSize // Data storage garbage collected since last commit
@@ -90,9 +88,9 @@ type Database struct {
 	flushnodes uint64             // Nodes flushed since last commit
 	flushsize  common.StorageSize // Data storage flushed since last commit
 
-	dirtiesSize   common.StorageSize // Storage size of the dirty node cache (exc. metadata)
-	childrenSize  common.StorageSize // Storage size of the external children tracking
-	preimagesSize common.StorageSize // Storage size of the preimages cache
+	dirtiesSize  common.StorageSize // Storage size of the dirty node cache (exc. metadata)
+	childrenSize common.StorageSize // Storage size of the external children tracking
+	preimages    *preimageStore     // The store for caching preimages
 
 	lock sync.RWMutex
 }
@@ -305,16 +303,18 @@ func NewDatabaseWithConfig(diskdb ethdb.KeyValueStore, config *Config) *Database
 			cleans = fastcache.LoadFromFileOrNew(config.Journal, config.Cache*1024*1024)
 		}
 	}
+	var preimage *preimageStore
+	if config != nil && config.Preimages {
+		preimage = newPreimageStore(diskdb)
+	}
 	db := &Database{
 		diskdb: diskdb,
 		cleans: cleans,
 		dirties: map[common.Hash]*cachedNode{{}: {
 			children: make(map[common.Hash]uint16),
 		}},
 		rawDirties: make(KvMap),
-	}
-	if config == nil || config.Preimages { // TODO(karalabe): Flip to default off in the future
-		db.preimages = make(map[common.Hash][]byte)
+		preimages:  preimage,
 	}
 	return db
 }
@@ -357,24 +357,6 @@ func (db *Database) insert(hash common.Hash, size int, node node) {
 	db.dirtiesSize += common.StorageSize(common.HashLength + entry.size)
 }
 
-// insertPreimage writes a new trie node pre-image to the memory database if it's
-// yet unknown. The method will NOT make a copy of the slice,
-// only use if the preimage will NOT be changed later on.
-//
-// Note, this method assumes that the database's lock is held!
-func (db *Database) insertPreimage(hash common.Hash, preimage []byte) {
-	// Short circuit if preimage collection is disabled
-	if db.preimages == nil {
-		return
-	}
-	// Track the preimage if a yet unknown one
-	if _, ok := db.preimages[hash]; ok {
-		return
-	}
-	db.preimages[hash] = preimage
-	db.preimagesSize += common.StorageSize(common.HashLength + len(preimage))
-}
-
 // node retrieves a cached trie node from memory, or returns nil if none can be
 // found in the memory cache.
 func (db *Database) node(hash common.Hash) node {
@@ -451,24 +433,6 @@ func (db *Database) Node(hash common.Hash) ([]byte, error) {
 	return nil, errors.New("not found")
 }
 
-// preimage retrieves a cached trie node pre-image from memory. If it cannot be
-// found cached, the method queries the persistent database for the content.
-func (db *Database) preimage(hash common.Hash) []byte {
-	// Short circuit if preimage collection is disabled
-	if db.preimages == nil {
-		return nil
-	}
-	// Retrieve the node from cache if available
-	db.lock.RLock()
-	preimage := db.preimages[hash]
-	db.lock.RUnlock()
-
-	if preimage != nil {
-		return preimage
-	}
-	return rawdb.ReadPreimage(db.diskdb, hash)
-}
-
 // Nodes retrieves the hashes of all the nodes cached within the memory database.
 // This method is extremely expensive and should only be used to validate internal
 // states in test code.
@@ -613,19 +577,8 @@ func (db *Database) Cap(limit common.StorageSize) error {
 
 	// If the preimage cache got large enough, push to disk. If it's still small
 	// leave for later to deduplicate writes.
-	flushPreimages := db.preimagesSize > 4*1024*1024
-	if flushPreimages {
-		if db.preimages == nil {
-			log.Error("Attempted to write preimages whilst disabled")
-		} else {
-			rawdb.WritePreimages(batch, db.preimages)
-			if batch.ValueSize() > ethdb.IdealBatchSize {
-				if err := batch.Write(); err != nil {
-					return err
-				}
-				batch.Reset()
-			}
-		}
+	if db.preimages != nil {
+		db.preimages.commit(false)
 	}
 	// Keep committing nodes from the flush-list until we're below allowance
 	oldest := db.oldest
@@ -660,13 +613,6 @@ func (db *Database) Cap(limit common.StorageSize) error {
 	db.lock.Lock()
 	defer db.lock.Unlock()
 
-	if flushPreimages {
-		if db.preimages == nil {
-			log.Error("Attempted to reset preimage cache whilst disabled")
-		} else {
-			db.preimages, db.preimagesSize = make(map[common.Hash][]byte), 0
-		}
-	}
 	for db.oldest != oldest {
 		node := db.dirties[db.oldest]
 		delete(db.dirties, db.oldest)
@@ -727,13 +673,7 @@ func (db *Database) Commit(node common.Hash, report bool, callback func(common.H
 
 	// Move all of the accumulated preimages into a write batch
 	if db.preimages != nil {
-		rawdb.WritePreimages(batch, db.preimages)
-		// Since we're going to replay trie node writes into the clean cache, flush out
-		// any batched pre-images before continuing.
-		if err := batch.Write(); err != nil {
-			return err
-		}
-		batch.Reset()
+		db.preimages.commit(true)
 	}
 	// Move the trie itself into the batch, flushing if enough data is accumulated
 	nodes, storage := len(db.dirties), db.dirtiesSize
@@ -756,9 +696,6 @@ func (db *Database) Commit(node common.Hash, report bool, callback func(common.H
 	batch.Reset()
 
 	// Reset the storage counters and bumped metrics
-	if db.preimages != nil {
-		db.preimages, db.preimagesSize = make(map[common.Hash][]byte), 0
-	}
 	memcacheCommitTimeTimer.Update(time.Since(start))
 	memcacheCommitSizeMeter.Mark(int64(storage - db.dirtiesSize))
 	memcacheCommitNodesMeter.Mark(int64(nodes - len(db.dirties)))
@@ -870,7 +807,11 @@ func (db *Database) Size() (common.StorageSize, common.StorageSize) {
 	// counted.
 	var metadataSize = common.StorageSize((len(db.dirties) - 1) * cachedNodeSize)
 	var metarootRefs = common.StorageSize(len(db.dirties[common.Hash{}].children) * (common.HashLength + 2))
-	return db.dirtiesSize + db.childrenSize + metadataSize - metarootRefs, db.preimagesSize
+	var preimageSize common.StorageSize
+	if db.preimages != nil {
+		preimageSize = db.preimages.size()
+	}
+	return db.dirtiesSize + db.childrenSize + metadataSize - metarootRefs, preimageSize
 }
 
 // saveCache saves clean state cache to given directory path

trie/preimages.go

@@ -0,0 +1,95 @@
+// Copyright 2022 The go-ethereum Authors
+// This file is part of the go-ethereum library.
+//
+// The go-ethereum library is free software: you can redistribute it and/or modify
+// it under the terms of the GNU Lesser General Public License as published by
+// the Free Software Foundation, either version 3 of the License, or
+// (at your option) any later version.
+//
+// The go-ethereum library is distributed in the hope that it will be useful,
+// but WITHOUT ANY WARRANTY; without even the implied warranty of
+// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
+// GNU Lesser General Public License for more details.
+//
+// You should have received a copy of the GNU Lesser General Public License
+// along with the go-ethereum library. If not, see <http://www.gnu.org/licenses/>.
+
+package trie
+
+import (
+	"sync"
+
+	"github.com/scroll-tech/go-ethereum/common"
+	"github.com/scroll-tech/go-ethereum/core/rawdb"
+	"github.com/scroll-tech/go-ethereum/ethdb"
+)
+
+// preimageStore is the store for caching preimages of node key.
+type preimageStore struct {
+	lock          sync.RWMutex
+	disk          ethdb.KeyValueStore
+	preimages     map[common.Hash][]byte // Preimages of nodes from the secure trie
+	preimagesSize common.StorageSize     // Storage size of the preimages cache
+}
+
+// newPreimageStore initializes the store for caching preimages.
+func newPreimageStore(disk ethdb.KeyValueStore) *preimageStore {
+	return &preimageStore{
+		disk:      disk,
+		preimages: make(map[common.Hash][]byte),
+	}
+}
+
+// insertPreimage writes a new trie node pre-image to the memory database if it's
+// yet unknown. The method will NOT make a copy of the slice, only use if the
+// preimage will NOT be changed later on.
+func (store *preimageStore) insertPreimage(preimages map[common.Hash][]byte) {
+	store.lock.Lock()
+	defer store.lock.Unlock()
+
+	for hash, preimage := range preimages {
+		if _, ok := store.preimages[hash]; ok {
+			continue
+		}
+		store.preimages[hash] = preimage
+		store.preimagesSize += common.StorageSize(common.HashLength + len(preimage))
+	}
+}
+
+// preimage retrieves a cached trie node pre-image from memory. If it cannot be
+// found cached, the method queries the persistent database for the content.
+func (store *preimageStore) preimage(hash common.Hash) []byte {
+	store.lock.RLock()
+	preimage := store.preimages[hash]
+	store.lock.RUnlock()
+
+	if preimage != nil {
+		return preimage
+	}
+	return rawdb.ReadPreimage(store.disk, hash)
+}
+
+// commit flushes the cached preimages into the disk.
+func (store *preimageStore) commit(force bool) error {
+	store.lock.Lock()
+	defer store.lock.Unlock()
+
+	if store.preimagesSize <= 4*1024*1024 && !force {
+		return nil
+	}
+	batch := store.disk.NewBatch()
+	rawdb.WritePreimages(batch, store.preimages)
+	if err := batch.Write(); err != nil {
+		return err
+	}
+	store.preimages, store.preimagesSize = make(map[common.Hash][]byte), 0
+	return nil
+}
+
+// size returns the current storage size of accumulated preimages.
+func (store *preimageStore) size() common.StorageSize {
+	store.lock.RLock()
+	defer store.lock.RUnlock()
+
+	return store.preimagesSize
+}

trie/secure_trie.go

@@ -37,6 +37,7 @@ import (
 // SecureTrie is not safe for concurrent use.
 type SecureTrie struct {
 	trie             Trie
+	preimages        *preimageStore
 	hashKeyBuf       [common.HashLength]byte
 	secKeyCache      map[string][]byte
 	secKeyCacheOwner *SecureTrie // Pointer to self, replace the key cache on mismatch
@@ -61,7 +62,7 @@ func NewSecure(root common.Hash, db *Database) (*SecureTrie, error) {
 	if err != nil {
 		return nil, err
 	}
-	return &SecureTrie{trie: *trie}, nil
+	return &SecureTrie{trie: *trie, preimages: db.preimages}, nil
 }
 
 // Get returns the value for key stored in the trie.
@@ -153,7 +154,10 @@ func (t *SecureTrie) GetKey(shaKey []byte) []byte {
 	if key, ok := t.getSecKeyCache()[string(shaKey)]; ok {
 		return key
 	}
-	return t.trie.db.preimage(common.BytesToHash(shaKey))
+	if t.preimages == nil {
+		return nil
+	}
+	return t.preimages.preimage(common.BytesToHash(shaKey))
 }
 
 // Commit writes all nodes and the secure hash pre-images to the trie's database.
@@ -164,12 +168,12 @@ func (t *SecureTrie) GetKey(shaKey []byte) []byte {
 func (t *SecureTrie) Commit(onleaf LeafCallback) (common.Hash, int, error) {
 	// Write all the pre-images to the actual disk database
 	if len(t.getSecKeyCache()) > 0 {
-		if t.trie.db.preimages != nil { // Ugly direct check but avoids the below write lock
-			t.trie.db.lock.Lock()
+		if t.preimages != nil {
+			preimages := make(map[common.Hash][]byte)
 			for hk, key := range t.secKeyCache {
-				t.trie.db.insertPreimage(common.BytesToHash([]byte(hk)), key)
+				preimages[common.BytesToHash([]byte(hk))] = key
 			}
-			t.trie.db.lock.Unlock()
+			t.preimages.insertPreimage(preimages)
 		}
 		t.secKeyCache = make(map[string][]byte)
 	}

trie/zk_trie.go

@@ -112,9 +112,10 @@ func (t *ZkTrie) GetKey(kHashBytes []byte) []byte {
 	if err != nil {
 		log.Error(fmt.Sprintf("Unhandled trie error: %v", err))
 	}
-
-	return t.db.db.preimage(common.BytesToHash(k.Bytes()))
-
+	if t.db.db.preimages != nil {
+		return t.db.db.preimages.preimage(common.BytesToHash(k.Bytes()))
+	}
+	return nil
 }
 
 // Commit writes all nodes and the secure hash pre-images to the trie's database.

trie/zk_trie_database.go

@@ -66,10 +66,8 @@ func (l *ZktrieDatabase) Get(key []byte) ([]byte, error) {
 func (l *ZktrieDatabase) UpdatePreimage(preimage []byte, hashField *big.Int) {
 	db := l.db
 	if db.preimages != nil { // Ugly direct check but avoids the below write lock
-		db.lock.Lock()
 		// we must copy the input key
-		db.insertPreimage(common.BytesToHash(hashField.Bytes()), common.CopyBytes(preimage))
-		db.lock.Unlock()
+		db.preimages.insertPreimage(map[common.Hash][]byte{common.BytesToHash(hashField.Bytes()): common.CopyBytes(preimage)})
 	}
 }
 

trie/zk_trie_test.go

@@ -35,7 +35,14 @@ import (
 )
 
 func newEmptyZkTrie() *ZkTrie {
-	trie, _ := NewZkTrie(common.Hash{}, NewZktrieDatabase(memorydb.New()))
+	trie, _ := NewZkTrie(
+		common.Hash{},
+		&ZktrieDatabase{
+			db: NewDatabaseWithConfig(memorydb.New(),
+				&Config{Preimages: true}),
+			prefix: []byte{},
+		},
+	)
 	return trie
 }