PR comments: added javadoc, fixed unlock issue (+ test for it)

Author: ldematte

x-pack/plugin/gpu/src/main/java/org/elasticsearch/xpack/gpu/codec/CuVSResourceManager.java

@@ -290,6 +290,7 @@ void lock(Runnable unlockAction) {
         }
 
         void unlock() {
+            unlockAction.run();
             unlockAction = NOT_LOCKED;
         }
 

x-pack/plugin/gpu/src/main/java/org/elasticsearch/xpack/gpu/codec/GPUMemoryService.java

@@ -9,6 +9,9 @@
 
 import com.nvidia.cuvs.CuVSResources;
 
+/**
+ * Abstracts GPU memory tracking (total vs available)
+ */
 interface GPUMemoryService {
 
     long totalMemoryInBytes(CuVSResources res);

x-pack/plugin/gpu/src/main/java/org/elasticsearch/xpack/gpu/codec/RealGPUMemoryService.java

@@ -10,6 +10,10 @@
 import com.nvidia.cuvs.CuVSResources;
 import com.nvidia.cuvs.GPUInfoProvider;
 
+/**
+ * A {@link GPUMemoryService} that tracks how much memory is currently used/available on a GPU by using the GPU free/total memory APIs
+ * (via a {@link GPUInfoProvider})
+ */
 class RealGPUMemoryService implements GPUMemoryService {
     private final GPUInfoProvider gpuInfoProvider;
 

x-pack/plugin/gpu/src/main/java/org/elasticsearch/xpack/gpu/codec/TrackingGPUMemoryService.java

@@ -9,6 +9,12 @@
 
 import com.nvidia.cuvs.CuVSResources;
 
+/**
+ * A {@link GPUMemoryService} that tracks manually how much memory is currently estimated to be used/available on a GPU.
+ * This implementation is useful when we are not able to use a "Real memory" measurement; for example, if we are using pooled RMM memory,
+ * the pool will permanently occupy most of the GPU RAM, allocations will happen inside the pool, and the "Real memory" measurement API
+ * will always report a (tiny) fixed amount of free memory (whatever is not in the pool).
+ */
 class TrackingGPUMemoryService implements GPUMemoryService {
 
     private final long totalMemoryInBytes;

x-pack/plugin/gpu/src/test/java/org/elasticsearch/xpack/gpu/codec/CuVSResourceManagerTests.java

@@ -24,6 +24,7 @@
 import static org.hamcrest.Matchers.anyOf;
 import static org.hamcrest.Matchers.containsString;
 import static org.hamcrest.Matchers.equalTo;
+import static org.hamcrest.Matchers.lessThan;
 import static org.hamcrest.Matchers.not;
 
 public class CuVSResourceManagerTests extends ESTestCase {
@@ -57,6 +58,27 @@ public void testBasicWithIvfPq() throws InterruptedException {
         testBasic(createIvfPqParams());
     }
 
+    public void testMultipleAcquireRelease() throws InterruptedException {
+        var mgr = new MockPoolingCuVSResourceManager(2);
+        var res1 = mgr.acquire(16 * 1024, 1024, CuVSMatrix.DataType.FLOAT, createNnDescentParams());
+        var res2 = mgr.acquire(16 * 1024, 1024, CuVSMatrix.DataType.FLOAT, createIvfPqParams());
+        assertThat(res1.toString(), containsString("id=0"));
+        assertThat(res2.toString(), containsString("id=1"));
+        assertThat(mgr.availableMemory(), lessThan(TOTAL_DEVICE_MEMORY_IN_BYTES / 2));
+        mgr.release(res1);
+        mgr.release(res2);
+        assertThat(mgr.availableMemory(), equalTo(TOTAL_DEVICE_MEMORY_IN_BYTES));
+        res1 = mgr.acquire(16 * 1024, 1024, CuVSMatrix.DataType.FLOAT, createNnDescentParams());
+        res2 = mgr.acquire(16 * 1024, 1024, CuVSMatrix.DataType.FLOAT, createIvfPqParams());
+        assertThat(res1.toString(), containsString("id=0"));
+        assertThat(res2.toString(), containsString("id=1"));
+        assertThat(mgr.availableMemory(), lessThan(TOTAL_DEVICE_MEMORY_IN_BYTES / 2));
+        mgr.release(res1);
+        mgr.release(res2);
+        assertThat(mgr.availableMemory(), equalTo(TOTAL_DEVICE_MEMORY_IN_BYTES));
+        mgr.shutdown();
+    }
+
     private static void testBlocking(CagraIndexParams params) throws Exception {
         var mgr = new MockPoolingCuVSResourceManager(2);
         var res1 = mgr.acquire(0, 0, CuVSMatrix.DataType.FLOAT, params);
@@ -185,13 +207,23 @@ private static CagraIndexParams createIvfPqParams() {
     static class MockPoolingCuVSResourceManager extends CuVSResourceManager.PoolingCuVSResourceManager {
 
         private final AtomicInteger idGenerator = new AtomicInteger();
+        private final GPUMemoryService gpuMemoryService;
 
         MockPoolingCuVSResourceManager(int capacity) {
             this(capacity, TOTAL_DEVICE_MEMORY_IN_BYTES);
         }
 
         MockPoolingCuVSResourceManager(int capacity, long totalMemoryInBytes) {
-            super(capacity, new TrackingGPUMemoryService(totalMemoryInBytes));
+            this(capacity, new TrackingGPUMemoryService(totalMemoryInBytes));
+        }
+
+        private MockPoolingCuVSResourceManager(int capacity, GPUMemoryService gpuMemoryService) {
+            super(capacity, gpuMemoryService);
+            this.gpuMemoryService = gpuMemoryService;
+        }
+
+        long availableMemory() {
+            return gpuMemoryService.availableMemoryInBytes(null);
         }
 
         @Override