Refactor ClassFinder into interface

5 files changed, 149 insertions, 23 deletions

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);
     }
 
 }