diff --git a/pkg/collector/process_test.go b/pkg/collector/process_test.go index 6b8b3346..26b17f8d 100644 --- a/pkg/collector/process_test.go +++ b/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() {