graph/db: version reject cache for policies

diff --git a/graph/db/kv_store.go b/graph/db/kv_store.go
index 2bbfe89..5e71ad7 100644
--- a/graph/db/kv_store.go
+++ b/graph/db/kv_store.go
@@ -1179,7 +1179,7 @@ func (c *KVStore) AddChannelEdge(ctx context.Context,
 			case alreadyExists:
 				return ErrEdgeAlreadyExist
 			default:
-				c.rejectCache.remove(edge.ChannelID)
+				c.rejectCache.remove(lnwire.GossipVersion1, edge.ChannelID)
 				c.chanCache.remove(edge.ChannelID)
 				return nil
 			}
@@ -1302,7 +1302,7 @@ func (c *KVStore) HasChannelEdge(
 	// We'll query the cache with the shared lock held to allow multiple
 	// readers to access values in the cache concurrently if they exist.
 	c.cacheMu.RLock()
-	if entry, ok := c.rejectCache.get(chanID); ok {
+	if entry, ok := c.rejectCache.get(lnwire.GossipVersion1, chanID); ok {
 		c.cacheMu.RUnlock()
 		upd1Time = time.Unix(entry.upd1Time, 0)
 		upd2Time = time.Unix(entry.upd2Time, 0)
@@ -1318,7 +1318,7 @@ func (c *KVStore) HasChannelEdge(
 	// The item was not found with the shared lock, so we'll acquire the
 	// exclusive lock and check the cache again in case another method added
 	// the entry to the cache while no lock was held.
-	if entry, ok := c.rejectCache.get(chanID); ok {
+	if entry, ok := c.rejectCache.get(lnwire.GossipVersion1, chanID); ok {
 		upd1Time = time.Unix(entry.upd1Time, 0)
 		upd2Time = time.Unix(entry.upd2Time, 0)
 		exists, isZombie = entry.flags.unpack()
@@ -1385,7 +1385,7 @@ func (c *KVStore) HasChannelEdge(
 		return time.Time{}, time.Time{}, exists, isZombie, err
 	}
 
-	c.rejectCache.insert(chanID, rejectCacheEntry{
+	c.rejectCache.insert(lnwire.GossipVersion1, chanID, rejectCacheEntry{
 		upd1Time: upd1Time.Unix(),
 		upd2Time: upd2Time.Unix(),
 		flags:    packRejectFlags(exists, isZombie),
@@ -1564,7 +1564,7 @@ func (c *KVStore) PruneGraph(spentOutputs []*wire.OutPoint,
 	}
 
 	for _, channel := range chansClosed {
-		c.rejectCache.remove(channel.ChannelID)
+		c.rejectCache.remove(lnwire.GossipVersion1, channel.ChannelID)
 		c.chanCache.remove(channel.ChannelID)
 	}
 
@@ -1831,7 +1831,7 @@ func (c *KVStore) DisconnectBlockAtHeight(height uint32) (
 	}
 
 	for _, channel := range removedChans {
-		c.rejectCache.remove(channel.ChannelID)
+		c.rejectCache.remove(lnwire.GossipVersion1, channel.ChannelID)
 		c.chanCache.remove(channel.ChannelID)
 	}
 
@@ -1950,7 +1950,7 @@ func (c *KVStore) DeleteChannelEdges(v lnwire.GossipVersion,
 	}
 
 	for _, chanID := range chanIDs {
-		c.rejectCache.remove(chanID)
+		c.rejectCache.remove(lnwire.GossipVersion1, chanID)
 		c.chanCache.remove(chanID)
 	}
 
@@ -3265,13 +3265,13 @@ func (c *KVStore) updateEdgeCache(e *models.ChannelEdgePolicy,
 	// the entry with the updated timestamp for the direction that was just
 	// written. If the edge doesn't exist, we'll load the cache entry lazily
 	// during the next query for this edge.
-	if entry, ok := c.rejectCache.get(e.ChannelID); ok {
+	if entry, ok := c.rejectCache.get(lnwire.GossipVersion1, e.ChannelID); ok {
 		if isUpdate1 {
 			entry.upd1Time = e.LastUpdate.Unix()
 		} else {
 			entry.upd2Time = e.LastUpdate.Unix()
 		}
-		c.rejectCache.insert(e.ChannelID, entry)
+		c.rejectCache.insert(lnwire.GossipVersion1, e.ChannelID, entry)
 	}
 
 	// If an entry for this channel is found in channel cache, we'll modify
@@ -4185,7 +4185,7 @@ func (c *KVStore) MarkEdgeZombie(chanID uint64,
 		return err
 	}
 
-	c.rejectCache.remove(chanID)
+	c.rejectCache.remove(lnwire.GossipVersion1, chanID)
 	c.chanCache.remove(chanID)
 
 	return nil
@@ -4254,7 +4254,7 @@ func (c *KVStore) markEdgeLiveUnsafe(tx kvdb.RwTx, chanID uint64) error {
 		return err
 	}
 
-	c.rejectCache.remove(chanID)
+	c.rejectCache.remove(lnwire.GossipVersion1, chanID)
 	c.chanCache.remove(chanID)
 
 	return nil
diff --git a/graph/db/reject_cache.go b/graph/db/reject_cache.go
index 2a27219..bfe3c47 100644
--- a/graph/db/reject_cache.go
+++ b/graph/db/reject_cache.go
@@ -1,5 +1,11 @@
 package graphdb
 
+import (
+	"time"
+
+	"github.com/lightningnetwork/lnd/lnwire"
+)
+
 // rejectFlags is a compact representation of various metadata stored by the
 // reject cache about a particular channel.
 type rejectFlags uint8
@@ -41,9 +47,68 @@ func (f rejectFlags) unpack() (bool, bool) {
 // including the timestamps of its latest edge policies and whether or not the
 // channel exists in the graph.
 type rejectCacheEntry struct {
+	// upd{1,2}Time are Unix timestamps for v1 policies.
 	upd1Time int64
 	upd2Time int64
-	flags    rejectFlags
+
+	// upd{1,2}BlockHeight are the last known block heights for v2
+	// policies.
+	upd1BlockHeight uint32
+	upd2BlockHeight uint32
+
+	flags rejectFlags
+}
+
+// newRejectCacheEntryV1 constructs a reject cache entry for v1 policies.
+func newRejectCacheEntryV1(upd1, upd2 time.Time, exists,
+	isZombie bool) rejectCacheEntry {
+
+	return rejectCacheEntry{
+		upd1Time: upd1.Unix(),
+		upd2Time: upd2.Unix(),
+		flags:    packRejectFlags(exists, isZombie),
+	}
+}
+
+// newRejectCacheEntryV2 constructs a reject cache entry for v2 policies.
+func newRejectCacheEntryV2(upd1, upd2 uint32, exists,
+	isZombie bool) rejectCacheEntry {
+
+	return rejectCacheEntry{
+		upd1BlockHeight: upd1,
+		upd2BlockHeight: upd2,
+		flags:           packRejectFlags(exists, isZombie),
+	}
+}
+
+// updateRejectCacheEntryV1 updates the cached v1 timestamps.
+func updateRejectCacheEntryV1(entry *rejectCacheEntry, isUpdate1 bool,
+	lastUpdate time.Time) {
+
+	if isUpdate1 {
+		entry.upd1Time = lastUpdate.Unix()
+	} else {
+		entry.upd2Time = lastUpdate.Unix()
+	}
+}
+
+// updateRejectCacheEntryV2 updates the cached v2 block heights.
+func updateRejectCacheEntryV2(entry *rejectCacheEntry, isUpdate1 bool,
+	blockHeight uint32) {
+
+	if isUpdate1 {
+		entry.upd1BlockHeight = blockHeight
+	} else {
+		entry.upd2BlockHeight = blockHeight
+	}
+}
+
+// rejectCacheKey uniquely identifies a channel entry in the reject cache by
+// gossip version and channel ID. This allows v1 and v2 policy state for the
+// same channel ID to be cached independently.
+type rejectCacheKey struct {
+	version lnwire.GossipVersion
+	chanID  uint64
 }
 
 // rejectCache is an in-memory cache used to improve the performance of
@@ -51,20 +116,25 @@ type rejectCacheEntry struct {
 // well as the most recent timestamps for each policy (if they exists).
 type rejectCache struct {
 	n     int
-	edges map[uint64]rejectCacheEntry
+	edges map[rejectCacheKey]rejectCacheEntry
 }
 
 // newRejectCache creates a new rejectCache with maximum capacity of n entries.
 func newRejectCache(n int) *rejectCache {
 	return &rejectCache{
 		n:     n,
-		edges: make(map[uint64]rejectCacheEntry, n),
+		edges: make(map[rejectCacheKey]rejectCacheEntry, n),
 	}
 }
 
 // get returns the entry from the cache for chanid, if it exists.
-func (c *rejectCache) get(chanid uint64) (rejectCacheEntry, bool) {
-	entry, ok := c.edges[chanid]
+func (c *rejectCache) get(version lnwire.GossipVersion, chanid uint64) (
+	rejectCacheEntry, bool) {
+
+	entry, ok := c.edges[rejectCacheKey{
+		version: version,
+		chanID:  chanid,
+	}]
 	return entry, ok
 }
 
@@ -72,10 +142,17 @@ func (c *rejectCache) get(chanid uint64) (rejectCacheEntry, bool) {
 // exists, it will be replaced with the new entry. If the entry doesn't exists,
 // it will be inserted to the cache, performing a random eviction if the cache
 // is at capacity.
-func (c *rejectCache) insert(chanid uint64, entry rejectCacheEntry) {
+func (c *rejectCache) insert(version lnwire.GossipVersion, chanid uint64,
+	entry rejectCacheEntry) {
+
+	key := rejectCacheKey{
+		version: version,
+		chanID:  chanid,
+	}
+
 	// If entry exists, replace it.
-	if _, ok := c.edges[chanid]; ok {
-		c.edges[chanid] = entry
+	if _, ok := c.edges[key]; ok {
+		c.edges[key] = entry
 		return
 	}
 
@@ -86,10 +163,13 @@ func (c *rejectCache) insert(chanid uint64, entry rejectCacheEntry) {
 			break
 		}
 	}
-	c.edges[chanid] = entry
+	c.edges[key] = entry
 }
 
 // remove deletes an entry for chanid from the cache, if it exists.
-func (c *rejectCache) remove(chanid uint64) {
-	delete(c.edges, chanid)
+func (c *rejectCache) remove(version lnwire.GossipVersion, chanid uint64) {
+	delete(c.edges, rejectCacheKey{
+		version: version,
+		chanID:  chanid,
+	})
 }
diff --git a/graph/db/reject_cache_test.go b/graph/db/reject_cache_test.go
index f64c39c..c526f48 100644
--- a/graph/db/reject_cache_test.go
+++ b/graph/db/reject_cache_test.go
@@ -3,6 +3,8 @@ package graphdb
 import (
 	"reflect"
 	"testing"
+
+	"github.com/lightningnetwork/lnd/lnwire"
 )
 
 // TestRejectCache checks the behavior of the rejectCache with respect to insertion,
@@ -15,14 +17,14 @@ func TestRejectCache(t *testing.T) {
 
 	// As a sanity check, assert that querying the empty cache does not
 	// return an entry.
-	_, ok := c.get(0)
+	_, ok := c.get(lnwire.GossipVersion1, 0)
 	if ok {
 		t.Fatalf("reject cache should be empty")
 	}
 
 	// Now, fill up the cache entirely.
 	for i := uint64(0); i < cacheSize; i++ {
-		c.insert(i, entryForInt(i))
+		c.insert(lnwire.GossipVersion1, i, entryForInt(i))
 	}
 
 	// Assert that the cache has all of the entries just inserted, since no
@@ -30,7 +32,10 @@ func TestRejectCache(t *testing.T) {
 	assertHasEntries(t, c, 0, cacheSize)
 
 	// Now, insert a new element that causes the cache to evict an element.
-	c.insert(cacheSize, entryForInt(cacheSize))
+	c.insert(
+		lnwire.GossipVersion1, cacheSize,
+		entryForInt(cacheSize),
+	)
 
 	// Assert that the cache has this last entry, as the cache should evict
 	// some prior element and not the newly inserted one.
@@ -40,7 +45,7 @@ func TestRejectCache(t *testing.T) {
 	// elements.
 	evicted := make(map[uint64]struct{})
 	for i := uint64(0); i < cacheSize+1; i++ {
-		_, ok := c.get(i)
+		_, ok := c.get(lnwire.GossipVersion1, i)
 		if !ok {
 			evicted[i] = struct{}{}
 		}
@@ -54,9 +59,9 @@ func TestRejectCache(t *testing.T) {
 
 	// Remove the highest item which initially caused the eviction and
 	// reinsert the element that was evicted prior.
-	c.remove(cacheSize)
+	c.remove(lnwire.GossipVersion1, cacheSize)
 	for i := range evicted {
-		c.insert(i, entryForInt(i))
+		c.insert(lnwire.GossipVersion1, i, entryForInt(i))
 	}
 
 	// Since the removal created an extra slot, the last insertion should
@@ -69,7 +74,7 @@ func TestRejectCache(t *testing.T) {
 	// happening on inserts for existing cache items, we expect this to fail
 	// with high probability.
 	for i := uint64(0); i < cacheSize; i++ {
-		c.insert(i, entryForInt(i))
+		c.insert(lnwire.GossipVersion1, i, entryForInt(i))
 	}
 	assertHasEntries(t, c, 0, cacheSize)
 
@@ -82,7 +87,7 @@ func assertHasEntries(t *testing.T, c *rejectCache, start, end uint64) {
 	t.Helper()
 
 	for i := start; i < end; i++ {
-		entry, ok := c.get(i)
+		entry, ok := c.get(lnwire.GossipVersion1, i)
 		if !ok {
 			t.Fatalf("reject cache should contain chan %d", i)
 		}
diff --git a/graph/db/sql_store.go b/graph/db/sql_store.go
index 4468ad3..d72c853 100644
--- a/graph/db/sql_store.go
+++ b/graph/db/sql_store.go
@@ -747,7 +747,7 @@ func (s *SQLStore) AddChannelEdge(ctx context.Context,
 			case alreadyExists:
 				return ErrEdgeAlreadyExist
 			default:
-				s.rejectCache.remove(edge.ChannelID)
+				s.rejectCache.remove(edge.Version, edge.ChannelID)
 				s.chanCache.remove(edge.ChannelID)
 				return nil
 			}
@@ -862,13 +862,18 @@ func (s *SQLStore) updateEdgeCache(e *models.ChannelEdgePolicy,
 	// the entry with the updated timestamp for the direction that was just
 	// written. If the edge doesn't exist, we'll load the cache entry lazily
 	// during the next query for this edge.
-	if entry, ok := s.rejectCache.get(e.ChannelID); ok {
-		if isUpdate1 {
-			entry.upd1Time = e.LastUpdate.Unix()
-		} else {
-			entry.upd2Time = e.LastUpdate.Unix()
+	if entry, ok := s.rejectCache.get(e.Version, e.ChannelID); ok {
+		switch e.Version {
+		case lnwire.GossipVersion1:
+			updateRejectCacheEntryV1(
+				&entry, isUpdate1, e.LastUpdate,
+			)
+		case lnwire.GossipVersion2:
+			updateRejectCacheEntryV2(
+				&entry, isUpdate1, e.LastBlockHeight,
+			)
 		}
-		s.rejectCache.insert(e.ChannelID, entry)
+		s.rejectCache.insert(e.Version, e.ChannelID, entry)
 	}
 
 	// If an entry for this channel is found in channel cache, we'll modify
@@ -1764,7 +1769,7 @@ func (s *SQLStore) MarkEdgeZombie(chanID uint64,
 			"(channel_id=%d): %w", chanID, err)
 	}
 
-	s.rejectCache.remove(chanID)
+	s.rejectCache.remove(lnwire.GossipVersion1, chanID)
 	s.chanCache.remove(chanID)
 
 	return nil
@@ -1813,7 +1818,7 @@ func (s *SQLStore) MarkEdgeLive(chanID uint64) error {
 			"(channel_id=%d): %w", chanID, err)
 	}
 
-	s.rejectCache.remove(chanID)
+	s.rejectCache.remove(lnwire.GossipVersion1, chanID)
 	s.chanCache.remove(chanID)
 
 	return err
@@ -1995,7 +2000,7 @@ func (s *SQLStore) DeleteChannelEdges(v lnwire.GossipVersion,
 	}
 
 	for _, chanID := range chanIDs {
-		s.rejectCache.remove(chanID)
+		s.rejectCache.remove(v, chanID)
 		s.chanCache.remove(chanID)
 	}
 
@@ -2222,7 +2227,7 @@ func (s *SQLStore) HasChannelEdge(chanID uint64) (time.Time, time.Time, bool,
 	// We'll query the cache with the shared lock held to allow multiple
 	// readers to access values in the cache concurrently if they exist.
 	s.cacheMu.RLock()
-	if entry, ok := s.rejectCache.get(chanID); ok {
+	if entry, ok := s.rejectCache.get(lnwire.GossipVersion1, chanID); ok {
 		s.cacheMu.RUnlock()
 		node1LastUpdate = time.Unix(entry.upd1Time, 0)
 		node2LastUpdate = time.Unix(entry.upd2Time, 0)
@@ -2238,7 +2243,7 @@ func (s *SQLStore) HasChannelEdge(chanID uint64) (time.Time, time.Time, bool,
 	// The item was not found with the shared lock, so we'll acquire the
 	// exclusive lock and check the cache again in case another method added
 	// the entry to the cache while no lock was held.
-	if entry, ok := s.rejectCache.get(chanID); ok {
+	if entry, ok := s.rejectCache.get(lnwire.GossipVersion1, chanID); ok {
 		node1LastUpdate = time.Unix(entry.upd1Time, 0)
 		node2LastUpdate = time.Unix(entry.upd2Time, 0)
 		exists, isZombie = entry.flags.unpack()
@@ -2309,11 +2314,13 @@ func (s *SQLStore) HasChannelEdge(chanID uint64) (time.Time, time.Time, bool,
 			fmt.Errorf("unable to fetch channel: %w", err)
 	}
 
-	s.rejectCache.insert(chanID, rejectCacheEntry{
-		upd1Time: node1LastUpdate.Unix(),
-		upd2Time: node2LastUpdate.Unix(),
-		flags:    packRejectFlags(exists, isZombie),
-	})
+	s.rejectCache.insert(
+		lnwire.GossipVersion1, chanID,
+		newRejectCacheEntryV1(
+			node1LastUpdate, node2LastUpdate, exists,
+			isZombie,
+		),
+	)
 
 	return node1LastUpdate, node2LastUpdate, exists, isZombie, nil
 }
@@ -2732,7 +2739,7 @@ func (s *SQLStore) PruneGraph(spentOutputs []*wire.OutPoint,
 	}
 
 	for _, channel := range closedChans {
-		s.rejectCache.remove(channel.ChannelID)
+		s.rejectCache.remove(channel.Version, channel.ChannelID)
 		s.chanCache.remove(channel.ChannelID)
 	}
 
@@ -3001,7 +3008,7 @@ func (s *SQLStore) DisconnectBlockAtHeight(height uint32) (
 
 	s.cacheMu.Lock()
 	for _, channel := range removedChans {
-		s.rejectCache.remove(channel.ChannelID)
+		s.rejectCache.remove(channel.Version, channel.ChannelID)
 		s.chanCache.remove(channel.ChannelID)
 	}
 	s.cacheMu.Unlock()