Test that requests and operations are tracked separately

Author: nicktindall

modules/repository-azure/src/internalClusterTest/java/org/elasticsearch/repositories/azure/AzureBlobStoreRepositoryMetricsTests.java

@@ -9,7 +9,6 @@
 
 package org.elasticsearch.repositories.azure;
 
-import com.sun.net.httpserver.Headers;
 import com.sun.net.httpserver.HttpExchange;
 import com.sun.net.httpserver.HttpHandler;
 
@@ -34,7 +33,6 @@
 
 import java.io.IOException;
 import java.nio.ByteBuffer;
-import java.nio.charset.StandardCharsets;
 import java.util.ArrayList;
 import java.util.List;
 import java.util.Map;
@@ -48,6 +46,7 @@
 import java.util.stream.IntStream;
 
 import static org.elasticsearch.repositories.azure.AbstractAzureServerTestCase.randomBlobContent;
+import static org.elasticsearch.repositories.azure.ResponseInjectingAzureHttpHandler.createFailNRequestsHandler;
 import static org.elasticsearch.repositories.blobstore.BlobStoreTestUtil.randomPurpose;
 import static org.hamcrest.Matchers.greaterThanOrEqualTo;
 import static org.hamcrest.Matchers.hasSize;
@@ -61,7 +60,7 @@ public class AzureBlobStoreRepositoryMetricsTests extends AzureBlobStoreReposito
     );
     private static final int MAX_RETRIES = 3;
 
-    private final Queue<RequestHandler> requestHandlers = new ConcurrentLinkedQueue<>();
+    private final Queue<ResponseInjectingAzureHttpHandler.RequestHandler> requestHandlers = new ConcurrentLinkedQueue<>();
 
     @Override
     protected Map<String, HttpHandler> createHttpHandlers() {
@@ -106,7 +105,8 @@ public void testThrottleResponsesAreCountedInMetrics() throws IOException {
 
         // Queue up some throttle responses
         final int numThrottles = randomIntBetween(1, MAX_RETRIES);
-        IntStream.range(0, numThrottles).forEach(i -> requestHandlers.offer(new FixedRequestHandler(RestStatus.TOO_MANY_REQUESTS)));
+        IntStream.range(0, numThrottles)
+            .forEach(i -> requestHandlers.offer(new ResponseInjectingAzureHttpHandler.FixedRequestHandler(RestStatus.TOO_MANY_REQUESTS)));
 
         // Check that the blob exists
         blobContainer.blobExists(purpose, blobName);
@@ -131,7 +131,13 @@ public void testRangeNotSatisfiedAreCountedInMetrics() throws IOException {
         clearMetrics(dataNodeName);
 
         // Queue up a range-not-satisfied error
-        requestHandlers.offer(new FixedRequestHandler(RestStatus.REQUESTED_RANGE_NOT_SATISFIED, null, GET_BLOB_REQUEST_PREDICATE));
+        requestHandlers.offer(
+            new ResponseInjectingAzureHttpHandler.FixedRequestHandler(
+                RestStatus.REQUESTED_RANGE_NOT_SATISFIED,
+                null,
+                GET_BLOB_REQUEST_PREDICATE
+            )
+        );
 
         // Attempt to read the blob
         assertThrows(RequestedRangeNotSatisfiedException.class, () -> blobContainer.readBlob(purpose, blobName));
@@ -163,7 +169,7 @@ public void testErrorResponsesAreCountedInMetrics() throws IOException {
             if (status == RestStatus.TOO_MANY_REQUESTS) {
                 throttles.incrementAndGet();
             }
-            requestHandlers.offer(new FixedRequestHandler(status));
+            requestHandlers.offer(new ResponseInjectingAzureHttpHandler.FixedRequestHandler(status));
         });
 
         // Check that the blob exists
@@ -259,7 +265,7 @@ public void testBatchDeleteFailure() throws IOException {
         clearMetrics(dataNodeName);
 
         // Handler will fail one or more of the batch requests
-        final RequestHandler failNRequestRequestHandler = createFailNRequestsHandler(failedBatches);
+        final ResponseInjectingAzureHttpHandler.RequestHandler failNRequestRequestHandler = createFailNRequestsHandler(failedBatches);
 
         // Exhaust the retries
         IntStream.range(0, (numberOfBatches - failedBatches) + (failedBatches * (MAX_RETRIES + 1)))
@@ -287,35 +293,6 @@ private long getLongCounterTotal(String dataNodeName, String metricKey) {
             .reduce(0L, Long::sum);
     }
 
-    /**
-     * Creates a {@link RequestHandler} that will persistently fail the first <code>numberToFail</code> distinct requests
-     * it sees. Any other requests are passed through to the delegate.
-     *
-     * @param numberToFail The number of requests to fail
-     * @return the handler
-     */
-    private static RequestHandler createFailNRequestsHandler(int numberToFail) {
-        final List<String> requestsToFail = new ArrayList<>(numberToFail);
-        return (exchange, delegate) -> {
-            final Headers requestHeaders = exchange.getRequestHeaders();
-            final String requestId = requestHeaders.get("X-ms-client-request-id").get(0);
-            boolean failRequest = false;
-            synchronized (requestsToFail) {
-                if (requestsToFail.contains(requestId)) {
-                    failRequest = true;
-                } else if (requestsToFail.size() < numberToFail) {
-                    requestsToFail.add(requestId);
-                    failRequest = true;
-                }
-            }
-            if (failRequest) {
-                exchange.sendResponseHeaders(500, -1);
-            } else {
-                delegate.handle(exchange);
-            }
-        };
-    }
-
     private void clearMetrics(String discoveryNode) {
         internalCluster().getInstance(PluginsService.class, discoveryNode)
             .filterPlugins(TestTelemetryPlugin.class)
@@ -480,80 +457,4 @@ private void assertMatchingMetricRecorded(MetricType metricType, String metricNa
             assertion.accept(measurement);
         }
     }
-
-    @SuppressForbidden(reason = "we use a HttpServer to emulate Azure")
-    private static class ResponseInjectingAzureHttpHandler implements DelegatingHttpHandler {
-
-        private final HttpHandler delegate;
-        private final Queue<RequestHandler> requestHandlerQueue;
-
-        ResponseInjectingAzureHttpHandler(Queue<RequestHandler> requestHandlerQueue, HttpHandler delegate) {
-            this.delegate = delegate;
-            this.requestHandlerQueue = requestHandlerQueue;
-        }
-
-        @Override
-        public void handle(HttpExchange exchange) throws IOException {
-            RequestHandler nextHandler = requestHandlerQueue.peek();
-            if (nextHandler != null && nextHandler.matchesRequest(exchange)) {
-                requestHandlerQueue.poll().writeResponse(exchange, delegate);
-            } else {
-                delegate.handle(exchange);
-            }
-        }
-
-        @Override
-        public HttpHandler getDelegate() {
-            return delegate;
-        }
-    }
-
-    @SuppressForbidden(reason = "we use a HttpServer to emulate Azure")
-    @FunctionalInterface
-    private interface RequestHandler {
-        void writeResponse(HttpExchange exchange, HttpHandler delegate) throws IOException;
-
-        default boolean matchesRequest(HttpExchange exchange) {
-            return true;
-        }
-    }
-
-    @SuppressForbidden(reason = "we use a HttpServer to emulate Azure")
-    private static class FixedRequestHandler implements RequestHandler {
-
-        private final RestStatus status;
-        private final String responseBody;
-        private final Predicate<HttpExchange> requestMatcher;
-
-        FixedRequestHandler(RestStatus status) {
-            this(status, null, req -> true);
-        }
-
-        /**
-         * Create a handler that only gets executed for requests that match the supplied predicate. Note
-         * that because the errors are stored in a queue this will prevent any subsequently queued errors from
-         * being returned until after it returns.
-         */
-        FixedRequestHandler(RestStatus status, String responseBody, Predicate<HttpExchange> requestMatcher) {
-            this.status = status;
-            this.responseBody = responseBody;
-            this.requestMatcher = requestMatcher;
-        }
-
-        @Override
-        public boolean matchesRequest(HttpExchange exchange) {
-            return requestMatcher.test(exchange);
-        }
-
-        @Override
-        public void writeResponse(HttpExchange exchange, HttpHandler delegateHandler) throws IOException {
-            if (responseBody != null) {
-                byte[] responseBytes = responseBody.getBytes(StandardCharsets.UTF_8);
-                exchange.sendResponseHeaders(status.getStatus(), responseBytes.length);
-                exchange.getResponseBody().write(responseBytes);
-            } else {
-                exchange.sendResponseHeaders(status.getStatus(), -1);
-            }
-        }
-    }
 }

modules/repository-azure/src/test/java/org/elasticsearch/repositories/azure/AzureBlobContainerStatsTests.java

@@ -16,19 +16,78 @@
 import org.elasticsearch.common.blobstore.OperationPurpose;
 import org.elasticsearch.common.bytes.BytesReference;
 import org.elasticsearch.core.SuppressForbidden;
+import org.elasticsearch.rest.RestStatus;
 import org.junit.Before;
 
 import java.io.IOException;
 import java.nio.ByteBuffer;
+import java.util.Arrays;
+import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
+import java.util.Queue;
+import java.util.concurrent.ConcurrentLinkedQueue;
+
+import static org.elasticsearch.repositories.azure.AzureBlobStore.Operation.BLOB_BATCH;
+import static org.elasticsearch.repositories.azure.AzureBlobStore.Operation.LIST_BLOBS;
+import static org.elasticsearch.repositories.azure.AzureBlobStore.Operation.PUT_BLOB;
 
 public class AzureBlobContainerStatsTests extends AbstractAzureServerTestCase {
 
+    private final Queue<ResponseInjectingAzureHttpHandler.RequestHandler> requestHandlers = new ConcurrentLinkedQueue<>();
+
     @SuppressForbidden(reason = "use a http server")
     @Before
     public void configureAzureHandler() {
-        httpServer.createContext("/", new AzureHttpHandler(ACCOUNT, CONTAINER, null, MockAzureBlobStore.LeaseExpiryPredicate.NEVER_EXPIRE));
+        httpServer.createContext(
+            "/",
+            new ResponseInjectingAzureHttpHandler(
+                requestHandlers,
+                new AzureHttpHandler(ACCOUNT, CONTAINER, null, MockAzureBlobStore.LeaseExpiryPredicate.NEVER_EXPIRE)
+            )
+        );
+    }
+
+    public void testRetriesAndOperationsAreTrackedSeparately() throws IOException {
+        serverlessMode = true;
+        final AzureBlobContainer blobContainer = asInstanceOf(AzureBlobContainer.class, createBlobContainer(between(1, 3)));
+        final AzureBlobStore blobStore = blobContainer.getBlobStore();
+        final OperationPurpose purpose = randomFrom(OperationPurpose.values());
+
+        // Just a sample of the easy operations to test
+        final List<AzureBlobStore.Operation> supportedOperations = Arrays.asList(PUT_BLOB, LIST_BLOBS, BLOB_BATCH);
+        final Map<AzureBlobStore.Operation, BlobStoreActionStats> expectedActionStats = new HashMap<>();
+        supportedOperations.forEach(operation -> expectedActionStats.put(operation, BlobStoreActionStats.ZERO));
+
+        for (int i = 0; i < randomIntBetween(10, 50); i++) {
+            final boolean triggerRetry = randomBoolean();
+            if (triggerRetry) {
+                requestHandlers.offer(new ResponseInjectingAzureHttpHandler.FixedRequestHandler(RestStatus.TOO_MANY_REQUESTS));
+            }
+            final AzureBlobStore.Operation operation = randomFrom(supportedOperations);
+            switch (operation) {
+                case PUT_BLOB -> blobStore.writeBlob(
+                    purpose,
+                    randomIdentifier(),
+                    BytesReference.fromByteBuffer(ByteBuffer.wrap(randomBlobContent())),
+                    false
+                );
+                case LIST_BLOBS -> blobStore.listBlobsByPrefix(purpose, randomIdentifier(), randomIdentifier());
+                case BLOB_BATCH -> blobStore.deleteBlobsIgnoringIfNotExists(
+                    purpose,
+                    List.of(randomIdentifier(), randomIdentifier(), randomIdentifier()).iterator()
+                );
+            }
+            expectedActionStats.computeIfPresent(operation, (op, stats) -> stats.add(new BlobStoreActionStats(1, triggerRetry ? 2 : 1)));
+        }
+
+        final Map<String, BlobStoreActionStats> stats = blobStore.stats();
+        expectedActionStats.forEach((operation, value) -> {
+            if (value.isZero() == false) {
+                String key = statsKey(purpose, operation);
+                assertEquals(key, stats.get(key), value);
+            }
+        });
     }
 
     public void testOperationPurposeIsReflectedInBlobStoreStats() throws IOException {