Allow thread names to be specified using regex

4 files changed, 115 insertions, 9 deletions

diff --git a/spark-common/src/main/java/me/lucko/spark/common/command/modules/SamplerModule.java b/spark-common/src/main/java/me/lucko/spark/common/command/modules/SamplerModule.java
index 5fd8b5b..fe719eb 100644
--- a/spark-common/src/main/java/me/lucko/spark/common/command/modules/SamplerModule.java
+++ b/spark-common/src/main/java/me/lucko/spark/common/command/modules/SamplerModule.java
@@ -56,6 +56,7 @@ public class SamplerModule<S> implements CommandModule<S> {
                 .aliases("start")
                 .argumentUsage("timeout", "timeout seconds")
                 .argumentUsage("thread", "thread name")
+                .argumentUsage("regex", null)
                 .argumentUsage("not-combined", null)
                 .argumentUsage("interval", "interval millis")
                 .argumentUsage("only-ticks-over", "tick length millis")
@@ -86,7 +87,12 @@ public class SamplerModule<S> implements CommandModule<S> {
                     } else if (threads.contains("*")) {
                         threadDumper = ThreadDumper.ALL;
                     } else {
-                        threadDumper = new ThreadDumper.Specific(threads);
+                        if (arguments.boolFlag("regex")) {
+                            threadDumper = new ThreadDumper.Regex(threads);
+                        } else {
+                            // specific matches
+                            threadDumper = new ThreadDumper.Specific(threads);
+                        }
                     }
 
                     ThreadGrouper threadGrouper;
diff --git a/spark-common/src/main/java/me/lucko/spark/sampler/Sampler.java b/spark-common/src/main/java/me/lucko/spark/sampler/Sampler.java
index 5758d85..d78b3fd 100644
--- a/spark-common/src/main/java/me/lucko/spark/sampler/Sampler.java
+++ b/spark-common/src/main/java/me/lucko/spark/sampler/Sampler.java
@@ -130,7 +130,7 @@ public class Sampler implements Runnable {
                 return;
             }
 
-            ThreadInfo[] threadDumps = this.threadDumper.dumpThreads(this.threadBean);
+            Iterable<ThreadInfo> threadDumps = this.threadDumper.dumpThreads(this.threadBean);
             for (ThreadInfo threadInfo : threadDumps) {
                 String threadName = threadInfo.getThreadName();
                 StackTraceElement[] stack = threadInfo.getStackTrace();
diff --git a/spark-common/src/main/java/me/lucko/spark/sampler/ThreadDumper.java b/spark-common/src/main/java/me/lucko/spark/sampler/ThreadDumper.java
index 940b1c6..61706f4 100644
--- a/spark-common/src/main/java/me/lucko/spark/sampler/ThreadDumper.java
+++ b/spark-common/src/main/java/me/lucko/spark/sampler/ThreadDumper.java
@@ -21,9 +21,15 @@
 
 package me.lucko.spark.sampler;
 
+import me.lucko.spark.util.Threads;
+
 import java.lang.management.ThreadInfo;
 import java.lang.management.ThreadMXBean;
+import java.util.Arrays;
+import java.util.Objects;
 import java.util.Set;
+import java.util.regex.Pattern;
+import java.util.regex.PatternSyntaxException;
 import java.util.stream.Collectors;
 
 /**
@@ -39,12 +45,12 @@ public interface ThreadDumper {
      * @param threadBean the thread bean instance to obtain the data from
      * @return an array of generated thread info instances
      */
-    ThreadInfo[] dumpThreads(ThreadMXBean threadBean);
+    Iterable<ThreadInfo> dumpThreads(ThreadMXBean threadBean);
 
     /**
      * Implementation of {@link ThreadDumper} that generates data for all threads.
      */
-    ThreadDumper ALL = threadBean -> threadBean.dumpAllThreads(false, false);
+    ThreadDumper ALL = threadBean -> Arrays.asList(threadBean.dumpAllThreads(false, false));
 
     /**
      * Implementation of {@link ThreadDumper} that generates data for a specific set of threads.
@@ -57,16 +63,52 @@ public interface ThreadDumper {
         }
 
         public Specific(Set<String> names) {
-            Set<String> threadNamesLower = names.stream().map(String::toLowerCase).collect(Collectors.toSet());
-            this.ids = Thread.getAllStackTraces().keySet().stream()
-                    .filter(t -> threadNamesLower.contains(t.getName().toLowerCase()))
+            Set<String> namesLower = names.stream().map(String::toLowerCase).collect(Collectors.toSet());
+            this.ids = Threads.getThreads()
+                    .filter(t -> namesLower.contains(t.getName().toLowerCase()))
                     .mapToLong(Thread::getId)
                     .toArray();
         }
 
         @Override
-        public ThreadInfo[] dumpThreads(ThreadMXBean threadBean) {
-            return threadBean.getThreadInfo(this.ids, Integer.MAX_VALUE);
+        public Iterable<ThreadInfo> dumpThreads(ThreadMXBean threadBean) {
+            return Arrays.asList(threadBean.getThreadInfo(this.ids, Integer.MAX_VALUE));
+        }
+    }
+
+    /**
+     * Implementation of {@link ThreadDumper} that generates data for a regex matched set of threads.
+     */
+    final class Regex implements ThreadDumper {
+        private final Set<Pattern> namePatterns;
+
+        public Regex(Set<String> namePatterns) {
+            this.namePatterns = namePatterns.stream()
+                    .map(regex -> {
+                        try {
+                            return Pattern.compile(regex);
+                        } catch (PatternSyntaxException e) {
+                            return null;
+                        }
+                    })
+                    .filter(Objects::nonNull)
+                    .collect(Collectors.toSet());
+        }
+
+        @Override
+        public Iterable<ThreadInfo> dumpThreads(ThreadMXBean threadBean) {
+            return Threads.getThreads()
+                    .filter(thread -> {
+                        for (Pattern pattern : this.namePatterns) {
+                            if (pattern.matcher(thread.getName()).matches()) {
+                                return true;
+                            }
+                        }
+                        return false;
+                    })
+                    .map(thread -> threadBean.getThreadInfo(thread.getId()))
+                    .filter(Objects::nonNull)
+                    .collect(Collectors.toList());
         }
     }
 
diff --git a/spark-common/src/main/java/me/lucko/spark/util/Threads.java b/spark-common/src/main/java/me/lucko/spark/util/Threads.java
new file mode 100644
index 0000000..d1efdff
--- /dev/null
+++ b/spark-common/src/main/java/me/lucko/spark/util/Threads.java
@@ -0,0 +1,58 @@
+/*
+ * This file is part of spark.
+ *
+ *  Copyright (c) lucko (Luck) <luck@lucko.me>
+ *  Copyright (c) contributors
+ *
+ *  This program is free software: you can redistribute it and/or modify
+ *  it under the terms of the GNU General Public License as published by
+ *  the Free Software Foundation, either version 3 of the License, or
+ *  (at your option) any later version.
+ *
+ *  This program is distributed in the hope that it will be useful,
+ *  but WITHOUT ANY WARRANTY; without even the implied warranty of
+ *  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+ *  GNU General Public License for more details.
+ *
+ *  You should have received a copy of the GNU General Public License
+ *  along with this program.  If not, see <http://www.gnu.org/licenses/>.
+ */
+
+package me.lucko.spark.util;
+
+import java.util.Arrays;
+import java.util.Objects;
+import java.util.stream.Stream;
+
+/**
+ * Utilities for working with {@link Thread}s.
+ */
+public final class Threads {
+
+    private Threads() {}
+
+    private static final ThreadGroup ROOT_THREAD_GROUP;
+
+    static {
+        ThreadGroup rootGroup = Thread.currentThread().getThreadGroup();
+        ThreadGroup parentGroup;
+        while ((parentGroup = rootGroup.getParent()) != null) {
+            rootGroup = parentGroup;
+        }
+        ROOT_THREAD_GROUP = rootGroup;
+    }
+
+    /**
+     * Gets a stream of all known active threads.
+     *
+     * @return a stream of threads
+     */
+    public static Stream<Thread> getThreads() {
+        Thread[] threads = new Thread[ROOT_THREAD_GROUP.activeCount()];
+        while (ROOT_THREAD_GROUP.enumerate(threads, true ) == threads.length) {
+            threads = new Thread[threads.length * 2];
+        }
+        return Arrays.stream(threads).filter(Objects::nonNull);
+    }
+
+}