Skip to content

Commit

Permalink
chore(lib/blocktree): cache nodes in map (#1633)
Browse files Browse the repository at this point in the history
  • Loading branch information
danigomez authored Jun 15, 2021
1 parent f5bd761 commit 3c18e47
Show file tree
Hide file tree
Showing 2 changed files with 97 additions and 8 deletions.
46 changes: 38 additions & 8 deletions lib/blocktree/blocktree.go
Original file line number Diff line number Diff line change
Expand Up @@ -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),
}
}

Expand All @@ -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),
}
}

Expand Down Expand Up @@ -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
}
Expand Down Expand Up @@ -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
}
Expand All @@ -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
Expand Down Expand Up @@ -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 {
Expand All @@ -368,5 +394,9 @@ func (bt *BlockTree) DeepCopy() *BlockTree {
}
}

for hash := range bt.nodeCache {
btCopy.nodeCache[hash] = btCopy.getNode(hash)
}

return btCopy
}
59 changes: 59 additions & 0 deletions lib/blocktree/blocktree_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -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{})
Expand Down Expand Up @@ -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")

Expand Down

0 comments on commit 3c18e47

Please sign in to comment.