Remove resolvers on db shutdown

neo4j · Jan 28, 2025 · c45d0bb · c45d0bb
1 parent 4ecc788
commit c45d0bb
Show file tree

Hide file tree

Showing 2 changed files with 24 additions and 2 deletions.
diff --git 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
@@ -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