p2p/discover: add RandomNodes iterator

p2p/discover/common.go

@@ -17,8 +17,10 @@
 package discover
 
 import (
+	"context"
 	"crypto/ecdsa"
 	"net"
+	"sync"
 
 	"github.com/ethereum/go-ethereum/log"
 	"github.com/ethereum/go-ethereum/p2p/enode"
@@ -55,3 +57,129 @@ type ReadPacket struct {
 	Data []byte
 	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.
+// It manages a set iterators which receive lookup results in real time.
+type lookupWalker struct {
+	lookup lookupFunc
+
+	newIterCh chan *lookupIterator
+	triggerCh chan struct{}
+	closeCh   chan struct{}
+	wg        sync.WaitGroup
+}
+
+func newLookupWalker(fn lookupFunc) *lookupWalker {
+	w := &lookupWalker{
+		lookup:    fn,
+		newIterCh: make(chan *lookupIterator),
+		triggerCh: make(chan struct{}),
+		closeCh:   make(chan struct{}),
+	}
+	w.wg.Add(1)
+	go w.loop()
+	return w
+}
+
+func (w *lookupWalker) close() {
+	close(w.closeCh)
+	w.wg.Wait()
+}
+
+func (w *lookupWalker) loop() {
+	var (
+		iters      = make(map[*lookupIterator]struct{})
+		foundNode  = make(chan *enode.Node)
+		lookupDone = make(chan struct{}, 1)
+		trigger    = w.triggerCh
+	)
+	for {
+		select {
+		case it := <-w.newIterCh:
+			iters[it] = struct{}{}
+
+		case <-trigger:
+			trigger = nil // stop listening to trigger until lookupDone
+			go w.runLookup(foundNode, lookupDone)
+
+		case <-lookupDone:
+			trigger = w.triggerCh
+
+		case n := <-foundNode:
+			for it := range iters {
+				it.deliver(n)
+			}
+
+		case <-w.closeCh:
+			w.wg.Done()
+			return
+		}
+	}
+}
+
+func (w *lookupWalker) runLookup(nodes chan<- *enode.Node, done chan struct{}) {
+	w.lookup(func(n *enode.Node) {
+		select {
+		case nodes <- n:
+		case <-w.closeCh:
+		}
+	})
+	done <- struct{}{}
+}
+
+type lookupIterator struct {
+	w   *lookupWalker
+	buf chan *enode.Node
+}
+
+const lookupIteratorBuffer = 100
+
+func (w *lookupWalker) newIterator() Iterator {
+	it := &lookupIterator{w, make(chan *enode.Node, lookupIteratorBuffer)}
+	select {
+	case w.newIterCh <- it:
+	case <-w.closeCh:
+	}
+	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 {
+	for {
+		select {
+		case it.w.triggerCh <- struct{}{}:
+			// lookup triggered
+		case n := <-it.buf:
+			return n
+		case <-ctx.Done():
+			return nil
+		}
+	}
+}
+
+// deliver sends n to the iterator buffer.
+func (it *lookupIterator) 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
+	// a single goroutine at a time.
+	remove := it.buf
+	if len(it.buf) < cap(it.buf) {
+		remove = nil
+	}
+	select {
+	case it.buf <- n:
+		return
+	case <-remove:
+	}
+}