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"
 )
@@ -156,50 +157,52 @@ func (f *filterIter) Next() bool {
 // 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
+	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 {
+		for f.it.Next() {
+			n := f.it.Node()
+			<-f.slots
 			// check the node async, in a separate goroutine
 			go func() {
-						if nn, err := f.check(n); err == nil {
+				if nn, err := check(ctx, n); err == nil {
 					select {
 					case f.passed <- nn:
-							case <-f.closed: // bale out if downstream is already closed and not calling Next
+					case <-ctx.Done(): // 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
 		}
-			} else {
 		// the iterator has ended
 		f.slots <- struct{}{}
-				break
-			}
-		}
 	}()
 
 	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.closeOnce.Do(func() {
 		f.it.Close()
-	close(f.closed) // override the passed channel
+		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