From 9d73954a3278ffa8285b98fbdf44db55f483b048 Mon Sep 17 00:00:00 2001
From: Felix Lange <fjl@twurst.com>
Date: Sat, 7 Dec 2019 11:32:21 +0100
Subject: [PATCH] p2p/dnsdisc: preserve linked trees across root updates

This improves the way links are handled when the link root changes.
Previously, sync would simply remove all links from the current tree and
garbage-collect all unreachable trees before syncing the new list of
links.

This behavior isn't great in certain cases: Consider a structure where
trees A, B, and C reference each other and D links to A. If D's link
root changed, the sync code would first remove trees A, B and C, only to
re-sync them later when the link to A was found again.

The fix for this problem is to track the current set of links in each
clientTree and removing old links only AFTER all links are synced.
---
 p2p/dnsdisc/sync.go      | 36 +++++++++++++++++++++++++-----------
 p2p/dnsdisc/sync_test.go |  4 ++--
 2 files changed, 27 insertions(+), 13 deletions(-)

diff --git a/p2p/dnsdisc/sync.go b/p2p/dnsdisc/sync.go
index d87ecb4344..53423527a5 100644
--- a/p2p/dnsdisc/sync.go
+++ b/p2p/dnsdisc/sync.go
@@ -18,7 +18,6 @@ package dnsdisc
 
 import (
 	"context"
-	"crypto/ecdsa"
 	"math/rand"
 	"time"
 
@@ -28,23 +27,23 @@ import (
 
 // clientTree is a full tree being synced.
 type clientTree struct {
-	c             *Client
-	lc            *linkCache
-	loc           *linkEntry
-	root          *rootEntry
+	c   *Client
+	loc *linkEntry // link to this tree
+
 	lastRootCheck mclock.AbsTime // last revalidation of root
+	root          *rootEntry
 	enrs          *subtreeSync
 	links         *subtreeSync
+
+	lc         *linkCache          // tracks all links between all trees
+	curLinks   map[string]struct{} // links contained in this tree
+	linkGCRoot string              // root on which last link GC has run
 }
 
 func newClientTree(c *Client, lc *linkCache, loc *linkEntry) *clientTree {
 	return &clientTree{c: c, lc: lc, loc: loc}
 }
 
-func keysEqual(k1, k2 *ecdsa.PublicKey) bool {
-	return k1.Curve == k2.Curve && k1.X.Cmp(k2.X) == 0 && k1.Y.Cmp(k2.Y) == 0
-}
-
 // syncAll retrieves all entries of the tree.
 func (ct *clientTree) syncAll(dest map[string]entry) error {
 	if err := ct.updateRoot(); err != nil {
@@ -72,6 +71,7 @@ func (ct *clientTree) syncRandom(ctx context.Context) (*enode.Node, error) {
 		err := ct.syncNextLink(ctx)
 		return nil, err
 	}
+	ct.gcLinks()
 
 	// Sync next random entry in ENR tree. Once every node has been visited, we simply
 	// start over. This is fine because entries are cached.
@@ -81,6 +81,16 @@ func (ct *clientTree) syncRandom(ctx context.Context) (*enode.Node, error) {
 	return ct.syncNextRandomENR(ctx)
 }
 
+// gcLinks removes outdated links from the global link cache. GC runs once
+// when the link sync finishes.
+func (ct *clientTree) gcLinks() {
+	if !ct.links.done() || ct.root.lroot == ct.linkGCRoot {
+		return
+	}
+	ct.lc.resetLinks(ct.loc.str, ct.curLinks)
+	ct.linkGCRoot = ct.root.lroot
+}
+
 func (ct *clientTree) syncNextLink(ctx context.Context) error {
 	hash := ct.links.missing[0]
 	e, err := ct.links.resolveNext(ctx, hash)
@@ -91,6 +101,7 @@ func (ct *clientTree) syncNextLink(ctx context.Context) error {
 
 	if dest, ok := e.(*linkEntry); ok {
 		ct.lc.addLink(ct.loc.str, dest.str)
+		ct.curLinks[dest.str] = struct{}{}
 	}
 	return nil
 }
@@ -140,7 +151,7 @@ func (ct *clientTree) updateRoot() error {
 	// Invalidate subtrees if changed.
 	if ct.links == nil || root.lroot != ct.links.root {
 		ct.links = newSubtreeSync(ct.c, ct.loc, root.lroot, true)
-		ct.lc.resetLinks(ct.loc.str)
+		ct.curLinks = make(map[string]struct{})
 	}
 	if ct.enrs == nil || root.eroot != ct.enrs.root {
 		ct.enrs = newSubtreeSync(ct.c, ct.loc, root.eroot, false)
@@ -231,13 +242,16 @@ func (lc *linkCache) addLink(from, to string) {
 }
 
 // resetLinks clears all links of the given tree.
-func (lc *linkCache) resetLinks(from string) {
+func (lc *linkCache) resetLinks(from string, keep map[string]struct{}) {
 	stk := []string{from}
 	for len(stk) > 0 {
 		item := stk[len(stk)-1]
 		stk = stk[:len(stk)-1]
 
 		for r, refs := range lc.backrefs {
+			if _, ok := keep[r]; ok {
+				continue
+			}
 			if _, ok := refs[item]; !ok {
 				continue
 			}
diff --git a/p2p/dnsdisc/sync_test.go b/p2p/dnsdisc/sync_test.go
index 1f1ea6ad7d..32af3656ef 100644
--- a/p2p/dnsdisc/sync_test.go
+++ b/p2p/dnsdisc/sync_test.go
@@ -47,7 +47,7 @@ func TestLinkCache(t *testing.T) {
 		t.Error("6 is referenced")
 	}
 
-	lc.resetLinks("1")
+	lc.resetLinks("1", nil)
 	if !lc.changed {
 		t.Error("changed flag not set")
 	}
@@ -74,7 +74,7 @@ func TestLinkCacheRandom(t *testing.T) {
 
 	// Remove all the links.
 	for _, s := range remove {
-		lc.resetLinks(s)
+		lc.resetLinks(s, nil)
 	}
 	if len(lc.backrefs) != 0 {
 		t.Logf("%+v", lc)