use thread safe collections for sets

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

@@ -15,7 +15,7 @@ public class ChunkListener implements Listener {
         List<LivingEntity> matchingEntities
                 = mcMMO.getTransientEntityTracker().getAllTransientEntitiesInChunk(event.getChunk());
         for(LivingEntity livingEntity : matchingEntities) {
-            mcMMO.getTransientEntityTracker().removeSummon(livingEntity, null, false);
+            mcMMO.getTransientEntityTracker().killSummonAndCleanMobFlags(livingEntity, null, false);
         }
     }
 }

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

@@ -682,7 +682,7 @@ public class EntityListener implements Listener {
         LivingEntity entity = event.getEntity();
 
         if (mcMMO.getTransientEntityTracker().isTransient(entity)) {
-            mcMMO.getTransientEntityTracker().removeSummon(entity, null, false);
+            mcMMO.getTransientEntityTracker().killSummonAndCleanMobFlags(entity, null, false);
         }
 
         /* WORLD BLACKLIST CHECK */

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

@@ -8,14 +8,18 @@ import org.bukkit.entity.LivingEntity;
 import org.bukkit.entity.Player;
 import org.jetbrains.annotations.NotNull;
 
+import java.util.UUID;
+
 public class TrackedTamingEntity extends CancellableRunnable {
     private final @NotNull LivingEntity livingEntity;
     private final @NotNull CallOfTheWildType callOfTheWildType;
     private final @NotNull Player player;
+    private final @NotNull UUID playerUUID;
 
     public TrackedTamingEntity(@NotNull LivingEntity livingEntity, @NotNull CallOfTheWildType callOfTheWildType,
                                @NotNull Player player) {
         this.player = player;
+        this.playerUUID = player.getUniqueId();
         this.callOfTheWildType = callOfTheWildType;
         this.livingEntity = livingEntity;
 
@@ -29,7 +33,8 @@ public class TrackedTamingEntity extends CancellableRunnable {
 
     @Override
     public void run() {
-        mcMMO.getTransientEntityTracker().removeSummon(this.getLivingEntity(), player, true);
+        mcMMO.getTransientEntityTracker().killSummonAndCleanMobFlags(this.getLivingEntity(), player, true);
+        mcMMO.getTransientEntityTracker().removeSummonFromTracker(playerUUID, this);
         this.cancel();
     }
 

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

@@ -12,18 +12,19 @@ import org.bukkit.persistence.PersistentDataType;
 import org.jetbrains.annotations.NotNull;
 
 import java.util.EnumMap;
-import java.util.HashSet;
+import java.util.Set;
+import java.util.concurrent.ConcurrentHashMap;
 import java.util.concurrent.ConcurrentMap;
 
 import static com.gmail.nossr50.util.MetadataService.*;
 
 public final class MobMetadataUtils {
-    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 final @NotNull ConcurrentMap<Entity, Set<MobMetaFlagType>> mobRegistry; // transient data
+    private static final @NotNull EnumMap<MobMetaFlagType, NamespacedKey> mobFlagKeyMap; // used for persistent data
     private static boolean isUsingPersistentData = false;
 
     private MobMetadataUtils() {
-        // private ctor
+        // private constructor to prevent instantiation
     }
 
     static {
@@ -39,8 +40,10 @@ public final class MobMetadataUtils {
         initMobFlagKeyMap();
 
         for (MobMetaFlagType metaFlagType : MobMetaFlagType.values()) {
-            if (PersistentDataConfig.getInstance().isMobPersistent(metaFlagType))
+            if (PersistentDataConfig.getInstance().isMobPersistent(metaFlagType)) {
                 isUsingPersistentData = true;
+                break;
+            }
         }
     }
 
@@ -58,61 +61,58 @@ public final class MobMetadataUtils {
                 case PLAYER_BRED_MOB -> mobFlagKeyMap.put(mobMetaFlagType, NSK_PLAYER_BRED_MOB);
                 case EXPLOITED_ENDERMEN -> mobFlagKeyMap.put(mobMetaFlagType, NSK_EXPLOITED_ENDERMEN);
                 case PLAYER_TAMED_MOB -> mobFlagKeyMap.put(mobMetaFlagType, NSK_PLAYER_TAMED_MOB);
-                default -> throw new IncompleteNamespacedKeyRegister("missing namespaced key register for type: " + mobMetaFlagType);
+                default -> throw new IncompleteNamespacedKeyRegister("Missing namespaced key register for type: " + mobMetaFlagType);
             }
         }
     }
 
     /**
-     * Whether a target {@link LivingEntity} has a specific mcMMO mob flags
+     * Checks if a {@link LivingEntity} has a specific mcMMO mob flag.
      *
      * @param flag         the type of mob flag to check for
-     * @param livingEntity the living entity to check for metadata
-     *
-     * @return true if the mob has metadata values for target {@link MobMetaFlagType}
+     * @param livingEntity the living entity to check
+     * @return true if the mob has the specified metadata flag
      */
     public static boolean hasMobFlag(@NotNull MobMetaFlagType flag, @NotNull LivingEntity livingEntity) {
         if (PersistentDataConfig.getInstance().isMobPersistent(flag)) {
             return livingEntity.getPersistentDataContainer().has(mobFlagKeyMap.get(flag), PersistentDataType.BYTE);
         } else {
-            if (mobRegistry.containsKey(livingEntity)) {
-                return mobRegistry.get(livingEntity).contains(flag);
-            }
-
-            return false;
+            final Set<MobMetaFlagType> flags = mobRegistry.get(livingEntity);
+            return flags != null && flags.contains(flag);
         }
     }
 
     /**
-     * Whether a target {@link LivingEntity} has any mcMMO mob flags
-     *
-     * @param livingEntity the living entity to check for metadata
+     * Checks if a {@link LivingEntity} has any mcMMO mob flags.
      *
-     * @return true if the mob has any mcMMO mob related metadata values
+     * @param livingEntity the living entity to check
+     * @return true if the mob has any mcMMO mob-related metadata flags
      */
     public static boolean hasMobFlags(@NotNull LivingEntity livingEntity) {
         if (isUsingPersistentData) {
             for (MobMetaFlagType metaFlagType : MobMetaFlagType.values()) {
-                if (hasMobFlag(metaFlagType, livingEntity))
+                if (hasMobFlag(metaFlagType, livingEntity)) {
                     return true;
+                }
             }
-
             return false;
         } else {
-            return mobRegistry.containsKey(livingEntity) && !mobRegistry.get(livingEntity).isEmpty();
+            final Set<MobMetaFlagType> flags = mobRegistry.get(livingEntity);
+            return flags != null && !flags.isEmpty();
         }
     }
 
     /**
-     * Copies all mcMMO mob flags from one {@link LivingEntity} to another {@link LivingEntity}
-     * This does not clear existing mcMMO mob flags on the target
+     * Copies all mcMMO mob flags from one {@link LivingEntity} to another.
+     * This does not clear existing mcMMO mob flags on the target.
      *
      * @param sourceEntity entity to copy from
      * @param targetEntity entity to copy to
      */
     public static void addMobFlags(@NotNull LivingEntity sourceEntity, @NotNull LivingEntity targetEntity) {
-        if (!hasMobFlags(sourceEntity))
+        if (!hasMobFlags(sourceEntity)) {
             return;
+        }
 
         if (isUsingPersistentData) {
             for (MobMetaFlagType flag : MobMetaFlagType.values()) {
@@ -121,14 +121,17 @@ public final class MobMetadataUtils {
                 }
             }
         } else {
-            HashSet<MobMetaFlagType> flags = new HashSet<>(mobRegistry.get(sourceEntity));
-            mobRegistry.put(targetEntity, flags);
+            Set<MobMetaFlagType> sourceFlags = mobRegistry.get(sourceEntity);
+            if (sourceFlags != null) {
+                Set<MobMetaFlagType> targetFlags = mobRegistry.computeIfAbsent(targetEntity, k -> ConcurrentHashMap.newKeySet());
+                targetFlags.addAll(sourceFlags);
+            }
         }
     }
 
     /**
-     * Adds a mob flag to a {@link LivingEntity} which effectively acts a true/false boolean
-     * Existence of the flag can be considered a true value, non-existence can be considered false for all intents and purposes
+     * Adds a mob flag to a {@link LivingEntity}.
+     * The existence of the flag acts as a true value; non-existence is false.
      *
      * @param flag         the desired flag to assign
      * @param livingEntity the target living entity
@@ -140,16 +143,15 @@ public final class MobMetadataUtils {
                 persistentDataContainer.set(mobFlagKeyMap.get(flag), PersistentDataType.BYTE, MetadataConstants.SIMPLE_FLAG_VALUE);
             }
         } else {
-            HashSet<MobMetaFlagType> flags = mobRegistry.getOrDefault(livingEntity, new HashSet<>());
-            flags.add(flag); // add the new flag
-            mobRegistry.put(livingEntity, flags); //update registry
+            final Set<MobMetaFlagType> flags = mobRegistry.computeIfAbsent(livingEntity, k -> ConcurrentHashMap.newKeySet());
+            flags.add(flag);
         }
     }
 
     /**
-     * Removes a specific mob flag from target {@link LivingEntity}
+     * Removes a specific mob flag from a {@link LivingEntity}.
      *
-     * @param flag         desired flag to remove
+     * @param flag         the flag to remove
      * @param livingEntity the target living entity
      */
     public static void removeMobFlag(@NotNull MobMetaFlagType flag, @NotNull LivingEntity livingEntity) {
@@ -159,19 +161,20 @@ public final class MobMetadataUtils {
                 persistentDataContainer.remove(mobFlagKeyMap.get(flag));
             }
         } else {
-            if (mobRegistry.containsKey(livingEntity)) {
-                mobRegistry.get(livingEntity).remove(flag);
-
-                if (mobRegistry.get(livingEntity).size() == 0)
-                    mobRegistry.remove(livingEntity);
+            final Set<MobMetaFlagType> flags = mobRegistry.get(livingEntity);
+            if (flags != null) {
+                flags.remove(flag);
+                if (flags.isEmpty()) {
+                    mobRegistry.remove(livingEntity, flags);
+                }
             }
         }
     }
 
     /**
-     * Remove all mcMMO related mob flags from the target {@link LivingEntity}
+     * Removes all mcMMO-related mob flags from a {@link LivingEntity}.
      *
-     * @param livingEntity target entity
+     * @param livingEntity the target entity
      */
     public static void removeMobFlags(@NotNull LivingEntity livingEntity) {
         if (isUsingPersistentData) {

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

@@ -16,23 +16,17 @@ 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 {
-    private final @NotNull Map<UUID, HashSet<TrackedTamingEntity>>
-            playerSummonedEntityTracker;
+    private final @NotNull Map<UUID, Set<TrackedTamingEntity>> playerSummonedEntityTracker;
 
     public TransientEntityTracker() {
         playerSummonedEntityTracker = new ConcurrentHashMap<>();
     }
 
-    public @Nullable Set<TrackedTamingEntity> getPlayerSummonedEntities(@NotNull UUID playerUUID) {
-        return playerSummonedEntityTracker.get(playerUUID);
-    }
-
     public void initPlayer(@NotNull Player player) {
-        playerSummonedEntityTracker.computeIfAbsent(player.getUniqueId(), __ -> new HashSet<>());
+        playerSummonedEntityTracker.computeIfAbsent(player.getUniqueId(), __ -> ConcurrentHashMap.newKeySet());
     }
 
     public void cleanupPlayer(@NotNull Player player) {
@@ -52,14 +46,14 @@ public class TransientEntityTracker {
     }
 
     public void addSummon(@NotNull UUID playerUUID, @NotNull TrackedTamingEntity trackedTamingEntity) {
-        playerSummonedEntityTracker.computeIfAbsent(playerUUID, __ -> new HashSet<>())
+        playerSummonedEntityTracker.computeIfAbsent(playerUUID, __ -> ConcurrentHashMap.newKeySet())
                 .add(trackedTamingEntity);
     }
 
-    public void removeSummon(@NotNull LivingEntity livingEntity, @Nullable Player player, boolean timeExpired) {
-        //Kill the summon & remove it
+    public void killSummonAndCleanMobFlags(@NotNull LivingEntity livingEntity, @Nullable Player player,
+                                           boolean timeExpired) {
         if (livingEntity.isValid()) {
-            livingEntity.setHealth(0); //Should trigger entity death events
+            livingEntity.setHealth(0); // Should trigger entity death events
             livingEntity.remove();
 
             Location location = livingEntity.getLocation();
@@ -69,7 +63,7 @@ public class TransientEntityTracker {
                 ParticleEffectUtils.playCallOfTheWildEffect(livingEntity);
             }
 
-            //Inform player of summon death
+            // Inform player of summon death
             if (player != null && player.isOnline()) {
                 if (timeExpired) {
                     NotificationManager.sendPlayerInformationChatOnly(player, "Taming.Summon.COTW.TimeExpired",
@@ -80,17 +74,6 @@ public class TransientEntityTracker {
                 }
             }
         }
-
-        //Remove our metadata
-        removeMobFlags(livingEntity);
-
-        //Clean from trackers
-        remove(livingEntity);
-    }
-
-    private void cleanPlayer(@Nullable Player player, @NotNull UUID playerUUID) {
-        cleanupAllSummons(playerUUID, player);
-        playerSummonedEntityTracker.remove(playerUUID);
     }
 
     public boolean isTransient(@NotNull LivingEntity livingEntity) {
@@ -99,35 +82,42 @@ public class TransientEntityTracker {
                         .anyMatch(trackedTamingEntity -> trackedTamingEntity.getLivingEntity().equals(livingEntity)));
     }
 
-
     private @NotNull Set<TrackedTamingEntity> getTrackedEntities(@NotNull UUID playerUUID,
                                                                  @NotNull CallOfTheWildType callOfTheWildType) {
-        final HashSet<TrackedTamingEntity> entities
-                = playerSummonedEntityTracker.computeIfAbsent(playerUUID, __ -> new HashSet<>());
+        final Set<TrackedTamingEntity> entities =
+                playerSummonedEntityTracker.computeIfAbsent(playerUUID, __ -> ConcurrentHashMap.newKeySet());
         return entities.stream()
                 .filter(trackedTamingEntity -> trackedTamingEntity.getCallOfTheWildType() == callOfTheWildType)
                 .collect(toSet());
     }
 
-    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 cleanPlayer(@Nullable Player player, @NotNull UUID playerUUID) {
+        killAndCleanAllSummons(playerUUID, player);
+        playerSummonedEntityTracker.remove(playerUUID);
     }
 
-    private void cleanupAllSummons(@NotNull UUID playerUUID, @Nullable Player player) {
-        if (playerSummonedEntityTracker.get(playerUUID) == null) {
+    private void killAndCleanAllSummons(@NotNull UUID playerUUID, @Nullable Player player) {
+        final Set<TrackedTamingEntity> entities = playerSummonedEntityTracker.get(playerUUID);
+        if (entities == null) {
             return;
         }
 
-        playerSummonedEntityTracker.get(playerUUID).forEach(trackedTamingEntity -> {
-            removeSummon(trackedTamingEntity.getLivingEntity(), player, false);
-        });
+        // Copy the set to avoid concurrent modification during iteration
+        final Set<TrackedTamingEntity> playerSummonsToRemove = new HashSet<>(entities);
+
+        // Kill and clean all summons
+        playerSummonsToRemove.forEach(
+                trackedTamingEntity -> killAndCleanSummon(playerUUID, player, trackedTamingEntity));
+    }
+
+    public void killAndCleanSummon(@NotNull UUID playerUUID, @Nullable Player player,
+                                   @NotNull TrackedTamingEntity trackedTamingEntity) {
+        killSummonAndCleanMobFlags(trackedTamingEntity.getLivingEntity(), player, false);
+        removeSummonFromTracker(playerUUID, trackedTamingEntity);
+    }
+
+    public void removeSummonFromTracker(@NotNull UUID playerUUID, @NotNull TrackedTamingEntity trackedTamingEntity) {
+        playerSummonedEntityTracker.computeIfAbsent(playerUUID, __ -> ConcurrentHashMap.newKeySet())
+                .remove(trackedTamingEntity);
     }
 }