Ensure scheduler is closed at server stop (#75, #78)

Co-authored-by: Luck <git@lucko.me>

3 files changed, 38 insertions, 5 deletions

diff --git a/spark-fabric/src/main/java/me/lucko/spark/fabric/plugin/FabricClientSparkPlugin.java b/spark-fabric/src/main/java/me/lucko/spark/fabric/plugin/FabricClientSparkPlugin.java
index 21e92dc..5b63904 100644
--- a/spark-fabric/src/main/java/me/lucko/spark/fabric/plugin/FabricClientSparkPlugin.java
+++ b/spark-fabric/src/main/java/me/lucko/spark/fabric/plugin/FabricClientSparkPlugin.java
@@ -36,6 +36,7 @@ import me.lucko.spark.fabric.FabricSparkGameHooks;
 import me.lucko.spark.fabric.FabricSparkMod;
 import me.lucko.spark.fabric.FabricTickHook;
 import me.lucko.spark.fabric.FabricTickReporter;
+import net.fabricmc.fabric.api.client.event.lifecycle.v1.ClientLifecycleEvents;
 import net.minecraft.client.MinecraftClient;
 import net.minecraft.client.network.ClientPlayNetworkHandler;
 import net.minecraft.command.CommandSource;
@@ -51,7 +52,17 @@ public class FabricClientSparkPlugin extends FabricSparkPlugin implements Sugges
 
     public static void register(FabricSparkMod mod, MinecraftClient client) {
         FabricClientSparkPlugin plugin = new FabricClientSparkPlugin(mod, client);
+        plugin.enable();
+
+        // ensure commands are registered
         plugin.scheduler.scheduleWithFixedDelay(plugin::checkCommandRegistered, 10, 10, TimeUnit.SECONDS);
+
+        // register shutdown hook
+        ClientLifecycleEvents.CLIENT_STOPPING.register(stoppingClient -> {
+            if (stoppingClient == plugin.minecraft) {
+                plugin.disable();
+            }
+        });
     }
 
     private final MinecraftClient minecraft;
diff --git a/spark-fabric/src/main/java/me/lucko/spark/fabric/plugin/FabricServerSparkPlugin.java b/spark-fabric/src/main/java/me/lucko/spark/fabric/plugin/FabricServerSparkPlugin.java
index 9375668..3293853 100644
--- a/spark-fabric/src/main/java/me/lucko/spark/fabric/plugin/FabricServerSparkPlugin.java
+++ b/spark-fabric/src/main/java/me/lucko/spark/fabric/plugin/FabricServerSparkPlugin.java
@@ -35,6 +35,7 @@ import me.lucko.spark.fabric.FabricSparkMod;
 import me.lucko.spark.fabric.FabricTickHook;
 import me.lucko.spark.fabric.FabricTickReporter;
 import net.fabricmc.fabric.api.command.v1.CommandRegistrationCallback;
+import net.fabricmc.fabric.api.event.lifecycle.v1.ServerLifecycleEvents;
 import net.minecraft.entity.player.PlayerEntity;
 import net.minecraft.server.MinecraftServer;
 import net.minecraft.server.command.CommandOutput;
@@ -49,8 +50,18 @@ public class FabricServerSparkPlugin extends FabricSparkPlugin implements Comman
 
     public static void register(FabricSparkMod mod, MinecraftServer server) {
         FabricServerSparkPlugin plugin = new FabricServerSparkPlugin(mod, server);
+        plugin.enable();
+
+        // register commands
         registerCommands(server.getCommandManager().getDispatcher(), plugin, plugin, "spark");
         CommandRegistrationCallback.EVENT.register((dispatcher, isDedicated) -> registerCommands(dispatcher, plugin, plugin, "spark"));
+
+        // register shutdown hook
+        ServerLifecycleEvents.SERVER_STOPPING.register(stoppingServer -> {
+            if (stoppingServer == plugin.server) {
+                plugin.disable();
+            }
+        });
     }
 
     private final MinecraftServer server;
diff --git a/spark-fabric/src/main/java/me/lucko/spark/fabric/plugin/FabricSparkPlugin.java b/spark-fabric/src/main/java/me/lucko/spark/fabric/plugin/FabricSparkPlugin.java
index 3cc125e..70078c7 100644
--- a/spark-fabric/src/main/java/me/lucko/spark/fabric/plugin/FabricSparkPlugin.java
+++ b/spark-fabric/src/main/java/me/lucko/spark/fabric/plugin/FabricSparkPlugin.java
@@ -20,7 +20,6 @@
 
 package me.lucko.spark.fabric.plugin;
 
-import com.google.common.util.concurrent.ThreadFactoryBuilder;
 import com.mojang.brigadier.Command;
 import com.mojang.brigadier.CommandDispatcher;
 import com.mojang.brigadier.arguments.StringArgumentType;
@@ -46,13 +45,25 @@ public abstract class FabricSparkPlugin implements SparkPlugin {
 
     protected FabricSparkPlugin(FabricSparkMod mod) {
         this.mod = mod;
-        this.scheduler = Executors.newSingleThreadScheduledExecutor(
-                new ThreadFactoryBuilder().setNameFormat("spark-fabric-async-worker").build()
-        );
+        this.scheduler = Executors.newSingleThreadScheduledExecutor(r -> {
+            Thread thread = Executors.defaultThreadFactory().newThread(r);
+            thread.setName("spark-fabric-async-worker");
+            thread.setDaemon(true);
+            return thread;
+        });
         this.platform = new SparkPlatform(this);
         this.platform.enable();
     }
 
+    public void enable() {
+        this.platform.enable();
+    }
+
+    public void disable() {
+        this.platform.disable();
+        this.scheduler.shutdown();
+    }
+
     public abstract boolean hasPermission(CommandOutput sender, String permission);
 
     @Override
@@ -75,7 +86,7 @@ public abstract class FabricSparkPlugin implements SparkPlugin {
         return new ThreadDumper.Specific(new long[]{Thread.currentThread().getId()});
     }
 
-    public static <T> void registerCommands(CommandDispatcher<T> dispatcher, Command<T> executor, SuggestionProvider<T> suggestor, String... aliases) {
+    protected static <T> void registerCommands(CommandDispatcher<T> dispatcher, Command<T> executor, SuggestionProvider<T> suggestor, String... aliases) {
         if (aliases.length == 0) {
             return;
         }