Fixed concurrency issue with Folia and COTW summons

Fixes #5086

Changelog.txt

@@ -1,5 +1,6 @@
 Version 2.2.027
     Fixed concurrency issue with Folia regarding locale strings
+    Fixed concurrency issue with Folia regarding COTW summons
 
 Version 2.2.026
     Fixed NullPointerException on ChunkUnloadEvent

src/main/java/com/gmail/nossr50/listeners/EntityListener.java

@@ -681,7 +681,7 @@ public class EntityListener implements Listener {
     public void onEntityDeath(EntityDeathEvent event) {
         LivingEntity entity = event.getEntity();
 
-        if (mcMMO.getTransientEntityTracker().isTransientSummon(entity)) {
+        if (mcMMO.getTransientEntityTracker().isTransient(entity)) {
             mcMMO.getTransientEntityTracker().removeSummon(entity, null, false);
         }
 

src/main/java/com/gmail/nossr50/skills/taming/TamingManager.java

@@ -475,11 +475,12 @@ public class TamingManager extends SkillManager {
     }
 
     private int getAmountCurrentlySummoned(@NotNull CallOfTheWildType callOfTheWildType) {
-        return mcMMO.getTransientEntityTracker().getAmountCurrentlySummoned(getPlayer().getUniqueId(), callOfTheWildType);
+        return mcMMO.getTransientEntityTracker().summonCountForPlayerOfType(getPlayer().getUniqueId(), callOfTheWildType);
     }
 
     private void addToTracker(@NotNull LivingEntity livingEntity, @NotNull CallOfTheWildType callOfTheWildType) {
-        mcMMO.getTransientEntityTracker().registerEntity(getPlayer().getUniqueId(), new TrackedTamingEntity(livingEntity, callOfTheWildType, getPlayer()));
+        mcMMO.getTransientEntityTracker().addSummon(getPlayer().getUniqueId(),
+                new TrackedTamingEntity(livingEntity, callOfTheWildType, getPlayer()));
     }
 
     /**

src/main/java/com/gmail/nossr50/util/MobMetadataUtils.java

@@ -3,6 +3,7 @@ package com.gmail.nossr50.util;
 import com.gmail.nossr50.api.exceptions.IncompleteNamespacedKeyRegister;
 import com.gmail.nossr50.config.PersistentDataConfig;
 import com.gmail.nossr50.metadata.MobMetaFlagType;
+import com.google.common.collect.MapMaker;
 import org.bukkit.NamespacedKey;
 import org.bukkit.entity.Entity;
 import org.bukkit.entity.LivingEntity;
@@ -12,13 +13,13 @@ import org.jetbrains.annotations.NotNull;
 
 import java.util.EnumMap;
 import java.util.HashSet;
-import java.util.WeakHashMap;
+import java.util.concurrent.ConcurrentMap;
 
 import static com.gmail.nossr50.util.MetadataService.*;
 
 //TODO: Use SpawnReason where appropriate instead of MobMetaFlagType
 public final class MobMetadataUtils {
-    private static final @NotNull WeakHashMap<Entity, HashSet<MobMetaFlagType>> mobRegistry; //transient data
+    private static final @NotNull ConcurrentMap<Entity, HashSet<MobMetaFlagType>> mobRegistry; //transient data
     private static final @NotNull EnumMap<MobMetaFlagType, NamespacedKey> mobFlagKeyMap; //used for persistent data
     private static boolean isUsingPersistentData = false;
 
@@ -28,7 +29,14 @@ public final class MobMetadataUtils {
 
     static {
         mobFlagKeyMap = new EnumMap<>(MobMetaFlagType.class);
-        mobRegistry = new WeakHashMap<>();
+        // Using Guava for a concurrent weak hash map
+        // IMPORTANT: This type of map uses == for comparison over .equals(),
+        // which is a violation of map contract
+        mobRegistry = new MapMaker()
+                .weakKeys()
+                .concurrencyLevel(4)
+                .makeMap();
+
         initMobFlagKeyMap();
 
         for (MobMetaFlagType metaFlagType : MobMetaFlagType.values()) {
@@ -92,7 +100,7 @@ public final class MobMetadataUtils {
 
             return false;
         } else {
-            return mobRegistry.containsKey(livingEntity) && mobRegistry.get(livingEntity).size() > 0;
+            return mobRegistry.containsKey(livingEntity) && !mobRegistry.get(livingEntity).isEmpty();
         }
     }
 

src/main/java/com/gmail/nossr50/util/TransientEntityTracker.java

@@ -5,8 +5,6 @@ import com.gmail.nossr50.skills.taming.TrackedTamingEntity;
 import com.gmail.nossr50.util.player.NotificationManager;
 import com.gmail.nossr50.util.skills.ParticleEffectUtils;
 import com.gmail.nossr50.util.text.StringUtils;
-import com.google.common.collect.ImmutableSet;
-import org.bukkit.Bukkit;
 import org.bukkit.Chunk;
 import org.bukkit.Location;
 import org.bukkit.Sound;
@@ -16,241 +14,49 @@ import org.jetbrains.annotations.NotNull;
 import org.jetbrains.annotations.Nullable;
 
 import java.util.*;
+import java.util.concurrent.ConcurrentHashMap;
 
 import static com.gmail.nossr50.util.MobMetadataUtils.removeMobFlags;
+import static java.util.stream.Collectors.toSet;
 
 public class TransientEntityTracker {
-    //These two are updated in step with each other
-    private final @NotNull HashMap<UUID, HashMap<CallOfTheWildType, HashSet<TrackedTamingEntity>>> perPlayerTransientEntityMap;
-    private final @NotNull HashSet<LivingEntity> chunkLookupCache;
+    private final @NotNull Map<UUID, HashSet<TrackedTamingEntity>>
+            playerSummonedEntityTracker;
 
     public TransientEntityTracker() {
-        perPlayerTransientEntityMap = new HashMap<>();
-        chunkLookupCache = new HashSet<>();
+        playerSummonedEntityTracker = new ConcurrentHashMap<>();
     }
 
-    public synchronized @NotNull HashSet<LivingEntity> getChunkLookupCache() {
-        return chunkLookupCache;
+    public @Nullable Set<TrackedTamingEntity> getPlayerSummonedEntities(@NotNull UUID playerUUID) {
+        return playerSummonedEntityTracker.get(playerUUID);
     }
 
-    public synchronized @NotNull HashMap<UUID, HashMap<CallOfTheWildType, HashSet<TrackedTamingEntity>>> getPerPlayerTransientEntityMap() {
-        return perPlayerTransientEntityMap;
+    public void initPlayer(@NotNull Player player) {
+        playerSummonedEntityTracker.computeIfAbsent(player.getUniqueId(), __ -> new HashSet<>());
     }
 
-    public synchronized void initPlayer(@NotNull Player player) {
-        if (!isPlayerRegistered(player.getUniqueId())) {
-            registerPlayer(player.getUniqueId());
-        }
-    }
-
-    /**
-     * Removes a player from the tracker
-     *
-     * @param playerUUID target player
-     */
-    public synchronized void cleanupPlayer(@NotNull UUID playerUUID) {
-        cleanPlayer(null, playerUUID);
-    }
-
-    /**
-     * Removes a player from the tracker
-     *
-     * @param player target player
-     */
-    public synchronized void cleanupPlayer(@NotNull Player player) {
+    public void cleanupPlayer(@NotNull Player player) {
         cleanPlayer(player, player.getUniqueId());
     }
 
-    /**
-     * Removes a player from the tracker
-     *
-     * @param player target player
-     * @param playerUUID target player UUID
-     */
-    private void cleanPlayer(@Nullable Player player, @NotNull UUID playerUUID) {
-        cleanupAllSummons(player, player.getUniqueId());
-        removePlayerFromMap(playerUUID);
-    }
-
-    private void removePlayerFromMap(@NotNull UUID playerUUID) {
-        getPerPlayerTransientEntityMap().remove(playerUUID);
-    }
-
-    /**
-     * Checks if a player has already been registered
-     * Being registered constitutes having necessary values initialized in our per-player map
-     *
-     * @param playerUUID target player
-     * @return true if the player is registered
-     */
-    private synchronized boolean isPlayerRegistered(@NotNull UUID playerUUID) {
-        return getPerPlayerTransientEntityMap().get(playerUUID) != null;
-    }
-
-    /**
-     * Register a player to our tracker, which initializes the necessary values in our per-player map
-     *
-     * @param playerUUID player to register
-     */
-    private synchronized void registerPlayer(@NotNull UUID playerUUID) {
-        getPerPlayerTransientEntityMap().put(playerUUID, new HashMap<>());
-
-        for(CallOfTheWildType callOfTheWildType : CallOfTheWildType.values()) {
-            getPerPlayerTransientEntityMap().get(playerUUID).put(callOfTheWildType, new HashSet<>());
-        }
-    }
-
-    /**
-     * Get the tracked transient entities map for a specific player
-     *
-     * @param playerUUID the target uuid of the player
-     * @return the tracked entities map for the player, null if the player isn't registered
-     */
-    public synchronized @Nullable HashMap<CallOfTheWildType, HashSet<TrackedTamingEntity>> getPlayerTrackedEntityMap(@NotNull UUID playerUUID) {
-        return getPerPlayerTransientEntityMap().get(playerUUID);
+    public @NotNull List<LivingEntity> getAllTransientEntitiesInChunk(@NotNull Chunk chunk) {
+        return playerSummonedEntityTracker.values().stream()
+                .flatMap(Collection::stream)
+                .map(TrackedTamingEntity::getLivingEntity)
+                .filter(livingEntity -> livingEntity.getLocation().getChunk() == chunk)
+                .toList();
     }
 
-    /**
-     * Registers an entity to a player
-     * This includes registration to our per-player map and our chunk lookup cache
-     *
-     * @param playerUUID target player's UUID
-     * @param trackedTamingEntity target entity
-     */
-    public synchronized void registerEntity(@NotNull UUID playerUUID, @NotNull TrackedTamingEntity trackedTamingEntity) {
-        //Add to map entry
-        getTrackedEntities(playerUUID, trackedTamingEntity.getCallOfTheWildType()).add(trackedTamingEntity);
-
-        //Add to cache for chunk lookups
-        addToChunkLookupCache(trackedTamingEntity);
-    }
-
-    /**
-     * Checks if a living entity is a summon
-     *
-     * @param livingEntity target livinig entity
-     * @return true if target living entity is a summon
-     */
-    public synchronized boolean isTransientSummon(@NotNull LivingEntity livingEntity) {
-        return getChunkLookupCache().contains(livingEntity);
-    }
-
-    /**
-     * Get the tracked taming entities for a player
-     * If the player isn't registered this will return null
-     *
-     * @param playerUUID the target uuid of the player
-     * @param callOfTheWildType target type
-     * @return the set of tracked entities for the player, null if the player isn't registered, the set can be empty
-     */
-    private synchronized @Nullable HashSet<TrackedTamingEntity> getTrackedEntities(@NotNull UUID playerUUID, @NotNull CallOfTheWildType callOfTheWildType) {
-        HashMap<CallOfTheWildType, HashSet<TrackedTamingEntity>> playerEntityMap = getPlayerTrackedEntityMap(playerUUID);
-
-        if (playerEntityMap == null)
-            return null;
-
-        return playerEntityMap.get(callOfTheWildType);
-    }
-
-    /**
-     * Adds an entity to our chunk lookup cache
-     *
-     * @param trackedTamingEntity target tracked taming entity
-     */
-    private synchronized void addToChunkLookupCache(@NotNull TrackedTamingEntity trackedTamingEntity) {
-        getChunkLookupCache().add(trackedTamingEntity.getLivingEntity());
-    }
-
-    /**
-     * Removes an entity from our tracker
-     * This includes removal from our per-player map and our chunk lookup cache
-     *
-     * @param livingEntity target entity
-     */
-    private void unregisterEntity(@NotNull LivingEntity livingEntity) {
-        chunkLookupCacheCleanup(livingEntity);
-        perPlayerTransientMapCleanup(livingEntity);
+    public int summonCountForPlayerOfType(@NotNull UUID playerUUID, @NotNull CallOfTheWildType callOfTheWildType) {
+        return getTrackedEntities(playerUUID, callOfTheWildType).size();
     }
 
-    /**
-     * Removes an entity from our chunk lookup cache
-     *
-     * @param livingEntity target entity
-     */
-    private void chunkLookupCacheCleanup(@NotNull LivingEntity livingEntity) {
-        getChunkLookupCache().remove(livingEntity);
+    public void addSummon(@NotNull UUID playerUUID, @NotNull TrackedTamingEntity trackedTamingEntity) {
+        playerSummonedEntityTracker.computeIfAbsent(playerUUID, __ -> new HashSet<>())
+                .add(trackedTamingEntity);
     }
 
-    /**
-     * Clean a living entity from our tracker
-     * Iterates over all players and their registered entities
-     * Doesn't do any kind of failure checking, if it doesn't find any player with a registered entity nothing bad happens or is reported
-     * However it should never happen like that, so maybe we could consider adding some failure to execute checking in the future
-     *
-     * @param livingEntity
-     */
-    private void perPlayerTransientMapCleanup(@NotNull LivingEntity livingEntity) {
-        for(UUID uuid : getPerPlayerTransientEntityMap().keySet()) {
-            for(CallOfTheWildType callOfTheWildType : CallOfTheWildType.values()) {
-
-                HashSet<TrackedTamingEntity> trackedEntities = getTrackedEntities(uuid, callOfTheWildType);
-
-                if (trackedEntities == null)
-                    continue;
-
-                Iterator<TrackedTamingEntity> iterator = trackedEntities.iterator();
-                while (iterator.hasNext()) {
-                    if (iterator.next().getLivingEntity().equals(livingEntity)) {
-                        iterator.remove();
-                        return;
-                    }
-                }
-            }
-        }
-    }
-
-    /**
-     * Get all transient entities that exist in a specific chunk
-     *
-     * @param chunk the chunk to match
-     * @return a list of transient entities that are located in the provided chunk
-     */
-    public synchronized @NotNull List<LivingEntity> getAllTransientEntitiesInChunk(@NotNull Chunk chunk) {
-        ArrayList<LivingEntity> matchingEntities = new ArrayList<>();
-
-        for(LivingEntity livingEntity : getChunkLookupCache()) {
-            if (livingEntity.getLocation().getChunk().equals(chunk)) {
-                matchingEntities.add(livingEntity);
-            }
-        }
-
-        return matchingEntities;
-    }
-
-    /**
-     * Get the amount of a summon currently active for a player
-     *
-     * @param playerUUID target player
-     * @param callOfTheWildType summon type
-     * @return the amount of summons currently active for player of target type
-     */
-    public synchronized int getAmountCurrentlySummoned(@NotNull UUID playerUUID, @NotNull CallOfTheWildType callOfTheWildType) {
-        HashSet<TrackedTamingEntity> trackedEntities = getTrackedEntities(playerUUID, callOfTheWildType);
-
-        if (trackedEntities == null)
-            return 0;
-
-        return trackedEntities.size();
-    }
-
-    /**
-     * Kills a summon and removes its metadata
-     * Then it removes it from the tracker / chunk lookup cache
-     *
-     * @param livingEntity entity to remove
-     * @param player associated player
-     */
-    public synchronized void removeSummon(@NotNull LivingEntity livingEntity, @Nullable Player player, boolean timeExpired) {
+    public void removeSummon(@NotNull LivingEntity livingEntity, @Nullable Player player, boolean timeExpired) {
         //Kill the summon & remove it
         if (livingEntity.isValid()) {
             livingEntity.setHealth(0); //Should trigger entity death events
@@ -266,9 +72,11 @@ public class TransientEntityTracker {
             //Inform player of summon death
             if (player != null && player.isOnline()) {
                 if (timeExpired) {
-                    NotificationManager.sendPlayerInformationChatOnly(player, "Taming.Summon.COTW.TimeExpired", StringUtils.getPrettyEntityTypeString(livingEntity.getType()));
+                    NotificationManager.sendPlayerInformationChatOnly(player, "Taming.Summon.COTW.TimeExpired",
+                            StringUtils.getPrettyEntityTypeString(livingEntity.getType()));
                 } else {
-                    NotificationManager.sendPlayerInformationChatOnly(player, "Taming.Summon.COTW.Removed", StringUtils.getPrettyEntityTypeString(livingEntity.getType()));
+                    NotificationManager.sendPlayerInformationChatOnly(player, "Taming.Summon.COTW.Removed",
+                            StringUtils.getPrettyEntityTypeString(livingEntity.getType()));
                 }
             }
         }
@@ -277,41 +85,49 @@ public class TransientEntityTracker {
         removeMobFlags(livingEntity);
 
         //Clean from trackers
-        unregisterEntity(livingEntity);
+        remove(livingEntity);
     }
 
-    /**
-     * Remove all tracked entities from existence if they currently exist
-     * Clear the tracked entity lists afterwards
-     *
-     * @deprecated use {@link #cleanupAllSummons(Player, UUID)} instead
-     */
-    @Deprecated
-    private void cleanupAllSummons(@NotNull UUID playerUUID) {
-        cleanupAllSummons(Bukkit.getPlayer(playerUUID), playerUUID);
+    private void cleanPlayer(@Nullable Player player, @NotNull UUID playerUUID) {
+        cleanupAllSummons(playerUUID, player);
+        playerSummonedEntityTracker.remove(playerUUID);
     }
 
-    /**
-     * Kills and cleans up all data related to all summoned entities for a player
-     *
-     * @param player used to send messages, can be null
-     * @param playerUUID used to grab associated data, cannot be null
-     */
-    private void cleanupAllSummons(@Nullable Player player, @NotNull UUID playerUUID) {
-        for(CallOfTheWildType callOfTheWildType : CallOfTheWildType.values()) {
-            HashSet<TrackedTamingEntity> trackedEntities = getTrackedEntities(playerUUID, callOfTheWildType);
+    public boolean isTransient(@NotNull LivingEntity livingEntity) {
+        return playerSummonedEntityTracker.values().stream().anyMatch(
+                trackedEntities -> trackedEntities.stream()
+                        .anyMatch(trackedTamingEntity -> trackedTamingEntity.getLivingEntity().equals(livingEntity)));
+    }
 
-            if (trackedEntities == null) {
-                continue;
-            }
 
-            ImmutableSet<TrackedTamingEntity> immutableSet = ImmutableSet.copyOf(trackedEntities);
+    private @NotNull Set<TrackedTamingEntity> getTrackedEntities(@NotNull UUID playerUUID,
+                                                                 @NotNull CallOfTheWildType callOfTheWildType) {
+        final HashSet<TrackedTamingEntity> entities
+                = playerSummonedEntityTracker.computeIfAbsent(playerUUID, __ -> new HashSet<>());
+        return entities.stream()
+                .filter(trackedTamingEntity -> trackedTamingEntity.getCallOfTheWildType() == callOfTheWildType)
+                .collect(toSet());
+    }
 
-            for(TrackedTamingEntity trackedTamingEntity : immutableSet) {
-                //Remove from existence
-                removeSummon(trackedTamingEntity.getLivingEntity(), player, false);
-            }
+    private void remove(@NotNull LivingEntity livingEntity) {
+        playerSummonedEntityTracker.values().forEach(trackedEntities -> {
+            Iterator<TrackedTamingEntity> iterator = trackedEntities.iterator();
+                while (iterator.hasNext()) {
+                    if (iterator.next().getLivingEntity().equals(livingEntity)) {
+                        iterator.remove();
+                        return;
+                    }
+                }
+            });
+    }
 
+    private void cleanupAllSummons(@NotNull UUID playerUUID, @Nullable Player player) {
+        if (playerSummonedEntityTracker.get(playerUUID) == null) {
+            return;
         }
+
+        playerSummonedEntityTracker.get(playerUUID).forEach(trackedTamingEntity -> {
+            removeSummon(trackedTamingEntity.getLivingEntity(), player, false);
+        });
     }
 }