feat(pkg/trie): Add TrieCache for trie values and nodes

go.mod

@@ -24,6 +24,7 @@ require (
 	github.com/ipfs/go-ds-badger2 v0.1.3
 	github.com/jpillora/backoff v1.0.0
 	github.com/jpillora/ipfilter v1.2.9
+	github.com/karlseguin/ccache/v3 v3.0.5
 	github.com/klauspost/compress v1.17.8
 	github.com/libp2p/go-libp2p v0.33.2
 	github.com/libp2p/go-libp2p-kad-dht v0.25.2

go.sum

@@ -325,6 +325,8 @@ github.com/json-iterator/go v1.1.6/go.mod h1:+SdeFBvtyEkXs7REEP0seUULqWtbJapLOCV
 github.com/jstemmer/go-junit-report v0.0.0-20190106144839-af01ea7f8024/go.mod h1:6v2b51hI/fHJwM22ozAgKL4VKDeJcHhJFhtBdhmNjmU=
 github.com/jtolds/gls v4.20.0+incompatible h1:xdiiI2gbIgH/gLH7ADydsJ1uDOEzR8yvV7C0MuV77Wo=
 github.com/jtolds/gls v4.20.0+incompatible/go.mod h1:QJZ7F/aHp+rZTRtaJ1ow/lLfFfVYBRgL+9YlvaHOwJU=
+github.com/karlseguin/ccache/v3 v3.0.5 h1:hFX25+fxzNjsRlREYsoGNa2LoVEw5mPF8wkWq/UnevQ=
+github.com/karlseguin/ccache/v3 v3.0.5/go.mod h1:qxC372+Qn+IBj8Pe3KvGjHPj0sWwEF7AeZVhsNPZ6uY=
 github.com/kisielk/errcheck v1.2.0/go.mod h1:/BMXB+zMLi60iA8Vv6Ksmxu/1UDYcXs4uQLJ+jE2L00=
 github.com/kisielk/errcheck v1.5.0/go.mod h1:pFxgyoBC7bSaBwPgfKdkLd5X25qrDl4LWUI2bnpBCr8=
 github.com/kisielk/gotool v1.0.0/go.mod h1:XhKaO+MFFWcvkIS/tQcRk01m1F5IRFswLeQ+oQHNcck=

pkg/trie/cache/inmemory/maxbytes_lru_cache.go

@@ -0,0 +1,49 @@
+// Copyright 2024 ChainSafe Systems (ON)
+// SPDX-License-Identifier: LGPL-3.0-only
+
+package inmemory
+
+import (
+	"github.com/karlseguin/ccache/v3"
+)
+
+// ccache uses 350 bytes of overhead per entry
+const cacheValueOverheadSize = 350
+
+// cacheValue is a helper alias over []byte to implement ccache.Sized
+type cacheValue []byte
+
+// Size returns the size of the cacheValue taking the overhead into account
+func (cv cacheValue) Size() int64 {
+	return int64(len(cv) + cacheValueOverheadSize)
+}
+
+// maxBytesLRUCache is an in-memory lru cache
+// consider that the values are deleted asyncronously so there is a chance that
+// the maxSize can be exceeded
+// we can use lru.GC() to force the deletion of the items that should be deleted
+type maxBytesLRUCache struct {
+	lru *ccache.Cache[cacheValue]
+}
+
+// newlruCache creates a new lruCache
+// maxSize is the cache max size in bytes
+func newLruCache(maxSize int64) *maxBytesLRUCache {
+	cache := ccache.New(ccache.Configure[cacheValue]().MaxSize(maxSize))
+	return &maxBytesLRUCache{
+		lru: cache,
+	}
+}
+
+func (cache *maxBytesLRUCache) get(key string) []byte {
+	item := cache.lru.Get(key)
+	if item != nil {
+		return item.Value()
+
+	}
+	return nil
+}
+
+func (cache *maxBytesLRUCache) set(key string, value []byte) {
+	cache.lru.Set(key, cacheValue(value), 0)
+}

pkg/trie/cache/inmemory/maxbytes_lru_cache_test.go

@@ -0,0 +1,56 @@
+// Copyright 2024 ChainSafe Systems (ON)
+// SPDX-License-Identifier: LGPL-3.0-only
+
+package inmemory
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+)
+
+func Test_TrieValueCache_SetAndGet(t *testing.T) {
+	const maxCacheSize = 10
+	t.Run("set_and_get_value_successful", func(t *testing.T) {
+		cache := newLruCache(maxCacheSize)
+		key := "key"
+		value := []byte("value")
+
+		cache.set(key, value)
+		valueFromCache := cache.get(key)
+
+		assert.Equal(t, value, valueFromCache)
+	})
+
+	t.Run("get_value_not_found", func(t *testing.T) {
+		cache := newLruCache(maxCacheSize)
+		valueFromCache := cache.get("missing")
+		assert.Nil(t, valueFromCache)
+	})
+
+	t.Run("replace_value_when_size_exceeded", func(t *testing.T) {
+		cache := newLruCache(maxCacheSize)
+		key1 := "key1"
+		key2 := "key2"
+		value1 := make([]byte, maxCacheSize/2+1)
+		value2 := make([]byte, maxCacheSize/2+1)
+
+		// First value is inserted successfully
+		cache.set(key1, value1)
+		valueFromCache := cache.get(key1)
+
+		assert.Equal(t, value1, valueFromCache)
+
+		// Second value is inserted successfully
+		cache.set(key2, value2)
+		valueFromCache = cache.get(key2)
+
+		assert.Equal(t, value2, valueFromCache)
+
+		// First value has been removed
+		cache.lru.GC() // Force GC to remove items that should be deleted
+
+		valueFromCache = cache.get(key1)
+		assert.Nil(t, valueFromCache)
+	})
+}

pkg/trie/cache/inmemory/trie_cache.go

@@ -0,0 +1,49 @@
+// Copyright 2024 ChainSafe Systems (ON)
+// SPDX-License-Identifier: LGPL-3.0-only
+
+package inmemory
+
+import (
+	lrucache "github.com/ChainSafe/gossamer/lib/utils/lru-cache"
+	"github.com/ChainSafe/gossamer/pkg/trie/cache"
+)
+
+// https://github.com/paritytech/polkadot-sdk/blob/a8f4f4f00f8fc0da512a09e1450bf4cda954d70d/substrate/primitives/trie/src/cache/mod.rs#L98
+const defaultNodeCacheMaxElements = 10000
+const defaultValueCacheMaxSize = 2 * 1024 * 1024 // 2MB
+
+// TrieInMemoryCache is an in-memory cache for trie nodes
+type TrieInMemoryCache struct {
+	nodeCache  *lrucache.LRUCache[string, []byte]
+	valueCache *maxBytesLRUCache
+}
+
+// NewTrieInMemoryCache creates a new TrieInMemoryCache
+func NewTrieInMemoryCache() *TrieInMemoryCache {
+	return &TrieInMemoryCache{
+		nodeCache:  lrucache.NewLRUCache[string, []byte](defaultNodeCacheMaxElements),
+		valueCache: newLruCache(defaultValueCacheMaxSize),
+	}
+}
+
+// GetValue returns the value for the given key
+func (tc *TrieInMemoryCache) GetValue(key []byte) []byte {
+	return tc.valueCache.get(string(key))
+}
+
+// SetValue sets the value for the given key
+func (tc *TrieInMemoryCache) SetValue(key []byte, value []byte) {
+	tc.valueCache.set(string(key), value)
+}
+
+// GetNode returns the node for the given key
+func (tc *TrieInMemoryCache) GetNode(key []byte) []byte {
+	return tc.nodeCache.Get(string(key))
+}
+
+// SetNode sets the node for the given key
+func (tc *TrieInMemoryCache) SetNode(key, value []byte) {
+	tc.nodeCache.Put(string(key), value)
+}
+
+var _ cache.TrieCache = (*TrieInMemoryCache)(nil)

pkg/trie/cache/inmemory/trie_cache_test.go

@@ -0,0 +1,29 @@
+// Copyright 2024 ChainSafe Systems (ON)
+// SPDX-License-Identifier: LGPL-3.0-only
+
+package inmemory
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+)
+
+func Test_TrieCache_SetAndGet(t *testing.T) {
+	t.Run("set_and_get_value_successful", func(t *testing.T) {
+		cache := NewTrieInMemoryCache()
+		key := []byte("key")
+		value := []byte("value")
+
+		cache.SetValue(key, value)
+		valueFromCache := cache.GetValue(key)
+
+		assert.Equal(t, value, valueFromCache)
+	})
+
+	t.Run("get_value_not_found", func(t *testing.T) {
+		cache := NewTrieInMemoryCache()
+		valueFromCache := cache.GetValue([]byte("missing"))
+		assert.Nil(t, valueFromCache)
+	})
+}

pkg/trie/cache/trie_cache.go

@@ -0,0 +1,11 @@
+// Copyright 2024 ChainSafe Systems (ON)
+// SPDX-License-Identifier: LGPL-3.0-only
+
+package cache
+
+type TrieCache interface {
+	GetValue(key []byte) []byte
+	SetValue(key, value []byte)
+	GetNode(key []byte) []byte
+	SetNode(key, value []byte)
+}

pkg/trie/triedb/in_memory_to_triedb_migration_test.go

@@ -10,12 +10,11 @@ import (
 	"github.com/ChainSafe/gossamer/pkg/trie"
 	"github.com/ChainSafe/gossamer/pkg/trie/inmemory"
 	"github.com/stretchr/testify/assert"
-	"github.com/stretchr/testify/require"
 )
 
-func newTestDB(t *testing.T) database.Table {
+func newTestDB(t assert.TestingT) database.Table {
 	db, err := database.NewPebble("", true)
-	require.NoError(t, err)
+	assert.NoError(t, err)
 	return database.NewTable(db, "trie")
 }
 
@@ -44,11 +43,12 @@ func TestTrieDB_Migration(t *testing.T) {
 
 	root, err := inMemoryTrie.Hash()
 	assert.NoError(t, err)
-	trieDB := NewTrieDB(root, db)
+	trieDB := NewTrieDB(root, db, nil)
 
 	t.Run("read_successful_from_db_created_using_v1_trie", func(t *testing.T) {
 		for k, v := range entries {
 			value := trieDB.Get([]byte(k))
+			assert.NotNil(t, value)
 			assert.Equal(t, v, value)
 		}
 

pkg/trie/triedb/lookup.go

@@ -7,21 +7,25 @@ import (
 	"bytes"
 
 	"github.com/ChainSafe/gossamer/lib/common"
+	"github.com/ChainSafe/gossamer/pkg/trie/cache"
 	"github.com/ChainSafe/gossamer/pkg/trie/db"
 	"github.com/ChainSafe/gossamer/pkg/trie/triedb/codec"
 )
 
 type TrieLookup struct {
-	// Database to query from
+	// db to query from
 	db db.DBGetter
-	// Hash to start at
+	// hash to start at
 	hash common.Hash
+	// cache to speed up the db lookups
+	cache cache.TrieCache
 }
 
-func NewTrieLookup(db db.DBGetter, hash common.Hash) TrieLookup {
+func NewTrieLookup(db db.DBGetter, hash common.Hash, cache cache.TrieCache) TrieLookup {
 	return TrieLookup{
-		db:   db,
-		hash: hash,
+		db:    db,
+		hash:  hash,
+		cache: cache,
 	}
 }
 
@@ -33,9 +37,21 @@ func (l *TrieLookup) lookupNode(keyNibbles []byte) (codec.Node, error) {
 	// Iterates through non inlined nodes
 	for {
 		// Get node from DB
-		nodeData, err := l.db.Get(hash)
-		if err != nil {
-			return nil, ErrIncompleteDB
+		var nodeData []byte
+		if l.cache != nil {
+			nodeData = l.cache.GetNode(hash)
+		}
+
+		if nodeData == nil {
+			var err error
+			nodeData, err = l.db.Get(hash)
+			if err != nil {
+				return nil, ErrIncompleteDB
+			}
+
+			if l.cache != nil {
+				l.cache.SetNode(hash, nodeData)
+			}
 		}
 
 	InlinedChildrenIterator:
@@ -103,14 +119,29 @@ func (l *TrieLookup) lookupNode(keyNibbles []byte) (codec.Node, error) {
 	}
 }
 
-func (l *TrieLookup) lookupValue(keyNibbles []byte) ([]byte, error) {
+func (l *TrieLookup) lookupValue(keyNibbles []byte) (value []byte, err error) {
+	if l.cache != nil {
+		if value = l.cache.GetValue(keyNibbles); value != nil {
+			return value, nil
+		}
+	}
+
 	node, err := l.lookupNode(keyNibbles)
 	if err != nil {
 		return nil, err
 	}
 
-	if value := node.GetValue(); value != nil {
-		return l.fetchValue(node.GetPartialKey(), value)
+	if nodeValue := node.GetValue(); nodeValue != nil {
+		value, err = l.fetchValue(node.GetPartialKey(), nodeValue)
+		if err != nil {
+			return nil, err
+		}
+
+		if l.cache != nil {
+			l.cache.SetValue(keyNibbles, value)
+		}
+
+		return value, nil
 	}
 
 	return nil, nil
@@ -124,7 +155,22 @@ func (l *TrieLookup) fetchValue(prefix []byte, value codec.NodeValue) ([]byte, e
 		return v.Data, nil
 	case codec.HashedValue:
 		prefixedKey := bytes.Join([][]byte{prefix, v.Data}, nil)
-		return l.db.Get(prefixedKey)
+		if l.cache != nil {
+			if value := l.cache.GetValue(prefixedKey); value != nil {
+				return value, nil
+			}
+		}
+
+		nodeData, err := l.db.Get(prefixedKey)
+		if err != nil {
+			return nil, err
+		}
+
+		if l.cache != nil {
+			l.cache.SetValue(prefixedKey, nodeData)
+		}
+
+		return nodeData, nil
 	default:
 		panic("unreachable")
 	}

pkg/trie/triedb/lookup_test.go

@@ -13,7 +13,7 @@ import (
 func TestTrieDB_Lookup(t *testing.T) {
 	t.Run("root_not_exists_in_db", func(t *testing.T) {
 		db := newTestDB(t)
-		lookup := NewTrieLookup(db, trie.EmptyHash)
+		lookup := NewTrieLookup(db, trie.EmptyHash, nil)
 
 		value, err := lookup.lookupValue([]byte("test"))
 		assert.Nil(t, value)