Short circuit to avoid livelock + spotless

Author: ldematte

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

@@ -8,9 +8,7 @@
 package org.elasticsearch.xpack.gpu.codec;
 
 import com.nvidia.cuvs.CuVSResources;
-
 import com.nvidia.cuvs.GPUInfoProvider;
-
 import com.nvidia.cuvs.spi.CuVSProvider;
 
 import org.elasticsearch.core.Strings;
@@ -84,7 +82,7 @@ class PoolingCuVSResourceManager implements CuVSResourceManager {
         private int createdCount;
 
         ReentrantLock lock = new ReentrantLock();
-        Condition enoughMemoryCondition = lock.newCondition();
+        Condition enoughResourcesCondition = lock.newCondition();
 
         public PoolingCuVSResourceManager(int capacity, GPUInfoProvider gpuInfoProvider) {
             if (capacity < 1 || capacity > MAX_RESOURCES) {
@@ -110,6 +108,17 @@ private ManagedCuVSResources getResourceFromPool() {
             return null;
         }
 
+        private int numLockedResources() {
+            int lockedResources = 0;
+            for (int i = 0; i < createdCount; ++i) {
+                var res = pool[i];
+                if (res.locked) {
+                    lockedResources++;
+                }
+            }
+            return lockedResources;
+        }
+
         @Override
         public ManagedCuVSResources acquire(int numVectors, int dims) throws InterruptedException {
             try {
@@ -119,10 +128,14 @@ public ManagedCuVSResources acquire(int numVectors, int dims) throws Interrupted
                 ManagedCuVSResources res = null;
                 while (allConditionsMet == false) {
                     res = getResourceFromPool();
+
                     final boolean enoughMemory;
                     if (res != null) {
+                        // If no resource in the pool is locked, short circuit to avoid livelock
+                        if (numLockedResources() == 0) {
+                            break;
+                        }
                         // Check resources availability
-                        // Memory
                         long requiredMemoryInBytes = estimateRequiredMemory(numVectors, dims);
                         if (requiredMemoryInBytes > gpuInfoProvider.getCurrentInfo(res).totalDeviceMemoryInBytes()) {
                             throw new IllegalArgumentException(
@@ -138,12 +151,11 @@ public ManagedCuVSResources acquire(int numVectors, int dims) throws Interrupted
                     } else {
                         enoughMemory = false;
                     }
-                    if (enoughMemory == false) {
-                        enoughMemoryCondition.await();
-                    }
-
                     // TODO: add enoughComputation / enoughComputationCondition here
                     allConditionsMet = enoughMemory; // && enoughComputation
+                    if (allConditionsMet == false) {
+                        enoughResourcesCondition.await();
+                    }
                 }
                 res.locked = true;
                 return res;
@@ -153,7 +165,7 @@ public ManagedCuVSResources acquire(int numVectors, int dims) throws Interrupted
         }
 
         private long estimateRequiredMemory(int numVectors, int dims) {
-            return (long)(GPU_COMPUTATION_MEMORY_FACTOR * numVectors * dims * Float.BYTES);
+            return (long) (GPU_COMPUTATION_MEMORY_FACTOR * numVectors * dims * Float.BYTES);
         }
 
         // visible for testing
@@ -164,7 +176,7 @@ protected CuVSResources createNew() {
         @Override
         public void finishedComputation(ManagedCuVSResources resources) {
             // currently does nothing, but could allow acquire to return possibly blocked resources
-            // something like enoughComputationCondition.signalAll()?
+            // enoughResourcesCondition.signalAll()
         }
 
         @Override
@@ -173,7 +185,7 @@ public void release(ManagedCuVSResources resources) {
                 lock.lock();
                 assert resources.locked;
                 resources.locked = false;
-                enoughMemoryCondition.signalAll();
+                enoughResourcesCondition.signalAll();
             } finally {
                 lock.unlock();
             }

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

@@ -8,7 +8,6 @@
 package org.elasticsearch.xpack.gpu.codec;
 
 import com.nvidia.cuvs.CuVSResources;
-
 import com.nvidia.cuvs.CuVSResourcesInfo;
 import com.nvidia.cuvs.GPUInfo;
 import com.nvidia.cuvs.GPUInfoProvider;
@@ -161,8 +160,8 @@ protected CuVSResources createNew() {
         @Override
         public ManagedCuVSResources acquire(int numVectors, int dims) throws InterruptedException {
             var res = super.acquire(numVectors, dims);
-            long memory = (long)(numVectors * dims * Float.BYTES *
-                CuVSResourceManager.PoolingCuVSResourceManager.GPU_COMPUTATION_MEMORY_FACTOR);
+            long memory = (long) (numVectors * dims * Float.BYTES
+                * CuVSResourceManager.PoolingCuVSResourceManager.GPU_COMPUTATION_MEMORY_FACTOR);
             allocations.add(memory);
             log.info("Added [{}]", memory);
             return res;