refactor(triedb): add clean cache to the rawDB engine

core/trie2/triedb/database.go

@@ -10,6 +10,12 @@ import (
 	"github.com/NethermindEth/juno/db"
 )
 
+const (
+	PathScheme string = "path"
+	HashScheme string = "hash"
+	RawScheme  string = "raw"
+)
+
 type Config struct {
 	PathConfig *pathdb.Config
 	HashConfig *hashdb.Config
@@ -19,13 +25,13 @@ type Config struct {
 func New(disk db.KeyValueStore, config *Config) (database.TrieDB, error) {
 	// Default to raw config if not provided
 	if config == nil {
-		return rawdb.New(disk), nil
+		return rawdb.New(disk, nil), nil
 	} else if config.PathConfig != nil {
 		return pathdb.New(disk, config.PathConfig)
 	} else if config.HashConfig != nil {
 		return hashdb.New(disk, config.HashConfig), nil
 	} else if config.RawConfig != nil {
-		return rawdb.New(disk), nil
+		return rawdb.New(disk, config.RawConfig), nil
 	}
 	return nil, fmt.Errorf("invalid trie database configuration: %v", config)
 }

core/trie2/triedb/rawdb/cache.go

@@ -0,0 +1,60 @@
+package rawdb
+
+import (
+	"math"
+
+	"github.com/NethermindEth/juno/core/felt"
+	"github.com/NethermindEth/juno/core/trie2/trieutils"
+	"github.com/VictoriaMetrics/fastcache"
+)
+
+const nodeCacheSize = ownerSize + trieutils.PathSize + 1
+
+type trieType byte
+
+const (
+	contract trieType = iota
+	class
+)
+
+// Stores committed trie nodes in memory
+type cleanCache struct {
+	cache *fastcache.Cache // map[nodeKey]node
+}
+
+// Creates a new clean cache with the given size.
+// The size is the maximum size of the cache in bytes.
+func newCleanCache(size uint64) cleanCache {
+	if size > uint64(math.MaxInt) {
+		panic("cache size too large: uint64 to int conversion would overflow")
+	}
+	return cleanCache{
+		cache: fastcache.New(int(size)),
+	}
+}
+
+func (c *cleanCache) putNode(owner *felt.Address, path *trieutils.Path, isClass bool, blob []byte) {
+	c.cache.Set(nodeKey(owner, path, isClass), blob)
+}
+
+func (c *cleanCache) getNode(owner *felt.Address, path *trieutils.Path, isClass bool) []byte {
+	return c.cache.Get(nil, nodeKey(owner, path, isClass))
+}
+
+func (c *cleanCache) deleteNode(owner *felt.Address, path *trieutils.Path, isClass bool) {
+	c.cache.Del(nodeKey(owner, path, isClass))
+}
+
+// key = owner (32 bytes) + path (20 bytes) + trie type (1 byte)
+func nodeKey(owner *felt.Address, path *trieutils.Path, isClass bool) []byte {
+	key := make([]byte, nodeCacheSize)
+	ownerBytes := owner.Bytes()
+	copy(key[:felt.Bytes], ownerBytes[:])
+	copy(key[felt.Bytes:felt.Bytes+trieutils.PathSize], path.EncodedBytes())
+
+	if isClass {
+		key[nodeCacheSize-1] = byte(class)
+	}
+
+	return key
+}

core/trie2/triedb/rawdb/database.go

@@ -13,19 +13,32 @@ import (
 
 var _ database.TrieDB = (*Database)(nil)
 
-type Config struct{}
+type Config struct {
+	CleanCacheSize uint64 // Maximum size (in bytes) for caching clean nodes
+}
 
 type Database struct {
 	disk db.KeyValueStore
 
 	lock sync.RWMutex
 	log  utils.StructuredLogger
+
+	config     Config
+	cleanCache *cleanCache
 }
 
-func New(disk db.KeyValueStore) *Database {
+func New(disk db.KeyValueStore, config *Config) *Database {
+	if config == nil {
+		config = &Config{
+			CleanCacheSize: 16 * utils.Megabyte,
+		}
+	}
+	cleanCache := newCleanCache(config.CleanCacheSize)
 	return &Database{
-		disk: disk,
-		log:  utils.NewNopZapLogger(),
+		disk:       disk,
+		config:     *config,
+		cleanCache: &cleanCache,
+		log:        utils.NewNopZapLogger(),
 	}
 }
 
@@ -37,11 +50,19 @@ func (d *Database) readNode(
 ) ([]byte, error) {
 	d.lock.RLock()
 	defer d.lock.RUnlock()
+
+	isClass := id.Type() == trieutils.Class
+	blob := d.cleanCache.getNode(owner, path, isClass)
+	if blob != nil {
+		return blob, nil
+	}
+
 	blob, err := trieutils.GetNodeByPath(d.disk, id.Bucket(), owner, path, isLeaf)
 	if err != nil {
 		return nil, err
 	}
 
+	d.cleanCache.putNode(owner, path, isClass, blob)
 	return blob, nil
 }
 
@@ -79,22 +100,22 @@ func (d *Database) Update(
 	}
 
 	for path, n := range classNodes {
-		err := d.updateNode(batch, db.ClassTrie, &felt.Address{}, &path, n)
+		err := d.updateNode(batch, db.ClassTrie, &felt.Address{}, &path, n, true)
 		if err != nil {
 			return err
 		}
 	}
 
 	for path, n := range contractNodes {
-		err := d.updateNode(batch, db.ContractTrieContract, &felt.Address{}, &path, n)
+		err := d.updateNode(batch, db.ContractTrieContract, &felt.Address{}, &path, n, false)
 		if err != nil {
 			return err
 		}
 	}
 
 	for owner, nodes := range contractStorageNodes {
 		for path, n := range nodes {
-			err := d.updateNode(batch, db.ContractTrieStorage, &owner, &path, n)
+			err := d.updateNode(batch, db.ContractTrieStorage, &owner, &path, n, false)
 			if err != nil {
 				return err
 			}
@@ -110,9 +131,27 @@ func (d *Database) updateNode(
 	owner *felt.Address,
 	path *trieutils.Path,
 	n trienode.TrieNode,
+	isClass bool,
 ) error {
 	if _, deleted := n.(*trienode.DeletedNode); deleted {
-		return trieutils.DeleteNodeByPath(batch, bucket, owner, path, n.IsLeaf())
+		err := trieutils.DeleteNodeByPath(batch, bucket, owner, path, n.IsLeaf())
+		if err != nil {
+			return err
+		}
+		d.cleanCache.deleteNode(owner, path, isClass)
+	} else {
+		err := trieutils.WriteNodeByPath(
+			batch,
+			bucket,
+			owner,
+			path,
+			n.IsLeaf(),
+			n.Blob(),
+		)
+		if err != nil {
+			return err
+		}
+		d.cleanCache.putNode(owner, path, isClass, n.Blob())
 	}
 	return trieutils.WriteNodeByPath(
 		batch,

core/trie2/triedb/rawdb/database_test.go

@@ -108,13 +108,13 @@ func verifyNode(
 func TestRawDB(t *testing.T) {
 	t.Run("New creates database", func(t *testing.T) {
 		memDB := memory.New()
-		database := New(memDB)
+		database := New(memDB, nil)
 		require.NotNil(t, database)
 	})
 
 	t.Run("Update with all node types", func(t *testing.T) {
 		memDB := memory.New()
-		database := New(memDB)
+		database := New(memDB, nil)
 
 		contractHash := felt.NewFromUint64[felt.Felt](210)
 		contractPath := trieutils.NewBitArray(1, 0x01)
@@ -171,7 +171,7 @@ func TestRawDB(t *testing.T) {
 
 	t.Run("Update with deleted nodes", func(t *testing.T) {
 		memDB := memory.New()
-		database := New(memDB)
+		database := New(memDB, nil)
 
 		err := database.Update(
 			&felt.StateRootHash{},
@@ -215,7 +215,7 @@ func TestRawDB(t *testing.T) {
 
 	t.Run("NodeReader returns correct reader", func(t *testing.T) {
 		memDB := memory.New()
-		database := New(memDB)
+		database := New(memDB, nil)
 
 		err := database.Update(
 			&felt.StateRootHash{},
@@ -245,7 +245,7 @@ func TestRawDB(t *testing.T) {
 
 	t.Run("Multiple updates", func(t *testing.T) {
 		memDB := memory.New()
-		database := New(memDB)
+		database := New(memDB, nil)
 
 		err := database.Update(
 			&felt.StateRootHash{},
@@ -286,7 +286,7 @@ func TestRawDB(t *testing.T) {
 
 	t.Run("Concurrent reads", func(t *testing.T) {
 		memDB := memory.New()
-		database := New(memDB)
+		database := New(memDB, nil)
 
 		err := database.Update(
 			&felt.StateRootHash{},

core/trie2/triedb/rawdb/types.go

@@ -11,3 +11,5 @@ type (
 	contractNodesMap        = map[trieutils.Path]trienode.TrieNode
 	contractStorageNodesMap = map[felt.Address]map[trieutils.Path]trienode.TrieNode
 )
+
+const ownerSize = felt.Bytes