Fix tasks not being canceled properly + annotations

src/main/java/net/minestom/server/timer/SchedulerManager.java

@@ -5,6 +5,7 @@ import it.unimi.dsi.fastutil.ints.Int2ObjectOpenHashMap;
 import it.unimi.dsi.fastutil.objects.ObjectCollection;
 import net.minestom.server.MinecraftServer;
 import net.minestom.server.utils.thread.MinestomThread;
+import org.jetbrains.annotations.NotNull;
 
 import java.util.Collection;
 import java.util.concurrent.ExecutorService;
@@ -64,7 +65,8 @@ public final class SchedulerManager {
      * @param runnable The {@link Task} to run when scheduled
      * @return the {@link TaskBuilder}
      */
-    public TaskBuilder buildTask(Runnable runnable) {
+    @NotNull
+    public TaskBuilder buildTask(@NotNull Runnable runnable) {
         return new TaskBuilder(this, runnable);
     }
 
@@ -74,7 +76,8 @@ public final class SchedulerManager {
      * @param runnable The shutdown {@link Task} to run when scheduled
      * @return the {@link TaskBuilder}
      */
-    public TaskBuilder buildShutdownTask(Runnable runnable) {
+    @NotNull
+    public TaskBuilder buildShutdownTask(@NotNull Runnable runnable) {
         return new TaskBuilder(this, runnable, true);
     }
 
@@ -83,19 +86,8 @@ public final class SchedulerManager {
      *
      * @param task The {@link Task} to remove
      */
-    public void removeTask(Task task) {
+    public void removeTask(@NotNull Task task) {
-        synchronized (tasks) {
+        task.cancel();
-            this.tasks.remove(task.getId());
-        }
-    }
-
-    /**
-     * Removes/Forces the end of a {@link Task}.
-     *
-     * @param task The {@link Task} to remove
-     */
-    public void removeShutdownTask(Task task) {
-        this.shutdownTasks.remove(task.getId());
     }
 
     /**
@@ -141,6 +133,7 @@ public final class SchedulerManager {
      *
      * @return a {@link Collection} with all the registered {@link Task}
      */
+    @NotNull
     public ObjectCollection<Task> getTasks() {
         return tasks.values();
     }
@@ -150,6 +143,7 @@ public final class SchedulerManager {
      *
      * @return a {@link Collection} with all the registered shutdown {@link Task}
      */
+    @NotNull
     public ObjectCollection<Task> getShutdownTasks() {
         return shutdownTasks.values();
     }
@@ -159,6 +153,7 @@ public final class SchedulerManager {
      *
      * @return the execution service for all the registered {@link Task}
      */
+    @NotNull
     public ExecutorService getBatchesPool() {
         return batchesPool;
     }
@@ -168,6 +163,7 @@ public final class SchedulerManager {
      *
      * @return the scheduled execution service for all the registered {@link Task}
      */
+    @NotNull
     public ScheduledExecutorService getTimerExecutionService() {
         return timerExecutionService;
     }

src/main/java/net/minestom/server/timer/Task.java

@@ -1,5 +1,8 @@
 package net.minestom.server.timer;
 
+import it.unimi.dsi.fastutil.ints.Int2ObjectMap;
+import org.jetbrains.annotations.NotNull;
+
 import java.util.Objects;
 import java.util.concurrent.ScheduledFuture;
 import java.util.concurrent.TimeUnit;
@@ -37,7 +40,7 @@ public class Task implements Runnable {
      * @param delay            The time to delay
      * @param repeat           The time until the repetition
      */
-    public Task(SchedulerManager schedulerManager, Runnable runnable, boolean shutdown, long delay, long repeat) {
+    public Task(@NotNull SchedulerManager schedulerManager, @NotNull Runnable runnable, boolean shutdown, long delay, long repeat) {
         this.schedulerManager = schedulerManager;
         this.runnable = runnable;
         this.shutdown = shutdown;
@@ -84,6 +87,7 @@ public class Task implements Runnable {
      *
      * @return the current stats of the task
      */
+    @NotNull
     public TaskStatus getStatus() {
         if (this.future == null) return TaskStatus.SCHEDULED;
         if (this.future.isCancelled()) return TaskStatus.CANCELLED;
@@ -119,10 +123,13 @@ public class Task implements Runnable {
      * Removes the task.
      */
     private void finish() {
-        if (this.shutdown)
+        Int2ObjectMap<Task> taskMap = shutdown ?
-            this.schedulerManager.removeShutdownTask(this);
+                this.schedulerManager.shutdownTasks :
-        else
+                this.schedulerManager.tasks;
-            this.schedulerManager.removeTask(this);
+
+        synchronized (taskMap) {
+            taskMap.remove(getId());
+        }
     }
 
     @Override

src/main/java/net/minestom/server/timer/TaskBuilder.java

@@ -2,6 +2,7 @@ package net.minestom.server.timer;
 
 import it.unimi.dsi.fastutil.ints.Int2ObjectMap;
 import net.minestom.server.utils.time.TimeUnit;
+import org.jetbrains.annotations.NotNull;
 
 /**
  * A builder which represents a fluent Object to schedule tasks.
@@ -30,7 +31,7 @@ public class TaskBuilder {
      * @param schedulerManager The manager for the tasks
      * @param runnable         The task to run when scheduled
      */
-    public TaskBuilder(SchedulerManager schedulerManager, Runnable runnable) {
+    public TaskBuilder(@NotNull SchedulerManager schedulerManager, @NotNull Runnable runnable) {
         this(schedulerManager, runnable, false);
     }
 
@@ -41,7 +42,7 @@ public class TaskBuilder {
      * @param runnable         The task to run when scheduled
      * @param shutdown         Defines whether the task is a shutdown task
      */
-    public TaskBuilder(SchedulerManager schedulerManager, Runnable runnable, boolean shutdown) {
+    public TaskBuilder(@NotNull SchedulerManager schedulerManager, @NotNull Runnable runnable, boolean shutdown) {
         this.schedulerManager = schedulerManager;
         this.runnable = runnable;
         this.shutdown = shutdown;
@@ -54,7 +55,8 @@ public class TaskBuilder {
      * @param unit The unit of time for {@code time}
      * @return this builder, for chaining
      */
-    public TaskBuilder delay(long time, TimeUnit unit) {
+    @NotNull
+    public TaskBuilder delay(long time, @NotNull TimeUnit unit) {
         this.delay = unit.toMilliseconds(time);
         return this;
     }
@@ -66,7 +68,8 @@ public class TaskBuilder {
      * @param unit The {@link TimeUnit} for {@code time}
      * @return this builder, for chaining
      */
-    public TaskBuilder repeat(long time, TimeUnit unit) {
+    @NotNull
+    public TaskBuilder repeat(long time, @NotNull TimeUnit unit) {
         this.repeat = unit.toMilliseconds(time);
         return this;
     }
@@ -76,6 +79,7 @@ public class TaskBuilder {
      *
      * @return this builder, for chaining
      */
+    @NotNull
     public TaskBuilder clearDelay() {
         this.delay = 0L;
         return this;
@@ -86,16 +90,18 @@ public class TaskBuilder {
      *
      * @return this builder, for chaining
      */
+    @NotNull
     public TaskBuilder clearRepeat() {
         this.repeat = 0L;
         return this;
     }
 
     /**
-     * Schedule this {@link Task} for execution.
+     * Schedules this {@link Task} for execution.
      *
      * @return the built {@link Task}
      */
+    @NotNull
     public Task schedule() {
         Task task = new Task(
                 this.schedulerManager,