From 45a2dc80c84b46112f3e2ae6f0815b97bf62df5a Mon Sep 17 00:00:00 2001
From: bitromortac <bitromortac@protonmail.com>
Date: Sat, 4 Dec 2021 12:28:07 +0100
Subject: [PATCH] autopilot: implement cached ChannelGraph

Using the cached implementation of a ChannelGraph greatly speeds up the
conversion of the channel database to a SimpleGraph.
---
 autopilot/graph.go | 90 ++++++++++++++++++++++++++++++++++++++++++++++
 rpcserver.go       |  2 +-
 2 files changed, 91 insertions(+), 1 deletion(-)

diff --git a/autopilot/graph.go b/autopilot/graph.go
index e630f8d35..ced7ed2e9 100644
--- a/autopilot/graph.go
+++ b/autopilot/graph.go
@@ -455,6 +455,96 @@ func (m *memChannelGraph) addRandNode() (*btcec.PublicKey, error) {
 	return newPub, nil
 }
 
+// databaseChannelGraphCached wraps a channeldb.ChannelGraph instance with the
+// necessary API to properly implement the autopilot.ChannelGraph interface.
+type databaseChannelGraphCached struct {
+	db *channeldb.ChannelGraph
+}
+
+// A compile time assertion to ensure databaseChannelGraphCached meets the
+// autopilot.ChannelGraph interface.
+var _ ChannelGraph = (*databaseChannelGraphCached)(nil)
+
+// ChannelGraphFromCachedDatabase returns an instance of the
+// autopilot.ChannelGraph backed by a live, open channeldb instance.
+func ChannelGraphFromCachedDatabase(db *channeldb.ChannelGraph) ChannelGraph {
+	return &databaseChannelGraphCached{
+		db: db,
+	}
+}
+
+// dbNodeCached is a wrapper struct around a database transaction for a
+// channeldb.LightningNode. The wrapper methods implement the autopilot.Node
+// interface.
+type dbNodeCached struct {
+	node     route.Vertex
+	channels map[uint64]*channeldb.DirectedChannel
+}
+
+// A compile time assertion to ensure dbNodeCached meets the autopilot.Node
+// interface.
+var _ Node = (*dbNodeCached)(nil)
+
+// PubKey is the identity public key of the node.
+//
+// NOTE: Part of the autopilot.Node interface.
+func (nc dbNodeCached) PubKey() [33]byte {
+	return nc.node
+}
+
+// Addrs returns a slice of publicly reachable public TCP addresses that the
+// peer is known to be listening on.
+//
+// NOTE: Part of the autopilot.Node interface.
+func (nc dbNodeCached) Addrs() []net.Addr {
+	// TODO: Add addresses to be usable by autopilot.
+	return []net.Addr{}
+}
+
+// ForEachChannel is a higher-order function that will be used to iterate
+// through all edges emanating from/to the target node. For each active
+// channel, this function should be called with the populated ChannelEdge that
+// describes the active channel.
+//
+// NOTE: Part of the autopilot.Node interface.
+func (nc dbNodeCached) ForEachChannel(cb func(ChannelEdge) error) error {
+	for cid, channel := range nc.channels {
+		edge := ChannelEdge{
+			ChanID:   lnwire.NewShortChanIDFromInt(cid),
+			Capacity: channel.Capacity,
+			Peer: dbNodeCached{
+				node: channel.OtherNode,
+			},
+		}
+
+		if err := cb(edge); err != nil {
+			return err
+		}
+	}
+
+	return nil
+}
+
+// ForEachNode is a higher-order function that should be called once for each
+// connected node within the channel graph. If the passed callback returns an
+// error, then execution should be terminated.
+//
+// NOTE: Part of the autopilot.ChannelGraph interface.
+func (dc *databaseChannelGraphCached) ForEachNode(cb func(Node) error) error {
+	return dc.db.ForEachNodeCached(func(n route.Vertex,
+		channels map[uint64]*channeldb.DirectedChannel) error {
+
+		if len(channels) > 0 {
+			node := dbNodeCached{
+				node:     n,
+				channels: channels,
+			}
+			return cb(node)
+		}
+		return nil
+	})
+}
+
 // memNode is a purely in-memory implementation of the autopilot.Node
 // interface.
 type memNode struct {
diff --git a/rpcserver.go b/rpcserver.go
index 26b1d12ef..15f2fc04e 100644
--- a/rpcserver.go
+++ b/rpcserver.go
@@ -5905,7 +5905,7 @@ func (r *rpcServer) GetNetworkInfo(ctx context.Context,
 	}
 
 	// Graph diameter.
-	channelGraph := autopilot.ChannelGraphFromDatabase(graph)
+	channelGraph := autopilot.ChannelGraphFromCachedDatabase(graph)
 	simpleGraph, err := autopilot.NewSimpleGraph(channelGraph)
 	if err != nil {
 		return nil, err