From 9c5c2fe5c7300563f3307dbfcb6c3109b658b100 Mon Sep 17 00:00:00 2001
From: Igor Motov <igor@nvidia.com>
Date: Fri, 21 Aug 2026 16:08:48 -1000
Subject: [PATCH] Fix intermittent ThreadLeak failures in the concurrent IT
 tests

`awaitTermination()` returns once the pool reaches TERMINATED, which happens before the workers finish `processWorkerExit()`, so the threads are still alive  and `ThreadLeakLingering` defaults to no wait, leaving `RandomizedRunner` to report dying threads as leaks. Annotate `CuVSTestCase` with `@ThreadLeakLingering(linger = 5000)`, inherited by all four IT classes that create executors. Also make `runConcurrently` shut the pool down on every path rather than only on success, bound the wait at 60 seconds, keep the pooled-memory reset unconditional, and report a failing task by catching `ExecutionException` instead of throwing from `exceptionally()`.
---
 .../nvidia/cuvs/CagraBuildAndSearchIT.java    | 33 ++++++++++---------
 .../java/com/nvidia/cuvs/CuVSTestCase.java    | 15 ++++++++-
 2 files changed, 31 insertions(+), 17 deletions(-)

diff --git a/java/cuvs-java/src/test/java/com/nvidia/cuvs/CagraBuildAndSearchIT.java b/java/cuvs-java/src/test/java/com/nvidia/cuvs/CagraBuildAndSearchIT.java
index e2287c0a22..500ef451f3 100644
--- a/java/cuvs-java/src/test/java/com/nvidia/cuvs/CagraBuildAndSearchIT.java
+++ b/java/cuvs-java/src/test/java/com/nvidia/cuvs/CagraBuildAndSearchIT.java
@@ -55,7 +55,8 @@ public void setup() {
   private static void runConcurrently(
       boolean usePooledMemory, int nThreads, Function<Integer, Runnable> runnableSupplier)
       throws ExecutionException, InterruptedException, TimeoutException {
-    try (ExecutorService parallelExecutor = Executors.newFixedThreadPool(nThreads)) {
+    ExecutorService parallelExecutor = Executors.newFixedThreadPool(nThreads);
+    try {
       if (usePooledMemory) {
         CuVSProvider.provider().enableRMMPooledMemory(10, 60);
       }
@@ -64,22 +65,22 @@ private static void runConcurrently(
         futures[j] = CompletableFuture.runAsync(runnableSupplier.apply(j), parallelExecutor);
       }
 
-      CompletableFuture.allOf(futures)
-          .exceptionally(
-              t -> {
-                log.error("Exception while executing runnable", t);
-                fail("Exception while executing runnable: " + unwrap(t));
-                return null;
-              })
-          .get(2000, TimeUnit.SECONDS);
-
-      parallelExecutor.shutdown();
-      assertTrue(
-          "Timeout waiting for parallelExecutor to finish",
-          parallelExecutor.awaitTermination(10, TimeUnit.SECONDS));
+      try {
+        CompletableFuture.allOf(futures).get(2000, TimeUnit.SECONDS);
+      } catch (ExecutionException e) {
+        log.error("Exception while executing runnable", e);
+        fail("Exception while executing runnable: " + unwrap(e));
+      }
     } finally {
-      if (usePooledMemory) {
-        CuVSProvider.provider().resetRMMPooledMemory();
+      try {
+        parallelExecutor.shutdown();
+        assertTrue(
+            "Timeout waiting for parallelExecutor to finish",
+            parallelExecutor.awaitTermination(60, TimeUnit.SECONDS));
+      } finally {
+        if (usePooledMemory) {
+          CuVSProvider.provider().resetRMMPooledMemory();
+        }
       }
     }
   }
diff --git a/java/cuvs-java/src/test/java/com/nvidia/cuvs/CuVSTestCase.java b/java/cuvs-java/src/test/java/com/nvidia/cuvs/CuVSTestCase.java
index f2c2840b03..62e636eb82 100644
--- a/java/cuvs-java/src/test/java/com/nvidia/cuvs/CuVSTestCase.java
+++ b/java/cuvs-java/src/test/java/com/nvidia/cuvs/CuVSTestCase.java
@@ -1,5 +1,5 @@
 /*
- * SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION.
+ * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
  * SPDX-License-Identifier: Apache-2.0
  */
 package com.nvidia.cuvs;
@@ -9,6 +9,7 @@
 import static org.junit.Assert.assertTrue;
 
 import com.carrotsearch.randomizedtesting.RandomizedContext;
+import com.carrotsearch.randomizedtesting.annotations.ThreadLeakLingering;
 import java.lang.invoke.MethodHandles;
 import java.util.ArrayList;
 import java.util.BitSet;
@@ -19,6 +20,18 @@
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
+/**
+ * Base class for the cuVS integration tests.
+ *
+ * <p>The lingering is not cosmetic. A {@link java.util.concurrent.ThreadPoolExecutor} reaches
+ * TERMINATED as soon as its worker count drops to zero, which happens in {@code getTask()} before
+ * the workers have finished {@code processWorkerExit()}, so {@code awaitTermination()} returns
+ * while the worker threads are still alive. {@code ThreadLeakLingering} defaults to no wait at
+ * all, which leaves the leak check sampling straight into that window and reporting threads that
+ * are in the middle of dying. Waiting a few seconds for them costs nothing when there is no leak,
+ * and a real one still fails the test.
+ */
+@ThreadLeakLingering(linger = 5000)
 public abstract class CuVSTestCase {
   protected Random random;
   private static final Logger log = LoggerFactory.getLogger(MethodHandles.lookup().lookupClass());