perf: memory fixes for node/fast-node caches and pruning

cache/cache.go

@@ -23,6 +23,9 @@ type Cache interface {
 	// Returns Node for the key, if exists. nil otherwise.
 	Get(key []byte) Node
 
+	// Peek is like Get but does not refresh the entry's LRU position.
+	Peek(key []byte) Node
+
 	// Has returns true if node with key exists in cache, false otherwise.
 	Has(key []byte) bool
 
@@ -87,6 +90,13 @@ func (c *lruCache) Get(key []byte) Node {
 	return nil
 }
 
+func (c *lruCache) Peek(key []byte) Node {
+	if ele, hit := c.dict[string(key)]; hit {
+		return ele.Value.(Node)
+	}
+	return nil
+}
+
 func (c *lruCache) Has(key []byte) bool {
 	_, exists := c.dict[string(key)]
 	return exists

cache/cache_test.go

@@ -307,3 +307,14 @@ func randBytes(length int) []byte {
 	rand.Read(key) //nolint:errcheck
 	return key
 }
+
+func TestPeekDoesNotRefreshRecency(t *testing.T) {
+	c := cache.New(2)
+	c.Add(&testNode{key: []byte("a")})
+	c.Add(&testNode{key: []byte("b")})
+	require.NotNil(t, c.Peek([]byte("a")))
+	require.Nil(t, c.Peek([]byte("zz")))
+	// "a" is still the oldest, so adding "c" evicts it.
+	evicted := c.Add(&testNode{key: []byte("c")})
+	require.Equal(t, []byte("a"), evicted.GetKey())
+}

cache_opts_test.go

@@ -0,0 +1,130 @@
+package iavl
+
+import (
+	"bytes"
+	"fmt"
+	"testing"
+
+	"github.com/stretchr/testify/require"
+
+	dbm "github.com/cosmos/iavl/db"
+)
+
+func TestCacheValueSizeLimit(t *testing.T) {
+	small, big := []byte("small"), bytes.Repeat([]byte("x"), 64)
+	tree := NewMutableTree(dbm.NewMemDB(), 1000, false, NewNopLogger(), CacheValueSizeLimitOption(16))
+	_, err := tree.Set([]byte("a"), small)
+	require.NoError(t, err)
+	_, err = tree.Set([]byte("b"), big)
+	require.NoError(t, err)
+	_, _, err = tree.SaveVersion()
+	require.NoError(t, err)
+
+	ndb := tree.ndb
+	require.True(t, ndb.fastNodeCache.Has([]byte("a")))
+	require.False(t, ndb.fastNodeCache.Has([]byte("b")), "large fast node must not be cached")
+	var bigLeaf int
+	leaves, err := ndb.leafNodes()
+	require.NoError(t, err)
+	for _, l := range leaves {
+		cached := ndb.nodeCache.Has(l.GetKey())
+		if len(l.value) > 16 {
+			bigLeaf++
+			require.False(t, cached, "large leaf must not be cached")
+		} else {
+			require.True(t, cached)
+		}
+	}
+	require.Equal(t, 1, bigLeaf)
+
+	// reads of the large value still work and don't populate the caches.
+	v, err := tree.Get([]byte("b"))
+	require.NoError(t, err)
+	require.Equal(t, big, v)
+	require.False(t, ndb.fastNodeCache.Has([]byte("b")))
+
+	// growing a cached value past the limit must drop the stale cache entry.
+	_, err = tree.Set([]byte("a"), big)
+	require.NoError(t, err)
+	_, _, err = tree.SaveVersion()
+	require.NoError(t, err)
+	require.False(t, ndb.fastNodeCache.Has([]byte("a")))
+	v, err = tree.Get([]byte("a"))
+	require.NoError(t, err)
+	require.Equal(t, big, v)
+}
+
+func TestFastNodeCacheSizeOption(t *testing.T) {
+	db := dbm.NewMemDB()
+	require.Equal(t, defaultFastNodeCacheSize, fastCacheCap(NewMutableTree(db, 0, false, NewNopLogger())))
+
+	tree := NewMutableTree(db, 0, false, NewNopLogger(), FastNodeCacheSizeOption(3))
+	for i := 0; i < 10; i++ {
+		_, err := tree.Set([]byte(fmt.Sprintf("k%d", i)), []byte("v"))
+		require.NoError(t, err)
+	}
+	_, _, err := tree.SaveVersion()
+	require.NoError(t, err)
+	require.Equal(t, 3, tree.ndb.fastNodeCache.Len())
+}
+
+func fastCacheCap(tree *MutableTree) int {
+	for i := 0; ; i++ {
+		tree.ndb.fastNodeCache.Add(&testCacheNode{key: []byte(fmt.Sprintf("%d", i))})
+		if tree.ndb.fastNodeCache.Len() <= i {
+			return i
+		}
+		if i > 2*defaultFastNodeCacheSize {
+			return -1
+		}
+	}
+}
+
+type testCacheNode struct{ key []byte }
+
+func (n *testCacheNode) GetKey() []byte { return n.key }
+
+func TestPruningEvictsAndSkipsCache(t *testing.T) {
+	db := dbm.NewMemDB()
+	tree := NewMutableTree(db, 100000, true, NewNopLogger())
+	for v := 0; v < 5; v++ {
+		for i := 0; i < 10; i++ {
+			_, err := tree.Set([]byte(fmt.Sprintf("k%d", i)), []byte(fmt.Sprintf("v%d-%d", v, i)))
+			require.NoError(t, err)
+		}
+		_, _, err := tree.SaveVersion()
+		require.NoError(t, err)
+	}
+	ndb := tree.ndb
+
+	// nodes orphaned by the version update.
+	var orphans [][]byte
+	require.NoError(t, ndb.traverseOrphans(1, 2, func(o *Node) error {
+		orphans = append(orphans, o.GetKey())
+		return nil
+	}))
+	require.NotEmpty(t, orphans)
+	for _, k := range orphans {
+		require.True(t, ndb.nodeCache.Has(k), "setup: orphan should be cached")
+	}
+
+	require.NoError(t, ndb.DeleteVersionsTo(1))
+	require.NoError(t, ndb.Commit())
+	// version 1 pruned: its orphans are gone from disk and must be gone from the cache.
+	require.NoError(t, ndb.DeleteVersionsTo(2))
+	require.NoError(t, ndb.Commit())
+	for _, k := range orphans {
+		require.False(t, ndb.nodeCache.Has(k), "pruned node still cached: %x", k)
+		has, err := ndb.Has(k)
+		require.NoError(t, err)
+		require.False(t, has)
+	}
+
+	// a pruning read of an uncached node must not populate the cache.
+	latest, err := ndb.GetRoot(5)
+	require.NoError(t, err)
+	ndb.nodeCache.Remove(latest)
+	_, err = ndb.getNode(latest, true)
+	require.NoError(t, err)
+	require.False(t, ndb.nodeCache.Has(latest))
+}

iterator.go

@@ -268,26 +268,30 @@ type NodeIterator struct {
 	nodesToVisit []*Node
 	ndb          *nodeDB
 	err          error
+	pruning      bool // read through the node cache without refreshing or populating it
 }
 
 // NewNodeIterator returns a new NodeIterator to traverse the tree of the root node.
 func NewNodeIterator(rootKey []byte, ndb *nodeDB) (*NodeIterator, error) {
+	return newNodeIterator(rootKey, ndb, false)
+}
+
+func newNodeIterator(rootKey []byte, ndb *nodeDB, pruning bool) (*NodeIterator, error) {
+	iter := &NodeIterator{nodesToVisit: []*Node{}, ndb: ndb, pruning: pruning}
 	if len(rootKey) == 0 {
-		return &NodeIterator{
-			nodesToVisit: []*Node{},
-			ndb:          ndb,
-		}, nil
+		return iter, nil
 	}
 
-	node, err := ndb.GetNode(rootKey)
+	node, err := iter.getNode(rootKey)
 	if err != nil {
 		return nil, err
 	}
+	iter.nodesToVisit = append(iter.nodesToVisit, node)
+	return iter, nil
+}
 
-	return &NodeIterator{
-		nodesToVisit: []*Node{node},
-		ndb:          ndb,
-	}, nil
+func (iter *NodeIterator) getNode(nk []byte) (*Node, error) {
+	return iter.ndb.getNode(nk, iter.pruning)
 }
 
 // GetNode returns the current visiting node.
@@ -322,14 +326,14 @@ func (iter *NodeIterator) Next(isSkipped bool) {
 		return
 	}
 
-	rightNode, err := iter.ndb.GetNode(node.rightNodeKey)
+	rightNode, err := iter.getNode(node.rightNodeKey)
 	if err != nil {
 		iter.err = err
 		return
 	}
 	iter.nodesToVisit = append(iter.nodesToVisit, rightNode)
 
-	leftNode, err := iter.ndb.GetNode(node.leftNodeKey)
+	leftNode, err := iter.getNode(node.leftNodeKey)
 	if err != nil {
 		iter.err = err
 		return