Extract class sources (plugin/mod names) and include in data payload

3 files changed, 299 insertions, 0 deletions

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.java
new file mode 100644
index 0000000..4481786
--- /dev/null
+++ b/spark-common/src/main/java/me/lucko/spark/common/util/ClassFinder.java
@@ -0,0 +1,71 @@
+/*
+ * 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;
+
+import net.bytebuddy.agent.ByteBuddyAgent;
+
+import org.checkerframework.checker.nullness.qual.Nullable;
+
+import java.lang.instrument.Instrumentation;
+import java.util.HashMap;
+import java.util.Map;
+
+/**
+ * Uses {@link Instrumentation} to find a class reference for given class names.
+ *
+ * <p>This is necessary as we don't always have access to the classloader for a given class.</p>
+ */
+public class ClassFinder {
+
+    private final Map<String, Class<?>> classes = new HashMap<>();
+
+    public ClassFinder() {
+        Instrumentation instrumentation;
+        try {
+            instrumentation = ByteBuddyAgent.install();
+        } catch (Exception e) {
+            return;
+        }
+
+        // obtain and cache loaded classes
+        for (Class<?> loadedClass : instrumentation.getAllLoadedClasses()) {
+            this.classes.put(loadedClass.getName(), loadedClass);
+        }
+    }
+
+    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;
+    }
+
+}
diff --git a/spark-common/src/main/java/me/lucko/spark/common/util/ClassSourceLookup.java b/spark-common/src/main/java/me/lucko/spark/common/util/ClassSourceLookup.java
new file mode 100644
index 0000000..27e3ec6
--- /dev/null
+++ b/spark-common/src/main/java/me/lucko/spark/common/util/ClassSourceLookup.java
@@ -0,0 +1,223 @@
+/*
+ * 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;
+
+import me.lucko.spark.common.sampler.node.StackTraceNode;
+import me.lucko.spark.common.sampler.node.ThreadNode;
+
+import org.checkerframework.checker.nullness.qual.Nullable;
+
+import java.io.IOException;
+import java.net.URISyntaxException;
+import java.net.URL;
+import java.net.URLClassLoader;
+import java.nio.file.Paths;
+import java.security.CodeSource;
+import java.security.ProtectionDomain;
+import java.util.Collections;
+import java.util.HashMap;
+import java.util.Map;
+import java.util.Objects;
+
+/**
+ * A function which defines the source of given {@link Class}es.
+ */
+public interface ClassSourceLookup {
+
+    /**
+     * Identify the given class.
+     *
+     * @param clazz the class
+     * @return the source of the class
+     */
+    @Nullable String identify(Class<?> clazz) throws Exception;
+
+    /**
+     * A no-operation {@link ClassSourceLookup}.
+     */
+    ClassSourceLookup NO_OP = new ClassSourceLookup() {
+        @Override
+        public @Nullable String identify(Class<?> clazz) {
+            return null;
+        }
+    };
+
+    /**
+     * A {@link ClassSourceLookup} which identifies classes based on their {@link ClassLoader}.
+     */
+    abstract class ByClassLoader implements ClassSourceLookup {
+
+        public abstract @Nullable String identify(ClassLoader loader) throws Exception;
+
+        @Override
+        public final @Nullable String identify(Class<?> clazz) throws Exception {
+            ClassLoader loader = clazz.getClassLoader();
+            while (loader != null) {
+                String source = identify(loader);
+                if (source != null) {
+                    return source;
+                }
+                loader = loader.getParent();
+            }
+            return null;
+        }
+    }
+
+    /**
+     * A {@link ClassSourceLookup} which identifies classes based on the first URL in a {@link URLClassLoader}.
+     */
+    class ByFirstUrlSource extends ByClassLoader {
+        @Override
+        public @Nullable String identify(ClassLoader loader) throws IOException, URISyntaxException {
+            if (loader instanceof URLClassLoader) {
+                URLClassLoader urlClassLoader = (URLClassLoader) loader;
+                URL[] urls = urlClassLoader.getURLs();
+                if (urls.length == 0) {
+                    return null;
+                }
+                return identifyUrl(urls[0]);
+            }
+            return null;
+        }
+    }
+
+    /**
+     * A {@link ClassSourceLookup} which identifies classes based on their {@link ProtectionDomain#getCodeSource()}.
+     */
+    class ByCodeSource implements ClassSourceLookup {
+        @Override
+        public @Nullable String identify(Class<?> clazz) throws URISyntaxException {
+            ProtectionDomain protectionDomain = clazz.getProtectionDomain();
+            if (protectionDomain == null) {
+                return null;
+            }
+            CodeSource codeSource = protectionDomain.getCodeSource();
+            if (codeSource == null) {
+                return null;
+            }
+
+            URL url = codeSource.getLocation();
+            return url == null ? null : identifyUrl(url);
+        }
+    }
+
+    /**
+     * Attempts to identify a jar file from a URL.
+     *
+     * @param url the url
+     * @return the name of the file
+     * @throws URISyntaxException thrown by {@link URL#toURI()}
+     */
+    static String identifyUrl(URL url) throws URISyntaxException {
+        if (url.getProtocol().equals("file")) {
+            String jarName = Paths.get(url.toURI()).getFileName().toString();
+            if (jarName.endsWith(".jar")) {
+                return jarName.substring(0, jarName.length() - 4);
+            }
+        }
+        return null;
+    }
+    interface Visitor {
+        void visit(ThreadNode node);
+
+        boolean hasMappings();
+
+        Map<String, String> getMapping();
+    }
+
+    static Visitor createVisitor(ClassSourceLookup lookup) {
+        if (lookup == ClassSourceLookup.NO_OP) {
+            return NoOpVistitor.INSTANCE; // don't bother!
+        }
+        return new VisitorImpl(lookup);
+    }
+
+    enum NoOpVistitor implements Visitor {
+        INSTANCE;
+
+        @Override
+        public void visit(ThreadNode node) {
+
+        }
+
+        @Override
+        public boolean hasMappings() {
+            return false;
+        }
+
+        @Override
+        public Map<String, String> getMapping() {
+            return Collections.emptyMap();
+        }
+    }
+
+    /**
+     * Visitor which scans {@link StackTraceNode}s and accumulates class identities.
+     */
+    class VisitorImpl implements Visitor {
+        private final ClassSourceLookup lookup;
+        private final ClassFinder classFinder = new ClassFinder();
+
+        // class name --> identifier (plugin name)
+        private final Map<String, String> map = new HashMap<>();
+
+        VisitorImpl(ClassSourceLookup lookup) {
+            this.lookup = lookup;
+        }
+
+        @Override
+        public void visit(ThreadNode node) {
+            for (StackTraceNode child : node.getChildren()) {
+                visitStackNode(child);
+            }
+        }
+
+        @Override
+        public boolean hasMappings() {
+            return !this.map.isEmpty();
+        }
+
+        @Override
+        public Map<String, String> getMapping() {
+            this.map.values().removeIf(Objects::isNull);
+            return this.map;
+        }
+
+        private void visitStackNode(StackTraceNode node) {
+            String className = node.getClassName();
+            if (!this.map.containsKey(className)) {
+                try {
+                    Class<?> clazz = this.classFinder.findClass(className);
+                    Objects.requireNonNull(clazz);
+                    this.map.put(className, this.lookup.identify(clazz));
+                } catch (Throwable e) {
+                    this.map.put(className, null);
+                }
+            }
+
+            // recursively
+            for (StackTraceNode child : node.getChildren()) {
+                visitStackNode(child);
+            }
+        }
+    }
+
+}
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 2e113a9..a35bf08 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
@@ -48,6 +48,11 @@ public final class MethodDisambiguator {
     private final Map<String, ComputedClass> cache = new ConcurrentHashMap<>();
 
     public Optional<MethodDescription> disambiguate(StackTraceNode element) {
+        String desc = element.getMethodDescription();
+        if (desc != null) {
+            return Optional.of(new MethodDescription(element.getMethodName(), desc));
+        }
+
         return disambiguate(element.getClassName(), element.getMethodName(), element.getLineNumber());
     }