diff --git a/p2p/discover/common.go b/p2p/discover/common.go index 592f776ac9..1354c09a47 100644 --- a/p2p/discover/common.go +++ b/p2p/discover/common.go @@ -27,6 +27,7 @@ import ( "github.com/ethereum/go-ethereum/p2p/netutil" ) +// UDPConn is a network connection on which discovery can operate. type UDPConn interface { ReadFromUDP(b []byte) (n int, addr *net.UDPAddr, err error) WriteToUDP(b []byte, addr *net.UDPAddr) (n int, err error) @@ -34,7 +35,7 @@ type UDPConn interface { LocalAddr() net.Addr } -// Config holds Table-related settings. +// Config holds settings for the discovery listener. type Config struct { // These settings are required and configure the UDP listener: PrivateKey *ecdsa.PrivateKey @@ -52,7 +53,7 @@ func ListenUDP(c UDPConn, ln *enode.LocalNode, cfg Config) (*UDPv4, error) { } // ReadPacket is a packet that couldn't be handled. Those packets are sent to the unhandled -// channel if configured. +// channel if configured. This is exported for internal use, do not use this type. type ReadPacket struct { Data []byte Addr *net.UDPAddr @@ -61,7 +62,7 @@ type ReadPacket struct { 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. +// It manages a set iterators which receive lookup results as they are found. type lookupWalker struct { lookup lookupFunc @@ -119,7 +120,10 @@ func (w *lookupWalker) loop() { case <-w.closeCh: for it := range iters { - close(it.buf) + it.drainAndClose() + } + if trigger == nil { + <-lookupDone } w.wg.Done() return @@ -127,7 +131,7 @@ func (w *lookupWalker) loop() { } } -func (w *lookupWalker) runLookup(nodes chan<- *enode.Node, done chan struct{}) { +func (w *lookupWalker) runLookup(nodes chan<- *enode.Node, done chan<- struct{}) { w.lookup(func(n *enode.Node) { select { case nodes <- n: @@ -139,15 +143,16 @@ func (w *lookupWalker) runLookup(nodes chan<- *enode.Node, done chan struct{}) { // Iterator is a sequence of discovered nodes. type Iterator struct { - w *lookupWalker - buf chan *enode.Node - closed bool + w *lookupWalker + buf chan *enode.Node + closed bool + closeOnce sync.Once } const lookupIteratorBuffer = 100 func (w *lookupWalker) newIterator() *Iterator { - it := &Iterator{w, make(chan *enode.Node, lookupIteratorBuffer), false} + it := &Iterator{w: w, buf: make(chan *enode.Node, lookupIteratorBuffer)} select { case w.newIterCh <- it: case <-w.closeCh: @@ -168,42 +173,50 @@ func (it *Iterator) NextNode(ctx context.Context) (n *enode.Node, isLive bool) { select { case it.w.triggerCh <- struct{}{}: // lookup triggered - case <-it.w.closeCh: - it.closed = true - return nil, false case n, ok := <-it.buf: if !ok { it.closed = true } - return n, it.closed + return n, !it.closed case <-ctx.Done(): - return nil, it.closed // TODO: should be permanently closed if channel is closed once. + return nil, !it.closed } } } // 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: + it.closeOnce.Do(func() { + select { + case it.w.delIterCh <- it: + case <-it.w.closeCh: + } + it.drainAndClose() + }) +} + +// deliver sends a node to the iterator buffer. +func (it *Iterator) deliver(n *enode.Node) { + // 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. + // This is OK because there is only one writer. + var remove chan *enode.Node + if len(it.buf) == cap(it.buf) { + remove = it.buf + } + for { + select { + case it.buf <- n: + return + case <-remove: + remove = nil + } } } -// deliver sends n to the iterator buffer. -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 - // 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: +func (it *Iterator) drainAndClose() { + for len(it.buf) > 0 { + <-it.buf } + close(it.buf) } diff --git a/p2p/discover/table_util_test.go b/p2p/discover/table_util_test.go index 8e5fc7374b..2adfd0dc0e 100644 --- a/p2p/discover/table_util_test.go +++ b/p2p/discover/table_util_test.go @@ -17,11 +17,14 @@ package discover import ( + "bytes" "crypto/ecdsa" "encoding/hex" + "errors" "fmt" "math/rand" "net" + "reflect" "sort" "sync" @@ -169,6 +172,28 @@ func hasDuplicates(slice []*node) bool { return false } +func checkNodesEqual(got, want []*enode.Node) error { + if reflect.DeepEqual(got, want) { + return nil + } + output := new(bytes.Buffer) + fmt.Fprintf(output, "got %d nodes:\n", len(got)) + for _, n := range got { + fmt.Fprintf(output, " %v %v\n", n.ID(), n) + } + fmt.Fprintf(output, "want %d:\n", len(want)) + for _, n := range want { + fmt.Fprintf(output, " %v %v\n", n.ID(), n) + } + return errors.New(output.String()) +} + +func sortByID(nodes []*enode.Node) { + sort.Slice(nodes, func(i, j int) bool { + return string(nodes[i].ID().Bytes()) < string(nodes[j].ID().Bytes()) + }) +} + func sortedByDistanceTo(distbase enode.ID, slice []*node) bool { return sort.SliceIsSorted(slice, func(i, j int) bool { return enode.DistCmp(distbase, slice[i].ID(), slice[j].ID()) < 0 diff --git a/p2p/discover/v4_udp.go b/p2p/discover/v4_udp.go index c6d9c4ec22..fd90071bc3 100644 --- a/p2p/discover/v4_udp.go +++ b/p2p/discover/v4_udp.go @@ -320,6 +320,7 @@ func (t *UDPv4) randomLookupWithCallback(callback func(*enode.Node)) { t.lookup(target, callback) } +// LookupPubkey finds the closest nodes to the given public key. func (t *UDPv4) LookupPubkey(key *ecdsa.PublicKey) []*enode.Node { if t.tab.len() == 0 { // All nodes were dropped, refresh. The very first query will hit this @@ -329,12 +330,14 @@ func (t *UDPv4) LookupPubkey(key *ecdsa.PublicKey) []*enode.Node { return unwrapNodes(t.lookup(encodePubkey(key), nil)) } +// for Table func (t *UDPv4) lookupRandom() []*enode.Node { var target encPubkey crand.Read(target[:]) return unwrapNodes(t.lookup(target, nil)) } +// for Table func (t *UDPv4) lookupSelf() []*enode.Node { return unwrapNodes(t.lookup(encodePubkey(&t.priv.PublicKey), nil)) }