[NOD-1008] In utxoDiffStore, keep diffData in memory for blocks whose blueScore is at least virtualBlueScore - X (#726)

* [NOD-1008] Use *blockNode as keys in utxoDiffStore.loaded and .dirty.

* [NOD-1008] Implement clearOldEntries.

* [NOD-1008] Increase maxBlueScoreDifferenceToKeepLoaded to 100.

* [NOD-1008] Fix a typo.

* [NOD-1008] Add clearOldEntries to saveChangesFromBlock.

* [NOD-1008] Begin implementing TestClearOldEntries.

* [NOD-1008] Finish implementing TestClearOldEntries.

* [NOD-1008] Fix a comment.

* [NOD-1008] Rename diffDataByHash to diffDataByBlockNode.

* [NOD-1008] Use dag.TipHashes instead of tracking tips manually.

blockdag/dag.go

@@ -804,6 +804,7 @@ func (dag *BlockDAG) saveChangesFromBlock(block *util.Block, virtualUTXODiff *UT
 
 	dag.index.clearDirtyEntries()
 	dag.utxoDiffStore.clearDirtyEntries()
+	dag.utxoDiffStore.clearOldEntries()
 	dag.reachabilityStore.clearDirtyEntries()
 	dag.multisetStore.clearNewEntries()
 
@@ -910,9 +911,9 @@ func (dag *BlockDAG) finalizeNodesBelowFinalityPoint(deleteDiffData bool) {
 	for parent := range dag.lastFinalityPoint.parents {
 		queue = append(queue, parent)
 	}
-	var blockHashesToDelete []*daghash.Hash
+	var nodesToDelete []*blockNode
 	if deleteDiffData {
-		blockHashesToDelete = make([]*daghash.Hash, 0, dag.dagParams.FinalityInterval)
+		nodesToDelete = make([]*blockNode, 0, dag.dagParams.FinalityInterval)
 	}
 	for len(queue) > 0 {
 		var current *blockNode
@@ -920,7 +921,7 @@ func (dag *BlockDAG) finalizeNodesBelowFinalityPoint(deleteDiffData bool) {
 		if !current.isFinalized {
 			current.isFinalized = true
 			if deleteDiffData {
-				blockHashesToDelete = append(blockHashesToDelete, current.hash)
+				nodesToDelete = append(nodesToDelete, current)
 			}
 			for parent := range current.parents {
 				queue = append(queue, parent)
@@ -928,7 +929,7 @@ func (dag *BlockDAG) finalizeNodesBelowFinalityPoint(deleteDiffData bool) {
 		}
 	}
 	if deleteDiffData {
-		err := dag.utxoDiffStore.removeBlocksDiffData(dbaccess.NoTx(), blockHashesToDelete)
+		err := dag.utxoDiffStore.removeBlocksDiffData(dbaccess.NoTx(), nodesToDelete)
 		if err != nil {
 			panic(fmt.Sprintf("Error removing diff data from utxoDiffStore: %s", err))
 		}

blockdag/dag_test.go

@@ -8,6 +8,7 @@ import (
 	"fmt"
 	"github.com/kaspanet/kaspad/dbaccess"
 	"github.com/pkg/errors"
+	"math"
 	"os"
 	"path/filepath"
 	"testing"
@@ -953,6 +954,11 @@ func testFinalizeNodesBelowFinalityPoint(t *testing.T, deleteDiffData bool) {
 	// Manually set the last finality point
 	dag.lastFinalityPoint = nodes[finalityInterval-1]
 
+	// Don't unload diffData
+	currentDifference := maxBlueScoreDifferenceToKeepLoaded
+	maxBlueScoreDifferenceToKeepLoaded = math.MaxUint64
+	defer func() { maxBlueScoreDifferenceToKeepLoaded = currentDifference }()
+
 	dag.finalizeNodesBelowFinalityPoint(deleteDiffData)
 	flushUTXODiffStore()
 
@@ -960,7 +966,7 @@ func testFinalizeNodesBelowFinalityPoint(t *testing.T, deleteDiffData bool) {
 		if !node.isFinalized {
 			t.Errorf("Node with blue score %d expected to be finalized", node.blueScore)
 		}
-		if _, ok := dag.utxoDiffStore.loaded[*node.hash]; deleteDiffData && ok {
+		if _, ok := dag.utxoDiffStore.loaded[node]; deleteDiffData && ok {
 			t.Errorf("The diff data of node with blue score %d should have been unloaded if deleteDiffData is %T", node.blueScore, deleteDiffData)
 		} else if !deleteDiffData && !ok {
 			t.Errorf("The diff data of node with blue score %d shouldn't have been unloaded if deleteDiffData is %T", node.blueScore, deleteDiffData)
@@ -988,7 +994,7 @@ func testFinalizeNodesBelowFinalityPoint(t *testing.T, deleteDiffData bool) {
 		if node.isFinalized {
 			t.Errorf("Node with blue score %d wasn't expected to be finalized", node.blueScore)
 		}
-		if _, ok := dag.utxoDiffStore.loaded[*node.hash]; !ok {
+		if _, ok := dag.utxoDiffStore.loaded[node]; !ok {
 			t.Errorf("The diff data of node with blue score %d shouldn't have been unloaded", node.blueScore)
 		}
 		if diffData, err := dag.utxoDiffStore.diffDataFromDB(node.hash); err != nil {

blockdag/utxodiffstore.go

@@ -14,16 +14,16 @@ type blockUTXODiffData struct {
 
 type utxoDiffStore struct {
 	dag    *BlockDAG
-	dirty  map[daghash.Hash]struct{}
+	dirty  map[*blockNode]struct{}
-	loaded map[daghash.Hash]*blockUTXODiffData
+	loaded map[*blockNode]*blockUTXODiffData
 	mtx    *locks.PriorityMutex
 }
 
 func newUTXODiffStore(dag *BlockDAG) *utxoDiffStore {
 	return &utxoDiffStore{
 		dag:    dag,
-		dirty:  make(map[daghash.Hash]struct{}),
+		dirty:  make(map[*blockNode]struct{}),
-		loaded: make(map[daghash.Hash]*blockUTXODiffData),
+		loaded: make(map[*blockNode]*blockUTXODiffData),
 		mtx:    locks.NewPriorityMutex(),
 	}
 }
@@ -32,15 +32,15 @@ func (diffStore *utxoDiffStore) setBlockDiff(node *blockNode, diff *UTXODiff) er
 	diffStore.mtx.HighPriorityWriteLock()
 	defer diffStore.mtx.HighPriorityWriteUnlock()
 	// load the diff data from DB to diffStore.loaded
-	_, err := diffStore.diffDataByHash(node.hash)
+	_, err := diffStore.diffDataByBlockNode(node)
 	if dbaccess.IsNotFoundError(err) {
-		diffStore.loaded[*node.hash] = &blockUTXODiffData{}
+		diffStore.loaded[node] = &blockUTXODiffData{}
 	} else if err != nil {
 		return err
 	}
 
-	diffStore.loaded[*node.hash].diff = diff
+	diffStore.loaded[node].diff = diff
-	diffStore.setBlockAsDirty(node.hash)
+	diffStore.setBlockAsDirty(node)
 	return nil
 }
 
@@ -48,19 +48,19 @@ func (diffStore *utxoDiffStore) setBlockDiffChild(node *blockNode, diffChild *bl
 	diffStore.mtx.HighPriorityWriteLock()
 	defer diffStore.mtx.HighPriorityWriteUnlock()
 	// load the diff data from DB to diffStore.loaded
-	_, err := diffStore.diffDataByHash(node.hash)
+	_, err := diffStore.diffDataByBlockNode(node)
 	if err != nil {
 		return err
 	}
 
-	diffStore.loaded[*node.hash].diffChild = diffChild
+	diffStore.loaded[node].diffChild = diffChild
-	diffStore.setBlockAsDirty(node.hash)
+	diffStore.setBlockAsDirty(node)
 	return nil
 }
 
-func (diffStore *utxoDiffStore) removeBlocksDiffData(dbContext dbaccess.Context, blockHashes []*daghash.Hash) error {
+func (diffStore *utxoDiffStore) removeBlocksDiffData(dbContext dbaccess.Context, nodes []*blockNode) error {
-	for _, hash := range blockHashes {
+	for _, node := range nodes {
-		err := diffStore.removeBlockDiffData(dbContext, hash)
+		err := diffStore.removeBlockDiffData(dbContext, node)
 		if err != nil {
 			return err
 		}
@@ -68,37 +68,37 @@ func (diffStore *utxoDiffStore) removeBlocksDiffData(dbContext dbaccess.Context,
 	return nil
 }
 
-func (diffStore *utxoDiffStore) removeBlockDiffData(dbContext dbaccess.Context, blockHash *daghash.Hash) error {
+func (diffStore *utxoDiffStore) removeBlockDiffData(dbContext dbaccess.Context, node *blockNode) error {
 	diffStore.mtx.LowPriorityWriteLock()
 	defer diffStore.mtx.LowPriorityWriteUnlock()
-	delete(diffStore.loaded, *blockHash)
+	delete(diffStore.loaded, node)
-	err := dbaccess.RemoveDiffData(dbContext, blockHash)
+	err := dbaccess.RemoveDiffData(dbContext, node.hash)
 	if err != nil {
 		return err
 	}
 	return nil
 }
 
-func (diffStore *utxoDiffStore) setBlockAsDirty(blockHash *daghash.Hash) {
+func (diffStore *utxoDiffStore) setBlockAsDirty(node *blockNode) {
-	diffStore.dirty[*blockHash] = struct{}{}
+	diffStore.dirty[node] = struct{}{}
 }
 
-func (diffStore *utxoDiffStore) diffDataByHash(hash *daghash.Hash) (*blockUTXODiffData, error) {
+func (diffStore *utxoDiffStore) diffDataByBlockNode(node *blockNode) (*blockUTXODiffData, error) {
-	if diffData, ok := diffStore.loaded[*hash]; ok {
+	if diffData, ok := diffStore.loaded[node]; ok {
 		return diffData, nil
 	}
-	diffData, err := diffStore.diffDataFromDB(hash)
+	diffData, err := diffStore.diffDataFromDB(node.hash)
 	if err != nil {
 		return nil, err
 	}
-	diffStore.loaded[*hash] = diffData
+	diffStore.loaded[node] = diffData
 	return diffData, nil
 }
 
 func (diffStore *utxoDiffStore) diffByNode(node *blockNode) (*UTXODiff, error) {
 	diffStore.mtx.HighPriorityReadLock()
 	defer diffStore.mtx.HighPriorityReadUnlock()
-	diffData, err := diffStore.diffDataByHash(node.hash)
+	diffData, err := diffStore.diffDataByBlockNode(node)
 	if err != nil {
 		return nil, err
 	}
@@ -108,7 +108,7 @@ func (diffStore *utxoDiffStore) diffByNode(node *blockNode) (*UTXODiff, error) {
 func (diffStore *utxoDiffStore) diffChildByNode(node *blockNode) (*blockNode, error) {
 	diffStore.mtx.HighPriorityReadLock()
 	defer diffStore.mtx.HighPriorityReadUnlock()
-	diffData, err := diffStore.diffDataByHash(node.hash)
+	diffData, err := diffStore.diffDataByBlockNode(node)
 	if err != nil {
 		return nil, err
 	}
@@ -135,11 +135,10 @@ func (diffStore *utxoDiffStore) flushToDB(dbContext *dbaccess.TxContext) error {
 	// Allocate a buffer here to avoid needless allocations/grows
 	// while writing each entry.
 	buffer := &bytes.Buffer{}
-	for hash := range diffStore.dirty {
+	for node := range diffStore.dirty {
-		hash := hash // Copy hash to a new variable to avoid passing the same pointer
 		buffer.Reset()
-		diffData := diffStore.loaded[hash]
+		diffData := diffStore.loaded[node]
-		err := storeDiffData(dbContext, buffer, &hash, diffData)
+		err := storeDiffData(dbContext, buffer, node.hash, diffData)
 		if err != nil {
 			return err
 		}
@@ -148,7 +147,32 @@ func (diffStore *utxoDiffStore) flushToDB(dbContext *dbaccess.TxContext) error {
 }
 
 func (diffStore *utxoDiffStore) clearDirtyEntries() {
-	diffStore.dirty = make(map[daghash.Hash]struct{})
+	diffStore.dirty = make(map[*blockNode]struct{})
+}
+
+// maxBlueScoreDifferenceToKeepLoaded is the maximum difference
+// between the virtual's blueScore and a blockNode's blueScore
+// under which to keep diff data loaded in memory.
+var maxBlueScoreDifferenceToKeepLoaded uint64 = 100
+
+// clearOldEntries removes entries whose blue score is lower than
+// virtual.blueScore - maxBlueScoreDifferenceToKeepLoaded.
+func (diffStore *utxoDiffStore) clearOldEntries() {
+	virtualBlueScore := diffStore.dag.VirtualBlueScore()
+	minBlueScore := virtualBlueScore - maxBlueScoreDifferenceToKeepLoaded
+	if maxBlueScoreDifferenceToKeepLoaded > virtualBlueScore {
+		minBlueScore = 0
+	}
+
+	toRemove := make(map[*blockNode]struct{})
+	for node := range diffStore.loaded {
+		if node.blueScore < minBlueScore {
+			toRemove[node] = struct{}{}
+		}
+	}
+	for node := range toRemove {
+		delete(diffStore.loaded, node)
+	}
 }
 
 // storeDiffData stores the UTXO diff data to the database.
@@ -156,7 +180,7 @@ func (diffStore *utxoDiffStore) clearDirtyEntries() {
 func storeDiffData(dbContext dbaccess.Context, w *bytes.Buffer, hash *daghash.Hash, diffData *blockUTXODiffData) error {
 	// To avoid a ton of allocs, use the io.Writer
 	// instead of allocating one. We expect the buffer to
-	// already be initalized and, in most cases, to already
+	// already be initialized and, in most cases, to already
 	// be large enough to accommodate the serialized data
 	// without growing.
 	err := serializeBlockUTXODiffData(w, diffData)

blockdag/utxodiffstore_test.go

@@ -78,7 +78,7 @@ func TestUTXODiffStore(t *testing.T) {
 	if err != nil {
 		t.Fatalf("Failed to commit database transaction: %s", err)
 	}
-	delete(dag.utxoDiffStore.loaded, *node.hash)
+	delete(dag.utxoDiffStore.loaded, node)
 
 	if storeDiff, err := dag.utxoDiffStore.diffByNode(node); err != nil {
 		t.Fatalf("diffByNode: unexpected error: %s", err)
@@ -87,9 +87,79 @@ func TestUTXODiffStore(t *testing.T) {
 	}
 
 	// Check if getBlockDiff caches the result in dag.utxoDiffStore.loaded
-	if loadedDiffData, ok := dag.utxoDiffStore.loaded[*node.hash]; !ok {
+	if loadedDiffData, ok := dag.utxoDiffStore.loaded[node]; !ok {
 		t.Errorf("the diff data wasn't added to loaded map after requesting it")
 	} else if !reflect.DeepEqual(loadedDiffData.diff, diff) {
 		t.Errorf("Expected diff and loadedDiff to be equal")
 	}
 }
+
+func TestClearOldEntries(t *testing.T) {
+	// Create a new database and DAG instance to run tests against.
+	dag, teardownFunc, err := DAGSetup("TestClearOldEntries", true, Config{
+		DAGParams: &dagconfig.SimnetParams,
+	})
+	if err != nil {
+		t.Fatalf("TestClearOldEntries: Failed to setup DAG instance: %v", err)
+	}
+	defer teardownFunc()
+
+	// Set maxBlueScoreDifferenceToKeepLoaded to 10 to make this test fast to run
+	currentDifference := maxBlueScoreDifferenceToKeepLoaded
+	maxBlueScoreDifferenceToKeepLoaded = 10
+	defer func() { maxBlueScoreDifferenceToKeepLoaded = currentDifference }()
+
+	// Add 10 blocks
+	blockNodes := make([]*blockNode, 10)
+	for i := 0; i < 10; i++ {
+		processedBlock := PrepareAndProcessBlockForTest(t, dag, dag.TipHashes(), nil)
+
+		node := dag.index.LookupNode(processedBlock.BlockHash())
+		if node == nil {
+			t.Fatalf("TestClearOldEntries: missing blockNode for hash %s", processedBlock.BlockHash())
+		}
+		blockNodes[i] = node
+	}
+
+	// Make sure that all of them exist in the loaded set
+	for _, node := range blockNodes {
+		_, ok := dag.utxoDiffStore.loaded[node]
+		if !ok {
+			t.Fatalf("TestClearOldEntries: diffData for node %s is not in the loaded set", node.hash)
+		}
+	}
+
+	// Add 10 more blocks on top of the others
+	for i := 0; i < 10; i++ {
+		PrepareAndProcessBlockForTest(t, dag, dag.TipHashes(), nil)
+	}
+
+	// Make sure that all the old nodes no longer exist in the loaded set
+	for _, node := range blockNodes {
+		_, ok := dag.utxoDiffStore.loaded[node]
+		if ok {
+			t.Fatalf("TestClearOldEntries: diffData for node %s is in the loaded set", node.hash)
+		}
+	}
+
+	// Add a block on top of the genesis to force the retrieval of all diffData
+	processedBlock := PrepareAndProcessBlockForTest(t, dag, []*daghash.Hash{dag.genesis.hash}, nil)
+	node := dag.index.LookupNode(processedBlock.BlockHash())
+	if node == nil {
+		t.Fatalf("TestClearOldEntries: missing blockNode for hash %s", processedBlock.BlockHash())
+	}
+
+	// Make sure that the child-of-genesis node isn't in the loaded set
+	_, ok := dag.utxoDiffStore.loaded[node]
+	if ok {
+		t.Fatalf("TestClearOldEntries: diffData for node %s is in the loaded set", node.hash)
+	}
+
+	// Make sure that all the old nodes still do not exist in the loaded set
+	for _, node := range blockNodes {
+		_, ok := dag.utxoDiffStore.loaded[node]
+		if ok {
+			t.Fatalf("TestClearOldEntries: diffData for node %s is in the loaded set", node.hash)
+		}
+	}
+}