Refactor waitForInstrumentation

dd-java-agent/agent-debugger/src/main/java/com/datadog/debugger/util/ExceptionHelper.java

@@ -73,7 +73,6 @@ public static Throwable getInnerMostThrowable(Throwable t) {
 
   public static StackTraceElement[] flattenStackTrace(Throwable t) {
     List<StackTraceElement> result = new ArrayList<>();
-    // result.add(null); // add a stack frame representing the exception message
     result.addAll(Arrays.asList(t.getStackTrace()));
     if (t.getCause() != null) {
       internalFlattenStackTrace(t.getCause(), t.getStackTrace(), result);

dd-java-agent/agent-debugger/src/test/java/com/datadog/debugger/exception/DefaultExceptionDebuggerTest.java

@@ -1,7 +1,7 @@
 package com.datadog.debugger.exception;
 
 import static com.datadog.debugger.exception.DefaultExceptionDebugger.SNAPSHOT_ID_TAG_FMT;
-import static com.datadog.debugger.exception.ExceptionProbeManagerTest.waitForInstrumentation;
+import static com.datadog.debugger.util.TestHelper.assertWithTimeout;
 import static java.util.Collections.emptyList;
 import static java.util.stream.Collectors.toMap;
 import static org.junit.jupiter.api.Assertions.*;
@@ -31,6 +31,7 @@
 import java.io.PrintWriter;
 import java.io.StringReader;
 import java.io.StringWriter;
+import java.time.Duration;
 import java.util.ArrayList;
 import java.util.HashMap;
 import java.util.List;
@@ -62,7 +63,9 @@ public void simpleException() {
     String fingerprint = Fingerprinter.fingerprint(exception, classNameFiltering);
     AgentSpan span = mock(AgentSpan.class);
     exceptionDebugger.handleException(exception, span);
-    waitForInstrumentation(exceptionDebugger.getExceptionProbeManager(), fingerprint);
+    assertWithTimeout(
+        () -> exceptionDebugger.getExceptionProbeManager().isAlreadyInstrumented(fingerprint),
+        Duration.ofSeconds(30));
     exceptionDebugger.handleException(exception, span);
     verify(configurationUpdater).accept(eq(ConfigurationAcceptor.Source.EXCEPTION), any());
   }

dd-java-agent/agent-debugger/src/test/java/com/datadog/debugger/exception/ExceptionProbeInstrumentationTest.java

@@ -4,7 +4,7 @@
 import static com.datadog.debugger.exception.DefaultExceptionDebugger.DD_DEBUG_ERROR_EXCEPTION_ID;
 import static com.datadog.debugger.exception.DefaultExceptionDebugger.ERROR_DEBUG_INFO_CAPTURED;
 import static com.datadog.debugger.exception.DefaultExceptionDebugger.SNAPSHOT_ID_TAG_FMT;
-import static com.datadog.debugger.exception.ExceptionProbeManagerTest.waitForInstrumentation;
+import static com.datadog.debugger.util.TestHelper.assertWithTimeout;
 import static org.junit.jupiter.api.Assertions.assertEquals;
 import static org.junit.jupiter.api.Assertions.assertTrue;
 import static org.mockito.Mockito.mock;
@@ -32,6 +32,7 @@
 import datadog.trace.core.CoreTracer;
 import java.lang.instrument.ClassFileTransformer;
 import java.lang.instrument.Instrumentation;
+import java.time.Duration;
 import java.util.Arrays;
 import java.util.Map;
 import java.util.Set;
@@ -80,7 +81,8 @@ public void onlyInstrument() throws Exception {
     Class<?> testClass = compileAndLoadClass(CLASS_NAME);
     String fingerprint =
         callMethodThrowingRuntimeException(testClass); // instrument exception stacktrace
-    waitForInstrumentation(exceptionProbeManager, fingerprint);
+    assertWithTimeout(
+        () -> exceptionProbeManager.isAlreadyInstrumented(fingerprint), Duration.ofSeconds(30));
     assertEquals(2, exceptionProbeManager.getProbes().size());
     callMethodNoException(testClass);
     assertEquals(0, listener.snapshots.size());
@@ -96,7 +98,8 @@ public void instrumentAndCaptureSnapshots() throws Exception {
     Class<?> testClass = compileAndLoadClass(CLASS_NAME);
     String fingerprint =
         callMethodThrowingRuntimeException(testClass); // instrument exception stacktrace
-    waitForInstrumentation(exceptionProbeManager, fingerprint);
+    assertWithTimeout(
+        () -> exceptionProbeManager.isAlreadyInstrumented(fingerprint), Duration.ofSeconds(30));
     assertEquals(2, exceptionProbeManager.getProbes().size());
     callMethodThrowingRuntimeException(testClass); // generate snapshots
     Map<String, Set<String>> probeIdsByMethodName =
@@ -125,11 +128,13 @@ public void differentExceptionsSameStack() throws Exception {
     Class<?> testClass = compileAndLoadClass(CLASS_NAME);
     // instrument RuntimeException stacktrace
     String fingerprint0 = callMethodThrowingRuntimeException(testClass);
-    waitForInstrumentation(exceptionProbeManager, fingerprint0);
+    assertWithTimeout(
+        () -> exceptionProbeManager.isAlreadyInstrumented(fingerprint0), Duration.ofSeconds(30));
     assertEquals(2, exceptionProbeManager.getProbes().size());
     // instrument IllegalArgumentException stacktrace
     String fingerprint1 = callMethodThrowingIllegalArgException(testClass);
-    waitForInstrumentation(exceptionProbeManager, fingerprint1);
+    assertWithTimeout(
+        () -> exceptionProbeManager.isAlreadyInstrumented(fingerprint1), Duration.ofSeconds(30));
     assertEquals(4, exceptionProbeManager.getProbes().size());
     Map<String, Set<String>> probeIdsByMethodName =
         extractProbeIdsByMethodName(exceptionProbeManager);

dd-java-agent/agent-debugger/src/test/java/com/datadog/debugger/exception/ExceptionProbeManagerTest.java

@@ -4,7 +4,6 @@
 
 import com.datadog.debugger.probe.ExceptionProbe;
 import com.datadog.debugger.util.ClassNameFiltering;
-import java.time.Duration;
 import java.util.Arrays;
 import java.util.Collections;
 import org.junit.jupiter.api.Test;
@@ -29,7 +28,7 @@ void instrumentSingleFrame() {
     ExceptionProbeManager exceptionProbeManager = new ExceptionProbeManager(classNameFiltering);
 
     String fingerprint = Fingerprinter.fingerprint(exception, classNameFiltering);
-    assertEquals("ca4d9f3a1033d7262a89855f4b5cbdc225ed63c592c6cdf83fc5a88589e5fb", fingerprint);
+    assertEquals("aa4a4dd768f6ef0fcc2b39a3bdedcbe44baff2e9dd0a779228db7bd8bf58", fingerprint);
     exceptionProbeManager.createProbesForException(fingerprint, exception.getStackTrace());
     assertEquals(1, exceptionProbeManager.getProbes().size());
     ExceptionProbe exceptionProbe = exceptionProbeManager.getProbes().iterator().next();
@@ -53,19 +52,4 @@ void filterAllFrames() {
     exceptionProbeManager.createProbesForException("", exception.getStackTrace());
     assertEquals(0, exceptionProbeManager.getProbes().size());
   }
-
-  static void waitForInstrumentation(
-      ExceptionProbeManager exceptionProbeManager, String fingerprint) {
-    Duration timeout = Duration.ofSeconds(30);
-    Duration sleepTime = Duration.ofMillis(10);
-    long count = timeout.toMillis() / sleepTime.toMillis();
-    while (count-- > 0 && !exceptionProbeManager.isAlreadyInstrumented(fingerprint)) {
-      try {
-        Thread.sleep(sleepTime.toMillis());
-      } catch (InterruptedException e) {
-        throw new RuntimeException(e);
-      }
-    }
-    assertTrue(exceptionProbeManager.isAlreadyInstrumented(fingerprint));
-  }
 }

dd-java-agent/agent-debugger/src/test/java/com/datadog/debugger/util/TestHelper.java

@@ -1,7 +1,11 @@
 package com.datadog.debugger.util;
 
+import static org.junit.jupiter.api.Assertions.assertTrue;
+
 import datadog.trace.api.Config;
 import java.lang.reflect.Field;
+import java.time.Duration;
+import java.util.function.BooleanSupplier;
 
 public class TestHelper {
 
@@ -14,4 +18,17 @@ public static void setFieldInConfig(Config config, String fieldName, Object valu
       e.printStackTrace();
     }
   }
+
+  public static void assertWithTimeout(BooleanSupplier predicate, Duration timeout) {
+    Duration sleepTime = Duration.ofMillis(10);
+    long count = timeout.toMillis() / sleepTime.toMillis();
+    while (count-- > 0 && !predicate.getAsBoolean()) {
+      try {
+        Thread.sleep(sleepTime.toMillis());
+      } catch (InterruptedException e) {
+        throw new RuntimeException(e);
+      }
+    }
+    assertTrue(predicate.getAsBoolean());
+  }
 }