Avoid double free in index creation/destruction; add tests for specific parts

Author: ldematte

java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/BruteForceIndexImpl.java

@@ -28,6 +28,7 @@
 import static com.nvidia.cuvs.internal.common.Util.cudaMemcpy;
 import static com.nvidia.cuvs.internal.common.Util.CudaMemcpyKind.*;
 import static com.nvidia.cuvs.internal.common.Util.prepareTensor;
+import static com.nvidia.cuvs.internal.panama.headers_h.cudaStream_t;
 import static com.nvidia.cuvs.internal.panama.headers_h.cuvsBruteForceBuild;
 import static com.nvidia.cuvs.internal.panama.headers_h.cuvsBruteForceDeserialize;
 import static com.nvidia.cuvs.internal.panama.headers_h.cuvsBruteForceIndexCreate;
@@ -37,14 +38,10 @@
 import static com.nvidia.cuvs.internal.panama.headers_h.cuvsBruteForceSerialize;
 import static com.nvidia.cuvs.internal.panama.headers_h.cuvsRMMAlloc;
 import static com.nvidia.cuvs.internal.panama.headers_h.cuvsRMMFree;
-import static com.nvidia.cuvs.internal.panama.headers_h.cuvsResources_t;
 import static com.nvidia.cuvs.internal.panama.headers_h.cuvsStreamGet;
 import static com.nvidia.cuvs.internal.panama.headers_h.cuvsStreamSync;
 import static com.nvidia.cuvs.internal.panama.headers_h.omp_set_num_threads;
-import static com.nvidia.cuvs.internal.panama.headers_h.cudaStream_t;
 
-import java.io.FileInputStream;
-import java.io.FileOutputStream;
 import java.io.InputStream;
 import java.io.OutputStream;
 import java.lang.foreign.Arena;
@@ -64,7 +61,6 @@
 import com.nvidia.cuvs.Dataset;
 import com.nvidia.cuvs.SearchResults;
 import com.nvidia.cuvs.internal.common.Util;
-import com.nvidia.cuvs.internal.panama.cuvsBruteForceIndex;
 import com.nvidia.cuvs.internal.panama.cuvsFilter;
 
 /**
@@ -130,8 +126,13 @@ private void checkNotDestroyed() {
   public void destroyIndex() throws Throwable {
     checkNotDestroyed();
     try {
-      int returnValue = cuvsBruteForceIndexDestroy(bruteForceIndexReference.getMemorySegment());
+      int returnValue = cuvsBruteForceIndexDestroy(bruteForceIndexReference.indexPtr);
       checkCuVSError(returnValue, "cuvsBruteForceIndexDestroy");
+
+      if (bruteForceIndexReference.datasetBytes > 0) {
+        returnValue = cuvsRMMFree(resources.getHandle(), bruteForceIndexReference.datasetPtr, bruteForceIndexReference.datasetBytes);
+        checkCuVSError(returnValue, "cuvsRMMFree");
+      }
     } finally {
       destroyed = true;
     }
@@ -145,7 +146,7 @@ public void destroyIndex() throws Throwable {
    * @return an instance of {@link IndexReference} that holds the pointer to the
    *         index
    */
-  private IndexReference build() throws Throwable {
+  private IndexReference build() {
     try (var localArena = Arena.ofConfined()) {
       long rows = dataset != null? dataset.size(): vectors.length;
       long cols = dataset != null? dataset.dimensions(): (rows > 0 ? vectors[0].length : 0);
@@ -155,16 +156,13 @@ private IndexReference build() throws Throwable {
           Util.buildMemorySegment(resources.getArena(), vectors);
 
       long cuvsResources = resources.getHandle();
-      MemorySegment stream = arena.allocate(cudaStream_t);
-      var returnValue = cuvsStreamGet(cuvsResources, stream);
-      checkCuVSError(returnValue, "cuvsStreamGet");
 
       omp_set_num_threads(bruteForceIndexParams.getNumWriterThreads());
 
-      MemorySegment datasetMemorySegment = arena.allocate(C_POINTER);
+      MemorySegment datasetMemorySegment = localArena.allocate(C_POINTER);
 
       long datasetBytes = C_FLOAT_BYTE_SIZE * rows * cols;
-      returnValue = cuvsRMMAlloc(cuvsResources, datasetMemorySegment, datasetBytes);
+      var returnValue = cuvsRMMAlloc(cuvsResources, datasetMemorySegment, datasetBytes);
       checkCuVSError(returnValue, "cuvsRMMAlloc");
 
       // IMPORTANT: this should only come AFTER cuvsRMMAlloc call
@@ -175,23 +173,20 @@ private IndexReference build() throws Throwable {
       long[] datasetShape = { rows, cols };
       MemorySegment datasetTensor = prepareTensor(arena, datasetMemorySegmentP, datasetShape, 2, 32, 2, 2, 1);
 
-      MemorySegment index = arena.allocate(cuvsBruteForceIndex_t);
-
-      returnValue = cuvsBruteForceIndexCreate(index);
-      checkCuVSError(returnValue, "cuvsBruteForceIndexCreate");
+      var indexReference = new IndexReference(datasetMemorySegmentP, datasetBytes, createBruteForceIndex());
 
       returnValue = cuvsStreamSync(cuvsResources);
       checkCuVSError(returnValue, "cuvsStreamSync");
 
-      returnValue = cuvsBruteForceBuild(cuvsResources, datasetTensor, 0, 0.0f, index);
+      returnValue = cuvsBruteForceBuild(cuvsResources, datasetTensor, 0, 0.0f, indexReference.indexPtr);
       checkCuVSError(returnValue, "cuvsBruteForceBuild");
 
       returnValue = cuvsStreamSync(cuvsResources);
       checkCuVSError(returnValue, "cuvsStreamSync");
 
       omp_set_num_threads(1);
 
-      return new IndexReference(index);
+      return indexReference;
     }
   }
 
@@ -236,17 +231,14 @@ public SearchResults search(BruteForceQuery cuvsQuery) throws Throwable {
       var returnValue = cuvsStreamGet(cuvsResources, stream);
       checkCuVSError(returnValue, "cuvsStreamGet");
 
-      MemorySegment queriesD = arena.allocate(C_POINTER);
-      MemorySegment neighborsD = arena.allocate(C_POINTER);
-      MemorySegment distancesD = arena.allocate(C_POINTER);
-      MemorySegment prefilterD = arena.allocate(C_POINTER);
-      MemorySegment prefilterDP = MemorySegment.NULL;
-      long prefilterLen = 0;
+      MemorySegment queriesD = localArena.allocate(C_POINTER);
+      MemorySegment neighborsD = localArena.allocate(C_POINTER);
+      MemorySegment distancesD = localArena.allocate(C_POINTER);
 
       long queriesBytes = C_FLOAT_BYTE_SIZE * numQueries * vectorDimension;
       long neighborsBytes = C_LONG_BYTE_SIZE * numQueries * topk;
       long distanceBytes = C_FLOAT_BYTE_SIZE * numQueries * topk;
-      long prefilterBytes = 0; // size assigned later
+      long prefilterBytes = 0;
 
       returnValue = cuvsRMMAlloc(cuvsResources, queriesD, queriesBytes);
       checkCuVSError(returnValue, "cuvsRMMAlloc");
@@ -259,14 +251,15 @@ public SearchResults search(BruteForceQuery cuvsQuery) throws Throwable {
       MemorySegment queriesDP = queriesD.get(C_POINTER, 0);
       MemorySegment neighborsDP = neighborsD.get(C_POINTER, 0);
       MemorySegment distancesDP = distancesD.get(C_POINTER, 0);
+      MemorySegment prefilterDP = MemorySegment.NULL;
 
       cudaMemcpy(queriesDP, querySeg, queriesBytes, INFER_DIRECTION);
 
-      long queriesShape[] = { numQueries, vectorDimension };
+      long[] queriesShape = { numQueries, vectorDimension };
       MemorySegment queriesTensor = prepareTensor(arena, queriesDP, queriesShape, 2, 32, 2, 2, 1);
-      long neighborsShape[] = { numQueries, topk };
+      long[] neighborsShape = { numQueries, topk };
       MemorySegment neighborsTensor = prepareTensor(arena, neighborsDP, neighborsShape, 0, 64, 2, 2, 1);
-      long distancesShape[] = { numQueries, topk };
+      long[] distancesShape = { numQueries, topk };
       MemorySegment distancesTensor = prepareTensor(arena, distancesDP, distancesShape, 2, 32, 2, 2, 1);
 
       MemorySegment prefilter = cuvsFilter.allocate(arena);
@@ -276,13 +269,14 @@ public SearchResults search(BruteForceQuery cuvsQuery) throws Throwable {
         cuvsFilter.type(prefilter, 0); // NO_FILTER
         cuvsFilter.addr(prefilter, 0);
       } else {
-        long prefilterShape[] = { (prefilterDataLength + 31) / 32 };
-        prefilterLen = prefilterShape[0];
+        long[] prefilterShape = { (prefilterDataLength + 31) / 32 };
+
+        MemorySegment prefilterD = localArena.allocate(C_POINTER);
+        long prefilterLen = prefilterShape[0];
         prefilterBytes = C_INT_BYTE_SIZE * prefilterLen;
 
         returnValue = cuvsRMMAlloc(cuvsResources, prefilterD, prefilterBytes);
         checkCuVSError(returnValue, "cuvsRMMAlloc");
-
         prefilterDP = prefilterD.get(C_POINTER, 0);
 
         cudaMemcpy(prefilterDP, prefilterDataMemorySegment, prefilterBytes, HOST_TO_DEVICE);
@@ -296,7 +290,7 @@ public SearchResults search(BruteForceQuery cuvsQuery) throws Throwable {
       returnValue = cuvsStreamSync(cuvsResources);
       checkCuVSError(returnValue, "cuvsStreamSync");
 
-      returnValue = cuvsBruteForceSearch(cuvsResources, bruteForceIndexReference.getMemorySegment(), queriesTensor,
+      returnValue = cuvsBruteForceSearch(cuvsResources, bruteForceIndexReference.indexPtr, queriesTensor,
           neighborsTensor, distancesTensor, prefilter);
       checkCuVSError(returnValue, "cuvsBruteForceSearch");
 
@@ -312,8 +306,10 @@ public SearchResults search(BruteForceQuery cuvsQuery) throws Throwable {
       checkCuVSError(returnValue, "cuvsRMMFree");
       returnValue = cuvsRMMFree(cuvsResources, queriesDP, queriesBytes);
       checkCuVSError(returnValue, "cuvsRMMFree");
-      returnValue = cuvsRMMFree(cuvsResources, prefilterDP, prefilterBytes);
-      checkCuVSError(returnValue, "cuvsRMMFree");
+      if (prefilterBytes > 0) {
+        returnValue = cuvsRMMFree(cuvsResources, prefilterDP, prefilterBytes);
+        checkCuVSError(returnValue, "cuvsRMMFree");
+      }
 
       return new BruteForceSearchResults(neighborsSequenceLayout, distancesSequenceLayout, neighborsMemorySegment,
           distancesMemorySegment, cuvsQuery.getTopK(), cuvsQuery.getMapping(), numQueries);
@@ -332,17 +328,33 @@ public void serialize(OutputStream outputStream, Path tempFile) throws Throwable
     tempFile = tempFile.toAbsolutePath();
 
     long cuvsRes = resources.getHandle();
-    int returnValue = cuvsBruteForceSerialize(cuvsRes, resources.getArena().allocateFrom(tempFile.toString()),
-        bruteForceIndexReference.getMemorySegment());
-    checkCuVSError(returnValue, "cuvsBruteForceSerialize");
+    try (var localArena = Arena.ofConfined()) {
+      int returnValue = cuvsBruteForceSerialize(cuvsRes, localArena.allocateFrom(tempFile.toString()),
+          bruteForceIndexReference.indexPtr);
+      checkCuVSError(returnValue, "cuvsBruteForceSerialize");
+    }
 
-    try (FileInputStream fileInputStream = new FileInputStream(tempFile.toFile())) {
-      fileInputStream.transferTo(outputStream);
+    try (var inputStream = Files.newInputStream(tempFile)) {
+      inputStream.transferTo(outputStream);
     } finally {
       Files.deleteIfExists(tempFile);
     }
   }
 
+  private static MemorySegment createBruteForceIndex() {
+    try (var localArena = Arena.ofConfined()) {
+      MemorySegment indexPtrPtr = localArena.allocate(cuvsBruteForceIndex_t);
+      // cuvsBruteForceIndexCreate gets a pointer to a cuvsBruteForceIndex_t, which is defined as a pointer to
+      // cuvsBruteForceIndex.
+      // It's basically a "out" parameter: the C functions will create the index and "return back" a pointer to it.
+      // The "out parameter" pointer is needed only for the duration of the function invocation (it could be a stack
+      // pointer, in C) so we allocate it from our localArena, unwrap it and return it.
+      var returnValue = cuvsBruteForceIndexCreate(indexPtrPtr);
+      checkCuVSError(returnValue, "cuvsBruteForceIndexCreate");
+      return indexPtrPtr.get(cuvsBruteForceIndex_t, 0);
+    }
+  }
+
   /**
    * Gets an instance of {@link IndexReference} by deserializing a BRUTEFORCE
    * index using an {@link InputStream}.
@@ -352,16 +364,16 @@ public void serialize(OutputStream outputStream, Path tempFile) throws Throwable
    */
   private IndexReference deserialize(InputStream inputStream) throws Throwable {
     checkNotDestroyed();
-    Path tmpIndexFile = Files.createTempFile(resources.tempDirectory(), UUID.randomUUID().toString(), ".bf");
-    tmpIndexFile = tmpIndexFile.toAbsolutePath();
-    IndexReference indexReference = new IndexReference(resources);
+    Path tmpIndexFile = Files.createTempFile(resources.tempDirectory(), UUID.randomUUID().toString(), ".bf")
+        .toAbsolutePath();
+    IndexReference indexReference = new IndexReference(createBruteForceIndex());
 
-    try (var in = inputStream; FileOutputStream fileOutputStream = new FileOutputStream(tmpIndexFile.toFile())) {
-      in.transferTo(fileOutputStream);
+    try (inputStream; var outputStream = Files.newOutputStream(tmpIndexFile); var arena = Arena.ofConfined()) {
+      inputStream.transferTo(outputStream);
 
       long cuvsRes = resources.getHandle();
-      int returnValue = cuvsBruteForceDeserialize(cuvsRes, resources.getArena().allocateFrom(tmpIndexFile.toString()),
-          indexReference.getMemorySegment());
+      int returnValue = cuvsBruteForceDeserialize(cuvsRes, arena.allocateFrom(tmpIndexFile.toString()),
+          indexReference.indexPtr);
       checkCuVSError(returnValue, "cuvsBruteForceDeserialize");
 
     } finally {
@@ -464,37 +476,24 @@ public BruteForceIndexImpl build() throws Throwable {
   }
 
   /**
-   * Holds the memory reference to a BRUTEFORCE index.
+   * Holds the memory reference to a BRUTEFORCE index and its associated dataset
    */
-  protected static class IndexReference {
+  private static class IndexReference {
 
-    private final MemorySegment memorySegment;
-
-    /**
-     * Constructs CagraIndexReference and allocate the MemorySegment.
-     */
-    protected IndexReference(CuVSResourcesImpl resources) {
-      memorySegment = cuvsBruteForceIndex.allocate(resources.getArena());
-    }
+      private final MemorySegment datasetPtr;
+      private final long datasetBytes;
+      private final MemorySegment indexPtr;
 
-    /**
-     * Constructs BruteForceIndexReference with an instance of MemorySegment passed
-     * as a parameter.
-     *
-     * @param indexMemorySegment the MemorySegment instance to use for containing
-     *                           index reference
-     */
-    protected IndexReference(MemorySegment indexMemorySegment) {
-      this.memorySegment = indexMemorySegment;
+    private IndexReference(MemorySegment datasetPtr, long datasetBytes, MemorySegment indexPtr) {
+        this.datasetPtr = datasetPtr;
+        this.datasetBytes = datasetBytes;
+        this.indexPtr = indexPtr;
     }
 
-    /**
-     * Gets the instance of index MemorySegment.
-     *
-     * @return index MemorySegment
-     */
-    protected MemorySegment getMemorySegment() {
-      return memorySegment;
+    private IndexReference(MemorySegment indexPtr) {
+      this.datasetPtr = MemorySegment.NULL;
+      this.datasetBytes = 0;
+      this.indexPtr = indexPtr;
     }
   }
 }