DRILL-7317: Close ClassLoaders used for udf jars uploading when closi…

…ng FunctionImplementationRegistry

exec/java-exec/src/main/java/org/apache/drill/exec/expr/fn/FunctionImplementationRegistry.java

@@ -455,7 +455,7 @@ private ScanResult scan(ClassLoader classLoader, Path path, URL[] urls) throws I
           return RunTimeScan.dynamicPackageScan(drillConfig, Sets.newHashSet(urls));
         } finally {
           if (markerFileConnection instanceof JarURLConnection) {
-            ((JarURLConnection) markerFile.openConnection()).getJarFile().close();
+            ((JarURLConnection) markerFileConnection).getJarFile().close();
           }
         }
       }
@@ -592,6 +592,7 @@ private void deleteQuietlyLocalJar(Path jar) {
    */
   @Override
   public void close() {
+    localFunctionRegistry.close();
     if (deleteTmpDir) {
       FileUtils.deleteQuietly(tmpDir);
     } else {

exec/java-exec/src/main/java/org/apache/drill/exec/expr/fn/registry/FunctionRegistryHolder.java

@@ -151,6 +151,20 @@ public void removeJar(String jarName) {
     }
   }
 
+  /**
+   * Removes all jars from {@link #jars} map and all associated with jars functions
+   * from {@link #functions}.
+   * This is write operation, so one user at a time can call perform such action,
+   * others will wait till first user completes his action.
+   */
+  public void removeAllJars() {
+    try (@SuppressWarnings("unused") Closeable lock = writeLock.open()) {
+      jars.keySet().stream()
+          .filter(jarName -> !jarName.equals(LocalFunctionRegistry.BUILT_IN))
+          .forEach(this::removeAllByJar);
+    }
+  }
+
   /**
    * Retrieves list of all jars name present in {@link #jars}
    * This is read operation, so several users can get this data.

exec/java-exec/src/main/java/org/apache/drill/exec/expr/fn/registry/LocalFunctionRegistry.java

@@ -53,7 +53,7 @@
 /**
  * Registry of Drill functions.
  */
-public class LocalFunctionRegistry {
+public class LocalFunctionRegistry implements AutoCloseable {
 
   public static final String BUILT_IN = "built-in";
 
@@ -356,4 +356,10 @@ private void registerOperatorsWithoutInference(DrillOperatorTable operatorTable,
       }
     }
   }
+
+  @Override
+  public void close() {
+    // unregisters all jars to close ClassLoader used for jar uploading
+    registryHolder.removeAllJars();
+  }
 }