Fix flaky TestTCPCollectingProcess_ConcurrentClient

pkg/collector/process_test.go

@@ -340,39 +340,59 @@ func TestUDPCollectingProcess_ReceiveDataRecord(t *testing.T) {
 func TestTCPCollectingProcess_ConcurrentClient(t *testing.T) {
 	input := getCollectorInput(tcpTransport, false, false)
 	cp, _ := InitCollectingProcess(input)
+	go cp.Start()
+	// Stop the collector when the test returns, and never before: stopping it from one
+	// of the client goroutines below used to race with the other clients, which could
+	// then fail to connect.
+	defer cp.Stop()
+	// wait until collector is ready
+	waitForCollectorReady(t, cp)
+	collectorAddr := cp.GetAddress()
+
+	const numClients = 2
+	// The connections are kept open until the end of the test, so that all the clients
+	// are connected to the collector simultaneously. They also need to be referenced
+	// until then, as the garbage collector closes unreachable connections.
+	conns := make([]net.Conn, numClients)
+	sessionIDs := make([]string, numClients)
 	var wg sync.WaitGroup
-	wg.Add(2)
-	go func() {
-		defer wg.Done()
-		// wait until collector is ready
-		waitForCollectorReady(t, cp)
-		collectorAddr := cp.GetAddress()
-		_, err := net.Dial(collectorAddr.Network(), collectorAddr.String())
-		if err != nil {
-			t.Errorf("Cannot establish connection to %s", collectorAddr.String())
+	for i := range numClients {
+		wg.Add(1)
+		go func() {
+			defer wg.Done()
+			conn, err := net.Dial(collectorAddr.Network(), collectorAddr.String())
+			// Use assert and not require, which must not be called from a goroutine
+			// other than the one running the test.
+			if !assert.NoErrorf(t, err, "Cannot establish connection to %s", collectorAddr.String()) {
+				return
+			}
+			conns[i] = conn
+			sessionIDs[i] = localConnSessionID(conn)
+		}()
+	}
+	wg.Wait()
+	for _, conn := range conns {
+		if conn != nil {
+			defer conn.Close()
 		}
-	}()
-	go func() {
-		defer wg.Done()
-		// wait until collector is ready
-		waitForCollectorReady(t, cp)
-		collectorAddr := cp.GetAddress()
-		_, err := net.Dial(collectorAddr.Network(), collectorAddr.String())
-		if err != nil {
-			t.Errorf("Cannot establish connection to %s", collectorAddr.String())
+	}
+
+	// Check for these specific sessions, and not just the session count, so that the
+	// assertion cannot be satisfied by an unrelated session. Sessions are registered
+	// asynchronously by the accept loop, hence Eventually.
+	assert.Eventually(t, func() bool {
+		cp.mutex.RLock()
+		defer cp.mutex.RUnlock()
+		if len(cp.sessions) != numClients {
+			return false
 		}
-		// Poll until both connections are registered by the collector's accept loop,
-		// rather than relying on a fixed sleep that can be too short on slow CI runners.
-		assert.Eventually(t, func() bool {
-			return cp.GetNumConnToCollector() >= 2
-		}, 5*time.Second, 10*time.Millisecond, "There should be at least two tcp clients.")
-		cp.Stop()
-	}()
-	cp.Start()
-	// Ensure both goroutines (and their calls into t) have finished before the
-	// test returns, otherwise a slow goroutine can call t.Errorf after the
-	// test has already completed, causing a panic.
-	wg.Wait()
+		for _, id := range sessionIDs {
+			if _, ok := cp.sessions[id]; !ok {
+				return false
+			}
+		}
+		return true
+	}, 5*time.Second, 10*time.Millisecond, "The collector should have a session for each tcp client.")
 }
 
 func TestUDPCollectingProcess_ConcurrentClient(t *testing.T) {
@@ -1024,18 +1044,22 @@ func getCollectorInput(network string, isEncrypted bool, isIPv6 bool) CollectorI
 	}
 }
 
+// waitForCollectorReady waits until the collector's socket is bound and its listening
+// address is available. Both the TCP and the UDP code paths bind the socket before
+// calling updateAddress, so a non-nil address means that the kernel is already queuing
+// incoming connections / datagrams for the collector, even if its accept or read loop
+// has not been scheduled yet.
+//
+// It deliberately does not dial the collector to probe for readiness: such a connection
+// is registered as a session by the TCP accept loop, and is only unregistered
+// asynchronously after it is closed, which contaminates session assertions in tests.
 func waitForCollectorReady(t *testing.T, cp *CollectingProcess) {
-	checkConn := func(ctx context.Context) (bool, error) {
-		if conn, err := net.Dial(cp.GetAddress().Network(), cp.GetAddress().String()); err != nil {
-			return false, err
-		} else {
-			defer conn.Close()
-			return true, nil
-		}
-	}
-	if err := wait.PollUntilContextTimeout(context.Background(), 100*time.Millisecond, 500*time.Millisecond, false, checkConn); err != nil {
-		t.Errorf("Cannot establish connection to %s", cp.GetAddress().String())
+	t.Helper()
+	checkAddr := func(ctx context.Context) (bool, error) {
+		return cp.GetAddress() != nil, nil
 	}
+	err := wait.PollUntilContextTimeout(context.Background(), 10*time.Millisecond, 5*time.Second, true, checkAddr)
+	require.NoError(t, err, "Collecting process is not ready")
 }
 
 func disableLogToStderr() {