Implement Heap Map (#2137)

Signed-off-by: Joshua Kim <[email protected]> Co-authored-by: Alberto Benegiamo <[email protected]> Co-authored-by: Stephen Buttolph <[email protected]>
ava-labs · Oct 20, 2023 · 7ed450b · 7ed450b
1 parent a21d0cf
commit 7ed450b
Show file tree

Hide file tree

Showing 16 changed files with 451 additions and 713 deletions.
diff --git a/snow/engine/avalanche/bootstrap/bootstrapper.go b/snow/engine/avalanche/bootstrap/bootstrapper.go
@@ -15,8 +15,8 @@ import (
 	"github.com/ava-labs/avalanchego/snow"
 	"github.com/ava-labs/avalanchego/snow/choices"
 	"github.com/ava-labs/avalanchego/snow/consensus/avalanche"
-	"github.com/ava-labs/avalanchego/snow/engine/avalanche/vertex"
 	"github.com/ava-labs/avalanchego/snow/engine/common"
+	"github.com/ava-labs/avalanchego/utils/heap"
 	"github.com/ava-labs/avalanchego/utils/logging"
 	"github.com/ava-labs/avalanchego/utils/set"
 	"github.com/ava-labs/avalanchego/version"
@@ -417,14 +417,15 @@ func (b *bootstrapper) fetch(ctx context.Context, vtxIDs ...ids.ID) error {
 
 // Process the vertices in [vtxs].
 func (b *bootstrapper) process(ctx context.Context, vtxs ...avalanche.Vertex) error {
-	// Vertices that we need to process. Store them in a heap for deduplication
-	// and so we always process vertices further down in the DAG first. This helps
-	// to reduce the number of repeated DAG traversals.
-	toProcess := vertex.NewHeap()
+	// Vertices that we need to process prioritized by vertices that are unknown
+	// or the furthest down the DAG. Unknown vertices are prioritized to ensure
+	// that once we have made it below a certain height in DAG traversal we do
+	// not need to reset and repeat DAG traversals.
+	toProcess := heap.NewMap[ids.ID, avalanche.Vertex](vertexLess)
 	for _, vtx := range vtxs {
 		vtxID := vtx.ID()
 		if _, ok := b.processedCache.Get(vtxID); !ok { // only process a vertex if we haven't already
-			toProcess.Push(vtx)
+			_, _ = toProcess.Push(vtxID, vtx)
 		} else {
 			b.VtxBlocked.RemoveMissingID(vtxID)
 		}
@@ -433,13 +434,15 @@ func (b *bootstrapper) process(ctx context.Context, vtxs ...avalanche.Vertex) er
 	vtxHeightSet := set.Set[ids.ID]{}
 	prevHeight := uint64(0)
 
-	for toProcess.Len() > 0 { // While there are unprocessed vertices
+	for {
 		if b.Halted() {
 			return nil
 		}
 
-		vtx := toProcess.Pop() // Get an unknown vertex or one furthest down the DAG
-		vtxID := vtx.ID()
+		vtxID, vtx, ok := toProcess.Pop()
+		if !ok {
+			break
+		}
 
 		switch vtx.Status() {
 		case choices.Unknown:
@@ -511,7 +514,7 @@ func (b *bootstrapper) process(ctx context.Context, vtxs ...avalanche.Vertex) er
 				parentID := parent.ID()
 				if _, ok := b.processedCache.Get(parentID); !ok { // But only if we haven't processed the parent
 					if !vtxHeightSet.Contains(parentID) {
-						toProcess.Push(parent)
+						toProcess.Push(parentID, parent)
 					}
 				}
 			}
@@ -633,3 +636,26 @@ func (b *bootstrapper) checkFinish(ctx context.Context) error {
 	b.processedCache.Flush()
 	return b.OnFinished(ctx, b.Config.SharedCfg.RequestID)
 }
+
+// A vertex is less than another vertex if it is unknown. Ties are broken by
+// prioritizing vertices that have a greater height.
+func vertexLess(i, j avalanche.Vertex) bool {
+	if !i.Status().Fetched() {
+		return true
+	}
+	if !j.Status().Fetched() {
+		return false
+	}
+
+	// Treat errors on retrieving the height as if the vertex is not fetched
+	heightI, errI := i.Height()
+	if errI != nil {
+		return true
+	}
+	heightJ, errJ := j.Height()
+	if errJ != nil {
+		return false
+	}
+
+	return heightI > heightJ
+}
diff --git a/snow/engine/avalanche/vertex/heap.go b/snow/engine/avalanche/vertex/heap.go
diff --git a/snow/engine/avalanche/vertex/heap_test.go b/snow/engine/avalanche/vertex/heap_test.go