From b5a88dafae52c7126b185aed841a418943fc5cb9 Mon Sep 17 00:00:00 2001
From: Martin HS <martin@swende.se>
Date: Thu, 26 Sep 2024 08:12:12 +0200
Subject: [PATCH] p2p/discover: fix flaky tests writing to test.log after
 completion (#30506)

This PR fixes two tests, which had a tendency to sometimes write to the `*testing.T` `log` facility after the test function had completed, which is not allowed. This PR fixes it by using waitgroups to ensure that the handler/logwriter terminates before the test exits.

closes #30505
---
 p2p/discover/v4_lookup_test.go | 26 ++++++++++++++++++++++----
 p2p/enode/nodedb.go            |  6 +++++-
 2 files changed, 27 insertions(+), 5 deletions(-)

diff --git a/p2p/discover/v4_lookup_test.go b/p2p/discover/v4_lookup_test.go
index 70bd7056fba2..29a9dd6645e0 100644
--- a/p2p/discover/v4_lookup_test.go
+++ b/p2p/discover/v4_lookup_test.go
@@ -21,6 +21,7 @@ import (
 	"fmt"
 	"net/netip"
 	"slices"
+	"sync"
 	"testing"
 
 	"github.com/ethereum/go-ethereum/crypto"
@@ -67,7 +68,11 @@ func TestUDPv4_Lookup(t *testing.T) {
 func TestUDPv4_LookupIterator(t *testing.T) {
 	t.Parallel()
 	test := newUDPTest(t)
-	defer test.close()
+	var wg sync.WaitGroup
+	defer func() {
+		test.close()
+		wg.Wait()
+	}()
 
 	// Seed table with initial nodes.
 	bootnodes := make([]*enode.Node, len(lookupTestnet.dists[256]))
@@ -75,7 +80,11 @@ func TestUDPv4_LookupIterator(t *testing.T) {
 		bootnodes[i] = lookupTestnet.node(256, i)
 	}
 	fillTable(test.table, bootnodes, true)
-	go serveTestnet(test, lookupTestnet)
+	wg.Add(1)
+	go func() {
+		serveTestnet(test, lookupTestnet)
+		wg.Done()
+	}()
 
 	// Create the iterator and collect the nodes it yields.
 	iter := test.udp.RandomNodes()
@@ -102,7 +111,11 @@ func TestUDPv4_LookupIterator(t *testing.T) {
 func TestUDPv4_LookupIteratorClose(t *testing.T) {
 	t.Parallel()
 	test := newUDPTest(t)
-	defer test.close()
+	var wg sync.WaitGroup
+	defer func() {
+		test.close()
+		wg.Wait()
+	}()
 
 	// Seed table with initial nodes.
 	bootnodes := make([]*enode.Node, len(lookupTestnet.dists[256]))
@@ -110,7 +123,12 @@ func TestUDPv4_LookupIteratorClose(t *testing.T) {
 		bootnodes[i] = lookupTestnet.node(256, i)
 	}
 	fillTable(test.table, bootnodes, true)
-	go serveTestnet(test, lookupTestnet)
+
+	wg.Add(1)
+	go func() {
+		serveTestnet(test, lookupTestnet)
+		wg.Done()
+	}()
 
 	it := test.udp.RandomNodes()
 	if ok := it.Next(); !ok || it.Node() == nil {
diff --git a/p2p/enode/nodedb.go b/p2p/enode/nodedb.go
index 1f31c98d2256..51e554e68adc 100644
--- a/p2p/enode/nodedb.go
+++ b/p2p/enode/nodedb.go
@@ -496,6 +496,10 @@ func nextNode(it iterator.Iterator) *Node {
 
 // Close flushes and closes the database files.
 func (db *DB) Close() {
-	close(db.quit)
+	select {
+	case <-db.quit: // already closed
+	default:
+		close(db.quit)
+	}
 	db.lvl.Close()
 }