Add jar file index to improve mc command compilation speed (#2736)

memorycompiler/src/main/java/com/taobao/arthas/compiler/ClassUriWrapper.java

@@ -0,0 +1,42 @@
+package com.taobao.arthas.compiler;
+
+/*-
+ * #%L
+ * compiler
+ * %%
+ * Copyright (C) 2017 - 2018 SkaLogs
+ * %%
+ * Licensed 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.
+ * #L%
+ */
+
+import java.net.URI;
+
+public class ClassUriWrapper {
+    private final URI uri;
+
+    private final String className;
+
+    public ClassUriWrapper(String className, URI uri) {
+        this.className = className;
+        this.uri = uri;
+    }
+
+    public URI getUri() {
+        return uri;
+    }
+
+    public String getClassName() {
+        return className;
+    }
+}

memorycompiler/src/main/java/com/taobao/arthas/compiler/CustomJavaFileObject.java

@@ -23,18 +23,20 @@
 import javax.lang.model.element.Modifier;
 import javax.lang.model.element.NestingKind;
 import javax.tools.JavaFileObject;
-import java.io.*;
+import java.io.IOException;
+import java.io.InputStream;
+import java.io.OutputStream;
+import java.io.Reader;
+import java.io.Writer;
 import java.net.URI;
 
 public class CustomJavaFileObject implements JavaFileObject {
-    private final String binaryName;
+    private final String className;
     private final URI uri;
-    private final String name;
 
-    public CustomJavaFileObject(String binaryName, URI uri) {
+    public CustomJavaFileObject(String className, URI uri) {
         this.uri = uri;
-        this.binaryName = binaryName;
-        name = uri.getPath() == null ? uri.getSchemeSpecificPart() : uri.getPath(); // for FS based URI the path is not null, for JAR URI the scheme specific part is not null
+        this.className = className;
     }
 
     public URI toUri() {
@@ -50,7 +52,7 @@ public OutputStream openOutputStream() {
     }
 
     public String getName() {
-        return name;
+        return this.className;
     }
 
     public Reader openReader(boolean ignoreEncodingErrors) {
@@ -78,10 +80,8 @@ public Kind getKind() {
     }
 
     public boolean isNameCompatible(String simpleName, Kind kind) {
-        String baseName = simpleName + kind.extension;
-        return kind.equals(getKind())
-                && (baseName.equals(getName())
-                || getName().endsWith("/" + baseName));
+        return Kind.CLASS.equals(getKind())
+                && this.className.endsWith(simpleName);
     }
 
     public NestingKind getNestingKind() {
@@ -92,8 +92,8 @@ public Modifier getAccessLevel() {
         throw new UnsupportedOperationException();
     }
 
-    public String binaryName() {
-        return binaryName;
+    public String getClassName() {
+        return this.className;
     }
 
 

memorycompiler/src/main/java/com/taobao/arthas/compiler/DynamicJavaFileManager.java

@@ -24,8 +24,7 @@ public class DynamicJavaFileManager extends ForwardingJavaFileManager<JavaFileMa
     public DynamicJavaFileManager(JavaFileManager fileManager, DynamicClassLoader classLoader) {
         super(fileManager);
         this.classLoader = classLoader;
-
-        finder = new PackageInternalsFinder(classLoader);
+        this.finder = new PackageInternalsFinder(classLoader);
     }
 
     @Override
@@ -53,7 +52,7 @@ public ClassLoader getClassLoader(JavaFileManager.Location location) {
     @Override
     public String inferBinaryName(Location location, JavaFileObject file) {
         if (file instanceof CustomJavaFileObject) {
-            return ((CustomJavaFileObject) file).binaryName();
+            return ((CustomJavaFileObject) file).getClassName();
         } else {
             /**
              * if it's not CustomJavaFileObject, then it's coming from standard file manager

memorycompiler/src/main/java/com/taobao/arthas/compiler/PackageInternalsFinder.java

@@ -28,15 +28,23 @@
 import java.net.URL;
 import java.net.URLDecoder;
 import java.util.ArrayList;
+import java.util.Arrays;
 import java.util.Collection;
+import java.util.Collections;
 import java.util.Enumeration;
+import java.util.HashMap;
 import java.util.List;
+import java.util.Map;
+import java.util.concurrent.ConcurrentHashMap;
 import java.util.jar.JarEntry;
+import java.util.stream.Collectors;
 
 public class PackageInternalsFinder {
     private final ClassLoader classLoader;
     private static final String CLASS_FILE_EXTENSION = ".class";
 
+    private static final Map<String, JarFileIndex> INDEXS = new ConcurrentHashMap<>();
+
     public PackageInternalsFinder(ClassLoader classLoader) {
         this.classLoader = classLoader;
     }
@@ -60,11 +68,30 @@ private Collection<JavaFileObject> listUnder(String packageName, URL packageFold
         if (directory.isDirectory()) { // browse local .class files - useful for local execution
             return processDir(packageName, directory);
         } else { // browse a jar file
-            return processJar(packageFolderURL);
-        } // maybe there can be something else for more involved class loaders
+            return processJar(packageName, packageFolderURL);
+        }
+    }
+
+    private List<JavaFileObject> processJar(String packageName, URL packageFolderURL) {
+        try {
+            String jarUri = packageFolderURL.toExternalForm().substring(0, packageFolderURL.toExternalForm().lastIndexOf("!/"));
+            JarFileIndex jarFileIndex = INDEXS.get(jarUri);
+            if (jarFileIndex == null) {
+                jarFileIndex = new JarFileIndex(jarUri, URI.create(jarUri + "!/"));
+                INDEXS.put(jarUri, jarFileIndex);
+            }
+            List<JavaFileObject> result = jarFileIndex.search(packageName);
+            if (result != null) {
+                return result;
+            }
+        } catch (Exception e) {
+            // ignore
+        }
+        // 保底
+        return fuse(packageFolderURL);
     }
 
-    private List<JavaFileObject> processJar(URL packageFolderURL) {
+    private List<JavaFileObject> fuse(URL packageFolderURL) {
         List<JavaFileObject> result = new ArrayList<JavaFileObject>();
         try {
             String jarUri = packageFolderURL.toExternalForm().substring(0, packageFolderURL.toExternalForm().lastIndexOf("!/"));
@@ -92,24 +119,16 @@ private List<JavaFileObject> processJar(URL packageFolderURL) {
     }
 
     private List<JavaFileObject> processDir(String packageName, File directory) {
-        List<JavaFileObject> result = new ArrayList<JavaFileObject>();
-
-        File[] childFiles = directory.listFiles();
-        if (childFiles != null) {
-            for (File childFile : childFiles) {
-                if (childFile.isFile()) {
-                    // We only want the .class files.
-                    if (childFile.getName().endsWith(CLASS_FILE_EXTENSION)) {
-                        String binaryName = packageName + "." + childFile.getName();
-                        binaryName = binaryName.replaceAll(CLASS_FILE_EXTENSION + "$", "");
-
-                        result.add(new CustomJavaFileObject(binaryName, childFile.toURI()));
-                    }
-                }
-            }
+        File[] files = directory.listFiles(item ->
+                item.isFile() && getKind(item.getName()) == JavaFileObject.Kind.CLASS);
+        if (files != null) {
+            return Arrays.stream(files).map(item -> {
+                String className = packageName + "." + item.getName()
+                        .replaceAll(CLASS_FILE_EXTENSION + "$", "");
+                return new CustomJavaFileObject(className, item.toURI());
+            }).collect(Collectors.toList());
         }
-
-        return result;
+        return Collections.emptyList();
     }
 
     private String decode(String filePath) {
@@ -121,4 +140,69 @@ private String decode(String filePath) {
 
         return filePath;
     }
+
+    public static JavaFileObject.Kind getKind(String name) {
+        if (name.endsWith(JavaFileObject.Kind.CLASS.extension))
+            return JavaFileObject.Kind.CLASS;
+        else if (name.endsWith(JavaFileObject.Kind.SOURCE.extension))
+            return JavaFileObject.Kind.SOURCE;
+        else if (name.endsWith(JavaFileObject.Kind.HTML.extension))
+            return JavaFileObject.Kind.HTML;
+        else
+            return JavaFileObject.Kind.OTHER;
+    }
+
+    public static class JarFileIndex {
+        private String jarUri;
+        private URI uri;
+
+        private Map<String, List<ClassUriWrapper>> packages = new HashMap<>();
+
+        public JarFileIndex(String jarUri, URI uri) throws IOException {
+            this.jarUri = jarUri;
+            this.uri = uri;
+            loadIndex();
+        }
+
+        private void loadIndex() throws IOException {
+            JarURLConnection jarConn = (JarURLConnection) uri.toURL().openConnection();
+            String rootEntryName = jarConn.getEntryName() == null ? "" : jarConn.getEntryName();
+            Enumeration<JarEntry> entryEnum = jarConn.getJarFile().entries();
+            while (entryEnum.hasMoreElements()) {
+                JarEntry jarEntry = entryEnum.nextElement();
+                String entryName = jarEntry.getName();
+                if (entryName.startsWith(rootEntryName) && entryName.endsWith(CLASS_FILE_EXTENSION)) {
+                    String className = entryName
+                            .substring(0, entryName.length() - CLASS_FILE_EXTENSION.length())
+                            .replace(rootEntryName, "")
+                            .replace("/", ".");
+                    if (className.startsWith(".")) className = className.substring(1);
+                    if (className.equals("package-info")
+                            || className.equals("module-info")
+                            || className.lastIndexOf(".") == -1) {
+                        continue;
+                    }
+                    String packageName = className.substring(0, className.lastIndexOf("."));
+                    List<ClassUriWrapper> classes = packages.get(packageName);
+                    if (classes == null) {
+                        classes = new ArrayList<>();
+                        packages.put(packageName, classes);
+                    }
+                    classes.add(new ClassUriWrapper(className, URI.create(jarUri + "!/" + entryName)));
+                }
+            }
+        }
+
+        public List<JavaFileObject> search(String packageName) {
+            if (this.packages.isEmpty()) {
+                return null;
+            }
+            if (this.packages.containsKey(packageName)) {
+                return packages.get(packageName).stream().map(item -> {
+                    return new CustomJavaFileObject(item.getClassName(), item.getUri());
+                }).collect(Collectors.toList());
+            }
+            return Collections.emptyList();
+        }
+    }
 }