simplify AsyncFilter design

Signed-off-by: Csaba Kiraly <csaba.kiraly@gmail.com>

# Conflicts:
#	eth/backend.go

eth/backend.go

@@ -18,6 +18,7 @@
 package eth
 
 import (
+	"context"
 	"encoding/json"
 	"fmt"
 	"math/big"
@@ -503,10 +504,13 @@ func (s *Ethereum) setupDiscovery() error {
 
 	// Add DHT nodes from discv4.
 	if s.p2pServer.DiscoveryV4() != nil {
-		asyncFilter := s.p2pServer.DiscoveryV4().RequestENR
-		filter := eth.NewNodeFilter(s.blockchain)
-		iter := enode.AsyncFilter(s.p2pServer.DiscoveryV4().RandomNodes(), asyncFilter, maxParallelENRRequests)
-		iter = enode.Filter(iter, filter)
+		iter := s.p2pServer.DiscoveryV4().RandomNodes()
+		resolverFunc := func(ctx context.Context, enr *enode.Node) (*enode.Node, error) {
+			// RequestENR does not yet support context. It will simply time out.
+			return s.p2pServer.DiscoveryV4().RequestENR(enr)
+		}
+		iter = enode.AsyncFilter(iter, resolverFunc, maxParallelENRRequests)
+		iter = enode.Filter(iter, eth.NewNodeFilter(s.blockchain))
 		s.discmix.AddSource(iter)
 	}
 

p2p/enode/iter.go

@@ -17,6 +17,7 @@
 package enode
 
 import (
+	"context"
 	"sync"
 	"time"
 )
@@ -155,51 +156,53 @@ func (f *filterIter) Next() bool {
 // AsyncFilterIter wraps an iterator such that Next only returns nodes for which
 // the 'check' function returns a (possibly modified) node.
 type AsyncFilterIter struct {
-	it     Iterator                   // the iterator to filter
-	check  func(*Node) (*Node, error) // the blocking check function
-	slots  chan struct{}              // the slots for parallel checking
-	passed chan *Node                 // channel to collect passed nodes
-	closed chan struct{}              // channel to override passed when closing
-	buffer *Node                      // buffer to serve the Node call
+	it        Iterator      // the iterator to filter
+	slots     chan struct{} // the slots for parallel checking
+	passed    chan *Node    // channel to collect passed nodes
+	buffer    *Node         // buffer to serve the Node call
+	cancel    context.CancelFunc
+	closeOnce sync.Once
 }
+type AsyncFilterFunc func(context.Context, *Node) (*Node, error)
 
 // AsyncFilter creates an iterator which checks nodes in parallel.
-func AsyncFilter(it Iterator, check func(*Node) (*Node, error), workers int) Iterator {
-	f := &AsyncFilterIter{it, check, make(chan struct{}, workers), make(chan *Node), make(chan struct{}), nil}
-
-	// create slots
+func AsyncFilter(it Iterator, check AsyncFilterFunc, workers int) Iterator {
+	f := &AsyncFilterIter{
+		it:     it,
+		slots:  make(chan struct{}, workers+1),
+		passed: make(chan *Node),
+	}
 	for range cap(f.slots) {
 		f.slots <- struct{}{}
 	}
+	ctx, cancel := context.WithCancel(context.Background())
+	f.cancel = cancel
 
 	go func() {
+		select {
+		case <-ctx.Done():
+			return
+		case <-f.slots:
+		}
 		// read from the iterator and start checking nodes in parallel
 		// when a node is checked, it will be sent to the passed channel
 		// and the slot will be released
-		for range f.slots {
-			if f.it.Next() {
-				if n := f.it.Node(); n != nil {
-					// check the node async, in a separate goroutine
-					go func() {
-						if nn, err := f.check(n); err == nil {
-							select {
-							case f.passed <- nn:
-							case <-f.closed: // bale out if downstream is already closed and not calling Next
-							}
-						}
-						f.slots <- struct{}{}
-					}()
-				} else {
-					// this is not supposed to happen
-					f.slots <- struct{}{}
-					break
+		for f.it.Next() {
+			n := f.it.Node()
+			<-f.slots
+			// check the node async, in a separate goroutine
+			go func() {
+				if nn, err := check(ctx, n); err == nil {
+					select {
+					case f.passed <- nn:
+					case <-ctx.Done(): // bale out if downstream is already closed and not calling Next
+					}
 				}
-			} else {
-				// the iterator has ended
 				f.slots <- struct{}{}
-				break
-			}
+			}()
 		}
+		// the iterator has ended
+		f.slots <- struct{}{}
 	}()
 
 	return f
@@ -218,13 +221,15 @@ func (f *AsyncFilterIter) Node() *Node {
 
 // Close ends the iterator, also closing the wrapped iterator.
 func (f *AsyncFilterIter) Close() {
-	f.it.Close()
-	close(f.closed) // override the passed channel
-	for range cap(f.slots) {
-		<-f.slots
-	}
-	close(f.slots)
-	close(f.passed)
+	f.closeOnce.Do(func() {
+		f.it.Close()
+		f.cancel()
+		for range cap(f.slots) {
+			<-f.slots
+		}
+		close(f.slots)
+		close(f.passed)
+	})
 }
 
 // FairMix aggregates multiple node iterators. The mixer itself is an iterator which ends