From 0223039fb4845cdab801092d845e1fd7a4ad5d3a Mon Sep 17 00:00:00 2001
From: Conner Fromknecht <conner@lightning.engineering>
Date: Mon, 15 Jul 2019 16:59:14 -0700
Subject: [PATCH] channeldb/graph: add NumZombies method and tests

---
 channeldb/graph.go      | 25 +++++++++++++++++++++++++
 channeldb/graph_test.go | 19 +++++++++++++++++++
 2 files changed, 44 insertions(+)

diff --git a/channeldb/graph.go b/channeldb/graph.go
index df233f1d9..b587189bd 100644
--- a/channeldb/graph.go
+++ b/channeldb/graph.go
@@ -3195,6 +3195,31 @@ func isZombieEdge(zombieIndex *bbolt.Bucket,
 	return true, pubKey1, pubKey2
 }
 
+// NumZombies returns the current number of zombie channels in the graph.
+func (c *ChannelGraph) NumZombies() (uint64, error) {
+	var numZombies uint64
+	err := c.db.View(func(tx *bbolt.Tx) error {
+		edges := tx.Bucket(edgeBucket)
+		if edges == nil {
+			return nil
+		}
+		zombieIndex := edges.Bucket(zombieBucket)
+		if zombieIndex == nil {
+			return nil
+		}
+
+		return zombieIndex.ForEach(func(_, _ []byte) error {
+			numZombies++
+			return nil
+		})
+	})
+	if err != nil {
+		return 0, err
+	}
+
+	return numZombies, nil
+}
+
 func putLightningNode(nodeBucket *bbolt.Bucket, aliasBucket *bbolt.Bucket,
 	updateIndex *bbolt.Bucket, node *LightningNode) error {
 
diff --git a/channeldb/graph_test.go b/channeldb/graph_test.go
index fd471231d..3e4323748 100644
--- a/channeldb/graph_test.go
+++ b/channeldb/graph_test.go
@@ -2847,6 +2847,22 @@ func TestEdgePolicyMissingMaxHtcl(t *testing.T) {
 	assertEdgeInfoEqual(t, dbEdgeInfo, edgeInfo)
 }
 
+// assertNumZombies queries the provided ChannelGraph for NumZombies, and
+// asserts that the returned number is equal to expZombies.
+func assertNumZombies(t *testing.T, graph *ChannelGraph, expZombies uint64) {
+	t.Helper()
+
+	numZombies, err := graph.NumZombies()
+	if err != nil {
+		t.Fatalf("unable to query number of zombies: %v", err)
+	}
+
+	if numZombies != expZombies {
+		t.Fatalf("expected %d zombies, found %d",
+			expZombies, numZombies)
+	}
+}
+
 // TestGraphZombieIndex ensures that we can mark edges correctly as zombie/live.
 func TestGraphZombieIndex(t *testing.T) {
 	t.Parallel()
@@ -2885,6 +2901,7 @@ func TestGraphZombieIndex(t *testing.T) {
 	if isZombie {
 		t.Fatal("expected edge to not be marked as zombie")
 	}
+	assertNumZombies(t, graph, 0)
 
 	// If we delete the edge and mark it as a zombie, then we should expect
 	// to see it within the index.
@@ -2904,6 +2921,7 @@ func TestGraphZombieIndex(t *testing.T) {
 		t.Fatalf("expected pubKey2 %x, got %x", node2.PubKeyBytes,
 			pubKey2)
 	}
+	assertNumZombies(t, graph, 1)
 
 	// Similarly, if we mark the same edge as live, we should no longer see
 	// it within the index.
@@ -2914,6 +2932,7 @@ func TestGraphZombieIndex(t *testing.T) {
 	if isZombie {
 		t.Fatal("expected edge to not be marked as zombie")
 	}
+	assertNumZombies(t, graph, 0)
 }
 
 // compareNodes is used to compare two LightningNodes while excluding the