Fix memory leak in SingletonExtensionFactory (#487) (#490)

2 files changed, 26 insertions, 7 deletions

diff --git a/pf4j/src/main/java/org/pf4j/SingletonExtensionFactory.java b/pf4j/src/main/java/org/pf4j/SingletonExtensionFactory.java
index ff94652..0eef139 100644
--- a/pf4j/src/main/java/org/pf4j/SingletonExtensionFactory.java
+++ b/pf4j/src/main/java/org/pf4j/SingletonExtensionFactory.java
@@ -19,7 +19,6 @@ import java.util.Arrays;
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
-import java.util.WeakHashMap;
 
 /**
  * An {@link ExtensionFactory} that always returns a specific instance.
@@ -32,12 +31,18 @@ public class SingletonExtensionFactory extends DefaultExtensionFactory {
 
     private final List<String> extensionClassNames;
 
-    private Map<ClassLoader, Map<String, Object>> cache;
+    private final Map<ClassLoader, Map<String, Object>> cache;
 
-    public SingletonExtensionFactory(String... extensionClassNames) {
+    public SingletonExtensionFactory(PluginManager pluginManager, String... extensionClassNames) {
         this.extensionClassNames = Arrays.asList(extensionClassNames);
 
-        cache = new WeakHashMap<>(); // simple cache implementation
+        cache = new HashMap<>();
+
+        pluginManager.addPluginStateListener(event -> {
+            if (event.getPluginState() != PluginState.STARTED) {
+                cache.remove(event.getPlugin().getPluginClassLoader());
+            }
+        });
     }
 
     @Override
diff --git a/pf4j/src/test/java/org/pf4j/SingletonExtensionFactoryTest.java b/pf4j/src/test/java/org/pf4j/SingletonExtensionFactoryTest.java
index 739a76f..ddfd2a4 100644
--- a/pf4j/src/test/java/org/pf4j/SingletonExtensionFactoryTest.java
+++ b/pf4j/src/test/java/org/pf4j/SingletonExtensionFactoryTest.java
@@ -15,6 +15,8 @@
  */
 package org.pf4j;
 
+import org.junit.jupiter.api.AfterEach;
+import org.junit.jupiter.api.BeforeEach;
 import org.junit.jupiter.api.Test;
 import org.pf4j.test.FailTestExtension;
 import org.pf4j.test.TestExtension;
@@ -33,9 +35,21 @@ import static org.junit.jupiter.api.Assertions.assertSame;
  */
 public class SingletonExtensionFactoryTest {
 
+    private PluginManager pluginManager;
+
+    @BeforeEach
+    public void setUp() {
+        pluginManager = new DefaultPluginManager();
+    }
+
+    @AfterEach
+    public void tearDown() {
+        pluginManager = null;
+    }
+
     @Test
     public void create() {
-        ExtensionFactory extensionFactory = new SingletonExtensionFactory();
+        ExtensionFactory extensionFactory = new SingletonExtensionFactory(pluginManager);
         Object extensionOne = extensionFactory.create(TestExtension.class);
         Object extensionTwo = extensionFactory.create(TestExtension.class);
         assertSame(extensionOne, extensionTwo);
@@ -43,7 +57,7 @@ public class SingletonExtensionFactoryTest {
 
     @Test
     public void createNewEachTime() {
-        ExtensionFactory extensionFactory = new SingletonExtensionFactory(FailTestExtension.class.getName());
+        ExtensionFactory extensionFactory = new SingletonExtensionFactory(pluginManager, FailTestExtension.class.getName());
         Object extensionOne = extensionFactory.create(TestExtension.class);
         Object extensionTwo = extensionFactory.create(TestExtension.class);
         assertNotSame(extensionOne, extensionTwo);
@@ -52,7 +66,7 @@ public class SingletonExtensionFactoryTest {
     @Test
     @SuppressWarnings("unchecked")
     public void createNewEachTimeFromDifferentClassLoaders() throws Exception {
-        ExtensionFactory extensionFactory = new SingletonExtensionFactory();
+        ExtensionFactory extensionFactory = new SingletonExtensionFactory(pluginManager);
 
         // Get classpath locations
         URL[] classpathReferences = getClasspathReferences();