diff --git a/routing/heap.go b/routing/heap.go index 59b26110..294e3306 100644 --- a/routing/heap.go +++ b/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 diff --git a/routing/heap_test.go b/routing/heap_test.go index 2ada4232..4214e965 100644 --- a/routing/heap_test.go +++ b/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) { diff --git a/routing/pathfind.go b/routing/pathfind.go index cc277176..7e951532 100644 --- a/routing/pathfind.go +++ b/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