diff --git a/java/src/main/java/ai/rapids/cudf/NativeDepsLoader.java b/java/src/main/java/ai/rapids/cudf/NativeDepsLoader.java index 4b322c5e7730..2b0a53530113 100755 --- a/java/src/main/java/ai/rapids/cudf/NativeDepsLoader.java +++ b/java/src/main/java/ai/rapids/cudf/NativeDepsLoader.java @@ -45,6 +45,7 @@ */ public class NativeDepsLoader { private static final int COPY_BUFFER_SIZE = 1024 * 1024; + private static final long EXECUTOR_SHUTDOWN_TIMEOUT_SECONDS = 10; // Positional extraction uses one copy buffer per worker. private static final int MAX_CONCURRENT_CHUNK_READS = Math.max(1, Math.min(12, Runtime.getRuntime().availableProcessors())); @@ -289,51 +290,60 @@ private static void loadNativeDeps(String[][] loadOrder, boolean preserveDeps) t Map timings = libLogLoadTiming ? new ConcurrentHashMap<>() : null; ExecutorService executor = Executors.newCachedThreadPool(); - List>> allFileFutures = new ArrayList<>(); - - // Start unpacking and creating the temporary files for each dependency. - // Unpacking a dependency does not depend on stage order. - for (String[] stageDependencies : loadOrder) { - List> stageFileFutures = new ArrayList<>(); - allFileFutures.add(stageFileFutures); - for (String name : stageDependencies) { - stageFileFutures.add(executor.submit(() -> createFileTimed(os, arch, name, timings))); + try { + List>> allFileFutures = new ArrayList<>(); + + // Start unpacking and creating the temporary files for each dependency. + // Unpacking a dependency does not depend on stage order. + for (String[] stageDependencies : loadOrder) { + List> stageFileFutures = new ArrayList<>(); + allFileFutures.add(stageFileFutures); + for (String name : stageDependencies) { + stageFileFutures.add(executor.submit(() -> createFileTimed(os, arch, name, timings))); + } } - } - List> loadCompletionFutures = new ArrayList<>(); - - // Proceed stage-by-stage waiting for the dependency file to have been - // produced then submit them to the thread pool to be loaded. - for (int i = 0; i < allFileFutures.size(); i++) { - List> stageFileFutures = allFileFutures.get(i); - String[] stageNames = loadOrder[i]; - // Submit all dependencies in the stage to be loaded in parallel - loadCompletionFutures.clear(); - for (int j = 0; j < stageFileFutures.size(); j++) { - Future fileFuture = stageFileFutures.get(j); - String name = stageNames[j]; - loadCompletionFutures.add( - executor.submit(() -> loadDepTimed(fileFuture, preserveDeps, name, timings))); - } + List> loadCompletionFutures = new ArrayList<>(); + + // Proceed stage-by-stage waiting for the dependency file to have been + // produced then submit them to the thread pool to be loaded. + for (int i = 0; i < allFileFutures.size(); i++) { + List> stageFileFutures = allFileFutures.get(i); + String[] stageNames = loadOrder[i]; + // Submit all dependencies in the stage to be loaded in parallel + loadCompletionFutures.clear(); + for (int j = 0; j < stageFileFutures.size(); j++) { + Future fileFuture = stageFileFutures.get(j); + String name = stageNames[j]; + loadCompletionFutures.add( + executor.submit(() -> loadDepTimed(fileFuture, preserveDeps, name, timings))); + } - // Wait for all dependencies in this stage to have been loaded - for (Future loadCompletionFuture : loadCompletionFutures) { - try { - loadCompletionFuture.get(); - } catch (ExecutionException | InterruptedException e) { - throw new IOException("Error loading dependencies", e); + // Wait for all dependencies in this stage to have been loaded + for (Future loadCompletionFuture : loadCompletionFutures) { + awaitLoadCompletion(loadCompletionFuture); } } + } finally { + shutdownAndAwait(executor); } - executor.shutdownNow(); - if (libLogLoadTiming) { logLoadSummary(loadOrder, timings, System.currentTimeMillis() - t0); } } + static void awaitLoadCompletion(Future loadCompletionFuture) throws IOException { + try { + loadCompletionFuture.get(); + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + throw new IOException("Interrupted while loading dependencies", e); + } catch (ExecutionException e) { + throw new IOException("Error loading dependencies", e); + } + } + /** * Allows other libraries to reuse the same native deps loading logic. Library will be searched * for under ${os.arch}/${os.name}/ in the class path using the class loader for this class. @@ -384,7 +394,10 @@ private static void loadDepTimed(Future fileFuture, boolean preserveDep, File path; try { path = fileFuture.get(); - } catch (ExecutionException | InterruptedException e) { + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + throw new RuntimeException("Interrupted while loading dependencies", e); + } catch (ExecutionException e) { throw new RuntimeException("Error loading dependencies", e); } long t0 = System.currentTimeMillis(); @@ -612,12 +625,20 @@ private static void awaitChunks(List> futures, String mappedName) private static void shutdownAndAwait(ExecutorService executor) { executor.shutdownNow(); boolean interrupted = Thread.interrupted(); - while (!executor.isTerminated()) { + long deadline = System.nanoTime() + + TimeUnit.SECONDS.toNanos(EXECUTOR_SHUTDOWN_TIMEOUT_SECONDS); + long remainingNanos = deadline - System.nanoTime(); + while (!executor.isTerminated() && remainingNanos > 0) { try { - executor.awaitTermination(1, TimeUnit.SECONDS); + executor.awaitTermination(remainingNanos, TimeUnit.NANOSECONDS); } catch (InterruptedException e) { interrupted = true; } + remainingNanos = deadline - System.nanoTime(); + } + if (!executor.isTerminated()) { + Log.warn("Timed out after {} seconds waiting for native dependency tasks to stop", + EXECUTOR_SHUTDOWN_TIMEOUT_SECONDS); } if (interrupted) { Thread.currentThread().interrupt(); diff --git a/java/src/test/java/ai/rapids/cudf/NativeDepsLoaderTest.java b/java/src/test/java/ai/rapids/cudf/NativeDepsLoaderTest.java index 0fbad5a3a393..907a0f29a3a4 100644 --- a/java/src/test/java/ai/rapids/cudf/NativeDepsLoaderTest.java +++ b/java/src/test/java/ai/rapids/cudf/NativeDepsLoaderTest.java @@ -15,6 +15,9 @@ import java.nio.file.Files; import java.nio.file.Path; import java.util.Comparator; +import java.util.HashSet; +import java.util.Set; +import java.util.concurrent.CompletableFuture; import java.util.stream.Stream; import static org.junit.jupiter.api.Assertions.assertEquals; @@ -130,4 +133,39 @@ void noArgLoad_failsSilently_andLeavesLibraryNotLoaded() { assertFalse(NativeDepsLoader.getLoaded(), "loaded flag should remain false after a failed load"); } + + @Test + void stagedLoadFailure_doesNotLeakExecutorThreads() throws IOException { + Files.createFile(libDir.resolve("libcudf.so")); + Files.createFile(libDir.resolve("libcudfjni.so")); + Set threadsBefore = new HashSet<>(Thread.getAllStackTraces().keySet()); + + NativeDepsLoader.loadNativeDeps(); + + assertFalse(NativeDepsLoader.getLoaded(), + "loaded flag should remain false after a failed load"); + Set newNonDaemonThreads = new HashSet<>(); + for (Thread thread : Thread.getAllStackTraces().keySet()) { + if (thread.isAlive() && !thread.isDaemon() && !threadsBefore.contains(thread)) { + newNonDaemonThreads.add(thread.getName()); + } + } + assertTrue(newNonDaemonThreads.isEmpty(), + "native dependency loading leaked threads: " + newNonDaemonThreads); + } + + @Test + void awaitLoadCompletion_preservesInterrupt() { + CompletableFuture incomplete = new CompletableFuture<>(); + Thread.currentThread().interrupt(); + try { + IOException ex = assertThrows(IOException.class, + () -> NativeDepsLoader.awaitLoadCompletion(incomplete)); + assertTrue(ex.getCause() instanceof InterruptedException); + assertTrue(Thread.currentThread().isInterrupted(), + "interrupted status should be restored"); + } finally { + Thread.interrupted(); + } + } }