From c45d0bb2f1519056b9a42c31b525600c5ff928ee Mon Sep 17 00:00:00 2001
From: Gemma Lamont <gemma.lamont@neotechnology.com>
Date: Mon, 27 Jan 2025 12:01:37 +0100
Subject: [PATCH] Remove resolvers on db shutdown

---
 .../main/java/apoc/ApocExtensionFactory.java   |  8 ++++++++
 .../java/apoc/RegisterComponentFactory.java    | 18 ++++++++++++++++--
 2 files changed, 24 insertions(+), 2 deletions(-)

diff --git a/common/src/main/java/apoc/ApocExtensionFactory.java b/common/src/main/java/apoc/ApocExtensionFactory.java
index 507d9bed4..959b57ad7 100644
--- a/common/src/main/java/apoc/ApocExtensionFactory.java
+++ b/common/src/main/java/apoc/ApocExtensionFactory.java
@@ -168,6 +168,14 @@ public void stop() {
             registeredListeners.clear();
         }
 
+        @Override
+        public void shutdown() throws Exception {
+            String databaseName = db.databaseName();
+            services.values().forEach(lifecycle -> dependencies
+                    .registerComponentLifecycle()
+                    .cleanUpResolver(databaseName, lifecycle.getClass()));
+        }
+
         public Collection<AvailabilityListener> getRegisteredListeners() {
             return registeredListeners;
         }
diff --git a/common/src/main/java/apoc/RegisterComponentFactory.java b/common/src/main/java/apoc/RegisterComponentFactory.java
index 995eca539..a61c69001 100644
--- a/common/src/main/java/apoc/RegisterComponentFactory.java
+++ b/common/src/main/java/apoc/RegisterComponentFactory.java
@@ -61,9 +61,23 @@ public class RegisterComponentLifecycle extends LifecycleAdapter {
 
         private final Map<Class, Map<String, Object>> resolvers = new ConcurrentHashMap<>();
 
-        public void addResolver(String databaseNamme, Class clazz, Object instance) {
+        public void addResolver(String databaseName, Class clazz, Object instance) {
             Map<String, Object> classInstanceMap = resolvers.computeIfAbsent(clazz, s -> new ConcurrentHashMap<>());
-            classInstanceMap.put(databaseNamme, instance);
+            classInstanceMap.put(databaseName, instance);
+        }
+
+        public void cleanUpResolver(String databaseName, Class clazz) {
+            Map<String, Object> innerMap = resolvers.get(clazz);
+
+            if (innerMap != null) {
+                // Remove the database to instance value
+                innerMap.remove(databaseName);
+
+                // If the inner map is now empty, remove the key from the outer map
+                if (innerMap.isEmpty()) {
+                    resolvers.remove(clazz);
+                }
+            }
         }
 
         @SuppressWarnings("unused") // used from extended