p2p/discover: improve iterator

This commit is contained in:
Felix Lange 2025-09-05 11:15:09 +02:00
parent f4046b0cfb
commit 4ed8f5ee2b

View file

@ -49,11 +49,14 @@ func newLookup(ctx context.Context, tab *Table, target enode.ID, q queryFunc) *l
result: nodesByDistance{target: target},
replyCh: make(chan []*enode.Node, alpha),
cancelCh: ctx.Done(),
queries: -1,
}
// Don't query further if we hit ourself.
// Unlikely to happen often in practice.
it.asked[tab.self().ID()] = true
// Initialize the lookup with nodes from table.
closest := it.tab.findnodeByID(it.result.target, bucketSize, false)
it.addNodes(closest.entries)
return it
}
@ -64,22 +67,18 @@ func (it *lookup) run() []*enode.Node {
return it.result.entries
}
func (it *lookup) empty() bool {
return len(it.replyBuffer) == 0
}
// advance advances the lookup until any new nodes have been found.
// It returns false when the lookup has ended.
func (it *lookup) advance() bool {
for it.startQueries() {
select {
case nodes := <-it.replyCh:
it.replyBuffer = it.replyBuffer[:0]
for _, n := range nodes {
if n != nil && !it.seen[n.ID()] {
it.seen[n.ID()] = true
it.result.push(n, bucketSize)
it.replyBuffer = append(it.replyBuffer, n)
}
}
it.queries--
if len(it.replyBuffer) > 0 {
if it.addNodes(nodes) {
return true
}
case <-it.cancelCh:
@ -89,6 +88,18 @@ func (it *lookup) advance() bool {
return false
}
func (it *lookup) addNodes(nodes []*enode.Node) (done bool) {
it.replyBuffer = it.replyBuffer[:0]
for _, n := range nodes {
if n != nil && !it.seen[n.ID()] {
it.seen[n.ID()] = true
it.result.push(n, bucketSize)
it.replyBuffer = append(it.replyBuffer, n)
}
}
return len(it.replyBuffer) == 0
}
func (it *lookup) shutdown() {
for it.queries > 0 {
<-it.replyCh
@ -103,17 +114,6 @@ func (it *lookup) startQueries() bool {
return false
}
// The first query returns nodes from the local table.
if it.queries == -1 {
closest := it.tab.findnodeByID(it.result.target, bucketSize, false)
if len(closest.entries) == 0 {
return false
}
it.queries = 1
it.replyCh <- closest.entries
return true
}
// Ask the closest nodes that we haven't asked yet.
for i := 0; i < len(it.result.entries) && it.queries < alpha; i++ {
n := it.result.entries[i]
@ -180,10 +180,14 @@ func (it *lookupIterator) Next() bool {
}
if it.lookup == nil {
it.lookup = it.nextLookup(it.ctx)
if it.lookup.empty() {
// If the lookup is empty right after creation, it means the local table
// is in a degraded state, and we need to wait for it to fill again.
it.lookupFailed(it.lookup.tab)
}
continue
}
if !it.lookup.advance() {
it.lookupFailed(it.lookup.tab)
it.lookup = nil
continue
}