From b0d1c4295507799237dc28fd1c3b73dc9a86d9ee Mon Sep 17 00:00:00 2001
From: xGinko <rachel.audrey@mail.de>
Date: Wed, 24 Jul 2024 16:51:26 +0200
Subject: [PATCH] remove from cache on death

---
 .../villageroptimizer/VillagerOptimizer.java  |  1 +
 .../villageroptimizer/WrapperCache.java       | 20 ++++++++++++++++++-
 2 files changed, 20 insertions(+), 1 deletion(-)

diff --git a/src/main/java/me/xginko/villageroptimizer/VillagerOptimizer.java b/src/main/java/me/xginko/villageroptimizer/VillagerOptimizer.java
index 559327d..22d60f3 100644
--- a/src/main/java/me/xginko/villageroptimizer/VillagerOptimizer.java
+++ b/src/main/java/me/xginko/villageroptimizer/VillagerOptimizer.java
@@ -192,6 +192,7 @@ public final class VillagerOptimizer extends JavaPlugin {
             config = new Config();
             if (wrapperCache != null) wrapperCache.disable();
             wrapperCache = new WrapperCache(config.cache_keep_time);
+            wrapperCache.enable();
             VillagerOptimizerCommand.reloadCommands();
             VillagerOptimizerModule.reloadModules();
             config.saveConfig();
diff --git a/src/main/java/me/xginko/villageroptimizer/WrapperCache.java b/src/main/java/me/xginko/villageroptimizer/WrapperCache.java
index a658f80..ff8b2bb 100644
--- a/src/main/java/me/xginko/villageroptimizer/WrapperCache.java
+++ b/src/main/java/me/xginko/villageroptimizer/WrapperCache.java
@@ -3,14 +3,20 @@ package me.xginko.villageroptimizer;
 import com.github.benmanes.caffeine.cache.Cache;
 import com.github.benmanes.caffeine.cache.Caffeine;
 import me.xginko.villageroptimizer.utils.Disableable;
+import me.xginko.villageroptimizer.utils.Enableable;
 import me.xginko.villageroptimizer.wrapper.WrappedVillager;
 import org.bukkit.entity.Villager;
+import org.bukkit.event.EventHandler;
+import org.bukkit.event.EventPriority;
+import org.bukkit.event.HandlerList;
+import org.bukkit.event.Listener;
+import org.bukkit.event.entity.EntityDeathEvent;
 import org.jetbrains.annotations.NotNull;
 
 import java.time.Duration;
 import java.util.UUID;
 
-public final class WrapperCache implements Disableable {
+public final class WrapperCache implements Enableable, Disableable, Listener {
 
     private final @NotNull Cache<UUID, WrappedVillager> wrapperCache;
 
@@ -18,12 +24,24 @@ public final class WrapperCache implements Disableable {
         this.wrapperCache = Caffeine.newBuilder().expireAfterWrite(cacheDuration).build();
     }
 
+    @Override
+    public void enable() {
+        VillagerOptimizer plugin = VillagerOptimizer.getInstance();
+        plugin.getServer().getPluginManager().registerEvents(this, plugin);
+    }
+
     @Override
     public void disable() {
+        HandlerList.unregisterAll(this);
         this.wrapperCache.invalidateAll();
         this.wrapperCache.cleanUp();
     }
 
+    @EventHandler(priority = EventPriority.MONITOR, ignoreCancelled = true)
+    private void onEntityDeath(EntityDeathEvent event) {
+        this.wrapperCache.invalidate(event.getEntity().getUniqueId());
+    }
+
     @SuppressWarnings("DataFlowIssue")
     public @NotNull WrappedVillager get(@NotNull Villager villager) {
         return this.wrapperCache.get(villager.getUniqueId(), k -> new WrappedVillager(villager));