Fixing flaky test for request service

Author: jonathan-buttner

muted-tests.yml

@@ -304,9 +304,6 @@ tests:
 - class: org.elasticsearch.indices.mapping.UpdateMappingIntegrationIT
   method: testUpdateMappingConcurrently
   issue: https://github.com/elastic/elasticsearch/issues/137758
-- class: org.elasticsearch.xpack.inference.external.http.sender.RequestExecutorServiceTests
-  method: testChangingCapacity_DoesNotRejectsOverflowTasks_BecauseOfQueueFull
-  issue: https://github.com/elastic/elasticsearch/issues/137823
 - class: org.elasticsearch.xpack.inference.integration.AuthorizationTaskExecutorMultipleNodesIT
   method: testAuthorizationTaskGetsRelocatedToAnotherNode_WhenTheNodeThatIsRunningItShutsDown
   issue: https://github.com/elastic/elasticsearch/issues/137911

x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/external/http/sender/RequestExecutorService.java

@@ -26,6 +26,7 @@
 import org.elasticsearch.xpack.inference.services.settings.RateLimitSettings;
 
 import java.time.Clock;
+import java.time.Duration;
 import java.time.Instant;
 import java.util.ArrayList;
 import java.util.List;
@@ -125,7 +126,10 @@ interface RateLimiterCreator {
     private final AtomicInteger rateLimitDivisor = new AtomicInteger(1);
     private final ThreadPool threadPool;
     private final CountDownLatch startupLatch;
-    private final CountDownLatch terminationLatch = new CountDownLatch(1);
+    // Two latches because we have two threads of execution, one thread blocking on a queue for items to be sent immediately, and
+    // another threads that is scheduled on an interval that checks items that can be rate limited
+    private final CountDownLatch immediateRequestQueueTerminationLatch = new CountDownLatch(1);
+    private final CountDownLatch rateLimitedTerminationLatch = new CountDownLatch(1);
     private final RequestSender requestSender;
     private final RequestExecutorServiceSettings settings;
     private final Clock clock;
@@ -184,11 +188,25 @@ public boolean isShutdown() {
     }
 
     public boolean awaitTermination(long timeout, TimeUnit unit) throws InterruptedException {
-        return terminationLatch.await(timeout, unit);
+        var totalWait = Duration.ofMillis(unit.toMillis(timeout));
+
+        var firstAwaitStart = Instant.now();
+        var firstLatchResult = immediateRequestQueueTerminationLatch.await(timeout, unit);
+        var firstAwaitEnd = Instant.now();
+
+        var remainingWaitTime = totalWait.minus(Duration.between(firstAwaitStart, firstAwaitEnd));
+
+        // If the first latch await returns false, we've run out of time
+        // If the remaining wait time is negative or zero, we've run out of time
+        if (firstLatchResult == false || remainingWaitTime.isNegative() || remainingWaitTime.isZero()) {
+            return false;
+        }
+
+        return rateLimitedTerminationLatch.await(remainingWaitTime.toMillis(), TimeUnit.MILLISECONDS);
     }
 
     public boolean isTerminated() {
-        return terminationLatch.getCount() == 0;
+        return immediateRequestQueueTerminationLatch.getCount() == 0 && rateLimitedTerminationLatch.getCount() == 0;
     }
 
     public int queueSize() {
@@ -213,6 +231,7 @@ public void start() {
         } catch (Exception e) {
             logger.warn("Failed to start request executor", e);
             cleanup(CleanupStrategy.RATE_LIMITED_REQUEST_QUEUES_ONLY);
+            cleanup(CleanupStrategy.REQUEST_QUEUE_ONLY);
         }
     }
 
@@ -350,12 +369,16 @@ private void cleanup(CleanupStrategy cleanupStrategy) {
             shutdown();
 
             switch (cleanupStrategy) {
-                case RATE_LIMITED_REQUEST_QUEUES_ONLY -> notifyRateLimitedRequestsOfShutdown();
-                case REQUEST_QUEUE_ONLY -> rejectRequestsInRequestQueue();
+                case RATE_LIMITED_REQUEST_QUEUES_ONLY -> {
+                    notifyRateLimitedRequestsOfShutdown();
+                    rateLimitedTerminationLatch.countDown();
+                }
+                case REQUEST_QUEUE_ONLY -> {
+                    rejectRequestsInRequestQueue();
+                    immediateRequestQueueTerminationLatch.countDown();
+                }
                 default -> logger.error(Strings.format("Unknown clean up strategy for request executor: [%s]", cleanupStrategy.toString()));
             }
-
-            terminationLatch.countDown();
         } catch (Exception e) {
             logger.warn("Encountered an error while cleaning up", e);
         }

x-pack/plugin/inference/src/test/java/org/elasticsearch/xpack/inference/external/http/sender/RequestExecutorServiceTests.java

@@ -103,6 +103,7 @@ public void testIsTerminated_IsTrue() throws InterruptedException {
         service.shutdown();
         service.start();
         latch.await(TIMEOUT.getSeconds(), TimeUnit.SECONDS);
+        service.awaitTermination(TIMEOUT.getSeconds(), TimeUnit.SECONDS);
 
         assertTrue(service.isTerminated());
     }
@@ -115,6 +116,8 @@ public void testCallingStartTwice_ThrowsAssertionException() throws InterruptedE
         service.start();
         latch.await(TIMEOUT.getSeconds(), TimeUnit.SECONDS);
 
+        service.awaitTermination(TIMEOUT.getSeconds(), TimeUnit.SECONDS);
+
         assertTrue(service.isTerminated());
         var exception = expectThrows(AssertionError.class, service::start);
         assertThat(exception.getMessage(), is("start() can only be called once"));
@@ -235,13 +238,15 @@ public void testTaskThrowsError_CallsOnFailure() throws InterruptedException {
         assertTrue(service.isTerminated());
     }
 
-    public void testShutdown_AllowsMultipleCalls() {
+    public void testShutdown_AllowsMultipleCalls() throws InterruptedException {
         var service = createRequestExecutorServiceWithMocks();
 
         service.shutdown();
         service.shutdown();
         service.start();
 
+        service.awaitTermination(TIMEOUT.getSeconds(), TimeUnit.SECONDS);
+
         assertTrue(service.isTerminated());
         assertTrue(service.isShutdown());
     }