Add '/spark heapdump' command

5 files changed, 250 insertions, 140 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 f73e3e4..ef21d1c 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
@@ -24,15 +24,16 @@ import com.google.common.collect.ImmutableList;
 
 import me.lucko.spark.common.command.Arguments;
 import me.lucko.spark.common.command.Command;
-import me.lucko.spark.common.command.modules.HeapModule;
-import me.lucko.spark.common.command.modules.MonitoringModule;
+import me.lucko.spark.common.command.modules.MemoryModule;
 import me.lucko.spark.common.command.modules.SamplerModule;
+import me.lucko.spark.common.command.modules.TickMonitoringModule;
 import me.lucko.spark.common.command.tabcomplete.CompletionSupplier;
 import me.lucko.spark.common.command.tabcomplete.TabCompleter;
 import me.lucko.spark.sampler.ThreadDumper;
 import me.lucko.spark.sampler.TickCounter;
 import me.lucko.spark.util.BytebinClient;
 
+import java.nio.file.Path;
 import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.Collections;
@@ -57,8 +58,8 @@ public abstract class SparkPlatform<S> {
     private static <T> List<Command<T>> prepareCommands() {
         ImmutableList.Builder<Command<T>> builder = ImmutableList.builder();
         new SamplerModule<T>().registerCommands(builder::add);
-        new MonitoringModule<T>().registerCommands(builder::add);
-        new HeapModule<T>().registerCommands(builder::add);
+        new TickMonitoringModule<T>().registerCommands(builder::add);
+        new MemoryModule<T>().registerCommands(builder::add);
         return builder.build();
     }
 
@@ -66,6 +67,7 @@ public abstract class SparkPlatform<S> {
     
     // abstract methods implemented by each platform
     public abstract String getVersion();
+    public abstract Path getPluginFolder();
     public abstract String getLabel();
     public abstract void sendMessage(S sender, String message);
     public abstract void sendMessage(String message);
diff --git a/spark-common/src/main/java/me/lucko/spark/common/command/modules/HeapModule.java b/spark-common/src/main/java/me/lucko/spark/common/command/modules/MemoryModule.java
index 318ce25..405b3d3 100644
--- a/spark-common/src/main/java/me/lucko/spark/common/command/modules/HeapModule.java
+++ b/spark-common/src/main/java/me/lucko/spark/common/command/modules/MemoryModule.java
@@ -24,26 +24,37 @@ import me.lucko.spark.common.SparkPlatform;
 import me.lucko.spark.common.command.Command;
 import me.lucko.spark.common.command.CommandModule;
 import me.lucko.spark.memory.HeapDump;
+import me.lucko.spark.memory.HeapDumpSummary;
 
 import okhttp3.MediaType;
 
 import java.io.IOException;
+import java.nio.file.Files;
+import java.nio.file.Path;
+import java.time.LocalDateTime;
+import java.time.format.DateTimeFormatter;
 import java.util.function.Consumer;
 
-public class HeapModule<S> implements CommandModule<S> {
+public class MemoryModule<S> implements CommandModule<S> {
     private static final MediaType JSON_TYPE = MediaType.parse("application/json; charset=utf-8");
 
     @Override
     public void registerCommands(Consumer<Command<S>> consumer) {
         consumer.accept(Command.<S>builder()
-                .aliases("heap", "memory")
+                .aliases("heapsummary")
+                .argumentUsage("run-gc-before", null)
                 .executor((platform, sender, arguments) -> {
                     platform.runAsync(() -> {
-                        platform.sendPrefixedMessage("&7Creating a new heap dump, please wait...");
+                        if (arguments.boolFlag("run-gc-before")) {
+                            platform.sendPrefixedMessage("&7Running garbage collector...");
+                            System.gc();
+                        }
 
-                        HeapDump heapDump;
+                        platform.sendPrefixedMessage("&7Creating a new heap dump summary, please wait...");
+
+                        HeapDumpSummary heapDump;
                         try {
-                            heapDump = HeapDump.createNew();
+                            heapDump = HeapDumpSummary.createNew();
                         } catch (Exception e) {
                             platform.sendPrefixedMessage("&cAn error occurred whilst inspecting the heap.");
                             e.printStackTrace();
@@ -53,7 +64,7 @@ public class HeapModule<S> implements CommandModule<S> {
                         byte[] output = heapDump.formCompressedDataPayload();
                         try {
                             String key = SparkPlatform.BYTEBIN_CLIENT.postGzippedContent(output, JSON_TYPE);
-                            platform.sendPrefixedMessage("&bHeap dump output:");
+                            platform.sendPrefixedMessage("&bHeap dump summmary output:");
                             platform.sendLink(SparkPlatform.VIEWER_URL + key);
                         } catch (IOException e) {
                             platform.sendPrefixedMessage("&cAn error occurred whilst uploading the data.");
@@ -63,6 +74,43 @@ public class HeapModule<S> implements CommandModule<S> {
                 })
                 .build()
         );
+
+        consumer.accept(Command.<S>builder()
+                .aliases("heapdump")
+                .argumentUsage("run-gc-before", null)
+                .argumentUsage("include-non-live", null)
+                .executor((platform, sender, arguments) -> {
+                    platform.runAsync(() -> {
+                        Path pluginFolder = platform.getPluginFolder();
+                        try {
+                            Files.createDirectories(pluginFolder);
+                        } catch (IOException e) {
+                            // ignore
+                        }
+
+                        Path file = pluginFolder.resolve("heap-" + DateTimeFormatter.ofPattern("yyyy-MM-dd_HH.mm.ss").format(LocalDateTime.now()) + ".hprof");
+                        boolean liveOnly = !arguments.boolFlag("include-non-live");
+
+                        if (arguments.boolFlag("run-gc-before")) {
+                            platform.sendPrefixedMessage("&7Running garbage collector...");
+                            System.gc();
+                        }
+
+                        platform.sendPrefixedMessage("&7Creating a new heap dump, please wait...");
+
+                        try {
+                            HeapDump.dumpHeap(file, liveOnly);
+                        } catch (Exception e) {
+                            platform.sendPrefixedMessage("&cAn error occurred whilst creating a heap dump.");
+                            e.printStackTrace();
+                            return;
+                        }
+
+                        platform.sendPrefixedMessage("&bHeap dump written to: " + file.toString());
+                    });
+                })
+                .build()
+        );
     }
 
 }
diff --git a/spark-common/src/main/java/me/lucko/spark/common/command/modules/MonitoringModule.java b/spark-common/src/main/java/me/lucko/spark/common/command/modules/TickMonitoringModule.java
index 608d6b4..d0513ab 100644
--- a/spark-common/src/main/java/me/lucko/spark/common/command/modules/MonitoringModule.java
+++ b/spark-common/src/main/java/me/lucko/spark/common/command/modules/TickMonitoringModule.java
@@ -33,7 +33,7 @@ import java.util.Arrays;
 import java.util.List;
 import java.util.function.Consumer;
 
-public class MonitoringModule<S> implements CommandModule<S> {
+public class TickMonitoringModule<S> implements CommandModule<S> {
 
     /** The tick monitor instance currently running, if any */
     private ReportingTickMonitor activeTickMonitor = null;
diff --git a/spark-common/src/main/java/me/lucko/spark/memory/HeapDump.java b/spark-common/src/main/java/me/lucko/spark/memory/HeapDump.java
index 4007bad..d91c476 100644
--- a/spark-common/src/main/java/me/lucko/spark/memory/HeapDump.java
+++ b/spark-common/src/main/java/me/lucko/spark/memory/HeapDump.java
@@ -20,154 +20,41 @@
 
 package me.lucko.spark.memory;
 
-import com.google.gson.stream.JsonWriter;
-
-import me.lucko.spark.util.TypeDescriptors;
-
-import java.io.ByteArrayOutputStream;
 import java.io.IOException;
-import java.io.OutputStreamWriter;
-import java.io.Writer;
 import java.lang.management.ManagementFactory;
-import java.nio.charset.StandardCharsets;
-import java.util.Arrays;
-import java.util.List;
-import java.util.Objects;
-import java.util.regex.Matcher;
-import java.util.regex.Pattern;
-import java.util.stream.Collectors;
-import java.util.zip.GZIPOutputStream;
+import java.nio.file.Path;
 
 import javax.management.JMX;
 import javax.management.MBeanServer;
 import javax.management.ObjectName;
 
 /**
- * Represents a "heap dump" from the VM.
- *
- * <p>Contains a number of entries, corresponding to types of objects in the virtual machine
- * and their recorded impact on memory usage.</p>
+ * Utility for creating .hprof memory heap snapshots.
  */
-public class HeapDump {
+public final class HeapDump {
+
+    private HeapDump() {}
 
-    /** The object name of the com.sun.management.DiagnosticCommandMBean */
-    private static final String DIAGNOSTIC_BEAN = "com.sun.management:type=DiagnosticCommand";
-    /** A regex pattern representing the expected format of the raw heap output */
-    private static final Pattern OUTPUT_FORMAT = Pattern.compile("^\\s*(\\d+):\\s*(\\d+)\\s*(\\d+)\\s*([^\\s]+).*$");
+    /** The object name of the com.sun.management.HotSpotDiagnosticMXBean */
+    private static final String DIAGNOSTIC_BEAN = "com.sun.management:type=HotSpotDiagnostic";
 
     /**
-     * Obtains the raw heap data output from the DiagnosticCommandMBean.
+     * Creates a heap dump at the given output path.
      *
-     * @return the raw output
-     * @throws Exception lots could go wrong!
+     * @param outputPath the path to write the snapshot to
+     * @param live if true dump only live objects i.e. objects that are reachable from others
+     * @throws Exception catch all
      */
-    private static String getRawHeapData() throws Exception {
+    public static void dumpHeap(Path outputPath, boolean live) throws Exception {
         MBeanServer beanServer = ManagementFactory.getPlatformMBeanServer();
         ObjectName diagnosticBeanName = ObjectName.getInstance(DIAGNOSTIC_BEAN);
 
-        DiagnosticCommandMXBean proxy = JMX.newMXBeanProxy(beanServer, diagnosticBeanName, DiagnosticCommandMXBean.class);
-        return proxy.gcClassHistogram(new String[0]);
-    }
-
-    /**
-     * Creates a new heap dump based on the current VM.
-     *
-     * @return the created heap dump
-     * @throws RuntimeException if an error occurred whilst requesting a heap dump from the VM
-     */
-    public static HeapDump createNew() {
-        String rawOutput;
-        try {
-            rawOutput = getRawHeapData();
-        } catch (Exception e) {
-            throw new RuntimeException("Unable to get heap dump", e);
-        }
-
-        return new HeapDump(Arrays.stream(rawOutput.split("\n"))
-                .map(line -> {
-                    Matcher matcher = OUTPUT_FORMAT.matcher(line);
-                    if (!matcher.matches()) {
-                        return null;
-                    }
-
-                    return new Entry(
-                            Integer.parseInt(matcher.group(1)),
-                            Integer.parseInt(matcher.group(2)),
-                            Long.parseLong(matcher.group(3)),
-                            TypeDescriptors.getJavaType(matcher.group(4))
-                    );
-                })
-                .filter(Objects::nonNull)
-                .collect(Collectors.toList()));
-    }
-
-    /** The entries in this heap dump */
-    private final List<Entry> entries;
-
-    private HeapDump(List<Entry> entries) {
-        this.entries = entries;
-    }
-
-    private void writeOutput(JsonWriter writer) throws IOException {
-        writer.beginObject();
-        writer.name("type").value("heap");
-        writer.name("entries").beginArray();
-        for (Entry entry : this.entries) {
-            writer.beginObject();
-            writer.name("#").value(entry.getOrder());
-            writer.name("i").value(entry.getInstances());
-            writer.name("s").value(entry.getBytes());
-            writer.name("t").value(entry.getType());
-            writer.endObject();
-        }
-        writer.endArray();
-        writer.endObject();
-    }
-
-    public byte[] formCompressedDataPayload() {
-        ByteArrayOutputStream byteOut = new ByteArrayOutputStream();
-        try (Writer writer = new OutputStreamWriter(new GZIPOutputStream(byteOut), StandardCharsets.UTF_8)) {
-            try (JsonWriter jsonWriter = new JsonWriter(writer)) {
-                writeOutput(jsonWriter);
-            }
-        } catch (IOException e) {
-            throw new RuntimeException(e);
-        }
-        return byteOut.toByteArray();
-    }
-
-    public static final class Entry {
-        private final int order;
-        private final int instances;
-        private final long bytes;
-        private final String type;
-
-        Entry(int order, int instances, long bytes, String type) {
-            this.order = order;
-            this.instances = instances;
-            this.bytes = bytes;
-            this.type = type;
-        }
-
-        public int getOrder() {
-            return this.order;
-        }
-
-        public int getInstances() {
-            return this.instances;
-        }
-
-        public long getBytes() {
-            return this.bytes;
-        }
-
-        public String getType() {
-            return this.type;
-        }
+        HotSpotDiagnosticMXBean proxy = JMX.newMXBeanProxy(beanServer, diagnosticBeanName, HotSpotDiagnosticMXBean.class);
+        proxy.dumpHeap(outputPath.toAbsolutePath().normalize().toString(), live);
     }
 
-    public interface DiagnosticCommandMXBean {
-        String gcClassHistogram(String[] args);
+    public interface HotSpotDiagnosticMXBean {
+        void dumpHeap(String outputFile, boolean live) throws IOException;
     }
 
 }
diff --git a/spark-common/src/main/java/me/lucko/spark/memory/HeapDumpSummary.java b/spark-common/src/main/java/me/lucko/spark/memory/HeapDumpSummary.java
new file mode 100644
index 0000000..402b89e
--- /dev/null
+++ b/spark-common/src/main/java/me/lucko/spark/memory/HeapDumpSummary.java
@@ -0,0 +1,173 @@
+/*
+ * 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.memory;
+
+import com.google.gson.stream.JsonWriter;
+
+import me.lucko.spark.util.TypeDescriptors;
+
+import java.io.ByteArrayOutputStream;
+import java.io.IOException;
+import java.io.OutputStreamWriter;
+import java.io.Writer;
+import java.lang.management.ManagementFactory;
+import java.nio.charset.StandardCharsets;
+import java.util.Arrays;
+import java.util.List;
+import java.util.Objects;
+import java.util.regex.Matcher;
+import java.util.regex.Pattern;
+import java.util.stream.Collectors;
+import java.util.zip.GZIPOutputStream;
+
+import javax.management.JMX;
+import javax.management.MBeanServer;
+import javax.management.ObjectName;
+
+/**
+ * Represents a "heap dump summary" from the VM.
+ *
+ * <p>Contains a number of entries, corresponding to types of objects in the virtual machine
+ * and their recorded impact on memory usage.</p>
+ */
+public final class HeapDumpSummary {
+
+    /** The object name of the com.sun.management.DiagnosticCommandMBean */
+    private static final String DIAGNOSTIC_BEAN = "com.sun.management:type=DiagnosticCommand";
+    /** A regex pattern representing the expected format of the raw heap output */
+    private static final Pattern OUTPUT_FORMAT = Pattern.compile("^\\s*(\\d+):\\s*(\\d+)\\s*(\\d+)\\s*([^\\s]+).*$");
+
+    /**
+     * Obtains the raw heap data output from the DiagnosticCommandMBean.
+     *
+     * @return the raw output
+     * @throws Exception lots could go wrong!
+     */
+    private static String getRawHeapData() throws Exception {
+        MBeanServer beanServer = ManagementFactory.getPlatformMBeanServer();
+        ObjectName diagnosticBeanName = ObjectName.getInstance(DIAGNOSTIC_BEAN);
+
+        DiagnosticCommandMXBean proxy = JMX.newMXBeanProxy(beanServer, diagnosticBeanName, DiagnosticCommandMXBean.class);
+        return proxy.gcClassHistogram(new String[0]);
+    }
+
+    /**
+     * Creates a new heap dump based on the current VM.
+     *
+     * @return the created heap dump
+     * @throws RuntimeException if an error occurred whilst requesting a heap dump from the VM
+     */
+    public static HeapDumpSummary createNew() {
+        String rawOutput;
+        try {
+            rawOutput = getRawHeapData();
+        } catch (Exception e) {
+            throw new RuntimeException("Unable to get heap dump", e);
+        }
+
+        return new HeapDumpSummary(Arrays.stream(rawOutput.split("\n"))
+                .map(line -> {
+                    Matcher matcher = OUTPUT_FORMAT.matcher(line);
+                    if (!matcher.matches()) {
+                        return null;
+                    }
+
+                    return new Entry(
+                            Integer.parseInt(matcher.group(1)),
+                            Integer.parseInt(matcher.group(2)),
+                            Long.parseLong(matcher.group(3)),
+                            TypeDescriptors.getJavaType(matcher.group(4))
+                    );
+                })
+                .filter(Objects::nonNull)
+                .collect(Collectors.toList()));
+    }
+
+    /** The entries in this heap dump */
+    private final List<Entry> entries;
+
+    private HeapDumpSummary(List<Entry> entries) {
+        this.entries = entries;
+    }
+
+    private void writeOutput(JsonWriter writer) throws IOException {
+        writer.beginObject();
+        writer.name("type").value("heap");
+        writer.name("entries").beginArray();
+        for (Entry entry : this.entries) {
+            writer.beginObject();
+            writer.name("#").value(entry.getOrder());
+            writer.name("i").value(entry.getInstances());
+            writer.name("s").value(entry.getBytes());
+            writer.name("t").value(entry.getType());
+            writer.endObject();
+        }
+        writer.endArray();
+        writer.endObject();
+    }
+
+    public byte[] formCompressedDataPayload() {
+        ByteArrayOutputStream byteOut = new ByteArrayOutputStream();
+        try (Writer writer = new OutputStreamWriter(new GZIPOutputStream(byteOut), StandardCharsets.UTF_8)) {
+            try (JsonWriter jsonWriter = new JsonWriter(writer)) {
+                writeOutput(jsonWriter);
+            }
+        } catch (IOException e) {
+            throw new RuntimeException(e);
+        }
+        return byteOut.toByteArray();
+    }
+
+    public static final class Entry {
+        private final int order;
+        private final int instances;
+        private final long bytes;
+        private final String type;
+
+        Entry(int order, int instances, long bytes, String type) {
+            this.order = order;
+            this.instances = instances;
+            this.bytes = bytes;
+            this.type = type;
+        }
+
+        public int getOrder() {
+            return this.order;
+        }
+
+        public int getInstances() {
+            return this.instances;
+        }
+
+        public long getBytes() {
+            return this.bytes;
+        }
+
+        public String getType() {
+            return this.type;
+        }
+    }
+
+    public interface DiagnosticCommandMXBean {
+        String gcClassHistogram(String[] args);
+    }
+
+}