From a6c88ed5060694fd6149b804717669cb41179b2b Mon Sep 17 00:00:00 2001
From: arnavb <arnavb@uber.com>
Date: Sat, 17 Aug 2024 15:54:28 +0000
Subject: [PATCH] update

---
 .../gluten/vectorized/JniLibLoader.java       | 77 +++++++++++--------
 1 file changed, 44 insertions(+), 33 deletions(-)

diff --git a/gluten-core/src/main/java/org/apache/gluten/vectorized/JniLibLoader.java b/gluten-core/src/main/java/org/apache/gluten/vectorized/JniLibLoader.java
index 04258e935778..ddb5c46cdeda 100644
--- a/gluten-core/src/main/java/org/apache/gluten/vectorized/JniLibLoader.java
+++ b/gluten-core/src/main/java/org/apache/gluten/vectorized/JniLibLoader.java
@@ -66,7 +66,10 @@ public class JniLibLoader {
   }
 
   public static void forceUnloadAll() {
-    List<String> loaded = new ArrayList<>(REQUIRE_UNLOAD_LIBRARY_PATHS);
+    List<String> loaded;
+    synchronized (REQUIRE_UNLOAD_LIBRARY_PATHS) {
+      loaded = new ArrayList<>(REQUIRE_UNLOAD_LIBRARY_PATHS);
+    }
     Collections.reverse(loaded); // use reversed order to unload
     loaded.forEach(JniLibLoader::unloadFromPath);
   }
@@ -86,19 +89,21 @@ private static String toRealPath(String libPath) {
 
   private static void loadFromPath0(String libPath, boolean requireUnload) {
     libPath = toRealPath(libPath);
-    if (LOADED_LIBRARY_PATHS.contains(libPath)) {
-      LOG.debug("Library in path {} has already been loaded, skipping", libPath);
-    } else {
-      System.load(libPath);
-      LOADED_LIBRARY_PATHS.add(libPath);
-      LOG.info("Library {} has been loaded using path-loading method", libPath);
+    synchronized (LOADED_LIBRARY_PATHS) {
+      if (LOADED_LIBRARY_PATHS.contains(libPath)) {
+        LOG.debug("Library in path {} has already been loaded, skipping", libPath);
+      } else {
+        System.load(libPath);
+        LOADED_LIBRARY_PATHS.add(libPath);
+        LOG.info("Library {} has been loaded using path-loading method", libPath);
+      }
     }
     if (requireUnload) {
       REQUIRE_UNLOAD_LIBRARY_PATHS.add(libPath);
     }
   }
 
-  public static synchronized void loadFromPath(String libPath, boolean requireUnload) {
+  public static void loadFromPath(String libPath, boolean requireUnload) {
     final File file = new File(libPath);
     if (!file.isFile() || !file.exists()) {
       throw new GlutenException("library at path: " + libPath + " is not a file or does not exist");
@@ -113,7 +118,6 @@ public static void unloadFromPath(String libPath) {
     }
     LOG.info("Starting unload library path: {} ", libPath);
     REQUIRE_UNLOAD_LIBRARY_PATHS.remove(libPath);
-
     try {
       ClassLoader classLoader = JniLibLoader.class.getClassLoader();
       Field field = ClassLoader.class.getDeclaredField("nativeLibraries");
@@ -149,44 +153,51 @@ public static void unloadFromPath(String libPath) {
   }
 
   public void mapAndLoad(String unmappedLibName, boolean requireUnload) {
-    try {
-      final String mappedLibName = System.mapLibraryName(unmappedLibName);
-      load(mappedLibName, requireUnload);
-    } catch (Exception e) {
-      throw new GlutenException(e);
+    synchronized (loadedLibraries) {
+      try {
+        final String mappedLibName = System.mapLibraryName(unmappedLibName);
+        load(mappedLibName, requireUnload);
+      } catch (Exception e) {
+        throw new GlutenException(e);
+      }
     }
   }
 
   public void load(String libName, boolean requireUnload) {
-    try {
-      if (loadedLibraries.contains(libName)) {
-        LOG.debug("Library {} has already been loaded, skipping", libName);
-        return;
+    synchronized (loadedLibraries) {
+      try {
+        if (loadedLibraries.contains(libName)) {
+          LOG.debug("Library {} has already been loaded, skipping", libName);
+          return;
+        }
+        File file = moveToWorkDir(workDir, libName);
+        loadWithLink(file.getAbsolutePath(), null, requireUnload);
+        loadedLibraries.add(libName);
+        LOG.info("Successfully loaded library {}", libName);
+      } catch (IOException e) {
+        throw new GlutenException(e);
       }
-      File file = moveToWorkDir(workDir, libName);
-      loadWithLink(file.getAbsolutePath(), null, requireUnload);
-      loadedLibraries.add(libName);
-      LOG.info("Successfully loaded library {}", libName);
-    } catch (IOException e) {
-      throw new GlutenException(e);
     }
   }
 
   public void loadAndCreateLink(String libName, String linkName, boolean requireUnload) {
-    try {
-      if (loadedLibraries.contains(libName)) {
-        LOG.debug("Library {} has already been loaded, skipping", libName);
+    synchronized (loadedLibraries) {
+      try {
+        if (loadedLibraries.contains(libName)) {
+          LOG.debug("Library {} has already been loaded, skipping", libName);
+        }
+        File file = moveToWorkDir(workDir, System.mapLibraryName(libName));
+        loadWithLink(file.getAbsolutePath(), linkName, requireUnload);
+        loadedLibraries.add(libName);
+        LOG.info("Successfully loaded library {}", libName);
+      } catch (IOException e) {
+        throw new GlutenException(e);
       }
-      File file = moveToWorkDir(workDir, System.mapLibraryName(libName));
-      loadWithLink(file.getAbsolutePath(), linkName, requireUnload);
-      loadedLibraries.add(libName);
-      LOG.info("Successfully loaded library {}", libName);
-    } catch (IOException e) {
-      throw new GlutenException(e);
     }
   }
 
   private File moveToWorkDir(String workDir, String libraryToLoad) throws IOException {
+    // final File temp = File.createTempFile(workDir, libraryToLoad);
     final Path libPath = Paths.get(workDir + "/" + libraryToLoad);
     if (Files.exists(libPath)) {
       Files.delete(libPath);