graph/db: move cache update out of pruneGraphNodes

In preparation for moving the cache write completely out of KVStore, we
move the cache write up one layer.

graph/db/kv_store.go

@@ -1427,7 +1427,18 @@ func (c *KVStore) PruneGraph(spentOutputs []*wire.OutPoint,
 		// 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(nodes, edgeIndex)
+		prunedNodes, err := c.pruneGraphNodes(nodes, edgeIndex)
+		if err != nil {
+			return err
+		}
+
+		if c.graphCache != nil {
+			for _, nodePubKey := range prunedNodes {
+				c.graphCache.RemoveNode(nodePubKey)
+			}
+		}
+
+		return nil
 	}, func() {
 		chansClosed = nil
 	})
@@ -1467,7 +1478,18 @@ func (c *KVStore) PruneGraphNodes() error {
 			return ErrGraphNoEdgesFound
 		}
 
-		return c.pruneGraphNodes(nodes, edgeIndex)
+		prunedNodes, err := c.pruneGraphNodes(nodes, edgeIndex)
+		if err != nil {
+			return err
+		}
+
+		if c.graphCache != nil {
+			for _, nodePubKey := range prunedNodes {
+				c.graphCache.RemoveNode(nodePubKey)
+			}
+		}
+
+		return nil
 	}, func() {})
 }
 
@@ -1475,7 +1497,7 @@ func (c *KVStore) PruneGraphNodes() error {
 // 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 *KVStore) pruneGraphNodes(nodes kvdb.RwBucket,
-	edgeIndex kvdb.RwBucket) error {
+	edgeIndex kvdb.RwBucket) ([]route.Vertex, error) {
 
 	log.Trace("Pruning nodes from graph with no open channels")
 
@@ -1483,7 +1505,7 @@ func (c *KVStore) pruneGraphNodes(nodes kvdb.RwBucket,
 	// even if it no longer has any open channels.
 	sourceNode, err := c.sourceNode(nodes)
 	if err != nil {
-		return err
+		return nil, err
 	}
 
 	// We'll use this map to keep count the number of references to a node
@@ -1505,7 +1527,7 @@ func (c *KVStore) pruneGraphNodes(nodes kvdb.RwBucket,
 		return nil
 	})
 	if err != nil {
-		return err
+		return nil, err
 	}
 
 	// To ensure we never delete the source node, we'll start off by
@@ -1531,12 +1553,12 @@ func (c *KVStore) pruneGraphNodes(nodes kvdb.RwBucket,
 		return nil
 	})
 	if err != nil {
-		return err
+		return nil, err
 	}
 
 	// Finally, we'll make a second pass over the set of nodes, and delete
 	// any nodes that have a ref count of zero.
-	var numNodesPruned int
+	var pruned []route.Vertex
 	for nodePubKey, refCount := range nodeRefCounts {
 		// If the ref count of the node isn't zero, then we can safely
 		// skip it as it still has edges to or from it within the
@@ -1545,10 +1567,6 @@ func (c *KVStore) pruneGraphNodes(nodes kvdb.RwBucket,
 			continue
 		}
 
-		if c.graphCache != nil {
-			c.graphCache.RemoveNode(nodePubKey)
-		}
-
 		// If we reach this point, then there are no longer any edges
 		// that connect this node, so we can delete it.
 		err := c.deleteLightningNode(nodes, nodePubKey[:])
@@ -1561,21 +1579,21 @@ func (c *KVStore) pruneGraphNodes(nodes kvdb.RwBucket,
 				continue
 			}
 
-			return err
+			return nil, err
 		}
 
 		log.Infof("Pruned unconnected node %x from channel graph",
 			nodePubKey[:])
 
-		numNodesPruned++
+		pruned = append(pruned, nodePubKey)
 	}
 
-	if numNodesPruned > 0 {
+	if len(pruned) > 0 {
 		log.Infof("Pruned %v unconnected nodes from the channel graph",
-			numNodesPruned)
+			len(pruned))
 	}
 
-	return nil
+	return pruned, err
 }
 
 // DisconnectBlockAtHeight is used to indicate that the block specified