review

elastic · Oct 14, 2024 · e9f1723 · e9f1723
1 parent e649ffb
commit e9f1723
Show file tree

Hide file tree

Showing 3 changed files with 7 additions and 8 deletions.
diff --git a/...tTest/java/org/elasticsearch/xpack/ml/integration/AdaptiveAllocationsScaleFromZeroIT.java b/...tTest/java/org/elasticsearch/xpack/ml/integration/AdaptiveAllocationsScaleFromZeroIT.java
@@ -23,6 +23,7 @@
 import java.util.concurrent.CountDownLatch;
 import java.util.concurrent.TimeUnit;
 
+import static org.hamcrest.Matchers.empty;
 import static org.hamcrest.Matchers.hasSize;
 import static org.hamcrest.Matchers.is;
 import static org.hamcrest.Matchers.not;
@@ -93,9 +94,7 @@ public void onFailure(Exception exception) {
         }
 
         latch.await();
-        if (failures.isEmpty() == false) {
-            fail(failures.getFirst());
-        }
+        assertThat(failures, empty());
     }
 
     @SuppressWarnings("unchecked")

diff --git a/...ml/src/main/java/org/elasticsearch/xpack/ml/action/TransportInternalInferModelAction.java b/...ml/src/main/java/org/elasticsearch/xpack/ml/action/TransportInternalInferModelAction.java
@@ -71,7 +71,7 @@ public class TransportInternalInferModelAction extends HandledTransportAction<Re
     private final XPackLicenseState licenseState;
     private final TrainedModelProvider trainedModelProvider;
     private final AdaptiveAllocationsScalerService adaptiveAllocationsScalerService;
-    private final InferenceWaitForAllocation scalingInference;
+    private final InferenceWaitForAllocation waitForAllocation;
     private final ThreadPool threadPool;
 
     TransportInternalInferModelAction(
@@ -94,7 +94,7 @@ public class TransportInternalInferModelAction extends HandledTransportAction<Re
         this.licenseState = licenseState;
         this.trainedModelProvider = trainedModelProvider;
         this.adaptiveAllocationsScalerService = adaptiveAllocationsScalerService;
-        this.scalingInference = new InferenceWaitForAllocation(assignmentService, this::inferOnBlockedRequest);
+        this.waitForAllocation = new InferenceWaitForAllocation(assignmentService, this::inferOnBlockedRequest);
         this.threadPool = threadPool;
     }
 
@@ -280,7 +280,7 @@ private void inferAgainstAllocatedModel(
             if (starting) {
                 message += "; starting deployment of one allocation";
                 logger.info(message);
-                scalingInference.waitForAssignment(
+                waitForAllocation.waitForAssignment(
                     new InferenceWaitForAllocation.WaitingRequest(request, responseBuilder, parentTaskId, listener)
                 );
                 return;

diff --git a/...gin/ml/src/main/java/org/elasticsearch/xpack/ml/inference/InferenceWaitForAllocation.java b/...gin/ml/src/main/java/org/elasticsearch/xpack/ml/inference/InferenceWaitForAllocation.java
@@ -85,7 +85,8 @@ public InferenceWaitForAllocation(
      */
     public synchronized void waitForAssignment(WaitingRequest request) {
         logger.info("waitForAssignment will wait for condition");
-        if (pendingRequestCount.get() > MAX_PENDING_REQUEST_COUNT) {
+        if (pendingRequestCount.incrementAndGet() >= MAX_PENDING_REQUEST_COUNT) {
+            pendingRequestCount.decrementAndGet();
             request.listener.onFailure(
                 new ElasticsearchStatusException(
                     "Rejected inference request waiting for an allocation of deployment [{}]. Too many pending requests",
@@ -96,7 +97,6 @@ public synchronized void waitForAssignment(WaitingRequest request) {
             return;
         }
 
-        pendingRequestCount.incrementAndGet();
         var predicate = new DeploymentHasAtLeastOneAllocation(request.deploymentId());
 
         assignmentService.waitForAssignmentCondition(