diff --git a/p2p/discover/common.go b/p2p/discover/common.go index b8911e707e..592f776ac9 100644 --- a/p2p/discover/common.go +++ b/p2p/discover/common.go @@ -58,13 +58,6 @@ type ReadPacket struct { Addr *net.UDPAddr } -// Iterator is an iterator over nodes. -type Iterator interface { - // NextNode returns the next node if one could be discovered before the - // passed context was canceled. - NextNode(context.Context) *enode.Node -} - type lookupFunc func(func(*enode.Node)) // lookupWalker performs recursive lookups, walking the DHT. @@ -72,7 +65,8 @@ type lookupFunc func(func(*enode.Node)) type lookupWalker struct { lookup lookupFunc - newIterCh chan *lookupIterator + newIterCh chan *Iterator + delIterCh chan *Iterator triggerCh chan struct{} closeCh chan struct{} wg sync.WaitGroup @@ -81,7 +75,8 @@ type lookupWalker struct { func newLookupWalker(fn lookupFunc) *lookupWalker { w := &lookupWalker{ lookup: fn, - newIterCh: make(chan *lookupIterator), + newIterCh: make(chan *Iterator), + delIterCh: make(chan *Iterator), triggerCh: make(chan struct{}), closeCh: make(chan struct{}), } @@ -97,7 +92,7 @@ func (w *lookupWalker) close() { func (w *lookupWalker) loop() { var ( - iters = make(map[*lookupIterator]struct{}) + iters = make(map[*Iterator]struct{}) foundNode = make(chan *enode.Node) lookupDone = make(chan struct{}, 1) trigger = w.triggerCh @@ -107,6 +102,9 @@ func (w *lookupWalker) loop() { case it := <-w.newIterCh: iters[it] = struct{}{} + case it := <-w.delIterCh: + delete(iters, it) + case <-trigger: trigger = nil // stop listening to trigger until lookupDone go w.runLookup(foundNode, lookupDone) @@ -120,6 +118,9 @@ func (w *lookupWalker) loop() { } case <-w.closeCh: + for it := range iters { + close(it.buf) + } w.wg.Done() return } @@ -136,39 +137,62 @@ func (w *lookupWalker) runLookup(nodes chan<- *enode.Node, done chan struct{}) { done <- struct{}{} } -type lookupIterator struct { - w *lookupWalker - buf chan *enode.Node +// Iterator is a sequence of discovered nodes. +type Iterator struct { + w *lookupWalker + buf chan *enode.Node + closed bool } const lookupIteratorBuffer = 100 -func (w *lookupWalker) newIterator() Iterator { - it := &lookupIterator{w, make(chan *enode.Node, lookupIteratorBuffer)} +func (w *lookupWalker) newIterator() *Iterator { + it := &Iterator{w, make(chan *enode.Node, lookupIteratorBuffer), false} select { case w.newIterCh <- it: case <-w.closeCh: + it.closed = true + close(it.buf) } return it } -// NextNode gets the next node from the buffer. -// This keeps triggering new lookups until a node is delivered. -func (it *lookupIterator) NextNode(ctx context.Context) *enode.Node { +// NextNode retrieves the next node if one could be discovered before the passed context +// was canceled. This triggers a lookup operation if none is running. The isLive return +// value says whether the iterator is still open. NextNode returns (nil, false) after Close +// has been called. +// +// NextNode is not safe for concurrent use. +func (it *Iterator) NextNode(ctx context.Context) (n *enode.Node, isLive bool) { for { select { case it.w.triggerCh <- struct{}{}: // lookup triggered - case n := <-it.buf: - return n + case <-it.w.closeCh: + it.closed = true + return nil, false + case n, ok := <-it.buf: + if !ok { + it.closed = true + } + return n, it.closed case <-ctx.Done(): - return nil + return nil, it.closed // TODO: should be permanently closed if channel is closed once. } } } +// Close ends the iterator. This can be called concurrently with NextNode. +func (it *Iterator) Close() { + select { + case it.w.delIterCh <- it: + close(it.buf) + case <-it.w.closeCh: + } +} + // deliver sends n to the iterator buffer. -func (it *lookupIterator) deliver(n *enode.Node) { +func (it *Iterator) deliver(n *enode.Node) { // We don't want deliver to block and replacing stale results is OK if they're not // being read fast enough. Check whether the buffer is full and enable the select case // which removes an element if so. This doesn't race because deliver is only called by diff --git a/p2p/discover/v4_udp.go b/p2p/discover/v4_udp.go index 11a9c9379e..c6d9c4ec22 100644 --- a/p2p/discover/v4_udp.go +++ b/p2p/discover/v4_udp.go @@ -304,7 +304,7 @@ func (t *UDPv4) Close() { } // RandomNodes is an iterator yielding nodes from a random walk of the DHT. -func (t *UDPv4) RandomNodes() Iterator { +func (t *UDPv4) RandomNodes() *Iterator { return t.randomWalk.newIterator() }