Merge pull request #36 from TechnicallyCoded/dev

Fix SpigotImplementation runAtEntityTimer() & runAtEntityLater() & minor cleanups

Author: TechnicallyCoded

common/src/main/java/com/tcoded/folialib/type/Ref.java

@@ -0,0 +1,23 @@
+package com.tcoded.folialib.type;
+
+public class Ref<T> {
+
+    private T value;
+
+    public Ref() {
+        this.value = null;
+    }
+
+    public Ref(T value) {
+        this.value = value;
+    }
+
+    public T get() {
+        return value;
+    }
+
+    public void set(T value) {
+        this.value = value;
+    }
+
+}

platform/folia/src/main/java/com/tcoded/folialib/impl/FoliaImplementation.java

@@ -8,6 +8,7 @@
 import com.tcoded.folialib.wrapper.task.WrappedTask;
 import io.papermc.paper.threadedregions.scheduler.AsyncScheduler;
 import io.papermc.paper.threadedregions.scheduler.GlobalRegionScheduler;
+import io.papermc.paper.threadedregions.scheduler.RegionScheduler;
 import io.papermc.paper.threadedregions.scheduler.ScheduledTask;
 import it.unimi.dsi.fastutil.longs.Long2ObjectOpenHashMap;
 import org.bukkit.Location;
@@ -35,11 +36,13 @@ public class FoliaImplementation implements PlatformScheduler {
 
     private final Plugin plugin;
     private final GlobalRegionScheduler globalRegionScheduler;
+    private final RegionScheduler regionScheduler;
     private final AsyncScheduler asyncScheduler;
 
     public FoliaImplementation(FoliaLib foliaLib) {
         this.plugin = foliaLib.getPlugin();
         this.globalRegionScheduler = plugin.getServer().getGlobalRegionScheduler();
+        this.regionScheduler = plugin.getServer().getRegionScheduler();
         this.asyncScheduler = plugin.getServer().getAsyncScheduler();
     }
 
@@ -234,7 +237,7 @@ public void runTimerAsync(@NotNull Consumer<WrappedTask> consumer, long delay, l
     public @NotNull CompletableFuture<Void> runAtLocation(Location location, @NotNull Consumer<WrappedTask> consumer) {
         CompletableFuture<Void> future = new CompletableFuture<>();
 
-        this.plugin.getServer().getRegionScheduler().run(plugin, location, task -> {
+        this.regionScheduler.run(plugin, location, task -> {
             consumer.accept(this.wrapTask(task));
             future.complete(null);
         });
@@ -249,7 +252,7 @@ public WrappedTask runAtLocationLater(Location location, @NotNull Runnable runna
             delay = 1;
         }
         return this.wrapTask(
-                this.plugin.getServer().getRegionScheduler().runDelayed(plugin, location, task -> runnable.run(), delay)
+                this.regionScheduler.runDelayed(plugin, location, task -> runnable.run(), delay)
         );
     }
 
@@ -261,7 +264,7 @@ public WrappedTask runAtLocationLater(Location location, @NotNull Runnable runna
             InvalidTickDelayNotifier.notifyOnce(plugin.getLogger(), delay);
             delay = 1;
         }
-        this.plugin.getServer().getRegionScheduler().runDelayed(plugin, location, task -> {
+        this.regionScheduler.runDelayed(plugin, location, task -> {
             consumer.accept(this.wrapTask(task));
             future.complete(null);
         }, delay);
@@ -290,7 +293,7 @@ public WrappedTask runAtLocationTimer(Location location, @NotNull Runnable runna
             period = 1;
         }
         return this.wrapTask(
-                this.plugin.getServer().getRegionScheduler().runAtFixedRate(plugin, location, task -> runnable.run(), delay, period)
+                this.regionScheduler.runAtFixedRate(plugin, location, task -> runnable.run(), delay, period)
         );
     }
 
@@ -304,7 +307,7 @@ public void runAtLocationTimer(Location location, @NotNull Consumer<WrappedTask>
             InvalidTickDelayNotifier.notifyOnce(plugin.getLogger(), period);
             period = 1;
         }
-        this.plugin.getServer().getRegionScheduler().runAtFixedRate(plugin, location, task -> consumer.accept(this.wrapTask(task)), delay, period);
+        this.regionScheduler.runAtFixedRate(plugin, location, task -> consumer.accept(this.wrapTask(task)), delay, period);
     }
 
 	@Override

platform/spigot/src/main/java/com/tcoded/folialib/impl/SpigotImplementation.java

@@ -2,6 +2,7 @@
 
 import com.tcoded.folialib.FoliaLib;
 import com.tcoded.folialib.enums.EntityTaskResult;
+import com.tcoded.folialib.type.Ref;
 import com.tcoded.folialib.util.TimeConverter;
 import com.tcoded.folialib.wrapper.task.WrappedTask;
 import com.tcoded.folialib.wrapper.task.WrappedBukkitTask;
@@ -289,18 +290,29 @@ public WrappedTask runAtEntityLater(Entity entity, @NotNull Runnable runnable, R
     public @NotNull CompletableFuture<Void> runAtEntityLater(Entity entity, @NotNull Consumer<WrappedTask> consumer, @Nullable Runnable fallback, long delay) {
         CompletableFuture<Void> future = new CompletableFuture<>();
 
+        // Entity is invalid. Complete the future and run the fallback if provided
         if (!isValid(entity)) {
+            future.complete(null);
+
             if (fallback != null) {
                 fallback.run();
-                future.complete(null);
             }
+
+            return future;
         }
-        else {
-            this.runLater(task -> {
-                consumer.accept(this.wrapTask(task));
+
+        this.runLater(task -> {
+            // Handle the case where the entity becomes invalid between scheduling and task execution
+            if (!isValid(entity)) {
+                if (fallback != null) fallback.run();
                 future.complete(null);
-            }, delay);
-        }
+                return;
+            }
+
+            // Everything is healthy, run the consumer and complete the future
+            consumer.accept(task);
+            future.complete(null);
+        }, delay);
 
         return future;
     }
@@ -326,8 +338,31 @@ public WrappedTask runAtEntityTimer(Entity entity, @NotNull Runnable runnable, @
             if (fallback != null) fallback.run();
             return null;
         }
-        return this.runTimer(runnable, delay, period);
-//        return this.wrapTask(this.scheduler.runTaskTimer(plugin, runnable, delay, period)); // todo: remove old
+
+        // This is called a magic trick...
+        // Close your eyes!
+        Ref<WrappedTask> wtRef = new Ref<>();
+
+        WrappedTask task = this.runTimer(() -> {
+            // Handle case where entity becomes invalid. The task should no longer run.
+            // We first check if the WrappedTask is non-null, since we never want to call the fallback
+            // multiple times: we need to be able to cancel the recurring task,
+            WrappedTask wt = wtRef.get();
+            if (wt != null && !isValid(entity)) {
+                wt.cancel();
+                // Run the fallback if provided
+                if (fallback != null) fallback.run();
+                return;
+            }
+
+            // Everything is healthy, run the task.
+            runnable.run();
+        }, delay, period);
+
+        wtRef.set(task);
+        // Ok, you can open your eyes again.
+
+        return task;
     }
 
     @Override