channeldb: add nwe PruneGraphNodes method to prune unconnected nodes

channeldb/graph.go

@@ -10,12 +10,12 @@ import (
 	"net"
 	"time"
 
-	"github.com/coreos/bbolt"
-	"github.com/lightningnetwork/lnd/lnwire"
 	"github.com/btcsuite/btcd/btcec"
 	"github.com/btcsuite/btcd/chaincfg/chainhash"
 	"github.com/btcsuite/btcd/wire"
 	"github.com/btcsuite/btcutil"
+	"github.com/coreos/bbolt"
+	"github.com/lightningnetwork/lnd/lnwire"
 )
 
 var (
@@ -698,6 +698,41 @@ func (c *ChannelGraph) PruneGraph(spentOutputs []*wire.OutPoint,
 	return chansClosed, nil
 }
 
+// PruneGraphNodes is a garbage collection method which attempts to prune out
+// any nodes from the channel graph that are currently unconnected. This ensure
+// that we only maintain a graph of reachable nodes. In the event that a pruned
+// node gains more channels, it will be re-added back to the graph.
+func (c *ChannelGraph) PruneGraphNodes() error {
+	// We'll use this map to collect a
+	nodesToMaybePrune := make(map[[33]byte]struct{})
+
+	return c.db.Update(func(tx *bolt.Tx) error {
+		nodes, err := tx.CreateBucketIfNotExists(nodeBucket)
+		if err != nil {
+			return err
+		}
+
+		err = nodes.ForEach(func(pubKey, nodeBytes []byte) error {
+			// Skip anything that may actually be a sub-bucket.
+			if len(pubKey) != 33 {
+				return nil
+			}
+
+			var nodePub [33]byte
+			copy(nodePub[:], pubKey)
+
+			nodesToMaybePrune[nodePub] = struct{}{}
+
+			return nil
+		})
+		if err != nil {
+			return err
+		}
+
+		return c.pruneGraphNodes(tx, nodes, nodesToMaybePrune)
+	})
+}
+
 // pruneGraphNodes attempts to remove any nodes from the graph who have had a
 // channel closed within the current block. If the node still has existing
 // channels in the graph, this will act as a no-op.

channeldb/graph_test.go

@@ -14,12 +14,12 @@ import (
 	"testing"
 	"time"
 
-	"github.com/coreos/bbolt"
-	"github.com/davecgh/go-spew/spew"
-	"github.com/lightningnetwork/lnd/lnwire"
 	"github.com/btcsuite/btcd/btcec"
 	"github.com/btcsuite/btcd/chaincfg/chainhash"
 	"github.com/btcsuite/btcd/wire"
+	"github.com/coreos/bbolt"
+	"github.com/davecgh/go-spew/spew"
+	"github.com/lightningnetwork/lnd/lnwire"
 )
 
 var (
@@ -2058,6 +2058,55 @@ func TestChannelEdgePruningUpdateIndexDeletion(t *testing.T) {
 	}
 }
 
+// TestPruneGraphNodes tests that unconnected vertexes are pruned via the
+// PruneSyncState method.
+func TestPruneGraphNodes(t *testing.T) {
+	t.Parallel()
+
+	db, cleanUp, err := makeTestDB()
+	defer cleanUp()
+	if err != nil {
+		t.Fatalf("unable to make test database: %v", err)
+	}
+
+	// We'll start off by inserting our source node, to ensure that it's
+	// the only node left after we prune the graph.
+	graph := db.ChannelGraph()
+	sourceNode, err := createTestVertex(db)
+	if err != nil {
+		t.Fatalf("unable to create source node: %v", err)
+	}
+	if err := graph.SetSourceNode(sourceNode); err != nil {
+		t.Fatalf("unable to set source node: %v", err)
+	}
+
+	// With the source node inserted, we'll now add two nodes into the
+	// channel graph, as they don't have any channels they should be
+	// removed from the graph at the end.
+	node1, err := createTestVertex(db)
+	if err != nil {
+		t.Fatalf("unable to create test node: %v", err)
+	}
+	if err := graph.AddLightningNode(node1); err != nil {
+		t.Fatalf("unable to add node: %v", err)
+	}
+	node2, err := createTestVertex(db)
+	if err != nil {
+		t.Fatalf("unable to create test node: %v", err)
+	}
+	if err := graph.AddLightningNode(node2); err != nil {
+		t.Fatalf("unable to add node: %v", err)
+	}
+
+	if err := graph.PruneGraphNodes(); err != nil {
+		t.Fatalf("unable to prune graph nodes: %v", err)
+	}
+
+	// There should only be a single node left at this point, the source
+	// node.
+	assertNumNodes(t, graph, 1)
+}
+
 // compareNodes is used to compare two LightningNodes while excluding the
 // Features struct, which cannot be compared as the semantics for reserializing
 // the featuresMap have not been defined.