account for more edge cases and improve checknig for protected entities

src/main/java/me/youhavetrouble/preventstabby/PreventStabby.java

@@ -11,12 +11,16 @@ import me.youhavetrouble.preventstabby.listeners.PvpListener;
 import me.youhavetrouble.preventstabby.listeners.UtilListener;
 import me.youhavetrouble.preventstabby.util.*;
 import org.bstats.bukkit.Metrics;
-import org.bstats.charts.MultiLineChart;
+import org.bukkit.Bukkit;
+import org.bukkit.Chunk;
 import org.bukkit.command.CommandSender;
 import org.bukkit.command.PluginCommand;
+import org.bukkit.entity.Entity;
+import org.bukkit.entity.Tameable;
 import org.bukkit.plugin.java.JavaPlugin;
 
 import java.io.File;
+import java.util.UUID;
 
 public final class PreventStabby extends JavaPlugin {
 
@@ -58,6 +62,24 @@ public final class PreventStabby extends JavaPlugin {
         }
 
         Metrics metrics = new Metrics(this, 14074);
+
+        getServer().getOnlinePlayers().forEach(player -> playerManager.getPlayerData(player.getUniqueId()));
+        getServer().getWorlds().forEach(world -> {
+            for (Chunk chunk : world.getLoadedChunks()) {
+                if (!chunk.isEntitiesLoaded()) continue;
+                Bukkit.getRegionScheduler().run(plugin, chunk.getWorld(), chunk.getX(), chunk.getZ(), (task)  -> {
+                    for (Entity entity : chunk.getEntities()) {
+                        if (!(entity instanceof Tameable tameable)) continue;
+                        UUID ownerId = tameable.getOwnerUniqueId();
+                        if (ownerId == null) continue;
+                        getPlayerManager().getPlayerData(ownerId).thenAccept(playerData -> {
+                            if (playerData == null) return;
+                            playerData.addRelatedEntity(entity.getUniqueId());
+                        });
+                    }
+                });
+            }
+        });
     }
 
     @Override

src/main/java/me/youhavetrouble/preventstabby/data/PlayerData.java

@@ -1,7 +1,10 @@
 package me.youhavetrouble.preventstabby.data;
 
 import me.youhavetrouble.preventstabby.PreventStabby;
+import org.jetbrains.annotations.Nullable;
 
+import java.util.HashSet;
+import java.util.Set;
 import java.util.UUID;
 
 /**
@@ -13,6 +16,7 @@ public class PlayerData {
     private long lastAccessTimestamp;
     private Long combatStartTimestamp, teleportTimestamp, loginTimestamp;
     private boolean pvpEnabled, lastCombatCheck;
+    private final Set<UUID> relatedEntities = new HashSet<>();
 
     public PlayerData(UUID playerUuid, boolean pvpEnabled) {
         this.playerUuid = playerUuid;
@@ -106,6 +110,7 @@ public class PlayerData {
      * Retrieves the login timestamp for the player.
      * @return The login timestamp for the player.
      */
+    @Nullable
     public Long getLoginTimestamp() {
         return loginTimestamp;
     }
@@ -162,4 +167,28 @@ public class PlayerData {
     public boolean isProtected() {
         return  hasLoginProtection() || hasTeleportProtection();
     }
+
+    /**
+     * Get the list of entities related to the player that should keep player's data loaded
+     * @return The list of entities related to the player that should keep player's data loaded
+     */
+    public Set<UUID> getRelatedEntities() {
+        return relatedEntities;
+    }
+
+    /**
+     * @see #getRelatedEntities()
+     * @param entity The entity id to add to the list of related entities
+     */
+    public void addRelatedEntity(UUID entity) {
+        relatedEntities.add(entity);
+    }
+
+    /**
+     * @see #getRelatedEntities()
+     * @param entity The entity id to remove from the list of related entities
+     */
+    public void removeRelatedEntity(UUID entity) {
+        relatedEntities.remove(entity);
+    }
 }

src/main/java/me/youhavetrouble/preventstabby/data/PlayerListener.java

@@ -78,7 +78,10 @@ public class PlayerListener implements Listener {
         if (!(entity instanceof Tameable tameable)) return;
         UUID ownerId = tameable.getOwnerUniqueId();
         if (ownerId == null) return;
-        PreventStabby.getPlugin().getPlayerManager().getPlayerData(ownerId);
+        PreventStabby.getPlugin().getPlayerManager().getPlayerData(ownerId).thenAccept(playerData -> {
+            if (playerData == null) return;
+            playerData.addRelatedEntity(entity.getUniqueId());
+        });
     }
 
 }

src/main/java/me/youhavetrouble/preventstabby/data/PlayerManager.java

@@ -6,11 +6,12 @@ import me.youhavetrouble.preventstabby.api.event.PlayerLeaveCombatEvent;
 import me.youhavetrouble.preventstabby.util.PluginMessages;
 import me.youhavetrouble.preventstabby.util.PvpState;
 import org.bukkit.Bukkit;
-import org.bukkit.Chunk;
 import org.bukkit.entity.Entity;
 import org.bukkit.entity.Player;
-import org.bukkit.entity.Tameable;
 import org.jetbrains.annotations.NotNull;
+
+import java.util.ArrayList;
+import java.util.Collection;
 import java.util.UUID;
 import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.ConcurrentHashMap;
@@ -29,25 +30,29 @@ public class PlayerManager {
             playerList.values().removeIf(PlayerData::isCacheExpired);
         }, 250, 250, TimeUnit.MILLISECONDS);
 
-        Bukkit.getGlobalRegionScheduler().runAtFixedRate(plugin, (task) -> Bukkit.getWorlds().forEach((world -> {
+        Bukkit.getAsyncScheduler().runAtFixedRate(plugin, (task) -> {
-            for (Chunk chunk : world.getLoadedChunks()) {
+            Collection<? extends Player> players = new ArrayList<>(Bukkit.getOnlinePlayers());
-                if (!chunk.isEntitiesLoaded()) continue;
+            // Load data for all online players as a failsafe
-                Bukkit.getRegionScheduler().run(plugin, chunk.getWorld(), chunk.getX(), chunk.getZ(), (task1)  -> {
+            CompletableFuture.allOf(players.stream().map(player -> getPlayerData(player.getUniqueId())).toArray(CompletableFuture[]::new)).join();
-                    for (Entity entity : chunk.getEntities()) {
+            // Refresh cache time for all players
-                        if (!(entity instanceof Tameable tameable)) continue;
+            for (PlayerData playerData : playerList.values()) {
-                        UUID ownerId = tameable.getOwnerUniqueId();
+                if (playerData == null) continue;
-                        if (ownerId == null) continue;
+                playerData.refreshCacheTime();
-                        getPlayerData(ownerId);
-                    }
-                });
             }
-        })), 5, 20 * 15);
+        }, 1, 1, TimeUnit.SECONDS);
 
         Bukkit.getGlobalRegionScheduler().runAtFixedRate(plugin, (task) -> {
             for (PlayerData playerData : playerList.values()) {
                 if (playerData == null) continue;
                 Player player = Bukkit.getPlayer(playerData.getPlayerUuid());
-                if (player == null || !player.isOnline()) continue;
+                if (player == null || !player.isOnline()) {
+                    // player not online, so check for related entities
+                    playerData.getRelatedEntities().removeIf( uuid -> {
+                        Entity entity = Bukkit.getEntity(uuid);
+                        return entity == null;
+                    });
+                    if (playerData.getRelatedEntities().isEmpty()) continue;
+                }
                 playerData.refreshCacheTime(); // Refresh cache timer if player is online
                 // leaving combat logic
                 if (playerData.getLastCombatCheckState() && !playerData.isInCombat()) {
@@ -238,13 +243,15 @@ public class PlayerManager {
         });
     }
 
-    public void setPlayerPvpState(UUID uuid, boolean state) {
+    public CompletableFuture<Void> setPlayerPvpState(UUID uuid, boolean state) {
         // If player is in cache update that
         if (getPlayer(uuid) != null) {
             getPlayer(uuid).setPvpEnabled(state);
         }
-        // Update the database aswell
+        return CompletableFuture.runAsync(() -> {
-        plugin.getSqLite().updatePlayerInfo(uuid, new PlayerData(uuid, state));
+            // Update the database aswell
+            plugin.getSqLite().updatePlayerInfo(uuid, new PlayerData(uuid, state));
+        });
     }
 
     public CompletableFuture<Boolean> togglePlayerPvpState(UUID uuid) {