[Java][Fix] Multithreaded querying fails without synchronization (#1082)

Fix CuVSResources usage for multithreaded queries
    
    * Added resources as a constructor argument for Builder classes in CagraIndex, HNSWIndex, TieredIndex and BruteForceIndex
    * Added a multithreaded CAGRA stability test (which fails with improper resources usage). This test uses CheckedCuVSResources.

Authors:
  - Ishan Chattopadhyaya (https://github.com/chatman)

Approvers:
  - MithunR (https://github.com/mythrocks)

URL: #1082

Author: chatman

java/cuvs-java/src/main/java/com/nvidia/cuvs/BruteForceQuery.java

@@ -17,11 +17,16 @@
 
 import java.util.Arrays;
 import java.util.BitSet;
+import java.util.Objects;
 import java.util.function.LongToIntFunction;
 
 /**
  * BruteForceQuery holds the query vectors to be used while invoking search.
  *
+ * <p><strong>Thread Safety:</strong> Each BruteForceQuery instance should use its own
+ * CuVSResources object that is not shared with other threads. Sharing CuVSResources
+ * between threads can lead to memory allocation errors or JVM crashes.
+ *
  * @since 25.02
  */
 public class BruteForceQuery {
@@ -31,6 +36,7 @@ public class BruteForceQuery {
   private final BitSet[] prefilters;
   private int numDocs = -1;
   private final int topK;
+  private final CuVSResources resources;
 
   /**
    * Constructs an instance of {@link BruteForceQuery} using queryVectors,
@@ -43,18 +49,21 @@ public class BruteForceQuery {
    *                     index
    * @param numDocs      Maximum of bits in each prefilter, representing number of documents in this index.
    *                     Used only when prefilter(s) is/are passed.
+   * @param resources    CuVSResources instance to use for this query
    */
   public BruteForceQuery(
       float[][] queryVectors,
       LongToIntFunction mapping,
       int topK,
       BitSet[] prefilters,
-      int numDocs) {
+      int numDocs,
+      CuVSResources resources) {
     this.queryVectors = queryVectors;
     this.mapping = mapping;
     this.topK = topK;
     this.prefilters = prefilters;
     this.numDocs = numDocs;
+    this.resources = resources;
   }
 
   /**
@@ -100,6 +109,15 @@ public int getNumDocs() {
     return numDocs;
   }
 
+  /**
+   * Gets the CuVSResources instance for this query.
+   *
+   * @return the CuVSResources instance
+   */
+  public CuVSResources getResources() {
+    return resources;
+  }
+
   @Override
   public String toString() {
     return "BruteForceQuery [mapping="
@@ -123,6 +141,20 @@ public static class Builder {
     private int numDocs;
     private LongToIntFunction mapping = SearchResults.IDENTITY_MAPPING;
     private int topK = 2;
+    private final CuVSResources resources;
+
+    /**
+     * Constructor that requires CuVSResources.
+     *
+     * <p><strong>Important:</strong> The provided CuVSResources instance should not be
+     * shared with other threads. Each thread performing searches should create its own
+     * CuVSResources instance to avoid memory allocation conflicts and potential JVM crashes.
+     *
+     * @param resources the CuVSResources instance to use for this query (must not be shared between threads)
+     */
+    public Builder(CuVSResources resources) {
+      this.resources = Objects.requireNonNull(resources, "resources cannot be null");
+    }
 
     /**
      * Registers the query vectors to be passed in the search call.
@@ -176,7 +208,7 @@ public Builder withPrefilters(BitSet[] prefilters, int numDocs) {
      * @return an instance of {@link BruteForceQuery}
      */
     public BruteForceQuery build() {
-      return new BruteForceQuery(queryVectors, mapping, topK, prefilters, numDocs);
+      return new BruteForceQuery(queryVectors, mapping, topK, prefilters, numDocs, resources);
     }
   }
 }

java/cuvs-java/src/main/java/com/nvidia/cuvs/CagraQuery.java

@@ -17,12 +17,17 @@
 
 import java.util.Arrays;
 import java.util.BitSet;
+import java.util.Objects;
 import java.util.function.LongToIntFunction;
 
 /**
  * CagraQuery holds the CagraSearchParams and the query vectors to be used while
  * invoking search.
  *
+ * <p><strong>Thread Safety:</strong> Each CagraQuery instance should use its own
+ * CuVSResources object that is not shared with other threads. Sharing CuVSResources
+ * between threads can lead to memory allocation errors or JVM crashes.
+ *
  * @since 25.02
  */
 public class CagraQuery {
@@ -33,6 +38,7 @@ public class CagraQuery {
   private final int topK;
   private final BitSet prefilter;
   private final int numDocs;
+  private final CuVSResources resources;
 
   /**
    * Constructs an instance of {@link CagraQuery} using cagraSearchParameters,
@@ -45,21 +51,24 @@ public class CagraQuery {
    * @param topK                  the top k results to return
    * @param prefilter             A single BitSet to use as filter while searching the CAGRA index
    * @param numDocs               Total number of dataset vectors; used to align the prefilter correctly
+   * @param resources             CuVSResources instance to use for this query
    */
   public CagraQuery(
       CagraSearchParams cagraSearchParameters,
       float[][] queryVectors,
       LongToIntFunction mapping,
       int topK,
       BitSet prefilter,
-      int numDocs) {
+      int numDocs,
+      CuVSResources resources) {
     super();
     this.cagraSearchParameters = cagraSearchParameters;
     this.queryVectors = queryVectors;
     this.mapping = mapping;
     this.topK = topK;
     this.prefilter = prefilter;
     this.numDocs = numDocs;
+    this.resources = resources;
   }
 
   /**
@@ -114,6 +123,15 @@ public int getNumDocs() {
     return numDocs;
   }
 
+  /**
+   * Gets the CuVSResources instance for this query.
+   *
+   * @return the CuVSResources instance
+   */
+  public CuVSResources getResources() {
+    return resources;
+  }
+
   @Override
   public String toString() {
     return "CuVSQuery [cagraSearchParameters="
@@ -138,11 +156,20 @@ public static class Builder {
     private int topK = 2;
     private BitSet prefilter;
     private int numDocs;
+    private final CuVSResources resources;
 
     /**
-     * Default constructor.
+     * Constructor that requires CuVSResources.
+     *
+     * <p><strong>Important:</strong> The provided CuVSResources instance should not be
+     * shared with other threads. Each thread performing searches should create its own
+     * CuVSResources instance to avoid memory allocation conflicts and potential JVM crashes.
+     *
+     * @param resources the CuVSResources instance to use for this query (must not be shared between threads)
      */
-    public Builder() {}
+    public Builder(CuVSResources resources) {
+      this.resources = Objects.requireNonNull(resources, "resources cannot be null");
+    }
 
     /**
      * Sets the instance of configured CagraSearchParams to be passed for search.
@@ -211,7 +238,8 @@ public Builder withPrefilter(BitSet prefilter, int numDocs) {
      * @return an instance of CuVSQuery
      */
     public CagraQuery build() {
-      return new CagraQuery(cagraSearchParams, queryVectors, mapping, topK, prefilter, numDocs);
+      return new CagraQuery(
+          cagraSearchParams, queryVectors, mapping, topK, prefilter, numDocs, resources);
     }
   }
 }

java/cuvs-java/src/main/java/com/nvidia/cuvs/CagraSearchParams.java

@@ -98,7 +98,6 @@ private HashMapMode(int value) {
   /**
    * Constructs an instance of CagraSearchParams with passed search parameters.
    *
-   * @param resources          the resources instance to use
    * @param maxQueries         the maximum number of queries to search at the same
    *                           time (batch size)
    * @param iTopKSize          the number of intermediate search results retained
@@ -120,7 +119,6 @@ private HashMapMode(int value) {
    *                           selection
    */
   private CagraSearchParams(
-      CuVSResources resources,
       int maxQueries,
       int iTopKSize,
       int maxIterations,
@@ -303,7 +301,6 @@ public String toString() {
    */
   public static class Builder {
 
-    private CuVSResources resources;
     private int maxQueries;
     private int iTopKSize = 64;
     private int maxIterations;
@@ -319,13 +316,9 @@ public static class Builder {
     private HashMapMode hashMapMode;
 
     /**
-     * Constructs this Builder with an instance of Arena.
-     *
-     * @param resources the {@link CuVSResources} instance to use
+     * Default constructor.
      */
-    public Builder(CuVSResources resources) {
-      this.resources = resources;
-    }
+    public Builder() {}
 
     /**
      * Sets the maximum number of queries to search at the same time (batch size).
@@ -486,7 +479,6 @@ public Builder withRandXorMask(long randXORMask) {
      */
     public CagraSearchParams build() {
       return new CagraSearchParams(
-          resources,
           maxQueries,
           iTopKSize,
           maxIterations,

java/cuvs-java/src/main/java/com/nvidia/cuvs/HnswQuery.java

@@ -16,12 +16,17 @@
 package com.nvidia.cuvs;
 
 import java.util.Arrays;
+import java.util.Objects;
 import java.util.function.LongToIntFunction;
 
 /**
  * HnswQuery holds the query vectors to be used while invoking search on the
  * HNSW index.
  *
+ * <p><strong>Thread Safety:</strong> Each HnswQuery instance should use its own
+ * CuVSResources object that is not shared with other threads. Sharing CuVSResources
+ * between threads can lead to memory allocation errors or JVM crashes.
+ *
  * @since 25.02
  */
 public class HnswQuery {
@@ -30,6 +35,7 @@ public class HnswQuery {
   private final LongToIntFunction mapping;
   private final float[][] queryVectors;
   private final int topK;
+  private final CuVSResources resources;
 
   /**
    * Constructs an instance of {@link HnswQuery} using queryVectors, mapping, and
@@ -39,16 +45,19 @@ public class HnswQuery {
    * @param queryVectors     2D float query vector array
    * @param mapping          a function mapping ordinals (neighbor IDs) to custom user IDs
    * @param topK             the top k results to return
+   * @param resources        CuVSResources instance to use for this query
    */
   private HnswQuery(
       HnswSearchParams hnswSearchParams,
       float[][] queryVectors,
       LongToIntFunction mapping,
-      int topK) {
+      int topK,
+      CuVSResources resources) {
     this.hnswSearchParams = hnswSearchParams;
     this.queryVectors = queryVectors;
     this.mapping = mapping;
     this.topK = topK;
+    this.resources = resources;
   }
 
   /**
@@ -85,6 +94,15 @@ public int getTopK() {
     return topK;
   }
 
+  /**
+   * Gets the CuVSResources instance for this query.
+   *
+   * @return the CuVSResources instance
+   */
+  public CuVSResources getResources() {
+    return resources;
+  }
+
   @Override
   public String toString() {
     return "HnswQuery [mapping="
@@ -97,14 +115,28 @@ public String toString() {
   }
 
   /**
-   * Builder helps configure and create an instance of BruteForceQuery.
+   * Builder helps configure and create an instance of HnswQuery.
    */
   public static class Builder {
 
     private HnswSearchParams hnswSearchParams;
     private float[][] queryVectors;
     private LongToIntFunction mapping = SearchResults.IDENTITY_MAPPING;
     private int topK = 2;
+    private final CuVSResources resources;
+
+    /**
+     * Constructor that requires CuVSResources.
+     *
+     * <p><strong>Important:</strong> The provided CuVSResources instance should not be
+     * shared with other threads. Each thread performing searches should create its own
+     * CuVSResources instance to avoid memory allocation conflicts and potential JVM crashes.
+     *
+     * @param resources the CuVSResources instance to use for this query (must not be shared between threads)
+     */
+    public Builder(CuVSResources resources) {
+      this.resources = Objects.requireNonNull(resources, "resources cannot be null");
+    }
 
     /**
      * Sets the instance of configured HnswSearchParams to be passed for search.
@@ -157,7 +189,7 @@ public Builder withTopK(int topK) {
      * @return an instance of {@link HnswQuery}
      */
     public HnswQuery build() {
-      return new HnswQuery(hnswSearchParams, queryVectors, mapping, topK);
+      return new HnswQuery(hnswSearchParams, queryVectors, mapping, topK, resources);
     }
   }
 }