Refactor ClassFinder into interface

12 files changed, 183 insertions, 33 deletions

diff --git a/spark-common/src/main/java/me/lucko/spark/common/SparkPlatform.java b/spark-common/src/main/java/me/lucko/spark/common/SparkPlatform.java
index 21f9210..733510d 100644
--- a/spark-common/src/main/java/me/lucko/spark/common/SparkPlatform.java
+++ b/spark-common/src/main/java/me/lucko/spark/common/SparkPlatform.java
@@ -55,6 +55,7 @@ import me.lucko.spark.common.util.BytebinClient;
 import me.lucko.spark.common.util.Configuration;
 import me.lucko.spark.common.util.SparkStaticLogger;
 import me.lucko.spark.common.util.TemporaryFiles;
+import me.lucko.spark.common.util.classfinder.ClassFinder;
 import me.lucko.spark.common.ws.TrustedKeyStore;
 import net.kyori.adventure.text.Component;
 import net.kyori.adventure.text.event.ClickEvent;
@@ -285,6 +286,10 @@ public class SparkPlatform {
         return this.plugin.createClassSourceLookup();
     }
 
+    public ClassFinder createClassFinder() {
+        return this.plugin.createClassFinder();
+    }
+
     public TickStatistics getTickStatistics() {
         return this.tickStatistics;
     }
diff --git a/spark-common/src/main/java/me/lucko/spark/common/SparkPlugin.java b/spark-common/src/main/java/me/lucko/spark/common/SparkPlugin.java
index a3bdceb..bf745df 100644
--- a/spark-common/src/main/java/me/lucko/spark/common/SparkPlugin.java
+++ b/spark-common/src/main/java/me/lucko/spark/common/SparkPlugin.java
@@ -33,6 +33,9 @@ import me.lucko.spark.common.sampler.source.ClassSourceLookup;
 import me.lucko.spark.common.sampler.source.SourceMetadata;
 import me.lucko.spark.common.tick.TickHook;
 import me.lucko.spark.common.tick.TickReporter;
+import me.lucko.spark.common.util.classfinder.ClassFinder;
+import me.lucko.spark.common.util.classfinder.FallbackClassFinder;
+import me.lucko.spark.common.util.classfinder.InstrumentationClassFinder;
 
 import java.nio.file.Path;
 import java.util.Collection;
@@ -150,6 +153,18 @@ public interface SparkPlugin {
     }
 
     /**
+     * Creates a class finder for the platform.
+     *
+     * @return the class finder
+     */
+    default ClassFinder createClassFinder() {
+        return ClassFinder.combining(
+                new InstrumentationClassFinder(this),
+                FallbackClassFinder.INSTANCE
+        );
+    }
+
+    /**
      * Gets a list of known sources (plugins/mods) on the platform.
      *
      * @return a list of sources
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 9ce66dc..9e2647a 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
@@ -507,7 +507,7 @@ public class SamplerModule implements CommandModule {
                 .creator(resp.senderData())
                 .comment(Iterables.getFirst(arguments.stringFlag("comment"), null))
                 .mergeMode(() -> {
-                    MethodDisambiguator methodDisambiguator = new MethodDisambiguator();
+                    MethodDisambiguator methodDisambiguator = new MethodDisambiguator(platform.createClassFinder());
                     return arguments.boolFlag("separate-parent-calls")
                             ? MergeMode.separateParentCalls(methodDisambiguator)
                             : MergeMode.sameMethod(methodDisambiguator);
diff --git a/spark-common/src/main/java/me/lucko/spark/common/sampler/AbstractSampler.java b/spark-common/src/main/java/me/lucko/spark/common/sampler/AbstractSampler.java
index 445702e..8a9c05f 100644
--- a/spark-common/src/main/java/me/lucko/spark/common/sampler/AbstractSampler.java
+++ b/spark-common/src/main/java/me/lucko/spark/common/sampler/AbstractSampler.java
@@ -32,6 +32,7 @@ import me.lucko.spark.common.sampler.source.ClassSourceLookup;
 import me.lucko.spark.common.sampler.source.SourceMetadata;
 import me.lucko.spark.common.sampler.window.ProtoTimeEncoder;
 import me.lucko.spark.common.sampler.window.WindowStatisticsCollector;
+import me.lucko.spark.common.util.classfinder.ClassFinder;
 import me.lucko.spark.common.ws.ViewerSocket;
 import me.lucko.spark.proto.SparkProtos;
 import me.lucko.spark.proto.SparkSamplerProtos.SamplerData;
@@ -44,6 +45,7 @@ import java.util.Locale;
 import java.util.Map;
 import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.CopyOnWriteArrayList;
+import java.util.function.Supplier;
 
 /**
  * Base implementation class for {@link Sampler}s.
@@ -230,11 +232,11 @@ public abstract class AbstractSampler implements Sampler {
         proto.setMetadata(metadata);
     }
 
-    protected void writeDataToProto(SamplerData.Builder proto, DataAggregator dataAggregator, MergeMode mergeMode, ClassSourceLookup classSourceLookup) {
+    protected void writeDataToProto(SamplerData.Builder proto, DataAggregator dataAggregator, MergeMode mergeMode, ClassSourceLookup classSourceLookup, Supplier<ClassFinder> classFinderSupplier) {
         List<ThreadNode> data = dataAggregator.exportData();
         data.sort(Comparator.comparing(ThreadNode::getThreadLabel));
 
-        ClassSourceLookup.Visitor classSourceVisitor = ClassSourceLookup.createVisitor(classSourceLookup);
+        ClassSourceLookup.Visitor classSourceVisitor = ClassSourceLookup.createVisitor(classSourceLookup, classFinderSupplier);
 
         ProtoTimeEncoder timeEncoder = new ProtoTimeEncoder(getMode().valueTransformer(), data);
         int[] timeWindows = timeEncoder.getKeys();
diff --git a/spark-common/src/main/java/me/lucko/spark/common/sampler/async/AsyncSampler.java b/spark-common/src/main/java/me/lucko/spark/common/sampler/async/AsyncSampler.java
index e3470b4..5350558 100644
--- a/spark-common/src/main/java/me/lucko/spark/common/sampler/async/AsyncSampler.java
+++ b/spark-common/src/main/java/me/lucko/spark/common/sampler/async/AsyncSampler.java
@@ -221,7 +221,7 @@ public class AsyncSampler extends AbstractSampler {
             proto.setChannelInfo(exportProps.channelInfo());
         }
         writeMetadataToProto(proto, platform, exportProps.creator(), exportProps.comment(), this.dataAggregator);
-        writeDataToProto(proto, this.dataAggregator, exportProps.mergeMode().get(), exportProps.classSourceLookup().get());
+        writeDataToProto(proto, this.dataAggregator, exportProps.mergeMode().get(), exportProps.classSourceLookup().get(), platform::createClassFinder);
         return proto.build();
     }
 
diff --git a/spark-common/src/main/java/me/lucko/spark/common/sampler/java/JavaSampler.java b/spark-common/src/main/java/me/lucko/spark/common/sampler/java/JavaSampler.java
index bae7e19..e3ae73a 100644
--- a/spark-common/src/main/java/me/lucko/spark/common/sampler/java/JavaSampler.java
+++ b/spark-common/src/main/java/me/lucko/spark/common/sampler/java/JavaSampler.java
@@ -193,7 +193,7 @@ public class JavaSampler extends AbstractSampler implements Runnable {
             proto.setChannelInfo(exportProps.channelInfo());
         }
         writeMetadataToProto(proto, platform, exportProps.creator(), exportProps.comment(), this.dataAggregator);
-        writeDataToProto(proto, this.dataAggregator, exportProps.mergeMode().get(), exportProps.classSourceLookup().get());
+        writeDataToProto(proto, this.dataAggregator, exportProps.mergeMode().get(), exportProps.classSourceLookup().get(), platform::createClassFinder);
         return proto.build();
     }
 
diff --git a/spark-common/src/main/java/me/lucko/spark/common/sampler/source/ClassSourceLookup.java b/spark-common/src/main/java/me/lucko/spark/common/sampler/source/ClassSourceLookup.java
index a62f8d1..c856123 100644
--- a/spark-common/src/main/java/me/lucko/spark/common/sampler/source/ClassSourceLookup.java
+++ b/spark-common/src/main/java/me/lucko/spark/common/sampler/source/ClassSourceLookup.java
@@ -23,7 +23,7 @@ package me.lucko.spark.common.sampler.source;
 import me.lucko.spark.common.SparkPlatform;
 import me.lucko.spark.common.sampler.node.StackTraceNode;
 import me.lucko.spark.common.sampler.node.ThreadNode;
-import me.lucko.spark.common.util.ClassFinder;
+import me.lucko.spark.common.util.classfinder.ClassFinder;
 import org.checkerframework.checker.nullness.qual.Nullable;
 
 import java.io.IOException;
@@ -42,6 +42,7 @@ import java.util.Map;
 import java.util.Objects;
 import java.util.Queue;
 import java.util.function.Function;
+import java.util.function.Supplier;
 import java.util.stream.Collectors;
 
 /**
@@ -203,11 +204,11 @@ public interface ClassSourceLookup {
         Map<String, String> getLineSourceMapping();
     }
 
-    static Visitor createVisitor(ClassSourceLookup lookup) {
+    static Visitor createVisitor(ClassSourceLookup lookup, Supplier<ClassFinder> classFinderSupplier) {
         if (lookup == ClassSourceLookup.NO_OP) {
             return NoOpVisitor.INSTANCE; // don't bother!
         }
-        return new VisitorImpl(lookup);
+        return new VisitorImpl(lookup, classFinderSupplier.get());
     }
 
     enum NoOpVisitor implements Visitor {
@@ -254,14 +255,15 @@ public interface ClassSourceLookup {
      */
     class VisitorImpl implements Visitor {
         private final ClassSourceLookup lookup;
-        private final ClassFinder classFinder = new ClassFinder();
+        private final ClassFinder classFinder;
 
         private final SourcesMap<String> classSources = new SourcesMap<>(Function.identity());
         private final SourcesMap<MethodCall> methodSources = new SourcesMap<>(MethodCall::toString);
         private final SourcesMap<MethodCallByLine> lineSources = new SourcesMap<>(MethodCallByLine::toString);
 
-        VisitorImpl(ClassSourceLookup lookup) {
+        VisitorImpl(ClassSourceLookup lookup, ClassFinder classFinder) {
             this.lookup = lookup;
+            this.classFinder = classFinder;
         }
 
         @Override
diff --git a/spark-common/src/main/java/me/lucko/spark/common/util/MethodDisambiguator.java b/spark-common/src/main/java/me/lucko/spark/common/util/MethodDisambiguator.java
index 2b2e3c7..3b0ddc3 100644
--- a/spark-common/src/main/java/me/lucko/spark/common/util/MethodDisambiguator.java
+++ b/spark-common/src/main/java/me/lucko/spark/common/util/MethodDisambiguator.java
@@ -24,6 +24,7 @@ import com.google.common.collect.ImmutableListMultimap;
 import com.google.common.collect.ImmutableMap;
 import com.google.common.collect.ListMultimap;
 import me.lucko.spark.common.sampler.node.StackTraceNode;
+import me.lucko.spark.common.util.classfinder.ClassFinder;
 import org.objectweb.asm.ClassReader;
 import org.objectweb.asm.ClassVisitor;
 import org.objectweb.asm.Label;
@@ -43,8 +44,13 @@ import java.util.concurrent.ConcurrentHashMap;
  * to a method (method name + method description).
  */
 public final class MethodDisambiguator {
-    private final Map<String, ComputedClass> cache = new ConcurrentHashMap<>();
-    private final ClassFinder classFinder = new ClassFinder();
+    private final ClassFinder classFinder;
+    private final Map<String, ComputedClass> cache;
+
+    public MethodDisambiguator(ClassFinder classFinder) {
+        this.classFinder = classFinder;
+        this.cache = new ConcurrentHashMap<>();
+    }
 
     public Optional<MethodDescription> disambiguate(StackTraceNode element) {
         String desc = element.getMethodDescription();
diff --git a/spark-common/src/main/java/me/lucko/spark/common/util/classfinder/ClassFinder.java b/spark-common/src/main/java/me/lucko/spark/common/util/classfinder/ClassFinder.java
new file mode 100644
index 0000000..1ee75c6
--- /dev/null
+++ b/spark-common/src/main/java/me/lucko/spark/common/util/classfinder/ClassFinder.java
@@ -0,0 +1,46 @@
+/*
+ * 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.common.util.classfinder;
+
+import com.google.common.collect.ImmutableList;
+import org.checkerframework.checker.nullness.qual.Nullable;
+
+public interface ClassFinder {
+
+    /**
+     * Creates a ClassFinder that combines the results of multiple other finders.
+     *
+     * @param finders the other class finders
+     * @return the combined class finder
+     */
+    static ClassFinder combining(ClassFinder... finders) {
+        return new CombinedClassFinder(ImmutableList.copyOf(finders));
+    }
+
+    /**
+     * Attempts to find a class by name.
+     *
+     * @param className the name of the class
+     * @return the class, if found
+     */
+    @Nullable Class<?> findClass(String className);
+
+}
diff --git a/spark-common/src/main/java/me/lucko/spark/common/util/classfinder/CombinedClassFinder.java b/spark-common/src/main/java/me/lucko/spark/common/util/classfinder/CombinedClassFinder.java
new file mode 100644
index 0000000..ed63f36
--- /dev/null
+++ b/spark-common/src/main/java/me/lucko/spark/common/util/classfinder/CombinedClassFinder.java
@@ -0,0 +1,44 @@
+/*
+ * 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.common.util.classfinder;
+
+import org.checkerframework.checker.nullness.qual.Nullable;
+
+import java.util.List;
+
+class CombinedClassFinder implements ClassFinder {
+    private final List<ClassFinder> finders;
+
+    CombinedClassFinder(List<ClassFinder> finders) {
+        this.finders = finders;
+    }
+
+    @Override
+    public @Nullable Class<?> findClass(String className) {
+        for (ClassFinder finder : this.finders) {
+            Class<?> clazz = finder.findClass(className);
+            if (clazz != null) {
+                return clazz;
+            }
+        }
+        return null;
+    }
+}
diff --git a/spark-common/src/main/java/me/lucko/spark/common/util/classfinder/FallbackClassFinder.java b/spark-common/src/main/java/me/lucko/spark/common/util/classfinder/FallbackClassFinder.java
new file mode 100644
index 0000000..dd3c9f0
--- /dev/null
+++ b/spark-common/src/main/java/me/lucko/spark/common/util/classfinder/FallbackClassFinder.java
@@ -0,0 +1,40 @@
+/*
+ * 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.common.util.classfinder;
+
+import org.checkerframework.checker.nullness.qual.Nullable;
+
+/**
+ * Uses {@link Class#forName(String)} to find a class reference for given class names.
+ */
+public enum FallbackClassFinder implements ClassFinder {
+    INSTANCE;
+
+    @Override
+    public @Nullable Class<?> findClass(String className) {
+        try {
+            return Class.forName(className);
+        } catch (Throwable e) {
+            return null;
+        }
+    }
+
+}
diff --git a/spark-common/src/main/java/me/lucko/spark/common/util/ClassFinder.java b/spark-common/src/main/java/me/lucko/spark/common/util/classfinder/InstrumentationClassFinder.java
index cead938..5f06d64 100644
--- a/spark-common/src/main/java/me/lucko/spark/common/util/ClassFinder.java
+++ b/spark-common/src/main/java/me/lucko/spark/common/util/classfinder/InstrumentationClassFinder.java
@@ -18,8 +18,10 @@
  *  along with this program.  If not, see <http://www.gnu.org/licenses/>.
  */
 
-package me.lucko.spark.common.util;
+package me.lucko.spark.common.util.classfinder;
 
+import me.lucko.spark.common.SparkPlugin;
+import me.lucko.spark.common.util.JavaVersion;
 import net.bytebuddy.agent.ByteBuddyAgent;
 import org.checkerframework.checker.nullness.qual.Nullable;
 
@@ -33,18 +35,18 @@ import java.util.logging.Level;
  *
  * <p>This is necessary as we don't always have access to the classloader for a given class.</p>
  */
-public class ClassFinder {
+public class InstrumentationClassFinder implements ClassFinder {
 
     private static boolean warned = false;
 
-    private static Instrumentation loadInstrumentation() {
+    private static Instrumentation loadInstrumentation(SparkPlugin plugin) {
         Instrumentation instrumentation = null;
         try {
             instrumentation = ByteBuddyAgent.install();
             if (!warned && JavaVersion.getJavaVersion() >= 21) {
                 warned = true;
-                SparkStaticLogger.log(Level.INFO, "If you see a warning above that says \"WARNING: A Java agent has been loaded dynamically\", it can be safely ignored.");
-                SparkStaticLogger.log(Level.INFO, "See here for more information: https://spark.lucko.me/docs/misc/Java-agent-warning");
+                plugin.log(Level.INFO, "If you see a warning above that says \"WARNING: A Java agent has been loaded dynamically\", it can be safely ignored.");
+                plugin.log(Level.INFO, "See here for more information: https://spark.lucko.me/docs/misc/Java-agent-warning");
             }
         } catch (Exception e) {
             // ignored
@@ -54,8 +56,8 @@ public class ClassFinder {
 
     private final Map<String, Class<?>> classes = new HashMap<>();
 
-    public ClassFinder() {
-        Instrumentation instrumentation = loadInstrumentation();
+    public InstrumentationClassFinder(SparkPlugin plugin) {
+        Instrumentation instrumentation = loadInstrumentation(plugin);
         if (instrumentation == null) {
             return;
         }
@@ -66,21 +68,9 @@ public class ClassFinder {
         }
     }
 
+    @Override
     public @Nullable Class<?> findClass(String className) {
-        // try instrumentation
-        Class<?> clazz = this.classes.get(className);
-        if (clazz != null) {
-            return clazz;
-        }
-
-        // try Class.forName
-        try {
-            return Class.forName(className);
-        } catch (Throwable e) {
-            // ignore
-        }
-
-        return null;
+        return this.classes.get(className);
     }
 
 }