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

core/trie2/triedb/database.go

@@ -35,7 +35,7 @@ func New(disk db.KeyValueStore, config *Config) (*Database, error) {
 	var err error
 	// Default to raw config if not provided
 	if config == nil {
-		triedb = rawdb.New(disk)
+		triedb = rawdb.New(disk, nil)
 	} else if config.PathConfig != nil {
 		triedb, err = pathdb.New(disk, config.PathConfig)
 		if err != nil {
@@ -44,7 +44,7 @@ func New(disk db.KeyValueStore, config *Config) (*Database, error) {
 	} else if config.HashConfig != nil {
 		triedb = hashdb.New(disk, config.HashConfig)
 	} else if config.RawConfig != nil {
-		triedb = rawdb.New(disk)
+		triedb = rawdb.New(disk, config.RawConfig)
 	}
 
 	return &Database{

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.Felt, path *trieutils.Path, isClass bool, blob []byte) {
+	c.cache.Set(nodeKey(owner, path, isClass), blob)
+}
+
+func (c *cleanCache) getNode(owner *felt.Felt, path *trieutils.Path, isClass bool) []byte {
+	return c.cache.Get(nil, nodeKey(owner, path, isClass))
+}
+
+func (c *cleanCache) deleteNode(owner *felt.Felt, 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.Felt, 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

@@ -10,19 +10,32 @@ import (
 	"github.com/NethermindEth/juno/utils"
 )
 
-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.SimpleLogger
+
+	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(),
 	}
 }
 
@@ -34,11 +47,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
 }
 
@@ -80,26 +101,26 @@ func (d *Database) Update(
 	}
 
 	for path, n := range classNodes {
-		if err := d.updateNode(batch, db.ClassTrie, &felt.Zero, &path, n); err != nil {
+		if err := d.updateNode(batch, db.ClassTrie, &felt.Zero, &path, n, true); err != nil {
 			return err
 		}
 	}
 
 	for path, n := range contractNodes {
-		if err := d.updateNode(batch, db.ContractTrieContract, &felt.Zero, &path, n); err != nil {
+		if err := d.updateNode(batch, db.ContractTrieContract, &felt.Zero, &path, n, false); err != nil {
 			return err
 		}
 	}
 
 	for owner, nodes := range contractStorageNodes {
 		for path, n := range nodes {
-			if err := d.updateNode(batch, db.ContractTrieStorage, &owner, &path, n); err != nil {
+			if err := d.updateNode(batch, db.ContractTrieStorage, &owner, &path, n, false); err != nil {
 				return err
 			}
 		}
 	}
 
-	return batch.Write()
+	return nil
 }
 
 func (d *Database) updateNode(
@@ -108,11 +129,13 @@ func (d *Database) updateNode(
 	owner *felt.Felt,
 	path *trieutils.Path,
 	n trienode.TrieNode,
+	isClass bool,
 ) error {
 	if _, deleted := n.(*trienode.DeletedNode); deleted {
 		if err := trieutils.DeleteNodeByPath(batch, bucket, owner, path, n.IsLeaf()); err != nil {
 			return err
 		}
+		d.cleanCache.deleteNode(owner, path, isClass)
 	} else {
 		if err := trieutils.WriteNodeByPath(
 			batch,
@@ -124,6 +147,7 @@ func (d *Database) updateNode(
 		); err != nil {
 			return err
 		}
+		d.cleanCache.putNode(owner, path, isClass, n.Blob())
 	}
 	return nil
 }

core/trie2/triedb/rawdb/types.go

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