From 3c18e47b203cee017dd4f65349e8630cbf1a112a Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Daniel=20G=C3=B3mez?= <dani91.gomez@gmail.com>
Date: Tue, 15 Jun 2021 11:47:04 -0300
Subject: [PATCH] chore(lib/blocktree): cache nodes in map (#1633)

---
 lib/blocktree/blocktree.go      | 46 ++++++++++++++++++++-----
 lib/blocktree/blocktree_test.go | 59 +++++++++++++++++++++++++++++++++
 2 files changed, 97 insertions(+), 8 deletions(-)

diff --git a/lib/blocktree/blocktree.go b/lib/blocktree/blocktree.go
index 9f5cc2d5e8..dcaeaab644 100644
--- a/lib/blocktree/blocktree.go
+++ b/lib/blocktree/blocktree.go
@@ -38,14 +38,16 @@ type BlockTree struct {
 	leaves *leafMap
 	db     database.Database
 	sync.RWMutex
+	nodeCache map[Hash]*node
 }
 
 // NewEmptyBlockTree creates a BlockTree with a nil head
 func NewEmptyBlockTree(db database.Database) *BlockTree {
 	return &BlockTree{
-		head:   nil,
-		leaves: newEmptyLeafMap(),
-		db:     db,
+		head:      nil,
+		leaves:    newEmptyLeafMap(),
+		db:        db,
+		nodeCache: make(map[Hash]*node),
 	}
 }
 
@@ -61,9 +63,10 @@ func NewBlockTreeFromRoot(root *types.Header, db database.Database) *BlockTree {
 	}
 
 	return &BlockTree{
-		head:   head,
-		leaves: newLeafMap(head),
-		db:     db,
+		head:      head,
+		leaves:    newLeafMap(head),
+		db:        db,
+		nodeCache: make(map[Hash]*node),
 	}
 }
 
@@ -103,6 +106,7 @@ func (bt *BlockTree) AddBlock(header *types.Header, arrivalTime uint64) error {
 	}
 	parent.addChild(n)
 	bt.leaves.replace(parent, n)
+	bt.setInCache(n)
 
 	return nil
 }
@@ -149,8 +153,24 @@ func (bt *BlockTree) GetAllBlocksAtDepth(hash common.Hash) []common.Hash {
 	return bt.head.getNodesWithDepth(depth, hashes)
 }
 
+func (bt *BlockTree) setInCache(b *node) {
+	if b == nil {
+		return
+	}
+
+	if _, has := bt.nodeCache[b.hash]; !has {
+		bt.nodeCache[b.hash] = b
+	}
+}
+
 // getNode finds and returns a node based on its Hash. Returns nil if not found.
-func (bt *BlockTree) getNode(h Hash) *node {
+func (bt *BlockTree) getNode(h Hash) (ret *node) {
+	defer func() { bt.setInCache(ret) }()
+
+	if b, ok := bt.nodeCache[h]; ok {
+		return b
+	}
+
 	if bt.head.hash == h {
 		return bt.head
 	}
@@ -175,6 +195,11 @@ func (bt *BlockTree) getNode(h Hash) *node {
 func (bt *BlockTree) Prune(finalised Hash) (pruned []Hash) {
 	bt.Lock()
 	defer bt.Unlock()
+	defer func() {
+		for _, hash := range pruned {
+			delete(bt.nodeCache, hash)
+		}
+	}()
 
 	if finalised == bt.head.hash {
 		return pruned
@@ -350,7 +375,8 @@ func (bt *BlockTree) DeepCopy() *BlockTree {
 	defer bt.RUnlock()
 
 	btCopy := &BlockTree{
-		db: bt.db,
+		db:        bt.db,
+		nodeCache: make(map[Hash]*node),
 	}
 
 	if bt.head == nil {
@@ -368,5 +394,9 @@ func (bt *BlockTree) DeepCopy() *BlockTree {
 		}
 	}
 
+	for hash := range bt.nodeCache {
+		btCopy.nodeCache[hash] = btCopy.getNode(hash)
+	}
+
 	return btCopy
 }
diff --git a/lib/blocktree/blocktree_test.go b/lib/blocktree/blocktree_test.go
index 2e0b905613..434a955928 100644
--- a/lib/blocktree/blocktree_test.go
+++ b/lib/blocktree/blocktree_test.go
@@ -229,6 +229,31 @@ func TestBlockTree_GetNode(t *testing.T) {
 	}
 }
 
+func TestBlockTree_GetNodeCache(t *testing.T) {
+	bt, branches := createTestBlockTree(testHeader, 16, nil)
+
+	for _, branch := range branches {
+		header := &types.Header{
+			ParentHash: branch.hash,
+			Number:     branch.depth,
+			StateRoot:  Hash{0x1},
+		}
+
+		err := bt.AddBlock(header, 0)
+		require.Nil(t, err)
+	}
+
+	block := bt.getNode(branches[0].hash)
+
+	cachedBlock, ok := bt.nodeCache[block.hash]
+
+	require.True(t, len(bt.nodeCache) > 0)
+	require.True(t, ok)
+	require.NotNil(t, cachedBlock)
+	require.Equal(t, cachedBlock, block)
+
+}
+
 func TestBlockTree_GetAllBlocksAtDepth(t *testing.T) {
 	bt, _ := createTestBlockTree(testHeader, 8, nil)
 	hashes := bt.head.getNodesWithDepth(big.NewInt(10), []common.Hash{})
@@ -386,12 +411,46 @@ func TestBlockTree_Prune(t *testing.T) {
 	}
 }
 
+func TestBlockTree_PruneCache(t *testing.T) {
+	var bt *BlockTree
+	var branches []testBranch
+
+	for {
+		bt, branches = createTestBlockTree(testHeader, 5, nil)
+		if len(branches) > 0 && len(bt.getNode(branches[0].hash).children) > 1 {
+			break
+		}
+	}
+
+	// pick some block to finalise
+	finalised := bt.head.children[0].children[0].children[0]
+	pruned := bt.Prune(finalised.hash)
+
+	for _, prunedHash := range pruned {
+		block, ok := bt.nodeCache[prunedHash]
+
+		require.False(t, ok)
+		require.Nil(t, block)
+	}
+
+}
+
 func TestBlockTree_DeepCopy(t *testing.T) {
 	bt, _ := createFlatTree(t, 8)
 
 	btCopy := bt.DeepCopy()
 
 	require.Equal(t, bt.db, btCopy.db)
+	for hash := range bt.nodeCache {
+		b, ok := btCopy.nodeCache[hash]
+		b2 := bt.nodeCache[hash]
+
+		require.True(t, ok)
+		require.True(t, b != b2)
+
+		require.True(t, equalNodeValue(b, b2))
+
+	}
 	require.True(t, equalNodeValue(bt.head, btCopy.head), "BlockTree heads not equal")
 	require.True(t, equalLeave(bt.leaves, btCopy.leaves), "BlockTree leaves not equal")