routing: add index map to distanceHeap

This commit adds the pubkeyIndices map to the distanceHeap to avoid
duplicate entries on the heap. This happened in the earlier iteration
of the findPath algorithm and would cause the driving loop to
evaluate already evaluated entries when there was no need.

routing/heap.go

@@ -1,6 +1,8 @@
 package routing
 
 import (
+	"container/heap"
+
 	"github.com/lightningnetwork/lnd/channeldb"
 	"github.com/lightningnetwork/lnd/lnwire"
 	"github.com/lightningnetwork/lnd/routing/route"
@@ -40,6 +42,21 @@ type nodeWithDist struct {
 // algorithm to keep track of the "closest" node to our source node.
 type distanceHeap struct {
 	nodes []nodeWithDist
+
+	// pubkeyIndices maps public keys of nodes to their respective index in
+	// the heap. This is used as a way to avoid db lookups by using heap.Fix
+	// instead of having duplicate entries on the heap.
+	pubkeyIndices map[route.Vertex]int
+}
+
+// newDistanceHeap initializes a new distance heap. This is required because
+// we must initialize the pubkeyIndices map for path-finding optimizations.
+func newDistanceHeap() distanceHeap {
+	distHeap := distanceHeap{
+		pubkeyIndices: make(map[route.Vertex]int),
+	}
+
+	return distHeap
 }
 
 // Len returns the number of nodes in the priority queue.
@@ -60,13 +77,17 @@ func (d *distanceHeap) Less(i, j int) bool {
 // NOTE: This is part of the heap.Interface implementation.
 func (d *distanceHeap) Swap(i, j int) {
 	d.nodes[i], d.nodes[j] = d.nodes[j], d.nodes[i]
+	d.pubkeyIndices[d.nodes[i].node] = i
+	d.pubkeyIndices[d.nodes[j].node] = j
 }
 
 // Push pushes the passed item onto the priority queue.
 //
 // NOTE: This is part of the heap.Interface implementation.
 func (d *distanceHeap) Push(x interface{}) {
-	d.nodes = append(d.nodes, x.(nodeWithDist))
+	n := x.(nodeWithDist)
+	d.nodes = append(d.nodes, n)
+	d.pubkeyIndices[n.node] = len(d.nodes) - 1
 }
 
 // Pop removes the highest priority item (according to Less) from the priority
@@ -77,9 +98,29 @@ func (d *distanceHeap) Pop() interface{} {
 	n := len(d.nodes)
 	x := d.nodes[n-1]
 	d.nodes = d.nodes[0 : n-1]
+	delete(d.pubkeyIndices, x.node)
 	return x
 }
 
+// PushOrFix attempts to adjust the position of a given node in the heap.
+// If the vertex already exists in the heap, then we must call heap.Fix to
+// modify its position and reorder the heap. If the vertex does not already
+// exist in the heap, then it is pushed onto the heap. Otherwise, we will end
+// up performing more db lookups on the same node in the pathfinding algorithm.
+func (d *distanceHeap) PushOrFix(dist nodeWithDist) {
+	index, ok := d.pubkeyIndices[dist.node]
+	if !ok {
+		heap.Push(d, dist)
+		return
+	}
+
+	// Change the value at the specified index.
+	d.nodes[index] = dist
+
+	// Call heap.Fix to reorder the heap.
+	heap.Fix(d, index)
+}
+
 // path represents an ordered set of edges which forms an available path from a
 // given source node to our destination. During the process of computing the
 // KSP's from a source to destination, several path swill be considered in the

routing/heap_test.go

@@ -6,7 +6,8 @@ import (
 	"reflect"
 	"sort"
 	"testing"
-	"time"
+
+	"github.com/lightningnetwork/lnd/routing/route"
 )
 
 // TestHeapOrdering ensures that the items inserted into the heap are properly
@@ -16,20 +17,33 @@ func TestHeapOrdering(t *testing.T) {
 
 	// First, create a blank heap, we'll use this to push on randomly
 	// generated items.
-	var nodeHeap distanceHeap
+	nodeHeap := newDistanceHeap()
 
-	prand.Seed(time.Now().Unix())
+	prand.Seed(1)
 
 	// Create 100 random entries adding them to the heap created above, but
 	// also a list that we'll sort with the entries.
 	const numEntries = 100
 	sortedEntries := make([]nodeWithDist, 0, numEntries)
 	for i := 0; i < numEntries; i++ {
+		var pubKey [33]byte
+		prand.Read(pubKey[:])
+
 		entry := nodeWithDist{
+			node: route.Vertex(pubKey),
 			dist: prand.Int63(),
 		}
 
-		heap.Push(&nodeHeap, entry)
+		// Use the PushOrFix method for the initial push to test the scenario
+		// where entry doesn't exist on the heap.
+		nodeHeap.PushOrFix(entry)
+
+		// Re-generate this entry's dist field
+		entry.dist = prand.Int63()
+
+		// Reorder the heap with a PushOrFix call.
+		nodeHeap.PushOrFix(entry)
+
 		sortedEntries = append(sortedEntries, entry)
 	}
 
@@ -47,6 +61,13 @@ func TestHeapOrdering(t *testing.T) {
 		poppedEntries = append(poppedEntries, e)
 	}
 
+	// Assert that the pubkeyIndices map is empty after popping all of the
+	// items off of it.
+	if len(nodeHeap.pubkeyIndices) != 0 {
+		t.Fatalf("there are still %d pubkeys in the pubkeyIndices map",
+			len(nodeHeap.pubkeyIndices))
+	}
+
 	// Finally, ensure that the items popped from the heap and the items we
 	// sorted are identical at this rate.
 	if !reflect.DeepEqual(poppedEntries, sortedEntries) {

routing/pathfind.go

@@ -300,7 +300,7 @@ func findPath(g *graphParams, r *RestrictParams, source, target route.Vertex,
 	// First we'll initialize an empty heap which'll help us to quickly
 	// locate the next edge we should visit next during our graph
 	// traversal.
-	var nodeHeap distanceHeap
+	nodeHeap := newDistanceHeap()
 
 	// For each node in the graph, we create an entry in the distance map
 	// for the node set with a distance of "infinity". graph.ForEachNode
@@ -527,9 +527,10 @@ func findPath(g *graphParams, r *RestrictParams, source, target route.Vertex,
 
 		next[fromVertex] = edge
 
-		// Add this new node to our heap as we'd like to further
-		// explore backwards through this edge.
-		heap.Push(&nodeHeap, distance[fromVertex])
+		// Either push distance[fromVertex] onto the heap if the node
+		// represented by fromVertex is not already on the heap OR adjust
+		// its position within the heap via heap.Fix.
+		nodeHeap.PushOrFix(distance[fromVertex])
 	}
 
 	// TODO(roasbeef): also add path caching