p2p/discover: add context in waitForNodes

This commit is contained in:
Felix Lange 2025-08-29 15:47:09 +02:00
parent 46e4f0b5c1
commit f8e0e8dc55

View file

@ -708,7 +708,7 @@ func (tab *Table) deleteNode(n *enode.Node) {
} }
// waitForNodes blocks until the table contains at least n nodes. // waitForNodes blocks until the table contains at least n nodes.
func (tab *Table) waitForNodes(n int) bool { func (tab *Table) waitForNodes(ctx context.Context, n int) error {
getlength := func() (count int) { getlength := func() (count int) {
for _, b := range &tab.buckets { for _, b := range &tab.buckets {
count += len(b.entries) count += len(b.entries)
@ -721,7 +721,7 @@ func (tab *Table) waitForNodes(n int) bool {
tab.mutex.Lock() tab.mutex.Lock()
if getlength() >= n { if getlength() >= n {
tab.mutex.Unlock() tab.mutex.Unlock()
return true return nil
} }
if ch == nil { if ch == nil {
// Init subscription. // Init subscription.
@ -734,8 +734,10 @@ func (tab *Table) waitForNodes(n int) bool {
// Wait for a node add event. // Wait for a node add event.
select { select {
case <-ch: case <-ch:
case <-ctx.Done():
return ctx.Err()
case <-tab.closeReq: case <-tab.closeReq:
return false return errClosed
} }
} }
} }