[FLINK-23029] Extend PluginLoader to allow closing ClassLoader

akalash · Jun 22, 2021 · 6168118 · 6168118
1 parent 3fdcb40
commit 6168118
Show file tree

Hide file tree

Showing 2 changed files with 45 additions and 6 deletions.
diff --git a/flink-core/src/main/java/org/apache/flink/core/plugin/PluginLoader.java b/flink-core/src/main/java/org/apache/flink/core/plugin/PluginLoader.java
@@ -22,6 +22,9 @@
 import org.apache.flink.util.ArrayUtils;
 import org.apache.flink.util.TemporaryClassLoaderContext;
 
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
 import javax.annotation.concurrent.ThreadSafe;
 
 import java.io.IOException;
@@ -40,21 +43,26 @@
  * construction.
  */
 @ThreadSafe
-public class PluginLoader {
+public class PluginLoader implements AutoCloseable {
+
+    private static final Logger LOG = LoggerFactory.getLogger(PluginLoader.class);
+
+    private final String pluginId;
 
     /**
      * Classloader which is used to load the plugin classes. We expect this classloader is
      * thread-safe.
      */
-    private final ClassLoader pluginClassLoader;
+    private final URLClassLoader pluginClassLoader;
 
     @VisibleForTesting
-    public PluginLoader(ClassLoader pluginClassLoader) {
+    public PluginLoader(String pluginId, URLClassLoader pluginClassLoader) {
+        this.pluginId = pluginId;
         this.pluginClassLoader = pluginClassLoader;
     }
 
     @VisibleForTesting
-    public static ClassLoader createPluginClassLoader(
+    public static URLClassLoader createPluginClassLoader(
             PluginDescriptor pluginDescriptor,
             ClassLoader parentClassLoader,
             String[] alwaysParentFirstPatterns) {
@@ -70,6 +78,7 @@ public static PluginLoader create(
             ClassLoader parentClassLoader,
             String[] alwaysParentFirstPatterns) {
         return new PluginLoader(
+                pluginDescriptor.getPluginId(),
                 createPluginClassLoader(
                         pluginDescriptor, parentClassLoader, alwaysParentFirstPatterns));
     }
@@ -91,6 +100,15 @@ public <P> Iterator<P> load(Class<P> service) {
         }
     }
 
+    @Override
+    public void close() {
+        try {
+            pluginClassLoader.close();
+        } catch (IOException e) {
+            LOG.warn("An error occurred while closing the classloader for plugin {}.", pluginId);
+        }
+    }
+
     /**
      * Wrapper for the service iterator. The wrapper will set/unset the context classloader to the
      * plugin classloader around the point where elements are returned.

diff --git a/flink-tests/src/test/java/org/apache/flink/test/plugin/PluginLoaderTest.java b/flink-tests/src/test/java/org/apache/flink/test/plugin/PluginLoaderTest.java
@@ -25,7 +25,9 @@
 import org.junit.Assert;
 import org.junit.Test;
 
+import java.net.MalformedURLException;
 import java.net.URL;
+import java.net.URLClassLoader;
 import java.util.Iterator;
 
 /** Test for {@link PluginLoader}. */
@@ -39,11 +41,11 @@ public void testPluginLoading() throws Exception {
         String[] parentPatterns = {TestSpi.class.getName(), OtherTestSpi.class.getName()};
         PluginDescriptor pluginDescriptorA =
                 new PluginDescriptor("A", new URL[] {classpathA}, parentPatterns);
-        ClassLoader pluginClassLoaderA =
+        URLClassLoader pluginClassLoaderA =
                 PluginLoader.createPluginClassLoader(
                         pluginDescriptorA, PARENT_CLASS_LOADER, new String[0]);
         Assert.assertNotEquals(PARENT_CLASS_LOADER, pluginClassLoaderA);
-        final PluginLoader pluginLoaderA = new PluginLoader(pluginClassLoaderA);
+        final PluginLoader pluginLoaderA = new PluginLoader("test-plugin", pluginClassLoaderA);
 
         Iterator<TestSpi> testSpiIteratorA = pluginLoaderA.load(TestSpi.class);
 
@@ -80,4 +82,23 @@ public void testPluginLoading() throws Exception {
                 secondTestSpiA.getClass().getCanonicalName());
         Assert.assertNotEquals(testSpiA.getClass(), secondTestSpiA.getClass());
     }
+
+    @Test
+    public void testClose() throws MalformedURLException {
+        final URL classpathA = createPluginJarURLFromString(PLUGIN_A);
+
+        String[] parentPatterns = {TestSpi.class.getName()};
+        PluginDescriptor pluginDescriptorA =
+                new PluginDescriptor("A", new URL[] {classpathA}, parentPatterns);
+        URLClassLoader pluginClassLoaderA =
+                PluginLoader.createPluginClassLoader(
+                        pluginDescriptorA, PARENT_CLASS_LOADER, new String[0]);
+
+        final PluginLoader pluginLoaderA = new PluginLoader("test-plugin", pluginClassLoaderA);
+        pluginLoaderA.close();
+
+        Assert.assertThrows(
+                ClassNotFoundException.class,
+                () -> pluginClassLoaderA.loadClass(junit.framework.Test.class.getName()));
+    }
 }