Improved Folia support (fixes #892)

src/main/java/net/coreprotect/command/TeleportCommand.java

@@ -14,7 +14,6 @@ import net.coreprotect.utility.Chat;
 import net.coreprotect.utility.ChatMessage;
 import net.coreprotect.utility.Color;
 import net.coreprotect.utility.Teleport;
-import net.coreprotect.utility.Util;
 import net.coreprotect.utility.WorldUtils;
 
 public class TeleportCommand {
@@ -56,6 +55,10 @@ public class TeleportCommand {
         World world = location.getWorld();
         if (wid > -1) {
             world = Bukkit.getServer().getWorld(WorldUtils.getWorldName(wid));
+            if (world == null) {
+                Chat.sendMessage(player, new ChatMessage(Phrase.build(Phrase.WORLD_NOT_FOUND, worldName)).build());
+                return;
+            }
         }
 
         String x = null;

src/main/java/net/coreprotect/consumer/process/EntityKillProcess.java

@@ -10,16 +10,20 @@ import org.bukkit.entity.EntityType;
 import net.coreprotect.consumer.Consumer;
 import net.coreprotect.database.logger.EntityKillLogger;
 import net.coreprotect.utility.EntityUtils;
-import net.coreprotect.utility.Util;
 
 class EntityKillProcess {
 
     static void process(PreparedStatement preparedStmt, PreparedStatement preparedStmtEntities, int batchCount, int processId, int id, Object object, String user) {
         if (object instanceof Object[]) {
-            BlockState block = (BlockState) ((Object[]) object)[0];
+            Object[] values = (Object[]) object;
-            EntityType type = (EntityType) ((Object[]) object)[1];
+            if (values.length <= 1 || !(values[0] instanceof BlockState) || !(values[1] instanceof EntityType)) {
+                return;
+            }
+
+            BlockState block = (BlockState) values[0];
+            EntityType type = (EntityType) values[1];
             Map<Integer, List<Object>> objectLists = Consumer.consumerObjectList.get(processId);
-            if (objectLists.get(id) != null) {
+            if (objectLists != null && objectLists.get(id) != null) {
                 List<Object> objectList = objectLists.get(id);
                 int entityId = EntityUtils.getEntityId(type);
                 EntityKillLogger.log(preparedStmt, preparedStmtEntities, batchCount, user, block, objectList, entityId);

src/main/java/net/coreprotect/consumer/process/PlayerKillProcess.java

@@ -10,8 +10,13 @@ class PlayerKillProcess {
 
     static void process(PreparedStatement preparedStmt, int batchCount, int id, Object object, String user) {
         if (object instanceof Object[]) {
-            BlockState block = (BlockState) ((Object[]) object)[0];
+            Object[] values = (Object[]) object;
-            String player = (String) ((Object[]) object)[1];
+            if (values.length <= 1 || !(values[0] instanceof BlockState) || !(values[1] instanceof String)) {
+                return;
+            }
+
+            BlockState block = (BlockState) values[0];
+            String player = (String) values[1];
             PlayerKillLogger.log(preparedStmt, batchCount, user, block, player);
         }
     }

src/main/java/net/coreprotect/database/rollback/Rollback.java

@@ -10,6 +10,7 @@ import java.util.List;
 import java.util.Map;
 import java.util.Map.Entry;
 import java.util.TreeMap;
+import java.util.concurrent.CompletableFuture;
 
 import org.bukkit.Bukkit;
 import org.bukkit.Location;
@@ -247,19 +248,21 @@ public class Rollback extends RollbackUtil {
                 }
 
                 ConfigHandler.rollbackHash.put(finalUserString, new int[] { itemCount, blockCount, entityCount, 0, scannedWorldData });
+                List<CompletableFuture<Boolean>> chunkFutures = new ArrayList<>();
                 for (Entry<Integer, World> rollbackWorlds : worldMap.entrySet()) {
                     Integer rollbackWorldId = rollbackWorlds.getKey();
                     World bukkitRollbackWorld = rollbackWorlds.getValue();
-                    Location chunkLocation = new Location(bukkitRollbackWorld, (finalChunkX << 4), 0, (finalChunkZ << 4));
                     final HashMap<Long, ArrayList<Object[]>> finalBlockList = dataList.get(rollbackWorldId);
                     final HashMap<Long, ArrayList<Object[]>> finalItemList = itemDataList.get(rollbackWorldId);
 
-                    Scheduler.scheduleSyncDelayedTask(CoreProtect.getInstance(), () -> {
+                    chunkFutures.add(scheduleChunkTask(finalChunkX, finalChunkZ, chunkKey, finalBlockList, finalItemList, rollbackType, preview, finalUserString, finalUser, bukkitRollbackWorld, inventoryRollback));
-                        // Process this chunk using our new RollbackProcessor class
+                }
-                        ArrayList<Object[]> blockData = finalBlockList != null ? finalBlockList.getOrDefault(chunkKey, new ArrayList<>()) : new ArrayList<>();
+
-                        ArrayList<Object[]> itemData = finalItemList != null ? finalItemList.getOrDefault(chunkKey, new ArrayList<>()) : new ArrayList<>();
+                if (ConfigHandler.isFolia) {
-                        RollbackProcessor.processChunk(finalChunkX, finalChunkZ, chunkKey, blockData, itemData, rollbackType, preview, finalUserString, finalUser instanceof Player ? (Player) finalUser : null, bukkitRollbackWorld, inventoryRollback);
+                    if (!awaitChunkTasks(chunkFutures, preview)) {
-                    }, chunkLocation, 0);
+                        Chat.console(Phrase.build(Phrase.ROLLBACK_ABORTED));
+                        break;
+                    }
                 }
 
                 rollbackHashData = ConfigHandler.rollbackHash.get(finalUserString);
@@ -268,24 +271,26 @@ public class Rollback extends RollbackUtil {
                 int sleepTime = 0;
                 int abort = 0;
 
-                while (next == 0 || scannedWorlds < worldMap.size()) {
+                if (!ConfigHandler.isFolia) {
-                    if (preview == 1) {
+                    while (next == 0 || scannedWorlds < worldMap.size()) {
-                        // Not actually changing blocks, so less intensive.
+                        if (preview == 1) {
-                        sleepTime = sleepTime + 1;
+                            // Not actually changing blocks, so less intensive.
-                        Thread.sleep(1);
+                            sleepTime = sleepTime + 1;
-                    }
+                            Thread.sleep(1);
-                    else {
+                        }
-                        sleepTime = sleepTime + 5;
+                        else {
-                        Thread.sleep(5);
+                            sleepTime = sleepTime + 5;
-                    }
+                            Thread.sleep(5);
+                        }
 
-                    rollbackHashData = ConfigHandler.rollbackHash.get(finalUserString);
+                        rollbackHashData = ConfigHandler.rollbackHash.get(finalUserString);
-                    next = rollbackHashData[3];
+                        next = rollbackHashData[3];
-                    scannedWorlds = rollbackHashData[4];
+                        scannedWorlds = rollbackHashData[4];
 
-                    if (sleepTime > 300000) {
+                        if (sleepTime > 300000) {
-                        abort = 1;
+                            abort = 1;
-                        break;
+                            break;
+                        }
                     }
                 }
 
@@ -331,4 +336,69 @@ public class Rollback extends RollbackUtil {
         return null;
     }
 
+    private static CompletableFuture<Boolean> scheduleChunkTask(int chunkX, int chunkZ, long chunkKey, HashMap<Long, ArrayList<Object[]>> blockList, HashMap<Long, ArrayList<Object[]>> itemList, int rollbackType, int preview, String userString, CommandSender user, World world, boolean inventoryRollback) {
+        CompletableFuture<Boolean> future = new CompletableFuture<>();
+        Location chunkLocation = new Location(world, (chunkX << 4), 0, (chunkZ << 4));
+
+        Scheduler.scheduleSyncDelayedTask(CoreProtect.getInstance(), () -> {
+            try {
+                ArrayList<Object[]> blockData = blockList != null ? blockList.getOrDefault(chunkKey, new ArrayList<>()) : new ArrayList<>();
+                ArrayList<Object[]> itemData = itemList != null ? itemList.getOrDefault(chunkKey, new ArrayList<>()) : new ArrayList<>();
+                Player rollbackPlayer = user instanceof Player ? (Player) user : null;
+                boolean result = RollbackProcessor.processChunk(chunkX, chunkZ, chunkKey, blockData, itemData, rollbackType, preview, userString, rollbackPlayer, world, inventoryRollback);
+                future.complete(result);
+            }
+            catch (Exception e) {
+                e.printStackTrace();
+                future.complete(false);
+            }
+        }, chunkLocation, 0);
+
+        return future;
+    }
+
+    private static boolean awaitChunkTasks(List<CompletableFuture<Boolean>> futures, int preview) throws InterruptedException {
+        if (futures.isEmpty()) {
+            return true;
+        }
+
+        int sleepTime = 0;
+        while (true) {
+            boolean allDone = true;
+            for (CompletableFuture<Boolean> future : futures) {
+                if (!future.isDone()) {
+                    allDone = false;
+                    break;
+                }
+            }
+
+            if (allDone) {
+                break;
+            }
+
+            int delay = preview == 1 ? 1 : 5;
+            sleepTime += delay;
+            if (sleepTime > 300000) {
+                return false;
+            }
+            Thread.sleep(delay);
+        }
+
+        for (CompletableFuture<Boolean> future : futures) {
+            Boolean result;
+            try {
+                result = future.getNow(Boolean.FALSE);
+            }
+            catch (Exception e) {
+                return false;
+            }
+
+            if (!Boolean.TRUE.equals(result)) {
+                return false;
+            }
+        }
+
+        return true;
+    }
+
 }

src/main/java/net/coreprotect/database/rollback/RollbackEntityHandler.java

@@ -1,5 +1,7 @@
 package net.coreprotect.database.rollback;
 
+import java.util.Arrays;
+
 import org.bukkit.Bukkit;
 import org.bukkit.Location;
 import org.bukkit.World;
@@ -78,16 +80,8 @@ public class RollbackEntityHandler {
                 else if (rowTypeRaw <= 0) {
                     // Attempt to remove entity
                     if (rowRolledBack == 1) {
-                        boolean removed = false;
+                        int entityId = getCachedEntityId(rowX, rowY, rowZ, rowWorldId, oldTypeRaw);
-                        int entityId = -1;
+                        EntityType targetType = EntityUtils.getEntityType(oldTypeRaw);
-                        String entityName = EntityUtils.getEntityType(oldTypeRaw).name();
-                        String token = "" + rowX + "." + rowY + "." + rowZ + "." + rowWorldId + "." + entityName + "";
-                        Object[] cachedEntity = CacheHandler.entityCache.get(token);
-
-                        if (cachedEntity != null) {
-                            entityId = (Integer) cachedEntity[1];
-                        }
-
                         int xmin = rowX - 5;
                         int xmax = rowX + 5;
                         int ymin = rowY - 1;
@@ -95,43 +89,10 @@ public class RollbackEntityHandler {
                         int zmin = rowZ - 5;
                         int zmax = rowZ + 5;
 
-                        for (Entity entity : block.getChunk().getEntities()) {
+                        boolean removed = removeMatchingEntity(Arrays.asList(block.getChunk().getEntities()), finalUserString, entityId, targetType, xmin, xmax, ymin, ymax, zmin, zmax, true);
-                            if (entityId > -1) {
+                        if (!removed && entityId > -1 && !ConfigHandler.isFolia) {
-                                int id = entity.getEntityId();
+                            // On non-Folia, keep world-wide fallback for moved entities.
-                                if (id == entityId) {
+                            removed = removeMatchingEntity(block.getWorld().getLivingEntities(), finalUserString, entityId, targetType, xmin, xmax, ymin, ymax, zmin, zmax, false);
-                                    updateEntityCount(finalUserString, 1);
-                                    removed = true;
-                                    entity.remove();
-                                    break;
-                                }
-                            }
-                            else {
-                                if (entity.getType().equals(EntityUtils.getEntityType(oldTypeRaw))) {
-                                    Location entityLocation = entity.getLocation();
-                                    int entityx = entityLocation.getBlockX();
-                                    int entityY = entityLocation.getBlockY();
-                                    int entityZ = entityLocation.getBlockZ();
-
-                                    if (entityx >= xmin && entityx <= xmax && entityY >= ymin && entityY <= ymax && entityZ >= zmin && entityZ <= zmax) {
-                                        updateEntityCount(finalUserString, 1);
-                                        removed = true;
-                                        entity.remove();
-                                        break;
-                                    }
-                                }
-                            }
-                        }
-
-                        if (!removed && entityId > -1) {
-                            for (Entity entity : block.getWorld().getLivingEntities()) {
-                                int id = entity.getEntityId();
-                                if (id == entityId) {
-                                    updateEntityCount(finalUserString, 1);
-                                    removed = true;
-                                    entity.remove();
-                                    break;
-                                }
-                            }
                         }
 
                         if (removed) {
@@ -159,6 +120,48 @@ public class RollbackEntityHandler {
         return WorldUtils.getWorldName(worldId);
     }
 
+    private static int getCachedEntityId(int rowX, int rowY, int rowZ, int rowWorldId, int oldTypeRaw) {
+        String entityName = EntityUtils.getEntityType(oldTypeRaw).name();
+        String token = rowX + "." + rowY + "." + rowZ + "." + rowWorldId + "." + entityName;
+        Object[] cachedEntity = CacheHandler.entityCache.get(token);
+        if (cachedEntity != null && cachedEntity.length > 1 && cachedEntity[1] instanceof Integer) {
+            return (Integer) cachedEntity[1];
+        }
+        return -1;
+    }
+
+    private static boolean removeMatchingEntity(Iterable<? extends Entity> entities, String userString, int cachedEntityId, EntityType targetType, int xmin, int xmax, int ymin, int ymax, int zmin, int zmax, boolean useBounds) {
+        for (Entity entity : entities) {
+            if (cachedEntityId > -1) {
+                if (entity.getEntityId() == cachedEntityId) {
+                    updateEntityCount(userString, 1);
+                    entity.remove();
+                    return true;
+                }
+                continue;
+            }
+
+            if (!entity.getType().equals(targetType)) {
+                continue;
+            }
+            if (useBounds && !isWithinBounds(entity.getLocation(), xmin, xmax, ymin, ymax, zmin, zmax)) {
+                continue;
+            }
+
+            updateEntityCount(userString, 1);
+            entity.remove();
+            return true;
+        }
+        return false;
+    }
+
+    private static boolean isWithinBounds(Location location, int xmin, int xmax, int ymin, int ymax, int zmin, int zmax) {
+        int entityX = location.getBlockX();
+        int entityY = location.getBlockY();
+        int entityZ = location.getBlockZ();
+        return entityX >= xmin && entityX <= xmax && entityY >= ymin && entityY <= ymax && entityZ >= zmin && entityZ <= zmax;
+    }
+
     /**
      * Updates the entity count in the rollback hash for a specific user.
      *

src/main/java/net/coreprotect/database/rollback/RollbackProcessor.java

@@ -25,11 +25,13 @@ import org.bukkit.entity.LivingEntity;
 import org.bukkit.entity.Player;
 import org.bukkit.inventory.ItemStack;
 
+import net.coreprotect.CoreProtect;
 import net.coreprotect.bukkit.BukkitAdapter;
 import net.coreprotect.config.Config;
 import net.coreprotect.config.ConfigHandler;
 import net.coreprotect.database.logger.ItemLogger;
 import net.coreprotect.model.BlockGroup;
+import net.coreprotect.thread.Scheduler;
 import net.coreprotect.utility.BlockUtils;
 import net.coreprotect.utility.ItemUtils;
 import net.coreprotect.utility.MaterialUtils;
@@ -435,7 +437,7 @@ public class RollbackProcessor {
                     int chunkZ = playerLocation.getBlockZ() >> 4;
 
                     if (bukkitRollbackWorld.getName().equals(playerWorld) && chunkX == finalChunkX && chunkZ == finalChunkZ) {
-                        Teleport.performSafeTeleport(player, playerLocation, false);
+                        Scheduler.runTask(CoreProtect.getInstance(), () -> Teleport.performSafeTeleport(player, playerLocation, false), player);
                     }
                 }
             }

src/main/java/net/coreprotect/listener/block/BlockDispenseListener.java

@@ -84,7 +84,10 @@ public final class BlockDispenseListener extends Queue implements Listener {
                 }
 
                 if (!dispenseSuccess && material == Material.BONE_MEAL) {
-                    CacheHandler.redstoneCache.put(newBlock.getLocation(), new Object[] { System.currentTimeMillis(), user });
+                    String key = CacheHandler.locationKey(newBlock.getLocation());
+                    if (!key.isEmpty()) {
+                        CacheHandler.redstoneCache.put(key, new Object[] { System.currentTimeMillis(), user });
+                    }
                 }
 
                 if (type == Material.FIRE && (!config.BLOCK_IGNITE || !(newBlockData instanceof Lightable))) {

src/main/java/net/coreprotect/listener/block/BlockFertilizeListener.java

@@ -59,7 +59,8 @@ public final class BlockFertilizeListener extends Queue implements Listener {
             user = player.getName();
         }
         else {
-            Object[] data = CacheHandler.redstoneCache.get(location);
+            String key = CacheHandler.locationKey(location);
+            Object[] data = CacheHandler.redstoneCache.get(key);
             if (data != null) {
                 long newTime = System.currentTimeMillis();
                 long oldTime = (long) data[0];
@@ -67,7 +68,7 @@ public final class BlockFertilizeListener extends Queue implements Listener {
                     user = (String) data[1];
                 }
 
-                CacheHandler.redstoneCache.remove(location);
+                CacheHandler.redstoneCache.remove(key);
             }
         }
 

src/main/java/net/coreprotect/thread/CacheHandler.java

@@ -23,7 +23,14 @@ public class CacheHandler implements Runnable {
     public static ConcurrentHashMap<String, Object[]> flowDuplicateCache = new ConcurrentHashMap<>(16, 0.75f, 2);
     public static ConcurrentHashMap<String, Object[]> bonemealDuplicateCache = new ConcurrentHashMap<>(16, 0.75f, 2);
     public static ConcurrentHashMap<String, Object[]> entityKillDuplicateCache = new ConcurrentHashMap<>(16, 0.75f, 2);
-    public static ConcurrentHashMap<Location, Object[]> redstoneCache = new ConcurrentHashMap<>(16, 0.75f, 2);
+    public static ConcurrentHashMap<String, Object[]> redstoneCache = new ConcurrentHashMap<>(16, 0.75f, 2);
+
+    public static String locationKey(Location location) {
+        if (location == null || location.getWorld() == null) {
+            return "";
+        }
+        return location.getWorld().getUID() + ":" + location.getBlockX() + ":" + location.getBlockY() + ":" + location.getBlockZ();
+    }
 
     public static boolean shouldSuppressRepeat(ConcurrentHashMap<String, Object[]> cache, String cacheId, int threshold, int resetWindowSeconds) {
         if (cache == null || cacheId == null || cacheId.isEmpty() || threshold <= 0) {

src/main/java/net/coreprotect/thread/Scheduler.java

@@ -29,10 +29,10 @@ public class Scheduler {
             else if (regionData instanceof Entity) {
                 Entity entity = (Entity) regionData;
                 if (delay == 0) {
-                    entity.getScheduler().run(plugin, value -> task.run(), task);
+                    entity.getScheduler().run(plugin, value -> task.run(), null);
                 }
                 else {
-                    entity.getScheduler().runDelayed(plugin, value -> task.run(), task, delay);
+                    entity.getScheduler().runDelayed(plugin, value -> task.run(), null, delay);
                 }
             }
             else {
@@ -62,7 +62,7 @@ public class Scheduler {
             }
             else if (regionData instanceof Entity) {
                 Entity entity = (Entity) regionData;
-                return entity.getScheduler().runAtFixedRate(plugin, value -> task.run(), task, delay, period);
+                return entity.getScheduler().runAtFixedRate(plugin, value -> task.run(), null, delay, period);
             }
             else {
                 return plugin.getServer().getGlobalRegionScheduler().runAtFixedRate(plugin, value -> task.run(), delay, period);

src/main/java/net/coreprotect/utility/entity/EntityUtil.java

@@ -74,12 +74,17 @@ public class EntityUtil {
         if (type == null) {
             return;
         }
+        final Location blockLocation = block.getLocation();
+        if (blockLocation.getWorld() == null) {
+            return;
+        }
+
         Scheduler.runTask(CoreProtect.getInstance(), () -> {
             try {
-                Location location = block.getLocation();
+                Location location = blockLocation.clone();
                 location.setX(location.getX() + 0.50);
                 location.setZ(location.getZ() + 0.50);
-                Entity entity = block.getLocation().getWorld().spawnEntity(location, type);
+                Entity entity = blockLocation.getWorld().spawnEntity(location, type);
 
                 if (list.isEmpty()) {
                     return;
@@ -599,7 +604,7 @@ public class EntityUtil {
             catch (Exception e) {
                 e.printStackTrace();
             }
-        }, block.getLocation());
+        }, blockLocation);
     }
 
 }