Clean up when module deactivates or crashes

api/src/main/java/fr/atlasworld/network/api/AtlasNetworkServer.java

@@ -4,6 +4,7 @@
 import fr.atlasworld.network.api.command.CommandManager;
 import fr.atlasworld.network.api.concurrent.action.FutureAction;
 import fr.atlasworld.network.api.file.FileManager;
+import fr.atlasworld.network.api.module.Module;
 import fr.atlasworld.network.api.networking.Socket;
 import fr.atlasworld.network.api.security.SecurityManager;
 import fr.atlasworld.network.api.services.ServiceManager;
@@ -60,15 +61,15 @@ public interface AtlasNetworkServer {
      * @return future action to notify your code when your task finishes, of fails.
      */
     @CanIgnoreReturnValue
-    <V> FutureAction<V> createAsyncSupplier(Supplier<V> supplier);
+    <V> FutureAction<V> createAsyncSupplier(Module module, Supplier<V> supplier);
 
     /**
      * Creates an asynchronous task using the system's Executor Pool.
      * @param runnable task to run.
      * @return future action to notify your code when your task finishes, or fails.
      */
     @CanIgnoreReturnValue
-    FutureAction<Void> createAsyncTask(Runnable runnable);
+    FutureAction<Void> createAsyncTask(Module module, Runnable runnable);
 
     /**
      * Creates an asynchronous task that will run every specified time interval.
@@ -78,5 +79,5 @@ public interface AtlasNetworkServer {
      * @return future action, this action will never successfully complete, it can only fail or be cancelled.
      */
     @CanIgnoreReturnValue
-    FutureAction<Void> createSequenceTask(Runnable runnable, long interval, TimeUnit intervalUnit);
+    FutureAction<Void> createSequenceTask(Module module, Runnable runnable, long interval, TimeUnit intervalUnit);
 }

api/src/main/java/fr/atlasworld/network/api/NetworkModule.java

@@ -1,5 +1,7 @@
 package fr.atlasworld.network.api;
 
+import fr.atlasworld.network.api.concurrent.action.FutureAction;
+import fr.atlasworld.network.api.concurrent.action.SimpleFutureAction;
 import fr.atlasworld.network.api.loader.ModuleClassLoader;
 import fr.atlasworld.network.api.module.Module;
 import fr.atlasworld.network.api.module.ModuleMeta;
@@ -30,6 +32,7 @@ public abstract class NetworkModule implements Module {
     private ModuleStatus status = ModuleStatus.INACTIVE;
     private Logger logger = null;
     private boolean loaded = false;
+    private SimpleFutureAction<Void> deactivationFuture;
 
     @ApiStatus.Internal
     public NetworkModule() {
@@ -106,6 +109,7 @@ public final void activate(@NotNull ModuleActivationContext ctx) {
         if (this.status == ModuleStatus.ACTIVE)
             throw new IllegalStateException("Module is already activated!");
 
+        this.deactivationFuture = new SimpleFutureAction<>();
         this.status = ModuleStatus.ACTIVE;
         this.onActivate(ctx);
     }
@@ -118,19 +122,26 @@ public final void deactivate(@NotNull ModuleDeactivationContext ctx) {
 
         this.status = ModuleStatus.INACTIVE;
         this.onDeactivate(ctx);
+        this.deactivationFuture.complete(null); // Trigger success listeners.
     }
 
     @Override
     @ApiStatus.Internal
     public final void crash(Throwable throwable) {
         this.status = ModuleStatus.CRASHED;
         this.logger.error("Oh no! I just crashed!", throwable);
+        this.deactivationFuture.complete(null); // Trigger success listeners.
     }
 
     public final Logger getLogger() {
         return this.logger;
     }
 
+    @Override
+    public final FutureAction<Void> deactivationFuture() {
+        return this.deactivationFuture;
+    }
+
     @Override
     public final @Nullable InputStream getResource(@NotNull String resource) {
         if (this.archive == null)

api/src/main/java/fr/atlasworld/network/api/concurrent/action/FutureAction.java

@@ -5,6 +5,7 @@
 
 import java.util.concurrent.Future;
 import java.util.concurrent.TimeUnit;
+import java.util.function.BiConsumer;
 import java.util.function.Consumer;
 
 /**
@@ -23,6 +24,11 @@ public interface FutureAction<V> extends Future<V> {
      */
     FutureAction<V> onFailure(@NotNull Consumer<Throwable> listener);
 
+    /**
+     * Execute code when the action is done, whether is succeeded or failed.
+     */
+    FutureAction<V> whenDone(BiConsumer<V, Throwable> listener);
+
     /**
      * Check if action was successful.
      * @return true if the action is done and was successful.

api/src/main/java/fr/atlasworld/network/api/concurrent/action/SimpleFutureAction.java

@@ -1,21 +1,26 @@
 package fr.atlasworld.network.api.concurrent.action;
 
 import com.google.common.util.concurrent.UncheckedExecutionException;
+import com.google.errorprone.annotations.CanIgnoreReturnValue;
 import org.jetbrains.annotations.NotNull;
 import org.jetbrains.annotations.Nullable;
 
+import java.util.ArrayList;
+import java.util.List;
 import java.util.concurrent.CountDownLatch;
 import java.util.concurrent.ExecutionException;
 import java.util.concurrent.TimeUnit;
 import java.util.concurrent.TimeoutException;
+import java.util.function.BiConsumer;
 import java.util.function.Consumer;
 import java.util.function.Function;
 
 public class SimpleFutureAction<V> implements FutureAction<V> {
     private final CountDownLatch latch = new CountDownLatch(1);
 
-    private Consumer<V> onSuccess = v -> {};
-    private Consumer<Throwable> onFailure = cause -> {};
+    private final List<Consumer<V>> onSuccess;
+    private final List<Consumer<Throwable>> onFailure;
+    private final List<BiConsumer<V, Throwable>> whenDone;
     private Function<Boolean, Boolean> onCancellation = interrupt -> false;
 
     private boolean done = false;
@@ -25,15 +30,29 @@ public class SimpleFutureAction<V> implements FutureAction<V> {
     private V value;
     private Throwable cause;
 
+    public SimpleFutureAction() {
+        this.onSuccess = new ArrayList<>();
+        this.onFailure = new ArrayList<>();
+        this.whenDone = new ArrayList<>();
+    }
+
     @Override
+    @CanIgnoreReturnValue
     public SimpleFutureAction<V> onSuccess(@NotNull Consumer<V> listener) {
-        this.onSuccess = listener;
+        this.onSuccess.add(listener);
         return this;
     }
 
     @Override
+    @CanIgnoreReturnValue
     public SimpleFutureAction<V> onFailure(@NotNull Consumer<Throwable> listener) {
-        this.onFailure = listener;
+        this.onFailure.add(listener);
+        return this;
+    }
+
+    @Override
+    public FutureAction<V> whenDone(BiConsumer<V, Throwable> listener) {
+        this.whenDone.add(listener);
         return this;
     }
 
@@ -68,6 +87,7 @@ public V syncUninterruptibly() {
         }
     }
 
+    @CanIgnoreReturnValue
     public SimpleFutureAction<V> onCancellation(Function<Boolean, Boolean> listener) {
         this.onCancellation = listener;
         return this;
@@ -112,7 +132,14 @@ public void complete(V value) {
         this.success = true;
 
         this.latch.countDown();
-        this.onSuccess.accept(this.value);
+
+        for (Consumer<V> listener : this.onSuccess) {
+            listener.accept(this.value);
+        }
+
+        for (BiConsumer<V, Throwable> listener : this.whenDone) {
+            listener.accept(this.value, this.cause);
+        }
     }
 
     public void fail(Throwable cause) {
@@ -124,6 +151,13 @@ public void fail(Throwable cause) {
         this.done = true;
 
         this.latch.countDown();
-        this.onFailure.accept(this.cause);
+
+        for (Consumer<Throwable> listener : this.onFailure) {
+            listener.accept(this.cause);
+        }
+
+        for (BiConsumer<V, Throwable> listener : this.whenDone) {
+            listener.accept(this.value, this.cause);
+        }
     }
 }

api/src/main/java/fr/atlasworld/network/api/module/Module.java

@@ -1,11 +1,13 @@
 package fr.atlasworld.network.api.module;
 
+import fr.atlasworld.network.api.concurrent.action.FutureAction;
 import fr.atlasworld.network.api.module.lifecycle.ModuleActivationContext;
 import fr.atlasworld.network.api.module.lifecycle.ModuleDeactivationContext;
 import fr.atlasworld.network.api.module.lifecycle.ModuleLoadContext;
 import fr.atlasworld.network.api.util.archive.Archive;
 import org.jetbrains.annotations.ApiStatus;
 import org.jetbrains.annotations.NotNull;
+import org.jetbrains.annotations.Nullable;
 
 import java.io.File;
 
@@ -60,4 +62,13 @@ public interface Module extends Archive {
      */
     @ApiStatus.Internal
     void crash(Throwable throwable);
+
+    /**
+     * Retrieve the deactivation future of the module.
+     * Success listeners are called when the module deactivates normally.
+     * Fail listeners are called when the module crashes.
+     * @return null if the module is not yet activated.
+     */
+    @Nullable
+    FutureAction<Void> deactivationFuture();
 }

api/src/main/java/fr/atlasworld/network/api/services/ServiceManager.java

@@ -1,5 +1,6 @@
 package fr.atlasworld.network.api.services;
 
+import fr.atlasworld.network.api.module.Module;
 import org.jetbrains.annotations.Nullable;
 
 /**
@@ -9,10 +10,11 @@ public interface ServiceManager {
 
     /**
      * Register a service implementation to the service manager.
+     * @param module module of the service class.
      * @param serviceClass parent service class.
      * @param serviceImpl implementation of the service class.
      */
-    <T extends Service> void registerService(Class<T> serviceClass, T serviceImpl);
+    <T extends Service> void registerService(Module module, Class<T> serviceClass, T serviceImpl);
 
     /**
      * Get a service from the service manager.

api/src/test/java/fr/atlasworld/network/api/test/TestModule.java

@@ -38,7 +38,7 @@ protected void onActivate(@NotNull ModuleActivationContext ctx) {
         LOGGER.info("Registering services..");
         ctx.getServer().getServiceManager().registerService(DatabaseService.class, new LocalDatabaseService());
 
-        ctx.getServer().createSequenceTask(() -> {
+        ctx.getServer().createSequenceTask(this, () -> {
             LOGGER.info("Hi!");
         }, 10, TimeUnit.SECONDS);
 

src/main/java/fr/atlasworld/network/core/AtlasNetwork.java

@@ -5,6 +5,7 @@
 import fr.atlasworld.network.api.command.CommandManager;
 import fr.atlasworld.network.api.concurrent.action.FutureAction;
 import fr.atlasworld.network.api.file.FileManager;
+import fr.atlasworld.network.api.module.Module;
 import fr.atlasworld.network.api.networking.Socket;
 import fr.atlasworld.network.api.security.SecurityManager;
 import fr.atlasworld.network.api.services.ServiceManager;
@@ -170,18 +171,24 @@ public Socket getSocket() {
     }
 
     @Override
-    public <V> FutureAction<V> createAsyncSupplier(Supplier<V> supplier) {
-        return ActionUtilities.supply(supplier);
+    public <V> FutureAction<V> createAsyncSupplier(Module module, Supplier<V> supplier) {
+        SimpleFutureAction<V> action = ActionUtilities.supply(supplier);
+        module.deactivationFuture().whenDone((v, cause) -> action.cancel(cause != null)); // Interrupt if the module crashed.
+        return action;
     }
 
     @Override
-    public FutureAction<Void> createAsyncTask(Runnable runnable) {
-        return ActionUtilities.execute(runnable);
+    public FutureAction<Void> createAsyncTask(Module module, Runnable runnable) {
+        SimpleFutureAction<Void> action = ActionUtilities.execute(runnable);
+        module.deactivationFuture().whenDone((v, cause) -> action.cancel(cause != null)); // Interrupt if the module crashed.
+        return action;
     }
 
     @Override
-    public FutureAction<Void> createSequenceTask(Runnable runnable, long interval, TimeUnit intervalUnit) {
-        return ActionUtilities.sequential(runnable, interval, intervalUnit);
+    public FutureAction<Void> createSequenceTask(Module module, Runnable runnable, long interval, TimeUnit intervalUnit) {
+        SimpleFutureAction<Void> action = ActionUtilities.sequential(runnable, interval, intervalUnit);
+        module.deactivationFuture().whenDone((v, cause) -> action.cancel(cause != null)); // Interrupt if the module crashed.
+        return action;
     }
 
     // Static fields

src/main/java/fr/atlasworld/network/core/concurrent/action/ActionUtilities.java

@@ -9,7 +9,7 @@
 import java.util.function.Supplier;
 
 public class ActionUtilities {
-    public static <T> FutureAction<T> supply(Supplier<T> supplier) {
+    public static <T> SimpleFutureAction<T> supply(Supplier<T> supplier) {
         SimpleFutureAction<T> action = new SimpleFutureAction<>();
 
         Environment.SYSTEM_EXECUTOR.execute(() -> {
@@ -23,7 +23,7 @@ public static <T> FutureAction<T> supply(Supplier<T> supplier) {
         return action;
     }
 
-    public static FutureAction<Void> execute(Runnable runnable) {
+    public static SimpleFutureAction<Void> execute(Runnable runnable) {
         SimpleFutureAction<Void> action = new SimpleFutureAction<>();
 
         Environment.SYSTEM_EXECUTOR.execute(() -> {
@@ -38,7 +38,7 @@ public static FutureAction<Void> execute(Runnable runnable) {
         return action;
     }
 
-    public static FutureAction<Void> sequential(Runnable runnable, long timeInterval, TimeUnit unit) {
+    public static SimpleFutureAction<Void> sequential(Runnable runnable, long timeInterval, TimeUnit unit) {
         SimpleFutureAction<Void> action = new SimpleFutureAction<>();
 
         ScheduledFuture<?> future = Environment.SYSTEM_EXECUTOR.scheduleAtFixedRate(() -> {

src/main/java/fr/atlasworld/network/core/console/SystemConsole.java

@@ -98,6 +98,8 @@ public void register(Module module, Consumer<CommandStore<CommandSource>> storeB
 
         this.commands.putIfAbsent(module, new ArrayList<>());
         this.commands.get(module).addAll(store.getRegisteredCommands());
+
+        module.deactivationFuture().whenDone((v, cause) -> this.unregister(module)); // Unregister commands when the module deactivates.
     }
 
     @Override

src/main/java/fr/atlasworld/network/core/services/SystemServiceManager.java

@@ -1,5 +1,6 @@
 package fr.atlasworld.network.core.services;
 
+import fr.atlasworld.network.api.module.Module;
 import fr.atlasworld.network.api.services.Service;
 import fr.atlasworld.network.api.services.ServiceManager;
 import org.jetbrains.annotations.Nullable;
@@ -20,8 +21,9 @@ public SystemServiceManager() {
     }
 
     @Override
-    public <T extends Service> void registerService(Class<T> serviceClass, T serviceImpl) {
+    public <T extends Service> void registerService(Module module, Class<T> serviceClass, T serviceImpl) {
         this.registeredServices.put(serviceClass, serviceImpl);
+        module.deactivationFuture().whenDone((v, cause) -> this.registeredServices.remove(serviceClass)); // Unregister services when the module deactivates.
     }
 
     @Override