Improve thread leak detector to poll for threads to terminate

lhotari · Oct 17, 2023 · 7f5c05f · 7f5c05f
1 parent 1d411d1
commit 7f5c05f
Showing 1 changed file with 18 additions and 6 deletions.
diff --git a/buildtools/src/main/java/org/apache/pulsar/tests/ThreadLeakDetectorListener.java b/buildtools/src/main/java/org/apache/pulsar/tests/ThreadLeakDetectorListener.java
@@ -38,6 +38,8 @@ public class ThreadLeakDetectorListener extends BetweenTestClassesListenerAdapte
             Long.parseLong(System.getenv().getOrDefault("THREAD_LEAK_DETECTOR_WAIT_MILLIS", "0"));
 
     private Set<ThreadKey> capturedThreadKeys;
+    private static final long THREAD_TERMINATION_POLL_INTERVAL = 250L;
+
 
     @Override
     protected void onBetweenTestClasses(Class<?> endedTestClass, Class<?> startedTestClass) {
@@ -47,14 +49,24 @@ protected void onBetweenTestClasses(Class<?> endedTestClass, Class<?> startedTes
                 compareThreads(capturedThreadKeys, endedTestClass, WAIT_FOR_THREAD_TERMINATION_MILLIS <= 0,
                         differenceDetected);
         if (WAIT_FOR_THREAD_TERMINATION_MILLIS > 0 && differenceDetected.booleanValue()) {
-            LOG.info("Difference detected in active threads. Waiting {} ms for threads to terminate.",
+            LOG.info("Difference detected in active threads. Waiting up to {} ms for threads to terminate.",
                     WAIT_FOR_THREAD_TERMINATION_MILLIS);
-            try {
-                Thread.sleep(WAIT_FOR_THREAD_TERMINATION_MILLIS);
-            } catch (InterruptedException e) {
-                Thread.currentThread().interrupt();
+            long endTime = System.currentTimeMillis() + WAIT_FOR_THREAD_TERMINATION_MILLIS;
+            while (System.currentTimeMillis() < endTime) {
+                try {
+                    Thread.sleep(THREAD_TERMINATION_POLL_INTERVAL);
+                } catch (InterruptedException e) {
+                    Thread.currentThread().interrupt();
+                }
+                differenceDetected.setFalse();
+                currentThreadKeys = compareThreads(capturedThreadKeys, endedTestClass, false, differenceDetected);
+                if (!differenceDetected.booleanValue()) {
+                    break;
+                }
+            }
+            if (differenceDetected.booleanValue()) {
+                currentThreadKeys = compareThreads(capturedThreadKeys, endedTestClass, true, null);
             }
-            currentThreadKeys = compareThreads(capturedThreadKeys, endedTestClass, true, null);
         }
         capturedThreadKeys = currentThreadKeys;
     }