les: fix data races in tests (#23457)

les/ulc_test.go

@@ -20,6 +20,7 @@ import (
 	"crypto/rand"
 	"fmt"
 	"net"
+	"sync/atomic"
 	"testing"
 	"time"
 
@@ -65,7 +66,7 @@ func testULCAnnounceThreshold(t *testing.T, protocol int) {
 
 		// Connect all servers.
 		for i := 0; i < len(servers); i++ {
-			connect(servers[i].handler, nodes[i].ID(), c.handler, protocol)
+			connect(servers[i].handler, nodes[i].ID(), c.handler, protocol, false)
 		}
 		for i := 0; i < len(servers); i++ {
 			for j := 0; j < testcase.height[i]; j++ {
@@ -86,7 +87,7 @@ func testULCAnnounceThreshold(t *testing.T, protocol int) {
 	}
 }
 
-func connect(server *serverHandler, serverId enode.ID, client *clientHandler, protocol int) (*serverPeer, *clientPeer, error) {
+func connect(server *serverHandler, serverId enode.ID, client *clientHandler, protocol int, noInitAnnounce bool) (*serverPeer, *clientPeer, error) {
 	// Create a message pipe to communicate through
 	app, net := p2p.MsgPipe()
 
@@ -110,16 +111,22 @@ func connect(server *serverHandler, serverId enode.ID, client *clientHandler, pr
 		select {
 		case <-client.closeCh:
 			errc1 <- p2p.DiscQuitting
-		case errc1 <- client.handle(peer1):
+		case errc1 <- client.handle(peer1, noInitAnnounce):
 		}
 	}()
-
-	select {
-	case <-time.After(time.Millisecond * 100):
-	case err := <-errc1:
-		return nil, nil, fmt.Errorf("peerLight handshake error: %v", err)
-	case err := <-errc2:
-		return nil, nil, fmt.Errorf("peerFull handshake error: %v", err)
+	// Ensure the connection is established or exits when any error occurs
+	for {
+		select {
+		case err := <-errc1:
+			return nil, nil, fmt.Errorf("failed to establish protocol connection %v", err)
+		case err := <-errc2:
+			return nil, nil, fmt.Errorf("failed to establish protocol connection %v", err)
+		default:
+		}
+		if atomic.LoadUint32(&peer1.serving) == 1 && atomic.LoadUint32(&peer2.serving) == 1 {
+			break
+		}
+		time.Sleep(50 * time.Millisecond)
 	}
 	return peer1, peer2, nil
 }