Ensure things are cleaned up nicely when the plugin disables

4 files changed, 95 insertions, 75 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 1bc6b57..dd55157 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
@@ -23,6 +23,7 @@ package me.lucko.spark.common;
 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.CommandModule;
 import me.lucko.spark.common.command.CommandResponseHandler;
 import me.lucko.spark.common.command.modules.ActivityLogModule;
 import me.lucko.spark.common.command.modules.HealthModule;
@@ -59,6 +60,7 @@ public class SparkPlatform {
     public static final BytebinClient BYTEBIN_CLIENT = new BytebinClient(OK_HTTP_CLIENT, "https://bytebin.lucko.me/", "spark-plugin");
 
     private final SparkPlugin plugin;
+    private final List<CommandModule> commandModules;
     private final List<Command> commands;
     private final ActivityLog activityLog;
     private final TickCounter tickCounter;
@@ -67,12 +69,18 @@ public class SparkPlatform {
     public SparkPlatform(SparkPlugin plugin) {
         this.plugin = plugin;
 
+        this.commandModules = ImmutableList.of(
+                new SamplerModule(),
+                new HealthModule(),
+                new TickMonitoringModule(),
+                new MemoryModule(),
+                new ActivityLogModule()
+        );
+
         ImmutableList.Builder<Command> commandsBuilder = ImmutableList.builder();
-        new SamplerModule().registerCommands(commandsBuilder::add);
-        new HealthModule().registerCommands(commandsBuilder::add);
-        new TickMonitoringModule().registerCommands(commandsBuilder::add);
-        new MemoryModule().registerCommands(commandsBuilder::add);
-        new ActivityLogModule().registerCommands(commandsBuilder::add);
+        for (CommandModule module : this.commandModules) {
+            module.registerCommands(commandsBuilder::add);
+        }
         this.commands = commandsBuilder.build();
 
         this.activityLog = new ActivityLog(plugin.getPluginFolder().resolve("activity.json"));
@@ -93,6 +101,10 @@ public class SparkPlatform {
         if (this.tickCounter != null) {
             this.tickCounter.close();
         }
+
+        for (CommandModule module : this.commandModules) {
+            module.close();
+        }
     }
 
     public SparkPlugin getPlugin() {
diff --git a/spark-common/src/main/java/me/lucko/spark/common/command/CommandModule.java b/spark-common/src/main/java/me/lucko/spark/common/command/CommandModule.java
index 76c24b3..2b536cc 100644
--- a/spark-common/src/main/java/me/lucko/spark/common/command/CommandModule.java
+++ b/spark-common/src/main/java/me/lucko/spark/common/command/CommandModule.java
@@ -22,8 +22,12 @@ package me.lucko.spark.common.command;
 
 import java.util.function.Consumer;
 
-public interface CommandModule {
+public interface CommandModule extends AutoCloseable {
 
     void registerCommands(Consumer<Command> consumer);
 
+    @Override
+    default void close() {
+
+    }
 }
diff --git a/spark-common/src/main/java/me/lucko/spark/common/command/modules/SamplerModule.java b/spark-common/src/main/java/me/lucko/spark/common/command/modules/SamplerModule.java
index c4f88cb..11d23da 100644
--- a/spark-common/src/main/java/me/lucko/spark/common/command/modules/SamplerModule.java
+++ b/spark-common/src/main/java/me/lucko/spark/common/command/modules/SamplerModule.java
@@ -50,12 +50,18 @@ import java.util.function.Consumer;
 public class SamplerModule implements CommandModule {
     private static final MediaType JSON_TYPE = MediaType.parse("application/json; charset=utf-8");
 
-    /** Guards {@link #activeSampler} */
-    private final Object[] activeSamplerMutex = new Object[0];
     /** The WarmRoast instance currently running, if any */
     private Sampler activeSampler = null;
 
     @Override
+    public void close() {
+        if (this.activeSampler != null) {
+            this.activeSampler.cancel();
+            this.activeSampler = null;
+        }
+    }
+
+    @Override
     public void registerCommands(Consumer<Command> consumer) {
         consumer.accept(Command.builder()
                 .aliases("sampler")
@@ -72,48 +78,41 @@ public class SamplerModule implements CommandModule {
                 .argumentUsage("include-line-numbers", null)
                 .executor((platform, sender, resp, arguments) -> {
                     if (arguments.boolFlag("info")) {
-                        synchronized (this.activeSamplerMutex) {
-                            if (this.activeSampler == null) {
-                                resp.replyPrefixed(TextComponent.of("There isn't an active sampling task running."));
+                        if (this.activeSampler == null) {
+                            resp.replyPrefixed(TextComponent.of("There isn't an active sampling task running."));
+                        } else {
+                            long timeout = this.activeSampler.getEndTime();
+                            if (timeout == -1) {
+                                resp.replyPrefixed(TextComponent.of("There is an active sampler currently running, with no defined timeout."));
                             } else {
-                                long timeout = this.activeSampler.getEndTime();
-                                if (timeout == -1) {
-                                    resp.replyPrefixed(TextComponent.of("There is an active sampler currently running, with no defined timeout."));
-                                } else {
-                                    long timeoutDiff = (timeout - System.currentTimeMillis()) / 1000L;
-                                    resp.replyPrefixed(TextComponent.of("There is an active sampler currently running, due to timeout in " + timeoutDiff + " seconds."));
-                                }
-
-                                long runningTime = (System.currentTimeMillis() - this.activeSampler.getStartTime()) / 1000L;
-                                resp.replyPrefixed(TextComponent.of("It has been sampling for " + runningTime + " seconds so far."));
+                                long timeoutDiff = (timeout - System.currentTimeMillis()) / 1000L;
+                                resp.replyPrefixed(TextComponent.of("There is an active sampler currently running, due to timeout in " + timeoutDiff + " seconds."));
                             }
+
+                            long runningTime = (System.currentTimeMillis() - this.activeSampler.getStartTime()) / 1000L;
+                            resp.replyPrefixed(TextComponent.of("It has been sampling for " + runningTime + " seconds so far."));
                         }
                         return;
                     }
 
                     if (arguments.boolFlag("cancel")) {
-                        synchronized (this.activeSamplerMutex) {
-                            if (this.activeSampler == null) {
-                                resp.replyPrefixed(TextComponent.of("There isn't an active sampling task running."));
-                            } else {
-                                this.activeSampler.cancel();
-                                this.activeSampler = null;
-                                resp.broadcastPrefixed(TextComponent.of("The active sampling task has been cancelled.", TextColor.GOLD));
-                            }
+                        if (this.activeSampler == null) {
+                            resp.replyPrefixed(TextComponent.of("There isn't an active sampling task running."));
+                        } else {
+                            close();
+                            resp.broadcastPrefixed(TextComponent.of("The active sampling task has been cancelled.", TextColor.GOLD));
                         }
                         return;
                     }
 
                     if (arguments.boolFlag("stop") || arguments.boolFlag("upload")) {
-                        synchronized (this.activeSamplerMutex) {
-                            if (this.activeSampler == null) {
-                                resp.replyPrefixed(TextComponent.of("There isn't an active sampling task running."));
-                            } else {
-                                this.activeSampler.cancel();
-                                resp.broadcastPrefixed(TextComponent.of("The active sampling operation has been stopped! Uploading results..."));
-                                handleUpload(platform, resp, this.activeSampler);
-                                this.activeSampler = null;
-                            }
+                        if (this.activeSampler == null) {
+                            resp.replyPrefixed(TextComponent.of("There isn't an active sampling task running."));
+                        } else {
+                            this.activeSampler.cancel();
+                            resp.broadcastPrefixed(TextComponent.of("The active sampling operation has been stopped! Uploading results..."));
+                            handleUpload(platform, resp, this.activeSampler);
+                            this.activeSampler = null;
                         }
                         return;
                     }
@@ -172,37 +171,34 @@ public class SamplerModule implements CommandModule {
                         }
                     }
 
-                    Sampler sampler;
-                    synchronized (this.activeSamplerMutex) {
-                        if (this.activeSampler != null) {
-                            resp.replyPrefixed(TextComponent.of("An active sampler is already running."));
-                            return;
-                        }
+                    if (this.activeSampler != null) {
+                        resp.replyPrefixed(TextComponent.of("An active sampler is already running."));
+                        return;
+                    }
 
-                        resp.broadcastPrefixed(TextComponent.of("Initializing a new profiler, please wait..."));
+                    resp.broadcastPrefixed(TextComponent.of("Initializing a new profiler, please wait..."));
 
-                        SamplerBuilder builder = new SamplerBuilder();
-                        builder.threadDumper(threadDumper);
-                        builder.threadGrouper(threadGrouper);
-                        if (timeoutSeconds != -1) {
-                            builder.completeAfter(timeoutSeconds, TimeUnit.SECONDS);
-                        }
-                        builder.samplingInterval(intervalMillis);
-                        builder.includeLineNumbers(includeLineNumbers);
-                        if (ticksOver != -1) {
-                            builder.ticksOver(ticksOver, tickCounter);
-                        }
-                        sampler = this.activeSampler = builder.start();
+                    SamplerBuilder builder = new SamplerBuilder();
+                    builder.threadDumper(threadDumper);
+                    builder.threadGrouper(threadGrouper);
+                    if (timeoutSeconds != -1) {
+                        builder.completeAfter(timeoutSeconds, TimeUnit.SECONDS);
+                    }
+                    builder.samplingInterval(intervalMillis);
+                    builder.includeLineNumbers(includeLineNumbers);
+                    if (ticksOver != -1) {
+                        builder.ticksOver(ticksOver, tickCounter);
+                    }
+                    Sampler sampler = this.activeSampler = builder.start();
 
-                        resp.broadcastPrefixed(TextComponent.of("Profiler now active!", TextColor.GOLD));
-                        if (timeoutSeconds == -1) {
-                            resp.broadcastPrefixed(TextComponent.of("Use '/" + platform.getPlugin().getLabel() + " sampler --stop' to stop profiling and upload the results."));
-                        } else {
-                            resp.broadcastPrefixed(TextComponent.of("The results will be automatically returned after the profiler has been running for " + timeoutSeconds + " seconds."));
-                        }
+                    resp.broadcastPrefixed(TextComponent.of("Profiler now active!", TextColor.GOLD));
+                    if (timeoutSeconds == -1) {
+                        resp.broadcastPrefixed(TextComponent.of("Use '/" + platform.getPlugin().getLabel() + " sampler --stop' to stop profiling and upload the results."));
+                    } else {
+                        resp.broadcastPrefixed(TextComponent.of("The results will be automatically returned after the profiler has been running for " + timeoutSeconds + " seconds."));
                     }
 
-                    CompletableFuture<Sampler> future = sampler.getFuture();
+                    CompletableFuture<Sampler> future = activeSampler.getFuture();
 
                     // send message if profiling fails
                     future.whenCompleteAsync((s, throwable) -> {
@@ -214,10 +210,8 @@ public class SamplerModule implements CommandModule {
 
                     // set activeSampler to null when complete.
                     future.whenCompleteAsync((s, throwable) -> {
-                        synchronized (this.activeSamplerMutex) {
-                            if (sampler == this.activeSampler) {
-                                this.activeSampler = null;
-                            }
+                        if (sampler == this.activeSampler) {
+                            this.activeSampler = null;
                         }
                     });
 
diff --git a/spark-common/src/main/java/me/lucko/spark/common/command/modules/TickMonitoringModule.java b/spark-common/src/main/java/me/lucko/spark/common/command/modules/TickMonitoringModule.java
index cbcf557..75d5d36 100644
--- a/spark-common/src/main/java/me/lucko/spark/common/command/modules/TickMonitoringModule.java
+++ b/spark-common/src/main/java/me/lucko/spark/common/command/modules/TickMonitoringModule.java
@@ -35,17 +35,29 @@ import java.util.function.Consumer;
 public class TickMonitoringModule implements CommandModule {
 
     /** The tick monitor instance currently running, if any */
+    private TickCounter tickCounter = null;
     private ReportingTickMonitor activeTickMonitor = null;
 
     @Override
+    public void close() {
+        if (this.activeTickMonitor != null) {
+            this.tickCounter.removeTickTask(this.activeTickMonitor);
+            this.activeTickMonitor.close();
+            this.activeTickMonitor = null;
+        }
+    }
+
+    @Override
     public void registerCommands(Consumer<Command> consumer) {
         consumer.accept(Command.builder()
                 .aliases("tickmonitoring")
                 .argumentUsage("threshold", "percentage increase")
                 .argumentUsage("without-gc", null)
                 .executor((platform, sender, resp, arguments) -> {
-                    TickCounter tickCounter = platform.getTickCounter();
-                    if (tickCounter == null) {
+                    if (this.tickCounter == null) {
+                        this.tickCounter = platform.getTickCounter();
+                    }
+                    if (this.tickCounter == null) {
                         resp.replyPrefixed(TextComponent.of("Not supported!", TextColor.RED));
                         return;
                     }
@@ -56,12 +68,10 @@ public class TickMonitoringModule implements CommandModule {
                             threshold = 100;
                         }
 
-                        this.activeTickMonitor = new ReportingTickMonitor(resp, tickCounter, threshold, !arguments.boolFlag("without-gc"));
-                        tickCounter.addTickTask(this.activeTickMonitor);
+                        this.activeTickMonitor = new ReportingTickMonitor(resp, this.tickCounter, threshold, !arguments.boolFlag("without-gc"));
+                        this.tickCounter.addTickTask(this.activeTickMonitor);
                     } else {
-                        tickCounter.removeTickTask(this.activeTickMonitor);
-                        this.activeTickMonitor.close();
-                        this.activeTickMonitor = null;
+                        close();
                         resp.broadcastPrefixed(TextComponent.of("Tick monitor disabled."));
                     }
                 })