diff --git a/eth/backend.go b/eth/backend.go index e93d023d25..1feb8865a2 100644 --- a/eth/backend.go +++ b/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) } diff --git a/p2p/enode/iter.go b/p2p/enode/iter.go index b4570d1b85..ee0c12fdd4 100644 --- a/p2p/enode/iter.go +++ b/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