KAFKA-16791 Add thread detection to ClusterTestExtensions (#16499)

Reviewers: Chia-Ping Tsai <[email protected]>

Author: FrankYang0529

core/src/test/java/kafka/test/junit/ClusterTestExtensions.java

@@ -28,19 +28,27 @@
 import kafka.test.annotation.Type;
 
 import org.apache.kafka.server.common.Features;
+import org.apache.kafka.test.TestUtils;
 
+import org.junit.jupiter.api.extension.AfterEachCallback;
+import org.junit.jupiter.api.extension.BeforeEachCallback;
 import org.junit.jupiter.api.extension.ExtensionContext;
+import org.junit.jupiter.api.extension.ExtensionContext.Namespace;
+import org.junit.jupiter.api.extension.ExtensionContext.Store;
 import org.junit.jupiter.api.extension.TestTemplateInvocationContext;
 import org.junit.jupiter.api.extension.TestTemplateInvocationContextProvider;
 import org.junit.platform.commons.util.ReflectionUtils;
 
 import java.lang.reflect.Method;
 import java.util.ArrayList;
 import java.util.Arrays;
+import java.util.Collections;
 import java.util.HashSet;
 import java.util.List;
 import java.util.Map;
 import java.util.Optional;
+import java.util.Set;
+import java.util.concurrent.atomic.AtomicReference;
 import java.util.function.Function;
 import java.util.stream.Collectors;
 import java.util.stream.Stream;
@@ -82,7 +90,20 @@
  * SomeIntegrationTest will be instantiated, lifecycle methods (before/after) will be run, and "someTest" will be invoked.
  *
  */
-public class ClusterTestExtensions implements TestTemplateInvocationContextProvider {
+public class ClusterTestExtensions implements TestTemplateInvocationContextProvider, BeforeEachCallback, AfterEachCallback {
+    private static final String METRICS_METER_TICK_THREAD_PREFIX = "metrics-meter-tick-thread";
+    private static final String SCALA_THREAD_PREFIX = "scala-";
+    private static final String FORK_JOIN_POOL_THREAD_PREFIX = "ForkJoinPool";
+    private static final String JUNIT_THREAD_PREFIX = "junit-";
+    private static final String ATTACH_LISTENER_THREAD_PREFIX = "Attach Listener";
+    private static final String PROCESS_REAPER_THREAD_PREFIX = "process reaper";
+    private static final String RMI_THREAD_PREFIX = "RMI";
+    private static final String DETECT_THREAD_LEAK_KEY = "detectThreadLeak";
+    private static final Set<String> SKIPPED_THREAD_PREFIX = Collections.unmodifiableSet(Stream.of(
+            METRICS_METER_TICK_THREAD_PREFIX, SCALA_THREAD_PREFIX, FORK_JOIN_POOL_THREAD_PREFIX, JUNIT_THREAD_PREFIX,
+            ATTACH_LISTENER_THREAD_PREFIX, PROCESS_REAPER_THREAD_PREFIX, RMI_THREAD_PREFIX)
+            .collect(Collectors.toSet()));
+
     @Override
     public boolean supportsTestTemplate(ExtensionContext context) {
         return true;
@@ -119,7 +140,31 @@ public Stream<TestTemplateInvocationContext> provideTestTemplateInvocationContex
         return generatedContexts.stream();
     }
 
+    @Override
+    public void beforeEach(ExtensionContext context) {
+        DetectThreadLeak detectThreadLeak = DetectThreadLeak.of(thread ->
+                SKIPPED_THREAD_PREFIX.stream().noneMatch(prefix -> thread.getName().startsWith(prefix)));
+        getStore(context).put(DETECT_THREAD_LEAK_KEY, detectThreadLeak);
+    }
+
+    @Override
+    public void afterEach(ExtensionContext context) throws InterruptedException {
+        DetectThreadLeak detectThreadLeak = getStore(context).remove(DETECT_THREAD_LEAK_KEY, DetectThreadLeak.class);
+        if (detectThreadLeak == null) {
+            return;
+        }
+        AtomicReference<List<Thread>> lastThread = new AtomicReference<>(Collections.emptyList());
+        TestUtils.waitForCondition(() -> {
+            List<Thread> threads = detectThreadLeak.newThreads();
+            lastThread.set(threads);
+            return threads.isEmpty();
+        }, () -> "Thread leak detected: " +
+                lastThread.get().stream().map(Thread::getName).collect(Collectors.joining(", ")));
+    }
 
+    private Store getStore(ExtensionContext context) {
+        return context.getStore(Namespace.create(context.getUniqueId()));
+    }
 
     List<TestTemplateInvocationContext> processClusterTemplate(ExtensionContext context, ClusterTemplate annot) {
         if (annot.value().trim().isEmpty()) {

core/src/test/java/kafka/test/junit/DetectThreadLeak.java

@@ -0,0 +1,51 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package kafka.test.junit;
+
+import java.util.List;
+import java.util.Set;
+import java.util.function.Predicate;
+import java.util.stream.Collectors;
+
+public interface DetectThreadLeak {
+    /**
+     * @return the new threads after `DetectThreadLeak` is created
+     */
+    List<Thread> newThreads();
+
+    /**
+     * Creates an instance of {@link DetectThreadLeak} that filters threads based on a given predicate.
+     * This method captures the current state of threads that match the predicate at the time of invocation
+     * and provides a way to detect new threads that match the predicate but were not present at the initial capture.
+     *
+     * @param predicate A {@link Predicate<Thread>} used to filter threads. Only threads that satisfy
+     *                  the predicate are considered for detection.
+     * @return An instance of {@link DetectThreadLeak} that can be used to detect new threads matching
+     *         the predicate that were not present at the time of this method's invocation.
+     *         The {@link DetectThreadLeak#newThreads()} method of the returned instance will return a list
+     *         of new threads that match the predicate and have been started after this method was called.
+     */
+    static DetectThreadLeak of(Predicate<Thread> predicate) {
+        Set<Long> before = Thread.getAllStackTraces().keySet()
+                .stream().filter(predicate).map(Thread::getId).collect(Collectors.toSet());
+        return () -> Thread.getAllStackTraces().keySet()
+                .stream().filter(predicate)
+                .filter(t -> !before.contains(t.getId()))
+                .collect(Collectors.toList());
+    }
+}

core/src/test/java/kafka/test/junit/DetectThreadLeakTest.java

@@ -0,0 +1,67 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package kafka.test.junit;
+
+import org.junit.jupiter.api.Test;
+
+import static org.junit.jupiter.api.Assertions.assertFalse;
+import static org.junit.jupiter.api.Assertions.assertTrue;
+
+public class DetectThreadLeakTest {
+
+    private static class LeakThread implements Runnable {
+        @Override
+        public void run() {
+            try {
+                Thread.sleep(20000);
+            } catch (InterruptedException e) {
+                // this can be neglected
+            }
+        }
+    }
+
+    @Test
+    public void testThreadLeak() throws InterruptedException {
+        DetectThreadLeak detectThreadLeak = DetectThreadLeak.of(thread -> true);
+        Thread leakThread = new Thread(new LeakThread());
+        try {
+            leakThread.start();
+            assertTrue(detectThreadLeak.newThreads().contains(leakThread));
+            leakThread.interrupt();
+        } finally {
+            leakThread.join();
+        }
+        assertFalse(leakThread.isAlive(), "Can't interrupt the thread");
+        assertFalse(detectThreadLeak.newThreads().contains(leakThread));
+    }
+
+    @Test
+    public void testDetectThreadLeakWithOverrideExpectedThreadNames() throws InterruptedException {
+        String threadName = "test-thread";
+        DetectThreadLeak detectThreadLeak = DetectThreadLeak.of(thread -> !thread.getName().equals(threadName));
+        Thread leakThread = new Thread(new LeakThread(), threadName);
+        try {
+            leakThread.start();
+            assertFalse(detectThreadLeak.newThreads().contains(leakThread));
+            leakThread.interrupt();
+        } finally {
+            leakThread.join();
+        }
+        assertFalse(leakThread.isAlive(), "Can't interrupt the thread");
+    }
+}