channeldb: prune nodes with no open channels left from the graph

channeldb/graph.go

@@ -588,6 +588,13 @@ func (c *ChannelGraph) PruneGraph(spentOutputs []*wire.OutPoint,
 
 	var chansClosed []*ChannelEdgeInfo
 
+	// nodesWithChansClosed is the set of nodes, each identified by their
+	// compressed public key, who had a channel closed within the latest
+	// block. We'll use this later on to determine whether we should prune
+	// them from the channel graph due to no longer having any other open
+	// channels.
+	nodesWithChansClosed := make(map[[33]byte]struct{})
+
 	err := c.db.Update(func(tx *bolt.Tx) error {
 		// First grab the edges bucket which houses the information
 		// we'd like to delete
@@ -636,7 +643,6 @@ func (c *ChannelGraph) PruneGraph(spentOutputs []*wire.OutPoint,
 			if err != nil {
 				return err
 			}
-			chansClosed = append(chansClosed, &edgeInfo)
 
 			// Attempt to delete the channel, an ErrEdgeNotFound
 			// will be returned if that outpoint isn't known to be
@@ -648,6 +654,12 @@ func (c *ChannelGraph) PruneGraph(spentOutputs []*wire.OutPoint,
 			if err != nil && err != ErrEdgeNotFound {
 				return err
 			}
+
+			// Include this channel in our list of closed channels
+			// and collect the node public keys at each end.
+			chansClosed = append(chansClosed, &edgeInfo)
+			nodesWithChansClosed[edgeInfo.NodeKey1Bytes] = struct{}{}
+			nodesWithChansClosed[edgeInfo.NodeKey2Bytes] = struct{}{}
 		}
 
 		metaBucket, err := tx.CreateBucketIfNotExists(graphMetaBucket)
@@ -669,7 +681,15 @@ func (c *ChannelGraph) PruneGraph(spentOutputs []*wire.OutPoint,
 		var newTip [pruneTipBytes]byte
 		copy(newTip[:], blockHash[:])
 
-		return pruneBucket.Put(blockHeightBytes[:], newTip[:])
+		err = pruneBucket.Put(blockHeightBytes[:], newTip[:])
+		if err != nil {
+			return err
+		}
+
+		// Now that the graph has been pruned, we'll also attempt to
+		// prune any nodes that have had a channel closed within the
+		// latest block.
+		return c.pruneGraphNodes(tx, nodes, nodesWithChansClosed)
 	})
 	if err != nil {
 		return nil, err
@@ -678,6 +698,58 @@ func (c *ChannelGraph) PruneGraph(spentOutputs []*wire.OutPoint,
 	return chansClosed, nil
 }
 
+// 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.
+func (c *ChannelGraph) pruneGraphNodes(tx *bolt.Tx, nodes *bolt.Bucket,
+	nodePubKeys map[[33]byte]struct{}) error {
+
+	log.Trace("Pruning nodes from graph with no open channels")
+
+	// We'll retrieve the graph's source node to ensure we don't remove it
+	// even if it no longer has any open channels.
+	sourceNode, err := c.sourceNode(tx)
+	if err != nil {
+		return err
+	}
+
+	// We'll now iterate over every node which had a channel closed and
+	// check whether they have any other open channels left within the
+	// graph. If they don't, they'll be pruned from the channel graph.
+	for nodePubKey := range nodePubKeys {
+		if bytes.Equal(nodePubKey[:], sourceNode.PubKeyBytes[:]) {
+			continue
+		}
+
+		node, err := fetchLightningNode(nodes, nodePubKey[:])
+		if err != nil {
+			continue
+		}
+		node.db = c.db
+
+		numChansLeft := 0
+		err = node.ForEachChannel(tx, func(*bolt.Tx, *ChannelEdgeInfo,
+			*ChannelEdgePolicy, *ChannelEdgePolicy) error {
+
+			numChansLeft++
+			return nil
+		})
+		if err != nil {
+			continue
+		}
+
+		if numChansLeft == 0 {
+			err := c.deleteLightningNode(tx, nodePubKey[:])
+			if err != nil {
+				log.Tracef("Unable to prune node %x from the "+
+					"graph: %v", nodePubKey, err)
+			}
+		}
+	}
+
+	return nil
+}
+
 // DisconnectBlockAtHeight is used to indicate that the block specified
 // by the passed height has been disconnected from the main chain. This
 // will "rewind" the graph back to the height below, deleting channels

channeldb/graph_test.go

@@ -421,6 +421,13 @@ func TestDisconnectBlockAtHeight(t *testing.T) {
 	}
 
 	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)
+	}
 
 	// We'd like to test the insertion/deletion of edges, so we create two
 	// vertexes to connect.
@@ -964,7 +971,7 @@ func assertNumChans(t *testing.T, graph *ChannelGraph, n int) {
 		return nil
 	}); err != nil {
 		_, _, line, _ := runtime.Caller(1)
-		t.Fatalf("line %v:unable to scan channels: %v", line, err)
+		t.Fatalf("line %v: unable to scan channels: %v", line, err)
 	}
 	if numChans != n {
 		_, _, line, _ := runtime.Caller(1)
@@ -973,6 +980,23 @@ func assertNumChans(t *testing.T, graph *ChannelGraph, n int) {
 	}
 }
 
+func assertNumNodes(t *testing.T, graph *ChannelGraph, n int) {
+	numNodes := 0
+	err := graph.ForEachNode(nil, func(_ *bolt.Tx, _ *LightningNode) error {
+		numNodes++
+		return nil
+	})
+	if err != nil {
+		_, _, line, _ := runtime.Caller(1)
+		t.Fatalf("line %v: unable to scan nodes: %v", line, err)
+	}
+
+	if numNodes != n {
+		_, _, line, _ := runtime.Caller(1)
+		t.Fatalf("line %v: expected %v nodes, got %v", line, n, numNodes)
+	}
+}
+
 func assertChanViewEqual(t *testing.T, a []wire.OutPoint, b []*wire.OutPoint) {
 	if len(a) != len(b) {
 		_, _, line, _ := runtime.Caller(1)
@@ -1003,6 +1027,13 @@ func TestGraphPruning(t *testing.T) {
 	}
 
 	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)
+	}
 
 	// As initial set up for the test, we'll create a graph with 5 vertexes
 	// and enough edges to create a fully connected graph. The graph will
@@ -1137,9 +1168,11 @@ func TestGraphPruning(t *testing.T) {
 		t.Fatalf("channels were pruned but shouldn't have been")
 	}
 
-	// Once again, the prune tip should have been updated.
+	// Once again, the prune tip should have been updated. We should still
+	// see both channels and their participants, along with the source node.
 	assertPruneTip(t, graph, &blockHash, blockHeight)
 	assertNumChans(t, graph, 2)
+	assertNumNodes(t, graph, 4)
 
 	// Finally, create a block that prunes the remainder of the channels
 	// from the graph.
@@ -1159,10 +1192,11 @@ func TestGraphPruning(t *testing.T) {
 			"expected %v, got %v", 2, len(prunedChans))
 	}
 
-	// The prune tip should be updated, and no channels should be found
-	// within the current graph.
+	// The prune tip should be updated, no channels should be found, and
+	// only the source node should remain within the current graph.
 	assertPruneTip(t, graph, &blockHash, blockHeight)
 	assertNumChans(t, graph, 0)
+	assertNumNodes(t, graph, 1)
 
 	// Finally, the channel view at this point in the graph should now be
 	// completely empty.  Those channels should also be missing from the
@@ -1888,6 +1922,13 @@ func TestChannelEdgePruningUpdateIndexDeletion(t *testing.T) {
 	}
 
 	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)
+	}
 
 	// We'll first populate our graph with two nodes. All channels created
 	// below will be made between these two nodes.