[Java] Exception-safe RMM Allocations (#1215)

This commit introduces exception safety for RMM allocations.

Previously, device memory allocated through `cuvsRmmAlloc()` was freed manually using `cuvsRmmFree()`, in all the index impl classes.  The problem there is that if an exception is thrown in the intervening time between alloc and free, it would lead to a leak of device memory.

This commit extends the `CloseableHandle` class to encapsulate the allocation of device memory.  This new class is used in try-with-resources blocks, to make device memory allocations exception-safe.

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

Approvers:
  - Lorenzo Dematté (https://github.com/ldematte)
  - Corey J. Nolet (https://github.com/cjnolet)

URL: #1215

Author: mythrocks

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

@@ -15,14 +15,14 @@
  */
 package com.nvidia.cuvs.internal;
 
+import static com.nvidia.cuvs.internal.common.CloseableRMMAllocation.allocateRMMSegment;
 import static com.nvidia.cuvs.internal.common.LinkerHelper.C_FLOAT;
 import static com.nvidia.cuvs.internal.common.LinkerHelper.C_FLOAT_BYTE_SIZE;
 import static com.nvidia.cuvs.internal.common.LinkerHelper.C_INT_BYTE_SIZE;
 import static com.nvidia.cuvs.internal.common.LinkerHelper.C_LONG;
 import static com.nvidia.cuvs.internal.common.LinkerHelper.C_LONG_BYTE_SIZE;
 import static com.nvidia.cuvs.internal.common.Util.CudaMemcpyKind.HOST_TO_DEVICE;
 import static com.nvidia.cuvs.internal.common.Util.CudaMemcpyKind.INFER_DIRECTION;
-import static com.nvidia.cuvs.internal.common.Util.allocateRMMSegment;
 import static com.nvidia.cuvs.internal.common.Util.buildMemorySegment;
 import static com.nvidia.cuvs.internal.common.Util.checkCuVSError;
 import static com.nvidia.cuvs.internal.common.Util.concatenate;
@@ -35,7 +35,6 @@
 import static com.nvidia.cuvs.internal.panama.headers_h.cuvsBruteForceIndex_t;
 import static com.nvidia.cuvs.internal.panama.headers_h.cuvsBruteForceSearch;
 import static com.nvidia.cuvs.internal.panama.headers_h.cuvsBruteForceSerialize;
-import static com.nvidia.cuvs.internal.panama.headers_h.cuvsRMMFree;
 import static com.nvidia.cuvs.internal.panama.headers_h.cuvsStreamSync;
 import static com.nvidia.cuvs.internal.panama.headers_h.omp_set_num_threads;
 
@@ -45,6 +44,7 @@
 import com.nvidia.cuvs.CuVSMatrix;
 import com.nvidia.cuvs.CuVSResources;
 import com.nvidia.cuvs.SearchResults;
+import com.nvidia.cuvs.internal.common.CloseableRMMAllocation;
 import com.nvidia.cuvs.internal.panama.cuvsFilter;
 import java.io.InputStream;
 import java.io.OutputStream;
@@ -118,20 +118,7 @@ public void destroyIndex() {
     try {
       int returnValue = cuvsBruteForceIndexDestroy(bruteForceIndexReference.indexPtr);
       checkCuVSError(returnValue, "cuvsBruteForceIndexDestroy");
-
-      if (bruteForceIndexReference.datasetBytes > 0) {
-        try (var resourcesAccessor = resources.access()) {
-          checkCuVSError(
-              cuvsRMMFree(
-                  resourcesAccessor.handle(),
-                  bruteForceIndexReference.datasetPtr,
-                  bruteForceIndexReference.datasetBytes),
-              "cuvsRMMFree");
-        }
-      }
-      if (bruteForceIndexReference.tensorDataArena != null) {
-        bruteForceIndexReference.tensorDataArena.close();
-      }
+      bruteForceIndexReference.close(resources);
     } finally {
       destroyed = true;
     }
@@ -158,25 +145,31 @@ private IndexReference build(
 
     try (var resourcesAccessor = resources.access()) {
       long cuvsResources = resourcesAccessor.handle();
-      MemorySegment datasetMemorySegmentP = allocateRMMSegment(cuvsResources, datasetBytes);
+      try (var closeableDataMemorySegmentP = allocateRMMSegment(cuvsResources, datasetBytes)) {
+        MemorySegment datasetMemorySegmentP = closeableDataMemorySegmentP.handle();
 
-      cudaMemcpy(datasetMemorySegmentP, datasetMemSegment, datasetBytes, INFER_DIRECTION);
+        cudaMemcpy(datasetMemorySegmentP, datasetMemSegment, datasetBytes, INFER_DIRECTION);
 
-      long[] datasetShape = {rows, cols};
-      var tensorDataArena = Arena.ofShared();
-      MemorySegment datasetTensor =
-          prepareTensor(tensorDataArena, datasetMemorySegmentP, datasetShape, 2, 32, 2, 1);
+        long[] datasetShape = {rows, cols};
+        var tensorDataArena = Arena.ofShared();
+        MemorySegment datasetTensor =
+            prepareTensor(tensorDataArena, datasetMemorySegmentP, datasetShape, 2, 32, 2, 1);
 
-      var returnValue = cuvsStreamSync(cuvsResources);
-      checkCuVSError(returnValue, "cuvsStreamSync");
+        var returnValue = cuvsStreamSync(cuvsResources);
+        checkCuVSError(returnValue, "cuvsStreamSync");
 
-      returnValue = cuvsBruteForceBuild(cuvsResources, datasetTensor, 0, 0.0f, index);
-      checkCuVSError(returnValue, "cuvsBruteForceBuild");
+        returnValue = cuvsBruteForceBuild(cuvsResources, datasetTensor, 0, 0.0f, index);
+        checkCuVSError(returnValue, "cuvsBruteForceBuild");
 
-      returnValue = cuvsStreamSync(cuvsResources);
-      checkCuVSError(returnValue, "cuvsStreamSync");
+        returnValue = cuvsStreamSync(cuvsResources);
+        checkCuVSError(returnValue, "cuvsStreamSync");
 
-      return new IndexReference(datasetMemorySegmentP, datasetBytes, tensorDataArena, index);
+        return new IndexReference(
+            new CloseableRMMAllocation(closeableDataMemorySegmentP),
+            datasetBytes,
+            tensorDataArena,
+            index);
+      }
     } finally {
       omp_set_num_threads(1);
     }
@@ -205,15 +198,19 @@ public SearchResults search(BruteForceQuery cuvsQuery) throws Throwable {
 
       // prepare the prefiltering data
       final long prefilterDataLength;
+      final long prefilterBytes;
       final MemorySegment prefilterDataMemorySegment;
       BitSet[] prefilters = cuvsQuery.getPrefilters();
       if (prefilters != null && prefilters.length > 0) {
         BitSet concatenatedFilters = concatenate(prefilters, cuvsQuery.getNumDocs());
         long[] filters = concatenatedFilters.toLongArray();
         prefilterDataMemorySegment = buildMemorySegment(localArena, filters);
         prefilterDataLength = (long) cuvsQuery.getNumDocs() * prefilters.length;
+        long[] prefilterShape = {(prefilterDataLength + 31) / 32};
+        prefilterBytes = C_INT_BYTE_SIZE * prefilterShape[0];
       } else {
         prefilterDataLength = 0;
+        prefilterBytes = 0;
         prefilterDataMemorySegment = MemorySegment.NULL;
       }
 
@@ -223,77 +220,66 @@ public SearchResults search(BruteForceQuery cuvsQuery) throws Throwable {
       try (var resourcesAccessor = cuvsQuery.getResources().access()) {
         long cuvsResources = resourcesAccessor.handle();
 
-        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
-
-        MemorySegment queriesDP = allocateRMMSegment(cuvsResources, queriesBytes);
-        MemorySegment neighborsDP = allocateRMMSegment(cuvsResources, neighborsBytes);
-        MemorySegment distancesDP = allocateRMMSegment(cuvsResources, distanceBytes);
-        MemorySegment prefilterDP = MemorySegment.NULL;
-
-        cudaMemcpy(queriesDP, querySeg, queriesBytes, INFER_DIRECTION);
-
-        long[] queriesShape = {numQueries, vectorDimension};
-        MemorySegment queriesTensor =
-            prepareTensor(localArena, queriesDP, queriesShape, 2, 32, 2, 1);
-        long[] neighborsShape = {numQueries, topk};
-        MemorySegment neighborsTensor =
-            prepareTensor(localArena, neighborsDP, neighborsShape, 0, 64, 2, 1);
-        long[] distancesShape = {numQueries, topk};
-        MemorySegment distancesTensor =
-            prepareTensor(localArena, distancesDP, distancesShape, 2, 32, 2, 1);
-
-        MemorySegment prefilter = cuvsFilter.allocate(localArena);
-        MemorySegment prefilterTensor;
-
-        if (prefilterDataMemorySegment == MemorySegment.NULL) {
-          cuvsFilter.type(prefilter, 0); // NO_FILTER
-          cuvsFilter.addr(prefilter, 0);
-        } else {
-          long[] prefilterShape = {(prefilterDataLength + 31) / 32};
-          long prefilterLen = prefilterShape[0];
-          prefilterBytes = C_INT_BYTE_SIZE * prefilterLen;
-
-          prefilterDP = allocateRMMSegment(cuvsResources, prefilterBytes);
-
-          cudaMemcpy(prefilterDP, prefilterDataMemorySegment, prefilterBytes, HOST_TO_DEVICE);
-
-          prefilterTensor = prepareTensor(localArena, prefilterDP, prefilterShape, 1, 32, 2, 1);
-
-          cuvsFilter.type(prefilter, 2);
-          cuvsFilter.addr(prefilter, prefilterTensor.address());
-        }
-
-        var returnValue = cuvsStreamSync(cuvsResources);
-        checkCuVSError(returnValue, "cuvsStreamSync");
-
-        returnValue =
-            cuvsBruteForceSearch(
-                cuvsResources,
-                bruteForceIndexReference.indexPtr,
-                queriesTensor,
-                neighborsTensor,
-                distancesTensor,
-                prefilter);
-        checkCuVSError(returnValue, "cuvsBruteForceSearch");
-
-        returnValue = cuvsStreamSync(cuvsResources);
-        checkCuVSError(returnValue, "cuvsStreamSync");
-
-        cudaMemcpy(neighborsMemorySegment, neighborsDP, neighborsBytes, INFER_DIRECTION);
-        cudaMemcpy(distancesMemorySegment, distancesDP, distanceBytes, INFER_DIRECTION);
-
-        returnValue = cuvsRMMFree(cuvsResources, neighborsDP, neighborsBytes);
-        checkCuVSError(returnValue, "cuvsRMMFree");
-        returnValue = cuvsRMMFree(cuvsResources, distancesDP, distanceBytes);
-        checkCuVSError(returnValue, "cuvsRMMFree");
-        returnValue = cuvsRMMFree(cuvsResources, queriesDP, queriesBytes);
-        checkCuVSError(returnValue, "cuvsRMMFree");
-        if (prefilterBytes > 0) {
-          returnValue = cuvsRMMFree(cuvsResources, prefilterDP, prefilterBytes);
-          checkCuVSError(returnValue, "cuvsRMMFree");
+        final long queriesBytes = C_FLOAT_BYTE_SIZE * numQueries * vectorDimension;
+        final long neighborsBytes = C_LONG_BYTE_SIZE * numQueries * topk;
+        final long distanceBytes = C_FLOAT_BYTE_SIZE * numQueries * topk;
+
+        try (var queriesDP = allocateRMMSegment(cuvsResources, queriesBytes);
+            var neighborsDP = allocateRMMSegment(cuvsResources, neighborsBytes);
+            var distancesDP = allocateRMMSegment(cuvsResources, distanceBytes);
+            var prefilterDP =
+                prefilterBytes > 0
+                    ? allocateRMMSegment(cuvsResources, prefilterBytes)
+                    : CloseableRMMAllocation.EMPTY) {
+
+          cudaMemcpy(queriesDP.handle(), querySeg, queriesBytes, INFER_DIRECTION);
+
+          long[] queriesShape = {numQueries, vectorDimension};
+          MemorySegment queriesTensor =
+              prepareTensor(localArena, queriesDP.handle(), queriesShape, 2, 32, 2, 1);
+          long[] neighborsShape = {numQueries, topk};
+          MemorySegment neighborsTensor =
+              prepareTensor(localArena, neighborsDP.handle(), neighborsShape, 0, 64, 2, 1);
+          long[] distancesShape = {numQueries, topk};
+          MemorySegment distancesTensor =
+              prepareTensor(localArena, distancesDP.handle(), distancesShape, 2, 32, 2, 1);
+
+          MemorySegment prefilter = cuvsFilter.allocate(localArena);
+          MemorySegment prefilterTensor;
+
+          if (prefilterDataMemorySegment == MemorySegment.NULL) {
+            cuvsFilter.type(prefilter, 0); // NO_FILTER
+            cuvsFilter.addr(prefilter, 0);
+          } else {
+            long[] prefilterShape = {(prefilterDataLength + 31) / 32};
+            cudaMemcpy(
+                prefilterDP.handle(), prefilterDataMemorySegment, prefilterBytes, HOST_TO_DEVICE);
+
+            prefilterTensor =
+                prepareTensor(localArena, prefilterDP.handle(), prefilterShape, 1, 32, 2, 1);
+
+            cuvsFilter.type(prefilter, 2);
+            cuvsFilter.addr(prefilter, prefilterTensor.address());
+          }
+
+          var returnValue = cuvsStreamSync(cuvsResources);
+          checkCuVSError(returnValue, "cuvsStreamSync");
+
+          returnValue =
+              cuvsBruteForceSearch(
+                  cuvsResources,
+                  bruteForceIndexReference.indexPtr,
+                  queriesTensor,
+                  neighborsTensor,
+                  distancesTensor,
+                  prefilter);
+          checkCuVSError(returnValue, "cuvsBruteForceSearch");
+
+          returnValue = cuvsStreamSync(cuvsResources);
+          checkCuVSError(returnValue, "cuvsStreamSync");
+
+          cudaMemcpy(neighborsMemorySegment, neighborsDP.handle(), neighborsBytes, INFER_DIRECTION);
+          cudaMemcpy(distancesMemorySegment, distancesDP.handle(), distanceBytes, INFER_DIRECTION);
         }
       }
       return BruteForceSearchResults.create(
@@ -479,27 +465,39 @@ public BruteForceIndexImpl build() throws Throwable {
    */
   private static class IndexReference {
 
-    private final MemorySegment datasetPtr;
+    private final CloseableRMMAllocation datasetAllocationHandle;
     private final long datasetBytes;
     private final Arena tensorDataArena;
     private final MemorySegment indexPtr;
 
     private IndexReference(
-        MemorySegment datasetPtr,
+        CloseableRMMAllocation datasetAllocationHandle,
         long datasetBytes,
         Arena tensorDataArena,
         MemorySegment indexPtr) {
-      this.datasetPtr = datasetPtr;
+      this.datasetAllocationHandle = datasetAllocationHandle;
       this.datasetBytes = datasetBytes;
       this.tensorDataArena = tensorDataArena;
       this.indexPtr = indexPtr;
     }
 
     private IndexReference(MemorySegment indexPtr) {
-      this.datasetPtr = MemorySegment.NULL;
+      this.datasetAllocationHandle = CloseableRMMAllocation.EMPTY;
       this.datasetBytes = 0;
       this.tensorDataArena = null;
       this.indexPtr = indexPtr;
     }
+
+    /**
+     * Free up the memory used for dataset, tensor-data.
+     */
+    private void close(CuVSResources resources) {
+      try (var resourcesAccessor = resources.access()) {
+        datasetAllocationHandle.close();
+      }
+      if (tensorDataArena != null) {
+        tensorDataArena.close();
+      }
+    }
   }
 }