From c2d9998350b3a15ff7678c9053bca434463cf915 Mon Sep 17 00:00:00 2001
From: Ajith Kumar <ajithalbus@gmail.com>
Date: Tue, 20 Oct 2020 14:36:15 +0530
Subject: [PATCH] Bucketed caching in SingletonExtensionFactory (#402)

---
 .../org/pf4j/SingletonExtensionFactory.java   | 20 +++++++---
 .../pf4j/SingletonExtensionFactoryTest.java   | 40 +++++++++++++++++++
 2 files changed, 55 insertions(+), 5 deletions(-)

diff --git a/pf4j/src/main/java/org/pf4j/SingletonExtensionFactory.java b/pf4j/src/main/java/org/pf4j/SingletonExtensionFactory.java
index 725b9a7..ff94652 100644
--- a/pf4j/src/main/java/org/pf4j/SingletonExtensionFactory.java
+++ b/pf4j/src/main/java/org/pf4j/SingletonExtensionFactory.java
@@ -19,36 +19,46 @@ 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.
  * Optional, you can specify the extension classes for which you want singletons.
  *
  * @author Decebal Suiu
+ * @author Ajith Kumar
  */
 public class SingletonExtensionFactory extends DefaultExtensionFactory {
 
     private final List<String> extensionClassNames;
 
-    private Map<String, Object> cache;
+    private Map<ClassLoader, Map<String, Object>> cache;
 
     public SingletonExtensionFactory(String... extensionClassNames) {
         this.extensionClassNames = Arrays.asList(extensionClassNames);
 
-        cache = new HashMap<>(); // simple cache implementation
+        cache = new WeakHashMap<>(); // simple cache implementation
     }
 
     @Override
     @SuppressWarnings("unchecked")
     public <T> T create(Class<T> extensionClass) {
         String extensionClassName = extensionClass.getName();
-        if (cache.containsKey(extensionClassName)) {
-            return (T) cache.get(extensionClassName);
+        ClassLoader extensionClassLoader = extensionClass.getClassLoader();
+
+        if (!cache.containsKey(extensionClassLoader)) {
+            cache.put(extensionClassLoader, new HashMap<>());
+        }
+
+        Map<String, Object> classLoaderBucket = cache.get(extensionClassLoader);
+
+        if (classLoaderBucket.containsKey(extensionClassName)) {
+            return (T) classLoaderBucket.get(extensionClassName);
         }
 
         T extension = super.create(extensionClass);
         if (extensionClassNames.isEmpty() || extensionClassNames.contains(extensionClassName)) {
-            cache.put(extensionClassName, extension);
+            classLoaderBucket.put(extensionClassName, extension);
         }
 
         return extension;
diff --git a/pf4j/src/test/java/org/pf4j/SingletonExtensionFactoryTest.java b/pf4j/src/test/java/org/pf4j/SingletonExtensionFactoryTest.java
index 80e7ab0..ac8a3c5 100644
--- a/pf4j/src/test/java/org/pf4j/SingletonExtensionFactoryTest.java
+++ b/pf4j/src/test/java/org/pf4j/SingletonExtensionFactoryTest.java
@@ -19,11 +19,17 @@ import org.junit.jupiter.api.Test;
 import org.pf4j.plugin.FailTestExtension;
 import org.pf4j.plugin.TestExtension;
 
+import java.io.File;
+import java.net.MalformedURLException;
+import java.net.URL;
+import java.net.URLClassLoader;
+
 import static org.junit.jupiter.api.Assertions.assertNotSame;
 import static org.junit.jupiter.api.Assertions.assertSame;
 
 /**
  * @author Decebal Suiu
+ * @author Ajith Kumar
  */
 public class SingletonExtensionFactoryTest {
 
@@ -43,4 +49,38 @@ public class SingletonExtensionFactoryTest {
         assertNotSame(extensionOne, extensionTwo);
     }
 
+    @Test
+    @SuppressWarnings("unchecked")
+    public void createNewEachTimeFromDifferentClassLoaders() throws Exception {
+        ExtensionFactory extensionFactory = new SingletonExtensionFactory();
+
+        // Get classpath locations
+        URL[] classpathReferences = getClasspathReferences();
+
+        // Create different classloaders for the classpath references and load classes respectively
+        ClassLoader klassLoaderOne = new URLClassLoader(classpathReferences, null);
+        Class klassOne = klassLoaderOne.loadClass(TestExtension.class.getName());
+        ClassLoader klassLoaderTwo = new URLClassLoader(classpathReferences, null);
+        Class klassTwo = klassLoaderTwo.loadClass(TestExtension.class.getName());
+
+        // create instances
+        Object instanceOne = extensionFactory.create(klassOne);
+        Object instanceTwo = extensionFactory.create(klassTwo);
+
+        // assert that instances not same
+        assertNotSame(instanceOne, instanceTwo);
+    }
+
+    private URL[] getClasspathReferences() throws MalformedURLException {
+        String classpathProperty = System.getProperty("java.class.path");
+
+        String[] classpaths = classpathProperty.split(":");
+        URL[] uris = new URL[classpaths.length];
+
+        for (int index = 0; index < classpaths.length; index++) {
+            uris[index] = new File(classpaths[index]).toURI().toURL();
+        }
+        return uris;
+    }
+
 }
-- 
2.39.5