From 6e90770991d7130c67bd36f9b71e6dfac0f54685 Mon Sep 17 00:00:00 2001
From: rkfg <rkfg@rkfg.me>
Date: Sat, 27 Aug 2022 17:40:27 +0300
Subject: [PATCH] Optimize channel pair selection

---
 channels.go | 48 ++++++++++++++++++++++++++++++++++++------------
 main.go     | 16 ++++++++++------
 routes.go   | 15 ---------------
 3 files changed, 46 insertions(+), 33 deletions(-)

diff --git a/channels.go b/channels.go
index 531b279..7383bf5 100644
--- a/channels.go
+++ b/channels.go
@@ -2,6 +2,8 @@ package main
 
 import (
 	"context"
+	"errors"
+	"fmt"
 	"math"
 	"math/rand"
 	"time"
@@ -9,6 +11,10 @@ import (
 	"github.com/lightningnetwork/lnd/lnrpc"
 )
 
+func formatChannelPair(a, b uint64) string {
+	return fmt.Sprintf("%d-%d", a, b)
+}
+
 func (r *regolancer) getChannels(ctx context.Context) error {
 	ctx, cancel := context.WithTimeout(ctx, time.Second*30)
 	defer cancel()
@@ -50,6 +56,12 @@ func (r *regolancer) getChannelCandidates(fromPerc, toPerc, amount int64) error
 			}
 		}
 	}
+	for _, fc := range r.fromChannels {
+		for _, tc := range r.toChannels {
+			pair := [2]*lnrpc.Channel{fc, tc}
+			r.channelPairs[formatChannelPair(pair[0].ChanId, pair[1].ChanId)] = pair
+		}
+	}
 	return nil
 }
 
@@ -63,22 +75,21 @@ func min(args ...int64) (result int64) {
 	return
 }
 
-func (r *regolancer) pickChannelPair(ctx context.Context, amount int64) (from uint64, to uint64, maxAmount int64, err error) {
+func (r *regolancer) pickChannelPair(amount int64) (from uint64, to uint64, maxAmount int64, err error) {
+	if len(r.channelPairs) == 0 {
+		return 0, 0, 0, errors.New("no routes")
+	}
 	var fromChan, toChan *lnrpc.Channel
-	for {
-		select {
-		case <-ctx.Done():
-			return 0, 0, 0, ctx.Err()
-		default:
-		}
-		fromIdx := rand.Int31n(int32(len(r.fromChannels)))
-		toIdx := rand.Int31n(int32(len(r.toChannels)))
-		fromChan = r.fromChannels[fromIdx]
-		toChan = r.toChannels[toIdx]
-		if !r.isFailedRoute(fromChan.ChanId, toChan.ChanId) {
+	idx := rand.Int31n(int32(len(r.channelPairs)))
+	var pair [2]*lnrpc.Channel
+	for _, pair = range r.channelPairs {
+		if idx == 0 {
 			break
 		}
+		idx--
 	}
+	fromChan = pair[0]
+	toChan = pair[1]
 	maxFrom := fromChan.Capacity/2 - fromChan.RemoteBalance
 	maxTo := toChan.Capacity/2 - toChan.LocalBalance
 	if amount == 0 {
@@ -86,5 +97,18 @@ func (r *regolancer) pickChannelPair(ctx context.Context, amount int64) (from ui
 	} else {
 		maxAmount = min(maxFrom, maxTo, amount)
 	}
+	for k, v := range r.failureCache {
+		if v.expiration.Before(time.Now()) {
+			r.channelPairs[k] = v.channelPair
+			delete(r.failureCache, k)
+		}
+	}
 	return fromChan.ChanId, toChan.ChanId, maxAmount, nil
 }
+
+func (r *regolancer) addFailedRoute(from, to uint64) {
+	t := time.Now().Add(time.Hour)
+	k := formatChannelPair(from, to)
+	r.failureCache[k] = failedRoute{channelPair: r.channelPairs[k], expiration: &t}
+	delete(r.channelPairs, k)
+}
diff --git a/main.go b/main.go
index 66efe4f..5092f24 100644
--- a/main.go
+++ b/main.go
@@ -41,6 +41,11 @@ var params struct {
 	StatFilename       string   `short:"s" long:"stat" description:"save successful rebalance information to the specified CSV file" json:"stat"`
 }
 
+type failedRoute struct {
+	channelPair [2]*lnrpc.Channel
+	expiration  *time.Time
+}
+
 type regolancer struct {
 	lnClient      lnrpc.LightningClient
 	routerClient  routerrpc.RouterClient
@@ -50,9 +55,10 @@ type regolancer struct {
 	fromChannelId uint64
 	toChannels    []*lnrpc.Channel
 	toChannelId   uint64
+	channelPairs  map[string][2]*lnrpc.Channel
 	nodeCache     map[string]*lnrpc.NodeInfo
 	chanCache     map[uint64]*lnrpc.ChannelEdge
-	failureCache  map[string]*time.Time
+	failureCache  map[string]failedRoute
 	excludeIn     map[uint64]struct{}
 	excludeOut    map[uint64]struct{}
 	excludeBoth   map[uint64]struct{}
@@ -79,10 +85,7 @@ func loadConfig() {
 
 func tryRebalance(ctx context.Context, r *regolancer, invoice **lnrpc.AddInvoiceResponse,
 	attempt *int) (err error, repeat bool) {
-	// purely local code with no RPC requests, should never take long unless all routes already failed
-	pickCtx, pickCtxCancel := context.WithTimeout(ctx, time.Second*5)
-	defer pickCtxCancel()
-	from, to, amt, err := r.pickChannelPair(pickCtx, params.Amount)
+	from, to, amt, err := r.pickChannelPair(params.Amount)
 	if err != nil {
 		log.Printf(errColor("Error during picking channel: %s"), err)
 		return err, false
@@ -179,7 +182,8 @@ func main() {
 	r := regolancer{
 		nodeCache:    map[string]*lnrpc.NodeInfo{},
 		chanCache:    map[uint64]*lnrpc.ChannelEdge{},
-		failureCache: map[string]*time.Time{},
+		channelPairs: map[string][2]*lnrpc.Channel{},
+		failureCache: map[string]failedRoute{},
 		statFilename: params.StatFilename,
 	}
 	r.lnClient = lnrpc.NewLightningClient(conn)
diff --git a/routes.go b/routes.go
index 61301f8..2c396bd 100644
--- a/routes.go
+++ b/routes.go
@@ -224,21 +224,6 @@ func (r *regolancer) probeRoute(ctx context.Context, route *lnrpc.Route, goodAmo
 	return 0, fmt.Errorf("unknown error: %+v", result)
 }
 
-func (r *regolancer) addFailedRoute(from, to uint64) {
-	t := time.Now().Add(time.Hour)
-	r.failureCache[fmt.Sprintf("%d-%d", from, to)] = &t
-	for k, v := range r.failureCache {
-		if v.Before(time.Now()) {
-			delete(r.failureCache, k)
-		}
-	}
-}
-
-func (r *regolancer) isFailedRoute(from, to uint64) bool {
-	_, ok := r.failureCache[fmt.Sprintf("%d-%d", from, to)]
-	return ok
-}
-
 func (r *regolancer) makeNodeList(nodes []string) error {
 	for _, nid := range nodes {
 		pk, err := hex.DecodeString(nid)