p2p/discover: return discutil.Iterator from RandomNodes

This commit is contained in:
Felix Lange 2019-07-07 23:15:03 +02:00
parent be439f4221
commit a6b300dd74
2 changed files with 15 additions and 21 deletions

View file

@ -66,8 +66,8 @@ type lookupFunc func(func(*enode.Node))
type lookupWalker struct { type lookupWalker struct {
lookup lookupFunc lookup lookupFunc
newIterCh chan *Iterator newIterCh chan *lookupIterator
delIterCh chan *Iterator delIterCh chan *lookupIterator
triggerCh chan struct{} triggerCh chan struct{}
closeCh chan struct{} closeCh chan struct{}
wg sync.WaitGroup wg sync.WaitGroup
@ -76,8 +76,8 @@ type lookupWalker struct {
func newLookupWalker(fn lookupFunc) *lookupWalker { func newLookupWalker(fn lookupFunc) *lookupWalker {
w := &lookupWalker{ w := &lookupWalker{
lookup: fn, lookup: fn,
newIterCh: make(chan *Iterator), newIterCh: make(chan *lookupIterator),
delIterCh: make(chan *Iterator), delIterCh: make(chan *lookupIterator),
triggerCh: make(chan struct{}), triggerCh: make(chan struct{}),
closeCh: make(chan struct{}), closeCh: make(chan struct{}),
} }
@ -93,7 +93,7 @@ func (w *lookupWalker) close() {
func (w *lookupWalker) loop() { func (w *lookupWalker) loop() {
var ( var (
iters = make(map[*Iterator]struct{}) iters = make(map[*lookupIterator]struct{})
foundNode = make(chan *enode.Node) foundNode = make(chan *enode.Node)
lookupDone = make(chan struct{}, 1) lookupDone = make(chan struct{}, 1)
trigger = w.triggerCh trigger = w.triggerCh
@ -141,8 +141,8 @@ func (w *lookupWalker) runLookup(nodes chan<- *enode.Node, done chan<- struct{})
done <- struct{}{} done <- struct{}{}
} }
// Iterator is a sequence of discovered nodes. // lookupIterator is a sequence of discovered nodes.
type Iterator struct { type lookupIterator struct {
w *lookupWalker w *lookupWalker
buf chan *enode.Node buf chan *enode.Node
closed bool closed bool
@ -151,8 +151,8 @@ type Iterator struct {
const lookupIteratorBuffer = 100 const lookupIteratorBuffer = 100
func (w *lookupWalker) newIterator() *Iterator { func (w *lookupWalker) newIterator() *lookupIterator {
it := &Iterator{w: w, buf: make(chan *enode.Node, lookupIteratorBuffer)} it := &lookupIterator{w: w, buf: make(chan *enode.Node, lookupIteratorBuffer)}
select { select {
case w.newIterCh <- it: case w.newIterCh <- it:
case <-w.closeCh: case <-w.closeCh:
@ -162,13 +162,7 @@ func (w *lookupWalker) newIterator() *Iterator {
return it return it
} }
// NextNode retrieves the next node if one could be discovered before the passed context func (it *lookupIterator) NextNode(ctx context.Context) (n *enode.Node, isLive bool) {
// 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 { for {
select { select {
case it.w.triggerCh <- struct{}{}: case it.w.triggerCh <- struct{}{}:
@ -184,8 +178,7 @@ func (it *Iterator) NextNode(ctx context.Context) (n *enode.Node, isLive bool) {
} }
} }
// Close ends the iterator. This can be called concurrently with NextNode. func (it *lookupIterator) Close() {
func (it *Iterator) Close() {
it.closeOnce.Do(func() { it.closeOnce.Do(func() {
select { select {
case it.w.delIterCh <- it: case it.w.delIterCh <- it:
@ -196,7 +189,7 @@ func (it *Iterator) Close() {
} }
// deliver sends a node to the iterator buffer. // deliver sends a node to the iterator buffer.
func (it *Iterator) deliver(n *enode.Node) { func (it *lookupIterator) deliver(n *enode.Node) {
// We don't want deliver to block and replace stale results when they're not being // We don't want deliver to block and replace stale results when they're not being
// read. Check whether the buffer is full and allow one receive from the buffer if so. // read. Check whether the buffer is full and allow one receive from the buffer if so.
// This is OK because there is only one writer. // This is OK because there is only one writer.
@ -214,7 +207,7 @@ func (it *Iterator) deliver(n *enode.Node) {
} }
} }
func (it *Iterator) drainAndClose() { func (it *lookupIterator) drainAndClose() {
for len(it.buf) > 0 { for len(it.buf) > 0 {
<-it.buf <-it.buf
} }

View file

@ -30,6 +30,7 @@ import (
"github.com/ethereum/go-ethereum/crypto" "github.com/ethereum/go-ethereum/crypto"
"github.com/ethereum/go-ethereum/log" "github.com/ethereum/go-ethereum/log"
"github.com/ethereum/go-ethereum/p2p/discutil"
"github.com/ethereum/go-ethereum/p2p/enode" "github.com/ethereum/go-ethereum/p2p/enode"
"github.com/ethereum/go-ethereum/p2p/enr" "github.com/ethereum/go-ethereum/p2p/enr"
"github.com/ethereum/go-ethereum/p2p/netutil" "github.com/ethereum/go-ethereum/p2p/netutil"
@ -304,7 +305,7 @@ func (t *UDPv4) Close() {
} }
// RandomNodes is an iterator yielding nodes from a random walk of the DHT. // RandomNodes is an iterator yielding nodes from a random walk of the DHT.
func (t *UDPv4) RandomNodes() *Iterator { func (t *UDPv4) RandomNodes() discutil.Iterator {
return t.randomWalk.newIterator() return t.randomWalk.newIterator()
} }