fix(virtual-thread): Reduce test flakiness (#10834)

fix(virtual-thread): Reduce test flakiness

Introduce generic blockUntil method to wait for trace pattern before evaluating them

Co-authored-by: bruce.bujon <bruce.bujon@datadoghq.com>

Author: PerfectSlayer

dd-java-agent/instrumentation-testing/src/main/java/datadog/trace/agent/test/AbstractInstrumentationTest.java

@@ -22,10 +22,12 @@
 import datadog.trace.core.TraceCollector;
 import java.lang.instrument.ClassFileTransformer;
 import java.lang.instrument.Instrumentation;
+import java.util.List;
 import java.util.ServiceLoader;
 import java.util.concurrent.TimeUnit;
 import java.util.concurrent.TimeoutException;
 import java.util.function.Function;
+import java.util.function.Predicate;
 import net.bytebuddy.agent.ByteBuddyAgent;
 import org.junit.jupiter.api.AfterEach;
 import org.junit.jupiter.api.BeforeEach;
@@ -138,6 +140,26 @@ protected void assertTraces(
     TraceAssertions.assertTraces(this.writer, options, matchers);
   }
 
+  /**
+   * Blocks the current thread until the traces written match the given predicate or the timeout
+   * occurs.
+   *
+   * @param predicate the condition that must be satisfied by the list of traces
+   */
+  protected void blockUntilTracesMatch(Predicate<List<List<DDSpan>>> predicate) {
+    long deadline = System.currentTimeMillis() + TIMEOUT_MILLIS;
+    while (!predicate.test(this.writer)) {
+      if (System.currentTimeMillis() > deadline) {
+        throw new RuntimeException(new TimeoutException("Timed out waiting for traces/spans."));
+      }
+      try {
+        Thread.sleep(10);
+      } catch (InterruptedException e) {
+        Thread.currentThread().interrupt();
+      }
+    }
+  }
+
   protected void blockUntilChildSpansFinished(final int numberOfSpans) {
     blockUntilChildSpansFinished(this.tracer.activeSpan(), numberOfSpans);
   }
@@ -147,7 +169,7 @@ static void blockUntilChildSpansFinished(AgentSpan span, int numberOfSpans) {
       TraceCollector traceCollector = ((DDSpan) span).context().getTraceCollector();
       if (!(traceCollector instanceof PendingTrace)) {
         throw new IllegalStateException(
-            "Expected $PendingTrace.name trace collector, got $traceCollector.class.name");
+            "Expected PendingTrace trace collector, got " + traceCollector.getClass().getName());
       }
 
       PendingTrace pendingTrace = (PendingTrace) traceCollector;

dd-java-agent/instrumentation/java/java-lang/java-lang-21.0/src/test/java/testdog/trace/instrumentation/java/lang/jdk21/VirtualThreadApiInstrumentationTest.java

@@ -6,7 +6,6 @@
 
 import datadog.trace.agent.test.AbstractInstrumentationTest;
 import datadog.trace.api.Trace;
-import java.util.concurrent.CountDownLatch;
 import java.util.concurrent.ThreadFactory;
 import java.util.concurrent.TimeoutException;
 import org.junit.jupiter.api.DisplayName;
@@ -96,7 +95,6 @@ public void run() {
   @Test
   void testNestedVirtualThreads() throws InterruptedException, TimeoutException {
     Thread.Builder.OfVirtual threadBuilder = Thread.ofVirtual();
-    CountDownLatch latch = new CountDownLatch(3);
 
     new Runnable() {
       @Trace(operationName = "parent")
@@ -118,19 +116,17 @@ public void run() {
                               @Override
                               public void run() {
                                 System.out.println("complete");
-                                latch.countDown();
                               }
                             });
-                        latch.countDown();
                       }
                     });
-                latch.countDown();
               }
             });
       }
     }.run();
 
-    latch.await();
+    // Block test thread until child spans are reported
+    blockUntilTracesMatch(traces -> traces.size() == 1 && traces.get(0).size() == 4);
 
     assertTraces(
         trace(