From be439f42211daf1fe3af322d2b479a85d8416d1e Mon Sep 17 00:00:00 2001
From: Felix Lange <fjl@twurst.com>
Date: Sun, 7 Jul 2019 23:11:07 +0200
Subject: [PATCH] p2p/discutil: add Close to iterator interface

---
 p2p/discutil/iter.go      | 11 +++++++++--
 p2p/discutil/iter_test.go | 10 ++++++++++
 2 files changed, 19 insertions(+), 2 deletions(-)

diff --git a/p2p/discutil/iter.go b/p2p/discutil/iter.go
index 69c7968da0..8081a91a6e 100644
--- a/p2p/discutil/iter.go
+++ b/p2p/discutil/iter.go
@@ -31,10 +31,13 @@ import (
 // node could be found before the context was canceled. The isLive return value reports
 // whether the iterator is still open. Once closed, iterators should keep returning (nil, false).
 //
-// Implementations are not required to be safe for concurrent use. It is therefore unsafe
-// to call NextNode from multiple goroutines at the same time.
+// Implementations of NextNode are not required to be safe for concurrent use. It is
+// therefore unsafe to call NextNode from multiple goroutines at the same time.
+//
+// Close may be called concurrently with NextNode, and interrupts NextNode.
 type Iterator interface {
 	NextNode(ctx context.Context) (n *enode.Node, isLive bool)
+	Close()
 }
 
 // ReadNodes reads at most n nodes from the given iterator. The return value contains no
@@ -83,6 +86,10 @@ func (f *filterIter) NextNode(ctx context.Context) (*enode.Node, bool) {
 	return n, isLive
 }
 
+func (f *filterIter) Close() {
+	f.it.Close()
+}
+
 // FairMix aggregates multiple node iterators. The mixer itself is an iterator which ends
 // only when Close is called. Source iterators added via AddSource are removed from the mix
 // when they end.
diff --git a/p2p/discutil/iter_test.go b/p2p/discutil/iter_test.go
index bc3708af3b..42f6f213d4 100644
--- a/p2p/discutil/iter_test.go
+++ b/p2p/discutil/iter_test.go
@@ -227,6 +227,8 @@ func (s *genIter) NextNode(ctx context.Context) (*enode.Node, bool) {
 	return n, true
 }
 
+func (s *genIter) Close() { panic("called") }
+
 func testNode(id, seq uint64) *enode.Node {
 	var nodeID enode.ID
 	binary.BigEndian.PutUint64(nodeID[:], id)
@@ -253,6 +255,8 @@ func (s *blockedIter) NextNode(ctx context.Context) (*enode.Node, bool) {
 	}
 }
 
+func (s *blockedIter) Close() { panic("called") }
+
 // cycleNodes is a never-ending interator that cycles through the given slice.
 type cycleNodes []*enode.Node
 
@@ -267,6 +271,8 @@ func (s cycleNodes) NextNode(ctx context.Context) (*enode.Node, bool) {
 	return n, true
 }
 
+func (s cycleNodes) Close() { panic("called") }
+
 // callCountIter counts calls to NextNode.
 type callCountIter struct {
 	child Iterator
@@ -277,3 +283,7 @@ func (it *callCountIter) NextNode(ctx context.Context) (*enode.Node, bool) {
 	it.count++
 	return it.child.NextNode(ctx)
 }
+
+func (it *callCountIter) Close() {
+	it.child.Close()
+}