Get node cache

lib/blocktree/blocktree.go

@@ -38,14 +38,16 @@ type BlockTree struct {
 	leaves *leafMap
 	db     database.Database
 	sync.RWMutex
+	blockCache 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,
+		blockCache: 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,
+		blockCache: make(map[Hash]*node),
 	}
 }
 
@@ -150,7 +153,19 @@ func (bt *BlockTree) GetAllBlocksAtDepth(hash common.Hash) []common.Hash {
 }
 
 // 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() {
+		if ret != nil {
+			if _, ok := bt.blockCache[ret.hash]; !ok {
+				bt.blockCache[ret.hash] = ret
+			}
+		}
+	}()
+
+	if b, ok := bt.blockCache[h]; ok {
+		return b
+	}
+
 	if bt.head.hash == h {
 		return bt.head
 	}
@@ -175,6 +190,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.blockCache, hash)
+		}
+	}()
 
 	if finalised == bt.head.hash {
 		return pruned
@@ -350,7 +370,8 @@ func (bt *BlockTree) DeepCopy() *BlockTree {
 	defer bt.RUnlock()
 
 	btCopy := &BlockTree{
-		db: bt.db,
+		db:         bt.db,
+		blockCache: bt.blockCache,
 	}
 
 	if bt.head == nil {

lib/blocktree/blocktree_test.go

@@ -229,6 +229,29 @@ 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.blockCache[block.hash]
+	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 +409,37 @@ 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.blockCache[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)
+	require.Equal(t, bt.blockCache, btCopy.blockCache)
 	require.True(t, equalNodeValue(bt.head, btCopy.head), "BlockTree heads not equal")
 	require.True(t, equalLeave(bt.leaves, btCopy.leaves), "BlockTree leaves not equal")