diff --git a/p2p/discover/common.go b/p2p/discover/common.go index 3c080359fd..b8911e707e 100644 --- a/p2p/discover/common.go +++ b/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: + } +}