Improved Scheduler API (#696)

* Improved Scheduler API

- Added `Scheduler#builder(plugin)`
This method allows a more simplified builder while maintaining the main requirement of the executor plugin
- Added `Scheduler#taskByPlugin(plugin)`
Allows to obtain the tasks that a plugin has sent to execute and that are currently active
- Added `TaskBuilder#task(Consumer<SchuledTask>)`
Allows to specify a task with access to the task itself with the ability to cancel itself

* Applied requested changes

- Removed tasks builder method
- Added `Scheduler#buildTask(plugin, Consumer<ScheduledTask>)`

* Removed some unused imports

* Applied suggested change

* Fix possible test bug

* Applied more suggested changes

* Fixed tests inside tasks

api/src/main/java/com/velocitypowered/api/command/Command.java

@@ -7,11 +7,7 @@
 
 package com.velocitypowered.api.command;
 
-import com.google.common.collect.ImmutableList;
 import com.velocitypowered.api.proxy.Player;
-import java.util.List;
-import java.util.concurrent.CompletableFuture;
-import org.checkerframework.checker.nullness.qual.NonNull;
 
 /**
  * Represents a command that can be executed by a {@link CommandSource}

api/src/main/java/com/velocitypowered/api/permission/PermissionSubject.java

@@ -8,7 +8,6 @@
 package com.velocitypowered.api.permission;
 
 import net.kyori.adventure.permission.PermissionChecker;
-import net.kyori.adventure.util.TriState;
 
 /**
  * Represents a object that has a set of queryable permissions.

api/src/main/java/com/velocitypowered/api/proxy/ProxyServer.java

@@ -23,7 +23,6 @@ import java.util.Collection;
 import java.util.Optional;
 import java.util.UUID;
 import net.kyori.adventure.audience.Audience;
-import org.checkerframework.checker.nullness.qual.NonNull;
 
 /**
  * Provides an interface to a Minecraft server proxy.

api/src/main/java/com/velocitypowered/api/scheduler/ScheduledTask.java

@@ -7,6 +7,8 @@
 
 package com.velocitypowered.api.scheduler;
 
+import org.jetbrains.annotations.NotNull;
+
 /**
  * Represents a task that is scheduled to run on the proxy.
  */
@@ -17,7 +19,7 @@ public interface ScheduledTask {
    *
    * @return the plugin that scheduled this task
    */
-  Object plugin();
+  @NotNull Object plugin();
 
   /**
    * Returns the current status of this task.

api/src/main/java/com/velocitypowered/api/scheduler/Scheduler.java

@@ -8,8 +8,12 @@
 package com.velocitypowered.api.scheduler;
 
 import java.time.Duration;
+import java.util.Collection;
 import java.util.concurrent.TimeUnit;
+import java.util.function.Consumer;
+
 import org.checkerframework.common.value.qual.IntRange;
+import org.jetbrains.annotations.NotNull;
 
 /**
  * Represents a scheduler to execute tasks on the proxy.
@@ -23,7 +27,24 @@ public interface Scheduler {
    * @param runnable the task to run when scheduled
    * @return the task builder
    */
-  TaskBuilder buildTask(Object plugin, Runnable runnable);
+  TaskBuilder buildTask(@NotNull Object plugin, @NotNull Runnable runnable);
+
+  /**
+   * Initializes a new {@link TaskBuilder} for creating a task on the proxy.
+   *
+   * @param plugin the plugin to request the task for
+   * @param consumer the task to be run when scheduled with the capacity to cancel itself
+   * @return the task builder
+   */
+  TaskBuilder buildTask(@NotNull Object plugin, @NotNull Consumer<ScheduledTask> consumer);
+
+  /**
+   * Get the {@link ScheduledTask} for a specific plugin.
+   *
+   * @param plugin the plugin object
+   * @return the list of {@link ScheduledTask} corresponding to a specific plugin
+   */
+  @NotNull Collection<ScheduledTask> tasksByPlugin(@NotNull Object plugin);
 
   /**
    * Represents a fluent interface to schedule tasks on the proxy.
@@ -37,7 +58,7 @@ public interface Scheduler {
      * @param unit the unit of time for {@code time}
      * @return this builder, for chaining
      */
-    TaskBuilder delay(@IntRange(from = 0) long time, TimeUnit unit);
+    TaskBuilder delay(@IntRange(from = 0) long time, @NotNull TimeUnit unit);
 
     /**
      * Specifies that the task should delay its execution by the specified amount of time.
@@ -45,7 +66,7 @@ public interface Scheduler {
      * @param duration the duration of the delay
      * @return this builder, for chaining
      */
-    default TaskBuilder delay(Duration duration) {
+    default TaskBuilder delay(@NotNull Duration duration) {
       return delay(duration.toMillis(), TimeUnit.MILLISECONDS);
     }
 
@@ -57,7 +78,7 @@ public interface Scheduler {
      * @param unit the unit of time for {@code time}
      * @return this builder, for chaining
      */
-    TaskBuilder repeat(@IntRange(from = 0) long time, TimeUnit unit);
+    TaskBuilder repeat(@IntRange(from = 0) long time, @NotNull TimeUnit unit);
 
     /**
      * Specifies that the task should continue running after waiting for the specified amount, until
@@ -66,7 +87,7 @@ public interface Scheduler {
      * @param duration the duration of the delay
      * @return this builder, for chaining
      */
-    default TaskBuilder repeat(Duration duration) {
+    default TaskBuilder repeat(@NotNull Duration duration) {
       return repeat(duration.toMillis(), TimeUnit.MILLISECONDS);
     }
 

native/src/main/java/com/velocitypowered/natives/compression/JavaVelocityCompressor.java

@@ -20,7 +20,6 @@ package com.velocitypowered.natives.compression;
 import static com.google.common.base.Preconditions.checkArgument;
 import static com.google.common.base.Preconditions.checkState;
 import static com.velocitypowered.natives.compression.CompressorUtils.ZLIB_BUFFER_SIZE;
-import static com.velocitypowered.natives.compression.CompressorUtils.ensureMaxSize;
 
 import com.velocitypowered.natives.util.BufferPreference;
 import io.netty.buffer.ByteBuf;

native/src/test/java/com/velocitypowered/natives/compression/VelocityCompressorTest.java

@@ -34,9 +34,7 @@ import java.util.zip.DataFormatException;
 import java.util.zip.Deflater;
 import org.junit.jupiter.api.BeforeAll;
 import org.junit.jupiter.api.Test;
-import org.junit.jupiter.api.condition.EnabledOnJre;
 import org.junit.jupiter.api.condition.EnabledOnOs;
-import org.junit.jupiter.api.condition.JRE;
 
 class VelocityCompressorTest {
 

proxy/src/main/java/com/velocitypowered/proxy/scheduler/VelocityScheduler.java

@@ -24,7 +24,6 @@ import com.google.common.collect.ImmutableList;
 import com.google.common.collect.Multimap;
 import com.google.common.collect.Multimaps;
 import com.google.common.util.concurrent.ThreadFactoryBuilder;
-import com.velocitypowered.api.plugin.PluginContainer;
 import com.velocitypowered.api.plugin.PluginManager;
 import com.velocitypowered.api.scheduler.ScheduledTask;
 import com.velocitypowered.api.scheduler.Scheduler;
@@ -32,13 +31,17 @@ import com.velocitypowered.api.scheduler.TaskStatus;
 import java.util.Collection;
 import java.util.HashSet;
 import java.util.IdentityHashMap;
+import java.util.Set;
 import java.util.concurrent.ExecutorService;
 import java.util.concurrent.Executors;
 import java.util.concurrent.ScheduledExecutorService;
 import java.util.concurrent.ScheduledFuture;
 import java.util.concurrent.TimeUnit;
+import java.util.function.Consumer;
+
 import org.apache.logging.log4j.LogManager;
 import org.apache.logging.log4j.Logger;
+import org.checkerframework.checker.nullness.qual.NonNull;
 import org.checkerframework.checker.nullness.qual.Nullable;
 
 public class VelocityScheduler implements Scheduler {
@@ -68,7 +71,25 @@ public class VelocityScheduler implements Scheduler {
     checkNotNull(plugin, "plugin");
     checkNotNull(runnable, "runnable");
     checkArgument(pluginManager.fromInstance(plugin).isPresent(), "plugin is not registered");
-    return new TaskBuilderImpl(plugin, runnable);
+    return new TaskBuilderImpl(plugin, runnable, null);
+  }
+
+  @Override
+  public TaskBuilder buildTask(Object plugin, Consumer<ScheduledTask> consumer) {
+    checkNotNull(plugin, "plugin");
+    checkNotNull(consumer, "consumer");
+    checkArgument(pluginManager.fromInstance(plugin).isPresent(), "plugin is not registered");
+    return new TaskBuilderImpl(plugin, null, consumer);
+  }
+
+  @Override
+  public @NonNull Collection<ScheduledTask> tasksByPlugin(@NonNull Object plugin) {
+    checkNotNull(plugin, "plugin");
+    checkArgument(pluginManager.fromInstance(plugin).isPresent(), "plugin is not registered");
+    final Collection<ScheduledTask> tasks = tasksByPlugin.get(plugin);
+    synchronized (tasksByPlugin) {
+      return Set.copyOf(tasks);
+    }
   }
 
   /**
@@ -93,12 +114,14 @@ public class VelocityScheduler implements Scheduler {
 
     private final Object plugin;
     private final Runnable runnable;
+    private final Consumer<ScheduledTask> consumer;
     private long delay; // ms
     private long repeat; // ms
 
-    private TaskBuilderImpl(Object plugin, Runnable runnable) {
+    private TaskBuilderImpl(Object plugin, Runnable runnable, Consumer<ScheduledTask> consumer) {
       this.plugin = plugin;
       this.runnable = runnable;
+      this.consumer = consumer;
     }
 
     @Override
@@ -127,7 +150,7 @@ public class VelocityScheduler implements Scheduler {
 
     @Override
     public ScheduledTask schedule() {
-      VelocityTask task = new VelocityTask(plugin, runnable, delay, repeat);
+      VelocityTask task = new VelocityTask(plugin, runnable, consumer, delay, repeat);
       tasksByPlugin.put(plugin, task);
       task.schedule();
       return task;
@@ -138,14 +161,16 @@ public class VelocityScheduler implements Scheduler {
 
     private final Object plugin;
     private final Runnable runnable;
+    private final Consumer<ScheduledTask> consumer;
     private final long delay;
     private final long repeat;
     private @Nullable ScheduledFuture<?> future;
     private volatile @Nullable Thread currentTaskThread;
 
-    private VelocityTask(Object plugin, Runnable runnable, long delay, long repeat) {
+    private VelocityTask(Object plugin, Runnable runnable, Consumer<ScheduledTask> consumer, long delay, long repeat) {
       this.plugin = plugin;
       this.runnable = runnable;
+      this.consumer = consumer;
       this.delay = delay;
       this.repeat = repeat;
     }
@@ -200,7 +225,11 @@ public class VelocityScheduler implements Scheduler {
       taskService.execute(() -> {
         currentTaskThread = Thread.currentThread();
         try {
+          if (runnable != null) {
             runnable.run();
+          } else {
+            consumer.accept(this);
+          }
         } catch (Throwable e) {
           //noinspection ConstantConditions
           if (e instanceof InterruptedException) {

proxy/src/test/java/com/velocitypowered/proxy/scheduler/VelocitySchedulerTest.java

@@ -22,9 +22,13 @@ import static org.junit.jupiter.api.Assertions.assertEquals;
 import com.velocitypowered.api.scheduler.ScheduledTask;
 import com.velocitypowered.api.scheduler.TaskStatus;
 import com.velocitypowered.proxy.testutil.FakePluginManager;
+
+import java.time.Duration;
 import java.util.concurrent.CountDownLatch;
 import java.util.concurrent.TimeUnit;
 import java.util.concurrent.atomic.AtomicInteger;
+import java.util.concurrent.atomic.AtomicReference;
+
 import org.junit.jupiter.api.Test;
 
 class VelocitySchedulerTest {
@@ -65,4 +69,65 @@ class VelocitySchedulerTest {
     task.cancel();
   }
 
+  @Test
+  void obtainTasksFromPlugin() throws Exception {
+    VelocityScheduler scheduler = new VelocityScheduler(new FakePluginManager());
+    AtomicInteger i = new AtomicInteger(0);
+    CountDownLatch latch = new CountDownLatch(1);
+
+    scheduler.buildTask(FakePluginManager.PLUGIN_A, task -> {
+      if (i.getAndIncrement() >= 1) {
+        task.cancel();
+        latch.countDown();
+      }
+    }).delay(50, TimeUnit.MILLISECONDS)
+      .repeat(Duration.ofMillis(5))
+      .schedule();
+
+    assertEquals(scheduler.tasksByPlugin(FakePluginManager.PLUGIN_A).size(), 1);
+
+    latch.await();
+
+    assertEquals(scheduler.tasksByPlugin(FakePluginManager.PLUGIN_A).size(), 0);
+  }
+
+  @Test
+  void testConsumerCancel() throws Exception {
+    VelocityScheduler scheduler = new VelocityScheduler(new FakePluginManager());
+    CountDownLatch latch = new CountDownLatch(1);
+
+    ScheduledTask task = scheduler.buildTask(FakePluginManager.PLUGIN_B, actualTask -> {
+      actualTask.cancel();
+      latch.countDown();
+    })
+        .repeat(5, TimeUnit.MILLISECONDS)
+        .schedule();
+
+    assertEquals(TaskStatus.SCHEDULED, task.status());
+
+    latch.await();
+
+    assertEquals(TaskStatus.CANCELLED, task.status());
+  }
+
+  @Test
+  void testConsumerEquality() throws Exception {
+    VelocityScheduler scheduler = new VelocityScheduler(new FakePluginManager());
+    CountDownLatch latch = new CountDownLatch(1);
+
+    AtomicReference<ScheduledTask> consumerTask = new AtomicReference<>();
+    AtomicReference<ScheduledTask> initialTask = new AtomicReference<>();
+
+    ScheduledTask task = scheduler.buildTask(FakePluginManager.PLUGIN_A, scheduledTask -> {
+      consumerTask.set(scheduledTask);
+      latch.countDown();
+    }).delay(60, TimeUnit.MILLISECONDS).schedule();
+
+    initialTask.set(task);
+    latch.await();
+
+    assertEquals(consumerTask.get(), initialTask.get());
+
+  }
+
 }