graph/db: move some cache writes to ChannelGraph.

Here we move the cache writes for DisconnectBlockAtHeight and
DeleteChannelEdges to the ChannelGraph.

graph/db/graph.go

@@ -259,3 +259,66 @@ func (c *ChannelGraph) MarkEdgeLive(chanID uint64) error {
 
 	return nil
 }
+
+// DeleteChannelEdges removes edges with the given channel IDs from the
+// database and marks them as zombies. This ensures that we're unable to re-add
+// it to our database once again. If an edge does not exist within the
+// database, then ErrEdgeNotFound will be returned. If strictZombiePruning is
+// true, then when we mark these edges as zombies, we'll set up the keys such
+// that we require the node that failed to send the fresh update to be the one
+// that resurrects the channel from its zombie state. The markZombie bool
+// denotes whether to mark the channel as a zombie.
+func (c *ChannelGraph) DeleteChannelEdges(strictZombiePruning, markZombie bool,
+	chanIDs ...uint64) error {
+
+	c.cacheMu.Lock()
+	defer c.cacheMu.Unlock()
+
+	infos, err := c.KVStore.DeleteChannelEdges(
+		strictZombiePruning, markZombie, chanIDs...,
+	)
+	if err != nil {
+		return err
+	}
+
+	if c.graphCache != nil {
+		for _, info := range infos {
+			c.graphCache.RemoveChannel(
+				info.NodeKey1Bytes, info.NodeKey2Bytes,
+				info.ChannelID,
+			)
+		}
+	}
+
+	return err
+}
+
+// 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
+// that are no longer confirmed from the graph. The prune log will be
+// set to the last prune height valid for the remaining chain.
+// Channels that were removed from the graph resulting from the
+// disconnected block are returned.
+func (c *ChannelGraph) DisconnectBlockAtHeight(height uint32) (
+	[]*models.ChannelEdgeInfo, error) {
+
+	c.cacheMu.Lock()
+	defer c.cacheMu.Unlock()
+
+	edges, err := c.KVStore.DisconnectBlockAtHeight(height)
+	if err != nil {
+		return nil, err
+	}
+
+	if c.graphCache != nil {
+		for _, edge := range edges {
+			c.graphCache.RemoveChannel(
+				edge.NodeKey1Bytes, edge.NodeKey2Bytes,
+				edge.ChannelID,
+			)
+		}
+	}
+
+	return edges, nil
+}

graph/db/kv_store.go

@@ -1654,14 +1654,6 @@ func (c *KVStore) DisconnectBlockAtHeight(height uint32) (
 				return err
 			}
 
-			if c.graphCache != nil {
-				c.graphCache.RemoveChannel(
-					edgeInfo.NodeKey1Bytes,
-					edgeInfo.NodeKey2Bytes,
-					edgeInfo.ChannelID,
-				)
-			}
-
 			removedChans = append(removedChans, edgeInfo)
 		}
 
@@ -1768,7 +1760,7 @@ func (c *KVStore) PruneTip() (*chainhash.Hash, uint32, error) {
 // that resurrects the channel from its zombie state. The markZombie bool
 // denotes whether or not to mark the channel as a zombie.
 func (c *KVStore) DeleteChannelEdges(strictZombiePruning, markZombie bool,
-	chanIDs ...uint64) error {
+	chanIDs ...uint64) ([]*models.ChannelEdgeInfo, error) {
 
 	// TODO(roasbeef): possibly delete from node bucket if node has no more
 	// channels
@@ -1777,6 +1769,7 @@ func (c *KVStore) DeleteChannelEdges(strictZombiePruning, markZombie bool,
 	c.cacheMu.Lock()
 	defer c.cacheMu.Unlock()
 
+	var infos []*models.ChannelEdgeInfo
 	err := kvdb.Update(c.db, func(tx kvdb.RwTx) error {
 		edges := tx.ReadWriteBucket(edgeBucket)
 		if edges == nil {
@@ -1810,19 +1803,15 @@ func (c *KVStore) DeleteChannelEdges(strictZombiePruning, markZombie bool,
 				return err
 			}
 
-			if c.graphCache != nil {
+			infos = append(infos, edgeInfo)
-				c.graphCache.RemoveChannel(
-					edgeInfo.NodeKey1Bytes,
-					edgeInfo.NodeKey2Bytes,
-					edgeInfo.ChannelID,
-				)
-			}
 		}
 
 		return nil
-	}, func() {})
+	}, func() {
+		infos = nil
+	})
 	if err != nil {
-		return err
+		return nil, err
 	}
 
 	for _, chanID := range chanIDs {
@@ -1830,7 +1819,7 @@ func (c *KVStore) DeleteChannelEdges(strictZombiePruning, markZombie bool,
 		c.chanCache.remove(chanID)
 	}
 
-	return nil
+	return infos, nil
 }
 
 // ChannelID attempt to lookup the 8-byte compact channel ID which maps to the