Self-attachment agent for JOL

agent.patch

Index: jol-core/src/main/java/org/openjdk/jol/util/VMSupport.java
IDEA additional info:
Subsystem: com.intellij.openapi.diff.impl.patch.CharsetEP
<+>UTF-8
===================================================================
--- jol-core/src/main/java/org/openjdk/jol/util/VMSupport.java	(revision 31:9a546334aa57fa0e114db159b2a8a5345c249b3e)
+++ jol-core/src/main/java/org/openjdk/jol/util/VMSupport.java	(revision 31+:9a546334aa57+)
@@ -35,13 +35,22 @@
 import javax.management.openmbean.CompositeDataSupport;
 import java.io.PrintWriter;
 import java.io.StringWriter;
+import java.io.InputStream;
+import java.io.File;
+import java.io.FileOutputStream;
 import java.lang.instrument.Instrumentation;
 import java.lang.management.ManagementFactory;
 import java.lang.reflect.Field;
+import java.net.URL;
+import java.net.URLClassLoader;
 import java.security.AccessController;
 import java.security.PrivilegedAction;
 import java.util.Arrays;
 import java.util.Random;
+import java.util.jar.Attributes;
+import java.util.jar.JarEntry;
+import java.util.jar.JarOutputStream;
+import java.util.jar.Manifest;
 import java.util.regex.Matcher;
 import java.util.regex.Pattern;
 
@@ -55,6 +64,8 @@
 
     private static Instrumentation INSTRUMENTATION;
 
+    private static boolean installAttempt;
+
     public static final Unsafe U;
 
     public static final String VM_NAME;
@@ -253,6 +264,7 @@
         private final boolean exactSizeAvail;
 
         public SizeInfo(Object o, ClassLayout layout) {
+            VMSupport.installInstrumentation();
             exactSizeAvail = VMSupport.INSTRUMENTATION != null && o != null;
             size = exactSizeAvail ? (int) VMSupport.INSTRUMENTATION.getObjectSize(o) : layout.instanceSize();
         }
@@ -418,6 +430,7 @@
     }
 
     public static int sizeOf(Object o) {
+        installInstrumentation();
         if (VMSupport.INSTRUMENTATION != null) {
             return VMSupport.align((int) VMSupport.INSTRUMENTATION.getObjectSize(o));
         }
@@ -425,6 +438,79 @@
         return new CurrentLayouter().layout(ClassData.parseInstance(o)).instanceSize();
     }
 
+    private static void installInstrumentation() {
+        if (!installAttempt) {
+            try {
+                doInstall();
+            } catch (Exception ignored) {
+            }
+        }
+    }
+
+    private static synchronized void doInstall() throws Exception {
+        if (installAttempt) {
+            return;
+        }
+        installAttempt = true;
+        ClassLoader classLoader = new URLClassLoader(new URL[]{new File(System.getProperty("java.home")
+                .replace('\\', '/') + "/../lib/tools.jar").toURI().toURL()}, null);
+        Class<?> virtualMachine = classLoader.loadClass("com.sun.tools.attach.VirtualMachine");
+        String runtimeName = ManagementFactory.getRuntimeMXBean().getName();
+        Object virtualMachineInstance = virtualMachine.getDeclaredMethod("attach", String.class)
+                .invoke(null, runtimeName.substring(0, runtimeName.indexOf('@')));
+        try {
+            File agentFile = File.createTempFile("jolAgent", ".jar");
+            try {
+                saveAgentJar(agentFile);
+                virtualMachine.getDeclaredMethod("loadAgent", String.class, String.class)
+                        .invoke(virtualMachineInstance, agentFile.getAbsolutePath(), "");
+                INSTRUMENTATION = doGetInstrumentation();
+            } finally {
+                agentFile.delete();
+            }
+        } finally {
+            virtualMachine.getDeclaredMethod("detach").invoke(virtualMachineInstance);
+        }
+    }
+
+    private static Instrumentation doGetInstrumentation() {
+        try {
+            Field field = ClassLoader.getSystemClassLoader()
+                    .loadClass(Installer.class.getName())
+                    .getDeclaredField("instrumentation");
+            field.setAccessible(true);
+            return (Instrumentation) field.get(null);
+        } catch (Exception e) {
+            throw null;
+        }
+    }
+
+    private static void saveAgentJar(File agentFile) throws Exception {
+        InputStream inputStream = Installer.class.getResourceAsStream('/' + Installer.class.getName().replace('.', '/') + ".class");
+        if (inputStream == null) {
+            return;
+        }
+        try {
+            Manifest manifest = new Manifest();
+            manifest.getMainAttributes().put(Attributes.Name.MANIFEST_VERSION, "1.0");
+            manifest.getMainAttributes().put(new Attributes.Name("Agent-Class"), Installer.class.getName());
+            JarOutputStream jarOutputStream = new JarOutputStream(new FileOutputStream(agentFile), manifest);
+            try {
+                jarOutputStream.putNextEntry(new JarEntry('/' + Installer.class.getName().replace('.', '/') + ".class"));
+                byte[] buffer = new byte[1024];
+                int index;
+                while ((index = inputStream.read(buffer)) != -1) {
+                    jarOutputStream.write(buffer, 0, index);
+                }
+                jarOutputStream.closeEntry();
+            } finally {
+                jarOutputStream.close();
+            }
+        } finally {
+            inputStream.close();
+        }
+    }
+
     /**
      * Produces the toString string, only calling toString() on known types,
      * which do not mutate the instance.
@@ -526,6 +612,15 @@
 
     static class MyDoubles4 {
         private double f1, f2, f3, f4;
+    }
+
+    public static class Installer {
+
+        private static volatile Instrumentation instrumentation;
+
+        public static void agentmain(String agentArgs, Instrumentation inst) {
+            instrumentation = inst;
+        }
     }
 
 }