Merge branch 'main' into segrep-backpressure-stats

Signed-off-by: Rishikesh Pasham <62345295+Rishikesh1159@users.noreply.github.com>

.ci/bwcVersions

@@ -25,4 +25,5 @@ BWC_VERSION:
   - "2.10.0"
   - "2.10.1"
   - "2.11.0"
+  - "2.11.1"
   - "2.12.0"

.github/CODEOWNERS

@@ -1 +1 @@
-*   @reta @anasalkouz @andrross @Bukhtawar @CEHENKLE @dblock @gbbafna @setiah @kartg @kotwanikunal @mch2 @nknize @owaiskazi19 @peternied @Rishikesh1159 @ryanbogan @saratvemulapalli @shwetathareja @dreamer-89 @tlfeng @VachaShah @dbwiddis @sachinpkale @sohami @msfroh
+*   @abbashus @adnapibar @anasalkouz @andrross @Bukhtawar @CEHENKLE @dblock @dbwiddis @dreamer-89 @gbbafna @kartg @kotwanikunal @mch2 @msfroh @nknize @owaiskazi19 @peternied @reta @Rishikesh1159 @ryanbogan @sachinpkale @saratvemulapalli @setiah @shwetathareja @sohami @tlfeng @VachaShah

CHANGELOG.md

@@ -91,6 +91,7 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
 ### Added
 - Per request phase latency ([#10351](https://github.com/opensearch-project/OpenSearch/issues/10351))
 - [Remote Store] Add repository stats for remote store([#10567](https://github.com/opensearch-project/OpenSearch/pull/10567))
+- Add search query categorizer ([#10255](https://github.com/opensearch-project/OpenSearch/pull/10255))
 
 ### Dependencies
 - Bump `com.google.api.grpc:proto-google-common-protos` from 2.10.0 to 2.25.1 ([#10208](https://github.com/opensearch-project/OpenSearch/pull/10208), [#10298](https://github.com/opensearch-project/OpenSearch/pull/10298))

libs/core/src/main/java/org/opensearch/Version.java

@@ -96,6 +96,7 @@ public class Version implements Comparable<Version>, ToXContentFragment {
     public static final Version V_2_10_0 = new Version(2100099, org.apache.lucene.util.Version.LUCENE_9_7_0);
     public static final Version V_2_10_1 = new Version(2100199, org.apache.lucene.util.Version.LUCENE_9_7_0);
     public static final Version V_2_11_0 = new Version(2110099, org.apache.lucene.util.Version.LUCENE_9_7_0);
+    public static final Version V_2_11_1 = new Version(2110199, org.apache.lucene.util.Version.LUCENE_9_7_0);
     public static final Version V_2_12_0 = new Version(2120099, org.apache.lucene.util.Version.LUCENE_9_8_0);
     public static final Version V_3_0_0 = new Version(3000099, org.apache.lucene.util.Version.LUCENE_9_8_0);
     public static final Version CURRENT = V_3_0_0;

plugins/repository-s3/src/internalClusterTest/java/org/opensearch/repositories/s3/S3BlobStoreRepositoryTests.java

@@ -165,6 +165,7 @@ protected Settings nodeSettings(int nodeOrdinal) {
         return builder.build();
     }
 
+    @AwaitsFix(bugUrl = "https://github.com/opensearch-project/OpenSearch/issues/10735")
     @Override
     public void testRequestStats() throws Exception {
         final String repository = createRepository(randomName());

plugins/repository-s3/src/main/java/org/opensearch/repositories/s3/async/AsyncPartsHandler.java

@@ -29,6 +29,7 @@
 
 import java.io.BufferedInputStream;
 import java.io.IOException;
+import java.io.InputStream;
 import java.util.ArrayList;
 import java.util.List;
 import java.util.concurrent.CompletableFuture;
@@ -140,28 +141,39 @@ private static void uploadPart(
         ExecutorService streamReadExecutor = uploadRequest.getWritePriority() == WritePriority.HIGH
             ? priorityExecutorService
             : executorService;
+        // Buffered stream is needed to allow mark and reset ops during IO errors so that only buffered
+        // data can be retried instead of retrying whole file by the application.
+        InputStream inputStream = new BufferedInputStream(inputStreamContainer.getInputStream(), (int) (ByteSizeUnit.MB.toBytes(1) + 1));
         CompletableFuture<UploadPartResponse> uploadPartResponseFuture = SocketAccess.doPrivileged(
             () -> s3AsyncClient.uploadPart(
                 uploadPartRequest,
-                AsyncRequestBody.fromInputStream(
-                    // Buffered stream is needed to allow mark and reset ops during IO errors so that only buffered
-                    // data can be retried instead of retrying whole file by the application.
-                    new BufferedInputStream(inputStreamContainer.getInputStream(), (int) (ByteSizeUnit.MB.toBytes(1) + 1)),
-                    inputStreamContainer.getContentLength(),
-                    streamReadExecutor
-                )
+                AsyncRequestBody.fromInputStream(inputStream, inputStreamContainer.getContentLength(), streamReadExecutor)
             )
         );
 
-        CompletableFuture<CompletedPart> convertFuture = uploadPartResponseFuture.thenApply(
-            uploadPartResponse -> convertUploadPartResponse(
-                completedParts,
-                inputStreamContainers,
-                uploadPartResponse,
-                partNumber,
-                uploadRequest.doRemoteDataIntegrityCheck()
-            )
-        );
+        CompletableFuture<CompletedPart> convertFuture = uploadPartResponseFuture.whenComplete((resp, throwable) -> {
+            try {
+                inputStream.close();
+            } catch (IOException ex) {
+                log.error(
+                    () -> new ParameterizedMessage(
+                        "Failed to close stream while uploading a part of idx {} and file {}.",
+                        uploadPartRequest.partNumber(),
+                        uploadPartRequest.key()
+                    ),
+                    ex
+                );
+            }
+        })
+            .thenApply(
+                uploadPartResponse -> convertUploadPartResponse(
+                    completedParts,
+                    inputStreamContainers,
+                    uploadPartResponse,
+                    partNumber,
+                    uploadRequest.doRemoteDataIntegrityCheck()
+                )
+            );
         futures.add(convertFuture);
 
         CompletableFutureUtils.forwardExceptionTo(convertFuture, uploadPartResponseFuture);

plugins/repository-s3/src/main/java/org/opensearch/repositories/s3/async/AsyncTransferManager.java

@@ -40,6 +40,7 @@
 
 import java.io.BufferedInputStream;
 import java.io.IOException;
+import java.io.InputStream;
 import java.util.Arrays;
 import java.util.Base64;
 import java.util.List;
@@ -310,17 +311,22 @@ private void uploadInOneChunk(
         ExecutorService streamReadExecutor = uploadRequest.getWritePriority() == WritePriority.HIGH
             ? priorityExecutorService
             : executorService;
+        // Buffered stream is needed to allow mark and reset ops during IO errors so that only buffered
+        // data can be retried instead of retrying whole file by the application.
+        InputStream inputStream = new BufferedInputStream(inputStreamContainer.getInputStream(), (int) (ByteSizeUnit.MB.toBytes(1) + 1));
         CompletableFuture<Void> putObjectFuture = SocketAccess.doPrivileged(
             () -> s3AsyncClient.putObject(
                 putObjectRequestBuilder.build(),
-                AsyncRequestBody.fromInputStream(
-                    // Buffered stream is needed to allow mark and reset ops during IO errors so that only buffered
-                    // data can be retried instead of retrying whole file by the application.
-                    new BufferedInputStream(inputStreamContainer.getInputStream(), (int) (ByteSizeUnit.MB.toBytes(1) + 1)),
-                    inputStreamContainer.getContentLength(),
-                    streamReadExecutor
-                )
+                AsyncRequestBody.fromInputStream(inputStream, inputStreamContainer.getContentLength(), streamReadExecutor)
             ).handle((resp, throwable) -> {
+                try {
+                    inputStream.close();
+                } catch (IOException e) {
+                    log.error(
+                        () -> new ParameterizedMessage("Failed to close stream while uploading single file {}.", uploadRequest.getKey()),
+                        e
+                    );
+                }
                 if (throwable != null) {
                     Throwable unwrappedThrowable = ExceptionsHelper.unwrap(throwable, S3Exception.class);
                     if (unwrappedThrowable != null) {

plugins/repository-s3/src/test/java/org/opensearch/repositories/s3/async/AsyncTransferManagerTests.java

@@ -37,9 +37,14 @@
 import org.opensearch.test.OpenSearchTestCase;
 import org.junit.Before;
 
+import java.io.IOException;
+import java.io.InputStream;
+import java.util.ArrayList;
+import java.util.List;
 import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.ExecutionException;
 import java.util.concurrent.Executors;
+import java.util.concurrent.atomic.AtomicReference;
 
 import static org.mockito.ArgumentMatchers.any;
 import static org.mockito.Mockito.mock;
@@ -71,17 +76,16 @@ public void testOneChunkUpload() {
             putObjectResponseCompletableFuture
         );
 
+        AtomicReference<InputStream> streamRef = new AtomicReference<>();
         CompletableFuture<Void> resultFuture = asyncTransferManager.uploadObject(
             s3AsyncClient,
             new UploadRequest("bucket", "key", ByteSizeUnit.MB.toBytes(1), WritePriority.HIGH, uploadSuccess -> {
                 // do nothing
             }, false, null),
-            new StreamContext(
-                (partIdx, partSize, position) -> new InputStreamContainer(new ZeroInputStream(partSize), partSize, position),
-                ByteSizeUnit.MB.toBytes(1),
-                ByteSizeUnit.MB.toBytes(1),
-                1
-            ),
+            new StreamContext((partIdx, partSize, position) -> {
+                streamRef.set(new ZeroInputStream(partSize));
+                return new InputStreamContainer(streamRef.get(), partSize, position);
+            }, ByteSizeUnit.MB.toBytes(1), ByteSizeUnit.MB.toBytes(1), 1),
             new StatsMetricPublisher()
         );
 
@@ -92,6 +96,14 @@ public void testOneChunkUpload() {
         }
 
         verify(s3AsyncClient, times(1)).putObject(any(PutObjectRequest.class), any(AsyncRequestBody.class));
+
+        boolean closeError = false;
+        try {
+            streamRef.get().available();
+        } catch (IOException e) {
+            closeError = e.getMessage().equals("Stream closed");
+        }
+        assertTrue("InputStream was still open after upload", closeError);
     }
 
     public void testOneChunkUploadCorruption() {
@@ -162,17 +174,17 @@ public void testMultipartUpload() {
             abortMultipartUploadResponseCompletableFuture
         );
 
+        List<InputStream> streams = new ArrayList<>();
         CompletableFuture<Void> resultFuture = asyncTransferManager.uploadObject(
             s3AsyncClient,
             new UploadRequest("bucket", "key", ByteSizeUnit.MB.toBytes(5), WritePriority.HIGH, uploadSuccess -> {
                 // do nothing
             }, true, 3376132981L),
-            new StreamContext(
-                (partIdx, partSize, position) -> new InputStreamContainer(new ZeroInputStream(partSize), partSize, position),
-                ByteSizeUnit.MB.toBytes(1),
-                ByteSizeUnit.MB.toBytes(1),
-                5
-            ),
+            new StreamContext((partIdx, partSize, position) -> {
+                InputStream stream = new ZeroInputStream(partSize);
+                streams.add(stream);
+                return new InputStreamContainer(stream, partSize, position);
+            }, ByteSizeUnit.MB.toBytes(1), ByteSizeUnit.MB.toBytes(1), 5),
             new StatsMetricPublisher()
         );
 
@@ -182,6 +194,16 @@ public void testMultipartUpload() {
             fail("did not expect resultFuture to fail");
         }
 
+        streams.forEach(stream -> {
+            boolean closeError = false;
+            try {
+                stream.available();
+            } catch (IOException e) {
+                closeError = e.getMessage().equals("Stream closed");
+            }
+            assertTrue("InputStream was still open after upload", closeError);
+        });
+
         verify(s3AsyncClient, times(1)).createMultipartUpload(any(CreateMultipartUploadRequest.class));
         verify(s3AsyncClient, times(5)).uploadPart(any(UploadPartRequest.class), any(AsyncRequestBody.class));
         verify(s3AsyncClient, times(1)).completeMultipartUpload(any(CompleteMultipartUploadRequest.class));

server/src/main/java/org/opensearch/action/admin/cluster/node/stats/NodeStats.java

@@ -221,7 +221,7 @@ public NodeStats(StreamInput in) throws IOException {
         } else {
             segmentReplicationRejectionStats = null;
         }
-        if (in.getVersion().onOrAfter(Version.V_3_0_0)) {
+        if (in.getVersion().onOrAfter(Version.V_2_12_0)) {
             repositoriesStats = in.readOptionalWriteable(RepositoriesStats::new);
         } else {
             repositoriesStats = null;
@@ -485,7 +485,7 @@ public void writeTo(StreamOutput out) throws IOException {
         if (out.getVersion().onOrAfter(Version.V_3_0_0)) {
             out.writeOptionalWriteable(segmentReplicationRejectionStats);
         }
-        if (out.getVersion().onOrAfter(Version.V_3_0_0)) {
+        if (out.getVersion().onOrAfter(Version.V_2_12_0)) {
             out.writeOptionalWriteable(repositoriesStats);
         }
     }

server/src/main/java/org/opensearch/action/search/SearchQueryCategorizer.java

@@ -0,0 +1,81 @@
+/*
+ * SPDX-License-Identifier: Apache-2.0
+ *
+ * The OpenSearch Contributors require contributions made to
+ * this file be licensed under the Apache-2.0 license or a
+ * compatible open source license.
+ */
+
+package org.opensearch.action.search;
+
+import org.apache.logging.log4j.LogManager;
+import org.apache.logging.log4j.Logger;
+import org.opensearch.index.query.QueryBuilder;
+import org.opensearch.index.query.QueryBuilderVisitor;
+import org.opensearch.index.query.QueryShapeVisitor;
+import org.opensearch.search.aggregations.AggregatorFactories;
+import org.opensearch.search.builder.SearchSourceBuilder;
+import org.opensearch.search.sort.SortBuilder;
+import org.opensearch.telemetry.metrics.MetricsRegistry;
+import org.opensearch.telemetry.metrics.tags.Tags;
+
+import java.util.List;
+import java.util.ListIterator;
+
+/**
+ * Class to categorize the search queries based on the type and increment the relevant counters.
+ * Class also logs the query shape.
+ */
+final class SearchQueryCategorizer {
+
+    private static final Logger log = LogManager.getLogger(SearchQueryCategorizer.class);
+
+    final SearchQueryCounters searchQueryCounters;
+
+    public SearchQueryCategorizer(MetricsRegistry metricsRegistry) {
+        searchQueryCounters = new SearchQueryCounters(metricsRegistry);
+    }
+
+    public void categorize(SearchSourceBuilder source) {
+        QueryBuilder topLevelQueryBuilder = source.query();
+
+        logQueryShape(topLevelQueryBuilder);
+        incrementQueryTypeCounters(topLevelQueryBuilder);
+        incrementQueryAggregationCounters(source.aggregations());
+        incrementQuerySortCounters(source.sorts());
+    }
+
+    private void incrementQuerySortCounters(List<SortBuilder<?>> sorts) {
+        if (sorts != null && sorts.size() > 0) {
+            for (ListIterator<SortBuilder<?>> it = sorts.listIterator(); it.hasNext();) {
+                SortBuilder sortBuilder = it.next();
+                String sortOrder = sortBuilder.order().toString();
+                searchQueryCounters.sortCounter.add(1, Tags.create().addTag("sort_order", sortOrder));
+            }
+        }
+    }
+
+    private void incrementQueryAggregationCounters(AggregatorFactories.Builder aggregations) {
+        if (aggregations != null) {
+            searchQueryCounters.aggCounter.add(1);
+        }
+    }
+
+    private void incrementQueryTypeCounters(QueryBuilder topLevelQueryBuilder) {
+        if (topLevelQueryBuilder == null) {
+            return;
+        }
+        QueryBuilderVisitor searchQueryVisitor = new SearchQueryCategorizingVisitor(searchQueryCounters);
+        topLevelQueryBuilder.visit(searchQueryVisitor);
+    }
+
+    private void logQueryShape(QueryBuilder topLevelQueryBuilder) {
+        if (topLevelQueryBuilder == null) {
+            return;
+        }
+        QueryShapeVisitor shapeVisitor = new QueryShapeVisitor();
+        topLevelQueryBuilder.visit(shapeVisitor);
+        log.debug("Query shape : {}", shapeVisitor.prettyPrintTree("  "));
+    }
+
+}