pkg/util/traceevent: stabilize flaky TestFlightRecorder

Author: D3Hunter

pkg/util/traceevent/flightrecorder_test.go

@@ -333,9 +333,25 @@ func testFlightRecorderConfigBadCase(t *testing.T) {
 	}
 }
 
+func testSamplingCounterDeterministic(t *testing.T) {
+	recorder := &HTTPFlightRecorder{}
+	conf := &DumpTriggerConfig{
+		Type:     "sampling",
+		Sampling: 5,
+	}
+	sampledCount := 0
+	for i := 0; i < 10; i++ {
+		if recorder.CheckSampling(conf) {
+			sampledCount++
+		}
+	}
+	require.Equal(t, 2, sampledCount)
+}
+
 func TestFlightRecorderConfig(t *testing.T) {
 	testFlightRecorderConfigGoodCase(t)
 	testFlightRecorderConfigBadCase(t)
+	testSamplingCounterDeterministic(t)
 }
 
 func TestParseTraceCategory(t *testing.T) {

pkg/util/traceevent/test/integration_test.go

@@ -24,6 +24,7 @@ import (
 	"regexp"
 	"strings"
 	"testing"
+	"time"
 
 	"github.com/pingcap/tidb/pkg/config"
 	"github.com/pingcap/tidb/pkg/config/kerneltype"
@@ -48,6 +49,38 @@ func drainEvents(eventCh <-chan []traceevent.Event) {
 	}
 }
 
+func eventsBelongToTrace(events []tracing.Event, traceID []byte) bool {
+	if len(events) == 0 || len(traceID) == 0 {
+		return false
+	}
+	matched := false
+	for _, event := range events {
+		if len(event.TraceID) == 0 {
+			continue
+		}
+		if !bytes.Equal(event.TraceID, traceID) {
+			return false
+		}
+		matched = true
+	}
+	return matched
+}
+
+func waitTraceEvents(t *testing.T, eventCh <-chan []tracing.Event, traceID []byte) []tracing.Event {
+	t.Helper()
+	deadline := time.After(3 * time.Second)
+	for {
+		select {
+		case events := <-eventCh:
+			if eventsBelongToTrace(events, traceID) {
+				return events
+			}
+		case <-deadline:
+			require.FailNowf(t, "failed to find trace events", "trace_id=%s", hex.EncodeToString(traceID))
+		}
+	}
+}
+
 func TestPrevTraceIDPersistence(t *testing.T) {
 	if kerneltype.IsClassic() {
 		t.Skip("trace events only work for next-gen kernel")
@@ -254,10 +287,12 @@ func TestFlightRecorder(t *testing.T) {
 		}
 		flightRecorder, err := traceevent.StartHTTPFlightRecorder(eventCh, &config)
 		require.NoError(t, err)
+		drainEvents(eventCh)
 		tk.MustQueryWithContext(ctx, "select * from t").Check(testkit.Rows())
+		traceID := bytes.Clone(tk.Session().GetSessionVars().PrevTraceID)
+		require.NotEmpty(t, traceID)
 		sink.DiscardOrFlush(ctx)
-		require.NotEmpty(t, eventCh)
-		events := <-eventCh
+		events := waitTraceEvents(t, eventCh, traceID)
 		for _, event := range events {
 			require.Equal(t, event.Category, traceevent.KvRequest)
 		}
@@ -279,7 +314,10 @@ func TestFlightRecorder(t *testing.T) {
 			tk.MustQueryWithContext(ctx, "select * from t").Check(testkit.Rows())
 			sink.DiscardOrFlush(ctx)
 		}
-		require.Len(t, eventCh, 2)
+		// The recorder state is process-global in this package, so unrelated traces can
+		// race into this channel. Keep this integration check as a smoke test and verify
+		// exact sampling math in unit tests.
+		require.GreaterOrEqual(t, len(eventCh), 1)
 		flightRecorder.Close()
 		drainEvents(eventCh)
 	}