Skip to content

Commit

Permalink
update
Browse files Browse the repository at this point in the history
  • Loading branch information
@ArnavBalyan
ArnavBalyan committed Aug 13, 2024
1 parent 60b02a0 commit c929a21
Showing 1 changed file with 61 additions and 56 deletions.
117 changes: 61 additions & 56 deletions gluten-core/src/main/java/org/apache/gluten/vectorized/JniLibLoader.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -47,8 +47,8 @@
public class JniLibLoader {
private static final Logger LOG = LoggerFactory.getLogger(JniLibLoader.class);

private static final Set<String> LOADED_LIBRARY_PATHS = new HashSet<>();
private static final Set<String> REQUIRE_UNLOAD_LIBRARY_PATHS = new LinkedHashSet<>();
private static final Set<String> LOADED_LIBRARY_PATHS = Collections.synchronizedSet(new HashSet<>());
private static final Set<String> REQUIRE_UNLOAD_LIBRARY_PATHS = Collections.synchronizedSet(new LinkedHashSet<>());

static {
GlutenShutdownManager.addHookForLibUnloading(
Expand All @@ -59,15 +59,17 @@ public class JniLibLoader {
}

private final String workDir;
private final Set<String> loadedLibraries = new HashSet<>();
private final Lock sync = new ReentrantLock();
private final Set<String> loadedLibraries = Collections.synchronizedSet(new HashSet<>());

public JniLibLoader(String workDir) {
JniLibLoader(String workDir) {
this.workDir = workDir;
}

public static synchronized void forceUnloadAll() {
List<String> loaded = new ArrayList<>(REQUIRE_UNLOAD_LIBRARY_PATHS);
public static void forceUnloadAll() {
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);
}
Expand All @@ -85,37 +87,43 @@ private static String toRealPath(String libPath) {
}
}

private static synchronized void loadFromPath0(String libPath, boolean requireUnload) {
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);
synchronized (REQUIRE_UNLOAD_LIBRARY_PATHS) {
REQUIRE_UNLOAD_LIBRARY_PATHS.add(libPath);
}
}
}

public static void loadFromPath(String libPath, boolean requireUnload) {
public static synchronized 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");
}
loadFromPath0(file.getAbsolutePath(), requireUnload);
}

public static synchronized void unloadFromPath(String libPath) {
if (!LOADED_LIBRARY_PATHS.remove(libPath)) {
LOG.warn("Library {} was not loaded or already unloaded:", libPath);
return;
public static void unloadFromPath(String libPath) {
synchronized (LOADED_LIBRARY_PATHS) {
if (!LOADED_LIBRARY_PATHS.remove(libPath)) {
LOG.warn("Library {} was not loaded or already unloaded:", libPath);
return;
}
}

LOG.info("Starting unload library path: {} ", libPath);
REQUIRE_UNLOAD_LIBRARY_PATHS.remove(libPath);

synchronized (REQUIRE_UNLOAD_LIBRARY_PATHS) {
REQUIRE_UNLOAD_LIBRARY_PATHS.remove(libPath);
}
try {
ClassLoader classLoader = JniLibLoader.class.getClassLoader();
Field field = ClassLoader.class.getDeclaredField("nativeLibraries");
Expand Down Expand Up @@ -151,49 +159,46 @@ public static synchronized void unloadFromPath(String libPath) {
}

public void mapAndLoad(String unmappedLibName, boolean requireUnload) {
sync.lock();
try {
final String mappedLibName = System.mapLibraryName(unmappedLibName);
load(mappedLibName, requireUnload);
} catch (Exception e) {
throw new GlutenException(e);
} finally {
sync.unlock();
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) {
sync.lock();
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);
} finally {
sync.unlock();
}
}

public void loadAndCreateLink(String libName, String linkName, boolean requireUnload) {
sync.lock();
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);
} finally {
sync.unlock();
}
}

Expand Down

0 comments on commit c929a21

Please sign in to comment.