From 366675f7f8946672d1cfb0daec922dadd320676f Mon Sep 17 00:00:00 2001
From: Nassim Jahnke <nassim@njahnke.dev>
Date: Mon, 9 Dec 2024 15:47:42 +0100
Subject: [PATCH] Make pick item provider Folia-compatible

---
 .../viaversion/ViaVersionPlugin.java          | 21 ++++++++++++++++
 .../providers/BukkitPickItemProvider.java     | 24 +++++++++----------
 2 files changed, 33 insertions(+), 12 deletions(-)

diff --git a/bukkit/src/main/java/com/viaversion/viaversion/ViaVersionPlugin.java b/bukkit/src/main/java/com/viaversion/viaversion/ViaVersionPlugin.java
index ea69290e0..71fde4438 100644
--- a/bukkit/src/main/java/com/viaversion/viaversion/ViaVersionPlugin.java
+++ b/bukkit/src/main/java/com/viaversion/viaversion/ViaVersionPlugin.java
@@ -47,6 +47,7 @@ import java.util.UUID;
 import java.util.concurrent.TimeUnit;
 import org.bukkit.Bukkit;
 import org.bukkit.ChatColor;
+import org.bukkit.block.Block;
 import org.bukkit.entity.Player;
 import org.bukkit.event.Event;
 import org.bukkit.event.EventPriority;
@@ -56,6 +57,8 @@ import org.bukkit.plugin.java.JavaPlugin;
 
 public class ViaVersionPlugin extends JavaPlugin implements ViaPlatform<Player> {
     private static final boolean FOLIA = PaperViaInjector.hasClass("io.papermc.paper.threadedregions.RegionizedServer");
+    private static final Runnable DUMMY_RUNNABLE = () -> {
+    };
     private static ViaVersionPlugin instance;
     private final BukkitCommandHandler commandHandler = new BukkitCommandHandler();
     private final BukkitViaConfig conf;
@@ -179,6 +182,24 @@ public class ViaVersionPlugin extends JavaPlugin implements ViaPlatform<Player>
         return new BukkitViaTask(getServer().getScheduler().runTaskLater(this, runnable, delay));
     }
 
+    public PlatformTask<?> runSyncAt(Runnable runnable, Block block) {
+        if (FOLIA) {
+            return new FoliaViaTask(getServer().getRegionScheduler().run(this, block.getLocation(), (e) -> runnable.run()));
+        }
+        return runSync(runnable);
+    }
+
+    public PlatformTask<?> runSyncFor(Runnable runnable, Player player) {
+        if (FOLIA) {
+            return new FoliaViaTask(player.getScheduler().run(this, (e) -> runnable.run(), DUMMY_RUNNABLE));
+        }
+        return runSync(() -> {
+            if (player.isOnline()) {
+                runnable.run();
+            }
+        });
+    }
+
     @Override
     public PlatformTask runRepeatingSync(Runnable runnable, long period) {
         if (FOLIA) {
diff --git a/bukkit/src/main/java/com/viaversion/viaversion/bukkit/providers/BukkitPickItemProvider.java b/bukkit/src/main/java/com/viaversion/viaversion/bukkit/providers/BukkitPickItemProvider.java
index ce3659aaf..5df39a696 100644
--- a/bukkit/src/main/java/com/viaversion/viaversion/bukkit/providers/BukkitPickItemProvider.java
+++ b/bukkit/src/main/java/com/viaversion/viaversion/bukkit/providers/BukkitPickItemProvider.java
@@ -52,12 +52,12 @@ public final class BukkitPickItemProvider extends PickItemProvider {
     @Override
     public void pickItemFromBlock(final UserConnection connection, final BlockPosition blockPosition, final boolean includeData) {
         final UUID uuid = connection.getProtocolInfo().getUuid();
-        plugin.getServer().getScheduler().runTask(plugin, () -> {
-            final Player player = plugin.getServer().getPlayer(uuid);
-            if (player == null) {
-                return;
-            }
+        final Player player = plugin.getServer().getPlayer(uuid);
+        if (player == null) {
+            return;
+        }
 
+        plugin.runSyncFor(() -> {
             final Location playerLocation = player.getLocation();
             if (blockPosition.distanceFromCenterSquared(playerLocation.getX(), playerLocation.getY(), playerLocation.getZ()) > BLOCK_RANGE_SQUARED) {
                 return;
@@ -72,7 +72,7 @@ public final class BukkitPickItemProvider extends PickItemProvider {
             if (item != null) {
                 pickItem(player, item);
             }
-        });
+        }, player);
     }
 
     private @Nullable ItemStack blockToItem(final Block block, final boolean includeData) {
@@ -96,12 +96,12 @@ public final class BukkitPickItemProvider extends PickItemProvider {
         }
 
         final UUID uuid = connection.getProtocolInfo().getUuid();
-        plugin.getServer().getScheduler().runTask(plugin, () -> {
-            final Player player = plugin.getServer().getPlayer(uuid);
-            if (player == null) {
-                return;
-            }
+        final Player player = plugin.getServer().getPlayer(uuid);
+        if (player == null) {
+            return;
+        }
 
+        plugin.runSyncFor(() -> {
             final Entity entity = player.getWorld().getNearbyEntities(player.getLocation(), ENTITY_RANGE, ENTITY_RANGE, ENTITY_RANGE).stream()
                 .filter(e -> e.getEntityId() == entityId)
                 .findAny()
@@ -114,7 +114,7 @@ public final class BukkitPickItemProvider extends PickItemProvider {
             if (spawnEggType != null) {
                 pickItem(player, new ItemStack(spawnEggType, 1));
             }
-        });
+        }, player);
     }
 
     private void pickItem(final Player player, final ItemStack item) {