p2p/discover: port discv5 to iterator API

p2p/discover/v4_udp.go

@@ -210,7 +210,7 @@ type UDPv4 struct {
 	addReplyMatcher chan *replyMatcher
 	gotreply        chan reply
 	closeCtx        context.Context
-	cancelCloseCtx  func()
+	cancelCloseCtx  context.CancelFunc
 }
 
 // replyMatcher represents a pending reply.

p2p/discover/v5_udp.go

@@ -18,6 +18,7 @@ package discover
 
 import (
 	"bytes"
+	"context"
 	"crypto/ecdsa"
 	crand "crypto/rand"
 	"errors"
@@ -92,9 +93,10 @@ type UDPv5 struct {
 	callQueue        map[enode.ID][]*callV5
 
 	// shutdown stuff
-	closing   chan struct{}
+	closeOnce      sync.Once
-	closeOnce sync.Once
+	closeCtx       context.Context
-	wg        sync.WaitGroup
+	cancelCloseCtx context.CancelFunc
+	wg             sync.WaitGroup
 }
 
 // callV5 represents a remote procedure call against another node.
@@ -133,6 +135,7 @@ func ListenV5(conn UDPConn, ln *enode.LocalNode, cfg Config) (*UDPv5, error) {
 
 // newUDPv5 creates a UDPv5 transport, but doesn't start any goroutines.
 func newUDPv5(conn UDPConn, ln *enode.LocalNode, cfg Config) (*UDPv5, error) {
+	closeCtx, cancelCloseCtx := context.WithCancel(context.Background())
 	cfg = cfg.withDefaults()
 	t := &UDPv5{
 		// static fields
@@ -150,12 +153,14 @@ func newUDPv5(conn UDPConn, ln *enode.LocalNode, cfg Config) (*UDPv5, error) {
 		callCh:        make(chan *callV5),
 		callDoneCh:    make(chan *callV5),
 		respTimeoutCh: make(chan *callTimeout),
-		closing:       make(chan struct{}),
 		// state of dispatch
 		codec:            newWireCodec(ln, cfg.PrivateKey, cfg.Clock),
 		activeCallByNode: make(map[enode.ID]*callV5),
 		activeCallByAuth: make(map[string]*callV5),
 		callQueue:        make(map[enode.ID][]*callV5),
+		// shutdown
+		closeCtx:       closeCtx,
+		cancelCloseCtx: cancelCloseCtx,
 	}
 	tab, err := newTable(t, t.db, cfg.Bootnodes, cfg.Log)
 	if err != nil {
@@ -173,7 +178,7 @@ func (t *UDPv5) Self() *enode.Node {
 // Close shuts down packet processing.
 func (t *UDPv5) Close() {
 	t.closeOnce.Do(func() {
-		close(t.closing)
+		t.cancelCloseCtx()
 		t.conn.Close()
 		t.wg.Wait()
 		t.tab.close()
@@ -206,81 +211,48 @@ func (t *UDPv5) Resolve(n *enode.Node) *enode.Node {
 	return n
 }
 
-// LookupRandom finds random nodes in the network.
+func (t *UDPv5) RandomNodes() enode.Iterator {
-func (t *UDPv5) LookupRandom() []*enode.Node {
 	if t.tab.len() == 0 {
 		// All nodes were dropped, refresh. The very first query will hit this
 		// case and run the bootstrapping logic.
 		<-t.tab.refresh()
 	}
-	return t.lookupRandom()
-}
 
-// lookupRandom looks up a random target.
+	return newLookupIterator(t.closeCtx, t.newRandomLookup)
-// This is needed to satisfy the transport interface.
-func (t *UDPv5) lookupRandom() []*enode.Node {
-	var target enode.ID
-	crand.Read(target[:])
-	return t.Lookup(target)
-}
-
-// lookupSelf looks up our own node ID.
-// This is needed to satisfy the transport interface.
-func (t *UDPv5) lookupSelf() []*enode.Node {
-	return t.Lookup(t.Self().ID())
 }
 
 // Lookup performs a recursive lookup for the given target.
 // It returns the closest nodes to target.
 func (t *UDPv5) Lookup(target enode.ID) []*enode.Node {
-	var (
+	return t.newLookup(t.closeCtx, target).run()
-		asked          = make(map[enode.ID]bool)
+}
-		seen           = make(map[enode.ID]bool)
-		response       = make(chan []*node, alpha)
-		pendingQueries = 0
-		result         *nodesByDistance
-	)
-	// Don't query further if we hit ourself.
-	// Unlikely to happen often in practice.
-	asked[t.Self().ID()] = true
 
-	// Generate the initial result set.
+// lookupRandom looks up a random target.
-	t.tab.mutex.Lock()
+// This is needed to satisfy the transport interface.
-	result = t.tab.closest(target, bucketSize, false)
+func (t *UDPv5) lookupRandom() []*enode.Node {
-	t.tab.mutex.Unlock()
+	return t.newRandomLookup(t.closeCtx).run()
+}
 
-	for {
+// lookupSelf looks up our own node ID.
-		// Ask the closest nodes we haven't asked yet.
+// This is needed to satisfy the transport interface.
-		for i := 0; i < len(result.entries) && pendingQueries < alpha; i++ {
+func (t *UDPv5) lookupSelf() []*enode.Node {
-			n := result.entries[i]
+	return t.newLookup(t.closeCtx, t.Self().ID()).run()
-			if !asked[n.ID()] {
+}
-				asked[n.ID()] = true
+
-				pendingQueries++
+func (t *UDPv5) newRandomLookup(ctx context.Context) *lookup {
-				go t.lookupWorker(n, target, response)
+	var target enode.ID
-			}
+	crand.Read(target[:])
-		}
+	return t.newLookup(ctx, target)
-		if pendingQueries == 0 {
+}
-			// We have asked all closest nodes, stop the search.
+
-			break
+func (t *UDPv5) newLookup(ctx context.Context, target enode.ID) *lookup {
-		}
+	return newLookup(ctx, t.tab, target, func(n *node) ([]*node, error) {
-		select {
+		return t.lookupWorker(n, target)
-		case nodes := <-response:
+	})
-			for _, n := range nodes {
-				if n != nil && !seen[n.ID()] {
-					seen[n.ID()] = true
-					result.push(n, bucketSize)
-				}
-			}
-		case <-t.closing:
-			return nil // Shutdown, no need to continue.
-		}
-		pendingQueries--
-	}
-	return unwrapNodes(result.entries)
 }
 
 // lookupWorker performs FINDNODE calls against a single node during lookup.
-func (t *UDPv5) lookupWorker(destNode *node, target enode.ID, response chan<- []*node) {
+func (t *UDPv5) lookupWorker(destNode *node, target enode.ID) ([]*node, error) {
 	var (
 		dists = lookupDistances(target, destNode.ID())
 		nodes = nodesByDistance{target: target}
@@ -288,22 +260,8 @@ func (t *UDPv5) lookupWorker(destNode *node, target enode.ID, response chan<- []
 	for i := 0; i < lookupRequestLimit && len(nodes.entries) < findnodeResultLimit; i++ {
 		fails := t.db.FindFailsV5(destNode.ID())
 		r, err := t.findnode(unwrapNode(destNode), dists[i])
-		if err == errClosed {
-			// Avoid recording failures on shutdown.
-			nodes.entries = nil
-			break
-		}
 		if len(r) == 0 {
-			// The query failed. Record the failure and drop the node if it fails repeatedly.
-			fails++
 			t.log.Trace("FINDNODE/v5 call found no useful nodes", "id", destNode.ID(), "d", dists[i], "failcount", fails, "err", err)
-			if fails >= maxFindnodeFailures {
-				t.log.Trace("Too many findnode failures, dropping", "id", destNode.ID(), "failcount", fails)
-				t.tab.delete(destNode)
-				break
-			}
-		} else if fails > 0 {
-			t.db.UpdateFindFailsV5(destNode.ID(), fails-1)
 		}
 		for _, n := range r {
 			if n.ID() != t.Self().ID() {
@@ -311,13 +269,7 @@ func (t *UDPv5) lookupWorker(destNode *node, target enode.ID, response chan<- []
 			}
 		}
 	}
-
+	return nodes.entries, nil
-	// Add all result nodes to table. Some of them might not be alive anymore, but we'll
-	// just remove those again during revalidation.
-	for _, n := range nodes.entries {
-		t.tab.addSeenNode(n)
-	}
-	response <- nodes.entries
 }
 
 // lookupDistances computes the distance parameter for FINDNODE calls to dest.
@@ -453,7 +405,7 @@ func (t *UDPv5) call(node *enode.Node, responseType byte, packet packetV5) *call
 	// Send call to dispatch.
 	select {
 	case t.callCh <- c:
-	case <-t.closing:
+	case <-t.closeCtx.Done():
 		c.err <- errClosed
 	}
 	return c
@@ -463,7 +415,7 @@ func (t *UDPv5) call(node *enode.Node, responseType byte, packet packetV5) *call
 func (t *UDPv5) callDone(c *callV5) {
 	select {
 	case t.callDoneCh <- c:
-	case <-t.closing:
+	case <-t.closeCtx.Done():
 	}
 }
 
@@ -513,7 +465,7 @@ func (t *UDPv5) dispatch() {
 			// Arm next read.
 			t.readNextCh <- struct{}{}
 
-		case <-t.closing:
+		case <-t.closeCtx.Done():
 			close(t.readNextCh)
 			for id, queue := range t.callQueue {
 				for _, c := range queue {
@@ -544,7 +496,7 @@ func (t *UDPv5) startResponseTimeout(c *callV5) {
 		<-done
 		select {
 		case t.respTimeoutCh <- &callTimeout{c, timer}:
-		case <-t.closing:
+		case <-t.closeCtx.Done():
 		}
 	})
 	c.timeout = timer
@@ -619,7 +571,7 @@ func (t *UDPv5) readLoop() {
 		}
 		select {
 		case t.packetInCh <- ReadPacket{Data: buf[:nbytes], Addr: from}:
-		case <-t.closing:
+		case <-t.closeCtx.Done():
 			return
 		}
 	}