From b780dfacdbb46c0b3caa18e6b67c0eeabafe179b Mon Sep 17 00:00:00 2001
From: Wilmer Paulino <wilmer.paulino@gmail.com>
Date: Wed, 27 Mar 2019 13:06:12 -0700
Subject: [PATCH] channeldb: add zombie edge index

In this commit, we add a zombie edge index to the database. This allows
us to quickly determine across restarts whether we're attempting to
process an edge we've previously deemed as zombie.
---
 channeldb/db.go         |   3 ++
 channeldb/graph.go      | 114 ++++++++++++++++++++++++++++++++++++++++
 channeldb/graph_test.go |  61 +++++++++++++++++++++
 3 files changed, 178 insertions(+)

diff --git a/channeldb/db.go b/channeldb/db.go
index 51f3fe9a..4b9476c3 100644
--- a/channeldb/db.go
+++ b/channeldb/db.go
@@ -268,6 +268,9 @@ func createChannelDB(dbPath string) error {
 		if _, err := edges.CreateBucket(channelPointBucket); err != nil {
 			return err
 		}
+		if _, err := edges.CreateBucket(zombieBucket); err != nil {
+			return err
+		}
 
 		graphMeta, err := tx.CreateBucket(graphMetaBucket)
 		if err != nil {
diff --git a/channeldb/graph.go b/channeldb/graph.go
index 2d5a2491..60c49e4e 100644
--- a/channeldb/graph.go
+++ b/channeldb/graph.go
@@ -106,6 +106,17 @@ var (
 	// maps: outPoint -> chanID
 	channelPointBucket = []byte("chan-index")
 
+	// zombieBucket is a sub-bucket of the main edgeBucket bucket
+	// responsible for maintaining an index of zombie channels. Each entry
+	// exists within the bucket as follows:
+	//
+	// maps: chanID -> pubKey1 || pubKey2
+	//
+	// The chanID represents the channel ID of the edge that is marked as a
+	// zombie and is used as the key, which maps to the public keys of the
+	// edge's participants.
+	zombieBucket = []byte("zombie-index")
+
 	// graphMetaBucket is a top-level bucket which stores various meta-deta
 	// related to the on-disk channel graph. Data stored in this bucket
 	// includes the block to which the graph has been synced to, the total
@@ -2782,6 +2793,109 @@ func (c *ChannelGraph) NewChannelEdgePolicy() *ChannelEdgePolicy {
 	return &ChannelEdgePolicy{db: c.db}
 }
 
+// MarkEdgeZombie marks an edge as a zombie within the graph's zombie index.
+// The public keys should represent the node public keys of the two parties
+// involved in the edge.
+func (c *ChannelGraph) MarkEdgeZombie(chanID uint64, pubKey1,
+	pubKey2 [33]byte) error {
+
+	return c.db.Batch(func(tx *bbolt.Tx) error {
+		edges := tx.Bucket(edgeBucket)
+		if edges == nil {
+			return ErrGraphNoEdgesFound
+		}
+		zombieIndex, err := edges.CreateBucketIfNotExists(zombieBucket)
+		if err != nil {
+			return err
+		}
+		return markEdgeZombie(zombieIndex, chanID, pubKey1, pubKey2)
+	})
+}
+
+// markEdgeZombie marks an edge as a zombie within our zombie index. The public
+// keys should represent the node public keys of the two parties involved in the
+// edge.
+func markEdgeZombie(zombieIndex *bbolt.Bucket, chanID uint64, pubKey1,
+	pubKey2 [33]byte) error {
+
+	var k [8]byte
+	byteOrder.PutUint64(k[:], chanID)
+
+	var v [66]byte
+	copy(v[:33], pubKey1[:])
+	copy(v[33:], pubKey2[:])
+
+	return zombieIndex.Put(k[:], v[:])
+}
+
+// MarkEdgeLive clears an edge from our zombie index, deeming it as live.
+func (c *ChannelGraph) MarkEdgeLive(chanID uint64) error {
+	return c.db.Batch(func(tx *bbolt.Tx) error {
+		edges := tx.Bucket(edgeBucket)
+		if edges == nil {
+			return ErrGraphNoEdgesFound
+		}
+		zombieIndex := edges.Bucket(zombieBucket)
+		if zombieIndex == nil {
+			return nil
+		}
+
+		var k [8]byte
+		byteOrder.PutUint64(k[:], chanID)
+		return zombieIndex.Delete(k[:])
+	})
+}
+
+// IsZombieEdge returns whether the edge is considered zombie. If it is a
+// zombie, then the two node public keys corresponding to this edge are also
+// returned.
+func (c *ChannelGraph) IsZombieEdge(chanID uint64) (bool, [33]byte, [33]byte) {
+	var (
+		isZombie         bool
+		pubKey1, pubKey2 [33]byte
+	)
+
+	err := c.db.View(func(tx *bbolt.Tx) error {
+		edges := tx.Bucket(edgeBucket)
+		if edges == nil {
+			return ErrGraphNoEdgesFound
+		}
+		zombieIndex := edges.Bucket(zombieBucket)
+		if zombieIndex == nil {
+			return nil
+		}
+
+		isZombie, pubKey1, pubKey2 = isZombieEdge(zombieIndex, chanID)
+		return nil
+	})
+	if err != nil {
+		return false, [33]byte{}, [33]byte{}
+	}
+
+	return isZombie, pubKey1, pubKey2
+}
+
+// isZombieEdge returns whether an entry exists for the given channel in the
+// zombie index. If an entry exists, then the two node public keys corresponding
+// to this edge are also returned.
+func isZombieEdge(zombieIndex *bbolt.Bucket,
+	chanID uint64) (bool, [33]byte, [33]byte) {
+
+	var k [8]byte
+	byteOrder.PutUint64(k[:], chanID)
+
+	v := zombieIndex.Get(k[:])
+	if v == nil {
+		return false, [33]byte{}, [33]byte{}
+	}
+
+	var pubKey1, pubKey2 [33]byte
+	copy(pubKey1[:], v[:33])
+	copy(pubKey2[:], v[33:])
+
+	return true, pubKey1, pubKey2
+}
+
 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 378391ec..0c7edc0e 100644
--- a/channeldb/graph_test.go
+++ b/channeldb/graph_test.go
@@ -2786,6 +2786,67 @@ func TestEdgePolicyMissingMaxHtcl(t *testing.T) {
 	assertEdgeInfoEqual(t, dbEdgeInfo, edgeInfo)
 }
 
+// TestGraphZombieIndex ensures that we can mark edges correctly as zombie/live.
+func TestGraphZombieIndex(t *testing.T) {
+	t.Parallel()
+
+	// We'll start by creating our test graph along with a test edge.
+	db, cleanUp, err := makeTestDB()
+	defer cleanUp()
+	if err != nil {
+		t.Fatalf("unable to create test database: %v", err)
+	}
+	graph := db.ChannelGraph()
+
+	node1, err := createTestVertex(db)
+	if err != nil {
+		t.Fatalf("unable to create test vertex: %v", err)
+	}
+	node2, err := createTestVertex(db)
+	if err != nil {
+		t.Fatalf("unable to create test vertex: %v", err)
+	}
+	edge, _, _ := createChannelEdge(db, node1, node2)
+
+	// If the graph is not aware of the edge, then it should not be a
+	// zombie.
+	isZombie, _, _ := graph.IsZombieEdge(edge.ChannelID)
+	if isZombie {
+		t.Fatal("expected edge to not be marked as zombie")
+	}
+
+	// If we mark the edge as a zombie, then we should expect to see it
+	// within the index.
+	err = graph.MarkEdgeZombie(
+		edge.ChannelID, node1.PubKeyBytes, node2.PubKeyBytes,
+	)
+	if err != nil {
+		t.Fatalf("unable to mark edge as zombie: %v", err)
+	}
+	isZombie, pubKey1, pubKey2 := graph.IsZombieEdge(edge.ChannelID)
+	if !isZombie {
+		t.Fatal("expected edge to be marked as zombie")
+	}
+	if pubKey1 != node1.PubKeyBytes {
+		t.Fatalf("expected pubKey1 %x, got %x", node1.PubKeyBytes,
+			pubKey1)
+	}
+	if pubKey2 != node2.PubKeyBytes {
+		t.Fatalf("expected pubKey2 %x, got %x", node2.PubKeyBytes,
+			pubKey2)
+	}
+
+	// Similarly, if we mark the same edge as live, we should no longer see
+	// it within the index.
+	if err := graph.MarkEdgeLive(edge.ChannelID); err != nil {
+		t.Fatalf("unable to mark edge as live: %v", err)
+	}
+	isZombie, _, _ = graph.IsZombieEdge(edge.ChannelID)
+	if isZombie {
+		t.Fatal("expected edge to not be marked as zombie")
+	}
+}
+
 // 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.