From ab031e3e6349f425ab6a6e53c71f01603e5b4c88 Mon Sep 17 00:00:00 2001
From: Andreas Troelsen <garbagemule@gmail.com>
Date: Mon, 5 Aug 2019 16:04:44 +0200
Subject: [PATCH] Remove pets in case they die.

This technically can't happen because all pet damage is cancelled, but we never know, and we best be sure to clean up.
---
 .../java/com/garbagemule/MobArena/ArenaListener.java |  3 +++
 .../com/garbagemule/MobArena/MonsterManager.java     | 12 ++++++++++++
 2 files changed, 15 insertions(+)

diff --git a/src/main/java/com/garbagemule/MobArena/ArenaListener.java b/src/main/java/com/garbagemule/MobArena/ArenaListener.java
index 24f1aff..67ef4ed 100644
--- a/src/main/java/com/garbagemule/MobArena/ArenaListener.java
+++ b/src/main/java/com/garbagemule/MobArena/ArenaListener.java
@@ -494,6 +494,9 @@ public class ArenaListener
         if (event instanceof PlayerDeathEvent) {
             onPlayerDeath((PlayerDeathEvent) event, (Player) event.getEntity());
         }
+        else if (monsters.hasPet(event.getEntity())) {
+            monsters.removePet(event.getEntity());
+        }
         else if (monsters.removeMonster(event.getEntity())) {
             onMonsterDeath(event);
         }
diff --git a/src/main/java/com/garbagemule/MobArena/MonsterManager.java b/src/main/java/com/garbagemule/MobArena/MonsterManager.java
index 98b2ac6..26e467b 100644
--- a/src/main/java/com/garbagemule/MobArena/MonsterManager.java
+++ b/src/main/java/com/garbagemule/MobArena/MonsterManager.java
@@ -131,6 +131,18 @@ public class MonsterManager
         return petToPlayer.containsKey(e);
     }
 
+    public void removePet(Entity pet) {
+        pet.remove();
+
+        Player owner = petToPlayer.remove(pet);
+        if (owner != null) {
+            Set<Entity> pets = playerToPets.get(owner);
+            if (pets != null) {
+                pets.remove(pet);
+            }
+        }
+    }
+
     public Player getOwner(Entity pet) {
         return petToPlayer.get(pet);
     }