From 3aaf32dc2eb622a5555e92d1b23cc5124e96cfb1 Mon Sep 17 00:00:00 2001
From: Joost Jager <joost.jager@gmail.com>
Date: Sun, 1 Dec 2019 14:55:01 +0100
Subject: [PATCH] routing: improve equal cost route comparison

When the (virtual) payment attempt cost is set to zero, probabilities
are no longer a factor in determining the best route. In case of routes
with equal costs, we'd just go with the first one found. This commit
refines this behavior by picking the route with the highest probability.
So even though probability doesn't affect the route cost, it is still
used as a tie breaker.
---
 routing/heap.go          |  5 +++
 routing/pathfind.go      | 24 +++++++++++----
 routing/pathfind_test.go | 66 ++++++++++++++++++++++++++++++++++++++++
 3 files changed, 89 insertions(+), 6 deletions(-)

diff --git a/routing/heap.go b/routing/heap.go
index 62b3e3dd..39a6f4ed 100644
--- a/routing/heap.go
+++ b/routing/heap.go
@@ -73,6 +73,11 @@ func (d *distanceHeap) Len() int { return len(d.nodes) }
 //
 // NOTE: This is part of the heap.Interface implementation.
 func (d *distanceHeap) Less(i, j int) bool {
+	// If distances are equal, tie break on probability.
+	if d.nodes[i].dist == d.nodes[j].dist {
+		return d.nodes[i].probability > d.nodes[j].probability
+	}
+
 	return d.nodes[i].dist < d.nodes[j].dist
 }
 
diff --git a/routing/pathfind.go b/routing/pathfind.go
index dee06b0c..6d7c7f94 100644
--- a/routing/pathfind.go
+++ b/routing/pathfind.go
@@ -493,13 +493,25 @@ func findPath(g *graphParams, r *RestrictParams, cfg *PathFindingConfig,
 			int64(cfg.PaymentAttemptPenalty),
 		)
 
-		// If the current best route is better than this candidate
-		// route, return. It is important to also return if the distance
-		// is equal, because otherwise the algorithm could run into an
-		// endless loop.
+		// If there is already a best route stored, compare this
+		// candidate route with the best route so far.
 		current, ok := distance[fromVertex]
-		if ok && tempDist >= current.dist {
-			return
+		if ok {
+			// If this route is worse than what we already found,
+			// skip this route.
+			if tempDist > current.dist {
+				return
+			}
+
+			// If the route is equally good and the probability
+			// isn't better, skip this route. It is important to
+			// also return if both cost and probability are equal,
+			// because otherwise the algorithm could run into an
+			// endless loop.
+			probNotBetter := probability <= current.probability
+			if tempDist == current.dist && probNotBetter {
+				return
+			}
 		}
 
 		// Every edge should have a positive time lock delta. If we
diff --git a/routing/pathfind_test.go b/routing/pathfind_test.go
index 6cde1600..eb7f746a 100644
--- a/routing/pathfind_test.go
+++ b/routing/pathfind_test.go
@@ -2154,6 +2154,72 @@ func testProbabilityRouting(t *testing.T, p10, p11, p20, minProbability float64,
 	}
 }
 
+// TestEqualCostRouteSelection asserts that route probability will be used as a
+// tie breaker in case the path finding probabilities are equal.
+func TestEqualCostRouteSelection(t *testing.T) {
+	t.Parallel()
+
+	// Set up a test graph with two possible paths to the target: via a and
+	// via b. The routing fees and probabilities are chosen such that the
+	// algorithm will first explore target->a->source (backwards search).
+	// This route has fee 6 and a penality of 4 for the 25% success
+	// probability. The algorithm will then proceed with evaluating
+	// target->b->source, which has a fee of 8 and a penalty of 2 for the
+	// 50% success probability. Both routes have the same path finding cost
+	// of 10. It is expected that in that case, the highest probability
+	// route (through b) is chosen.
+	testChannels := []*testChannel{
+		symmetricTestChannel("source", "a", 100000, &testChannelPolicy{}),
+		symmetricTestChannel("source", "b", 100000, &testChannelPolicy{}),
+		symmetricTestChannel("a", "target", 100000, &testChannelPolicy{
+			Expiry:      144,
+			FeeBaseMsat: lnwire.NewMSatFromSatoshis(6),
+			MinHTLC:     1,
+		}, 1),
+		symmetricTestChannel("b", "target", 100000, &testChannelPolicy{
+			Expiry:      100,
+			FeeBaseMsat: lnwire.NewMSatFromSatoshis(8),
+			MinHTLC:     1,
+		}, 2),
+	}
+
+	ctx := newPathFindingTestContext(t, testChannels, "source")
+	defer ctx.cleanup()
+
+	alias := ctx.testGraphInstance.aliasMap
+
+	paymentAmt := lnwire.NewMSatFromSatoshis(100)
+	target := ctx.testGraphInstance.aliasMap["target"]
+
+	ctx.restrictParams.ProbabilitySource = func(fromNode, toNode route.Vertex,
+		amt lnwire.MilliSatoshi) float64 {
+
+		switch {
+		case fromNode == alias["source"] && toNode == alias["a"]:
+			return 0.25
+		case fromNode == alias["source"] && toNode == alias["b"]:
+			return 0.5
+		default:
+			return 1
+		}
+	}
+
+	ctx.pathFindingConfig = PathFindingConfig{
+		PaymentAttemptPenalty: lnwire.NewMSatFromSatoshis(1),
+	}
+
+	path, err := ctx.findPath(target, paymentAmt)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	if path[1].ChannelID != 2 {
+		t.Fatalf("expected route to pass through channel %v, "+
+			"but channel %v was selected instead", 2,
+			path[1].ChannelID)
+	}
+}
+
 // TestNoCycle tries to guide the path finding algorithm into reconstructing an
 // endless route. It asserts that the algorithm is able to handle this properly.
 func TestNoCycle(t *testing.T) {