Add support for heap dump compression (#16)

4 files changed, 180 insertions, 27 deletions

diff --git a/spark-common/build.gradle b/spark-common/build.gradle
index 27bb3af..5213b6b 100644
--- a/spark-common/build.gradle
+++ b/spark-common/build.gradle
@@ -1,6 +1,7 @@
 dependencies {
     compile 'com.squareup.okhttp3:okhttp:3.14.1'
     compile 'com.squareup.okio:okio:1.17.3'
+    compile 'org.tukaani:xz:1.8'
     compile('net.kyori:text-api:3.0.0') {
         exclude(module: 'checker-qual')
     }
diff --git a/spark-common/src/main/java/me/lucko/spark/common/command/modules/HealthModule.java b/spark-common/src/main/java/me/lucko/spark/common/command/modules/HealthModule.java
index 921495c..7ad9525 100644
--- a/spark-common/src/main/java/me/lucko/spark/common/command/modules/HealthModule.java
+++ b/spark-common/src/main/java/me/lucko/spark/common/command/modules/HealthModule.java
@@ -26,6 +26,7 @@ import me.lucko.spark.common.command.CommandModule;
 import me.lucko.spark.common.command.tabcomplete.TabCompleter;
 import me.lucko.spark.common.monitor.cpu.CpuMonitor;
 import me.lucko.spark.common.monitor.tick.TpsCalculator;
+import me.lucko.spark.common.util.FormatUtil;
 import net.kyori.text.Component;
 import net.kyori.text.TextComponent;
 import net.kyori.text.format.TextColor;
@@ -144,14 +145,14 @@ public class HealthModule implements CommandModule {
                                 .build()
                         );
                         report.add(TextComponent.builder("    ")
-                                .append(TextComponent.of(formatBytes(heapUsage.getUsed()), TextColor.WHITE))
+                                .append(TextComponent.of(FormatUtil.formatBytes(heapUsage.getUsed()), TextColor.WHITE))
                                 .append(TextComponent.space())
                                 .append(TextComponent.of("/", TextColor.GRAY))
                                 .append(TextComponent.space())
-                                .append(TextComponent.of(formatBytes(heapUsage.getMax()), TextColor.WHITE))
+                                .append(TextComponent.of(FormatUtil.formatBytes(heapUsage.getMax()), TextColor.WHITE))
                                 .append(TextComponent.of("   "))
                                 .append(TextComponent.of("(", TextColor.GRAY))
-                                .append(TextComponent.of(percent(heapUsage.getUsed(), heapUsage.getMax()), TextColor.GREEN))
+                                .append(TextComponent.of(FormatUtil.percent(heapUsage.getUsed(), heapUsage.getMax()), TextColor.GREEN))
                                 .append(TextComponent.of(")", TextColor.GRAY))
                                 .build()
                         );
@@ -167,7 +168,7 @@ public class HealthModule implements CommandModule {
                                     .build()
                             );
                             report.add(TextComponent.builder("    ")
-                                    .append(TextComponent.of(formatBytes(nonHeapUsage.getUsed()), TextColor.WHITE))
+                                    .append(TextComponent.of(FormatUtil.formatBytes(nonHeapUsage.getUsed()), TextColor.WHITE))
                                     .build()
                             );
                             report.add(TextComponent.empty());
@@ -192,14 +193,14 @@ public class HealthModule implements CommandModule {
                                         .build()
                                 );
                                 report.add(TextComponent.builder("    ")
-                                        .append(TextComponent.of(formatBytes(usage.getUsed()), TextColor.WHITE))
+                                        .append(TextComponent.of(FormatUtil.formatBytes(usage.getUsed()), TextColor.WHITE))
                                         .append(TextComponent.space())
                                         .append(TextComponent.of("/", TextColor.GRAY))
                                         .append(TextComponent.space())
-                                        .append(TextComponent.of(formatBytes(usage.getMax()), TextColor.WHITE))
+                                        .append(TextComponent.of(FormatUtil.formatBytes(usage.getMax()), TextColor.WHITE))
                                         .append(TextComponent.of("   "))
                                         .append(TextComponent.of("(", TextColor.GRAY))
-                                        .append(TextComponent.of(percent(usage.getUsed(), usage.getMax()), TextColor.GREEN))
+                                        .append(TextComponent.of(FormatUtil.percent(usage.getUsed(), usage.getMax()), TextColor.GREEN))
                                         .append(TextComponent.of(")", TextColor.GRAY))
                                         .build()
                                 );
@@ -211,7 +212,7 @@ public class HealthModule implements CommandModule {
                                             .append(TextComponent.space())
                                             .append(TextComponent.of("Usage at last GC:", TextColor.GRAY))
                                             .append(TextComponent.space())
-                                            .append(TextComponent.of(formatBytes(collectionUsage.getUsed()), TextColor.WHITE))
+                                            .append(TextComponent.of(FormatUtil.formatBytes(collectionUsage.getUsed()), TextColor.WHITE))
                                             .build()
                                     );
                                 }
@@ -230,14 +231,14 @@ public class HealthModule implements CommandModule {
                                     .build()
                             );
                             report.add(TextComponent.builder("    ")
-                                    .append(TextComponent.of(formatBytes(usedSpace), TextColor.WHITE))
+                                    .append(TextComponent.of(FormatUtil.formatBytes(usedSpace), TextColor.WHITE))
                                     .append(TextComponent.space())
                                     .append(TextComponent.of("/", TextColor.GRAY))
                                     .append(TextComponent.space())
-                                    .append(TextComponent.of(formatBytes(totalSpace), TextColor.WHITE))
+                                    .append(TextComponent.of(FormatUtil.formatBytes(totalSpace), TextColor.WHITE))
                                     .append(TextComponent.of("   "))
                                     .append(TextComponent.of("(", TextColor.GRAY))
-                                    .append(TextComponent.of(percent(usedSpace, totalSpace), TextColor.GREEN))
+                                    .append(TextComponent.of(FormatUtil.percent(usedSpace, totalSpace), TextColor.GREEN))
                                     .append(TextComponent.of(")", TextColor.GRAY))
                                     .build()
                             );
@@ -257,11 +258,6 @@ public class HealthModule implements CommandModule {
         );
     }
 
-    private static String percent(double value, double max) {
-        double percent = (value * 100d) / max;
-        return (int) percent + "%";
-    }
-
     private static TextComponent formatTps(double tps) {
         TextColor color;
         if (tps > 18.0) {
@@ -285,7 +281,7 @@ public class HealthModule implements CommandModule {
             color = TextColor.GREEN;
         }
 
-        return TextComponent.of(percent(usage, 1d), color);
+        return TextComponent.of(FormatUtil.percent(usage, 1d), color);
     }
 
     private static TextComponent generateMemoryUsageDiagram(MemoryUsage usage, int length) {
@@ -356,13 +352,4 @@ public class HealthModule implements CommandModule {
                 .build();
     }
 
-    private static String formatBytes(long bytes) {
-        if (bytes == 0) {
-            return "0 bytes";
-        }
-        String[] sizes = new String[]{"bytes", "KB", "MB", "GB", "TB", "PB", "EB", "ZB", "YB"};
-        int sizeIndex = (int) (Math.log(bytes) / Math.log(1024));
-        return String.format("%.1f", bytes / Math.pow(1024, sizeIndex)) + " " + sizes[sizeIndex];
-    }
-
 }
diff --git a/spark-common/src/main/java/me/lucko/spark/common/command/modules/MemoryModule.java b/spark-common/src/main/java/me/lucko/spark/common/command/modules/MemoryModule.java
index f5854f0..a52bf8c 100644
--- a/spark-common/src/main/java/me/lucko/spark/common/command/modules/MemoryModule.java
+++ b/spark-common/src/main/java/me/lucko/spark/common/command/modules/MemoryModule.java
@@ -24,20 +24,34 @@ import me.lucko.spark.common.ActivityLog.Activity;
 import me.lucko.spark.common.SparkPlatform;
 import me.lucko.spark.common.command.Command;
 import me.lucko.spark.common.command.CommandModule;
+import me.lucko.spark.common.command.tabcomplete.CompletionSupplier;
 import me.lucko.spark.common.command.tabcomplete.TabCompleter;
 import me.lucko.spark.common.heapdump.HeapDump;
 import me.lucko.spark.common.heapdump.HeapDumpSummary;
+import me.lucko.spark.common.util.FormatUtil;
 import net.kyori.text.TextComponent;
 import net.kyori.text.event.ClickEvent;
 import net.kyori.text.format.TextColor;
 import okhttp3.MediaType;
+import org.tukaani.xz.LZMA2Options;
+import org.tukaani.xz.LZMAOutputStream;
+import org.tukaani.xz.XZOutputStream;
 
 import java.io.IOException;
+import java.io.InputStream;
+import java.io.OutputStream;
 import java.nio.file.Files;
 import java.nio.file.Path;
 import java.time.LocalDateTime;
 import java.time.format.DateTimeFormatter;
+import java.util.ArrayList;
+import java.util.Arrays;
+import java.util.List;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicLong;
 import java.util.function.Consumer;
+import java.util.function.LongConsumer;
+import java.util.zip.GZIPOutputStream;
 
 public class MemoryModule implements CommandModule {
     private static final MediaType JSON_TYPE = MediaType.parse("application/json; charset=utf-8");
@@ -90,6 +104,9 @@ public class MemoryModule implements CommandModule {
 
         consumer.accept(Command.builder()
                 .aliases("heapdump")
+                .argumentUsage("xz", null)
+                .argumentUsage("lzma", null)
+                .argumentUsage("gzip", null)
                 .argumentUsage("run-gc-before", null)
                 .argumentUsage("include-non-live", null)
                 .executor((platform, sender, resp, arguments) -> {
@@ -124,11 +141,120 @@ public class MemoryModule implements CommandModule {
                                 .build()
                         );
                         platform.getActivityLog().addToLog(Activity.fileActivity(sender, System.currentTimeMillis(), "Heap dump", file.toString()));
+
+                        if (arguments.boolFlag("xz") || arguments.boolFlag("lzma") || arguments.boolFlag("gzip")) {
+                            resp.broadcastPrefixed(TextComponent.of("Compressing heap dump, please wait..."));
+                            try {
+                                long size = Files.size(file);
+                                AtomicLong lastReport = new AtomicLong(System.currentTimeMillis());
+
+                                LongConsumer progressHandler = progress -> {
+                                    long timeSinceLastReport = System.currentTimeMillis() - lastReport.get();
+                                    if (timeSinceLastReport > TimeUnit.SECONDS.toMillis(5)) {
+                                        lastReport.set(System.currentTimeMillis());
+
+                                        platform.getPlugin().runAsync(() -> {
+                                            resp.broadcastPrefixed(TextComponent.builder("").color(TextColor.GRAY)
+                                                    .append(TextComponent.of("Compressed "))
+                                                    .append(TextComponent.of(FormatUtil.formatBytes(progress), TextColor.GOLD))
+                                                    .append(TextComponent.of(" / "))
+                                                    .append(TextComponent.of(FormatUtil.formatBytes(size), TextColor.GOLD))
+                                                    .append(TextComponent.of(" so far... ("))
+                                                    .append(TextComponent.of(FormatUtil.percent(progress, size), TextColor.GREEN))
+                                                    .append(TextComponent.of(")"))
+                                                    .build()
+                                            );
+                                        });
+                                    }
+                                };
+
+                                Path compressedFile;
+                                if (arguments.boolFlag("xz")) {
+                                    compressedFile = file.getParent().resolve(file.getFileName().toString() + ".xz");
+                                    try (InputStream in = Files.newInputStream(file)) {
+                                        try (OutputStream out = Files.newOutputStream(compressedFile)) {
+                                            try (XZOutputStream compressionOut = new XZOutputStream(out, new LZMA2Options())) {
+                                                copy(in, compressionOut, progressHandler);
+                                            }
+                                        }
+                                    }
+                                } else if (arguments.boolFlag("lzma")) {
+                                    compressedFile = file.getParent().resolve(file.getFileName().toString() + ".lzma");
+                                    try (InputStream in = Files.newInputStream(file)) {
+                                        try (OutputStream out = Files.newOutputStream(compressedFile)) {
+                                            try (LZMAOutputStream compressionOut = new LZMAOutputStream(out, new LZMA2Options(), true)) {
+                                                copy(in, compressionOut, progressHandler);
+                                            }
+                                        }
+                                    }
+                                } else {
+                                    compressedFile = file.getParent().resolve(file.getFileName().toString() + ".gz");
+                                    try (InputStream in = Files.newInputStream(file)) {
+                                        try (OutputStream out = Files.newOutputStream(compressedFile)) {
+                                            try (GZIPOutputStream compressionOut = new GZIPOutputStream(out, 1024 * 64)) {
+                                                copy(in, compressionOut, progressHandler);
+                                            }
+                                        }
+                                    }
+                                }
+
+                                long compressedSize = Files.size(compressedFile);
+
+                                resp.broadcastPrefixed(TextComponent.builder("").color(TextColor.GRAY)
+                                        .append(TextComponent.of("Compression complete: "))
+                                        .append(TextComponent.of(FormatUtil.formatBytes(size), TextColor.GOLD))
+                                        .append(TextComponent.of(" --> "))
+                                        .append(TextComponent.of(FormatUtil.formatBytes(compressedSize), TextColor.GOLD))
+                                        .append(TextComponent.of(" ("))
+                                        .append(TextComponent.of(FormatUtil.percent(compressedSize, size), TextColor.GREEN))
+                                        .append(TextComponent.of(")"))
+                                        .build()
+                                );
+
+                                resp.broadcastPrefixed(TextComponent.builder("Compressed heap dump written to: ", TextColor.GOLD)
+                                        .append(TextComponent.of(compressedFile.toString(), TextColor.GRAY))
+                                        .build()
+                                );
+                            } catch (IOException e) {
+                                e.printStackTrace();
+                            }
+                        }
                     });
                 })
-                .tabCompleter((platform, sender, arguments) -> TabCompleter.completeForOpts(arguments, "--run-gc-before", "--include-non-live"))
+                .tabCompleter((platform, sender, arguments) -> {
+                    List<String> opts = new ArrayList<>(Arrays.asList("--run-gc-before", "--include-non-live"));
+                    opts.removeAll(arguments);
+
+                    if (!arguments.contains("--xz") && !arguments.contains("--lzma") && !arguments.contains("--gzip")) {
+                        opts.addAll(Arrays.asList("--xz", "--lzma", "--gzip"));
+                    }
+
+                    return TabCompleter.create()
+                            .from(0, CompletionSupplier.startsWith(opts))
+                            .complete(arguments);
+                })
                 .build()
         );
     }
 
+    public static long copy(InputStream from, OutputStream to, LongConsumer progress) throws IOException {
+        byte[] buf = new byte[1024 * 64];
+        long total = 0;
+        long iterations = 0;
+        while (true) {
+            int r = from.read(buf);
+            if (r == -1) {
+                break;
+            }
+            to.write(buf, 0, r);
+            total += r;
+
+            // report progress every 5MB
+            if (iterations++ % ((1024 / 64) * 5) == 0) {
+                progress.accept(total);
+            }
+        }
+        return total;
+    }
+
 }
diff --git a/spark-common/src/main/java/me/lucko/spark/common/util/FormatUtil.java b/spark-common/src/main/java/me/lucko/spark/common/util/FormatUtil.java
new file mode 100644
index 0000000..37165d1
--- /dev/null
+++ b/spark-common/src/main/java/me/lucko/spark/common/util/FormatUtil.java
@@ -0,0 +1,39 @@
+/*
+ * 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;
+
+public final class FormatUtil {
+    private FormatUtil() {}
+
+    public static String percent(double value, double max) {
+        double percent = (value * 100d) / max;
+        return (int) percent + "%";
+    }
+
+    public static String formatBytes(long bytes) {
+        if (bytes == 0) {
+            return "0 bytes";
+        }
+        String[] sizes = new String[]{"bytes", "KB", "MB", "GB", "TB", "PB", "EB", "ZB", "YB"};
+        int sizeIndex = (int) (Math.log(bytes) / Math.log(1024));
+        return String.format("%.1f", bytes / Math.pow(1024, sizeIndex)) + " " + sizes[sizeIndex];
+    }
+}