From 589b7f69c12b76a1bc9fbde1de723f99f9960561 Mon Sep 17 00:00:00 2001
From: kezhenxu94 <kezhenxu94@apache.org>
Date: Thu, 9 Nov 2023 06:47:57 +0000
Subject: [PATCH] Fix `limit` doesn't work for `findEndpoint` API in ES storage

---
 .../elasticsearch/ElasticSearchScroller.java  |  8 +++---
 .../elasticsearch/bulk/ElasticSearchIT.java   | 27 +++++++++++++++++++
 .../query/MetadataQueryEsDAO.java             |  2 +-
 3 files changed, 32 insertions(+), 5 deletions(-)

diff --git a/oap-server/server-library/library-client/src/main/java/org/apache/skywalking/oap/server/library/client/elasticsearch/ElasticSearchScroller.java b/oap-server/server-library/library-client/src/main/java/org/apache/skywalking/oap/server/library/client/elasticsearch/ElasticSearchScroller.java
index f414f4c502cc0..62504e83fc81e 100644
--- a/oap-server/server-library/library-client/src/main/java/org/apache/skywalking/oap/server/library/client/elasticsearch/ElasticSearchScroller.java
+++ b/oap-server/server-library/library-client/src/main/java/org/apache/skywalking/oap/server/library/client/elasticsearch/ElasticSearchScroller.java
@@ -61,12 +61,12 @@ public List<T> scroll() {
                 }
                 for (final var searchHit : response.getHits()) {
                     results.add(resultConverter.apply(searchHit));
+                    if (queryMaxSize > 0 && results.size() >= queryMaxSize) {
+                        return results;
+                    }
                 }
                 if (search.getSize() != null && response.getHits().getHits().size() < search.getSize()) {
-                    break;
-                }
-                if (queryMaxSize > 0 && results.size() >= queryMaxSize) {
-                    break;
+                    return results;
                 }
                 response = client.scroll(SCROLL_CONTEXT_RETENTION, scrollId);
             }
diff --git a/oap-server/server-library/library-client/src/test/java/org/apache/skywalking/library/elasticsearch/bulk/ElasticSearchIT.java b/oap-server/server-library/library-client/src/test/java/org/apache/skywalking/library/elasticsearch/bulk/ElasticSearchIT.java
index cdaca0a2b47cf..fa6fcb43e9b64 100644
--- a/oap-server/server-library/library-client/src/test/java/org/apache/skywalking/library/elasticsearch/bulk/ElasticSearchIT.java
+++ b/oap-server/server-library/library-client/src/test/java/org/apache/skywalking/library/elasticsearch/bulk/ElasticSearchIT.java
@@ -26,8 +26,10 @@
 import org.apache.skywalking.library.elasticsearch.response.Document;
 import org.apache.skywalking.library.elasticsearch.response.Index;
 import org.apache.skywalking.library.elasticsearch.response.Mappings;
+import org.apache.skywalking.library.elasticsearch.response.search.SearchHit;
 import org.apache.skywalking.library.elasticsearch.response.search.SearchResponse;
 import org.apache.skywalking.oap.server.library.client.elasticsearch.ElasticSearchClient;
+import org.apache.skywalking.oap.server.library.client.elasticsearch.ElasticSearchScroller;
 import org.apache.skywalking.oap.server.library.client.elasticsearch.IndexRequestWrapper;
 import org.apache.skywalking.oap.server.library.util.StringUtil;
 import org.junit.jupiter.api.Assertions;
@@ -41,6 +43,7 @@
 import java.util.HashMap;
 import java.util.Map;
 import java.util.Optional;
+import java.util.UUID;
 import java.util.function.Function;
 
 @Slf4j
@@ -171,6 +174,30 @@ public void documentOperate(final ElasticsearchContainer server,
                                                                       .get("message"));
         client.deleteById(indexName, id);
         Assertions.assertFalse(client.existDoc(indexName, id));
+
+        for (int i = 0; i < 100; i++) {
+            builder = ImmutableMap.<String, Object>builder()
+                                  .put("user", "sw")
+                                  .put("post_date", "2009-11-15T14:12:12")
+                                  .put("message", "trying out Elasticsearch")
+                                  .build();
+
+            indexName = "test_scroller";
+            client.forceInsert(indexName, UUID.randomUUID().toString(), builder);
+        }
+
+        SearchBuilder search = Search.builder().size(50);
+        search.query(Query.term("user", "sw"));
+        final var scroller = ElasticSearchScroller
+            .<SearchHit>builder()
+            .client(client)
+            .search(search.build())
+            .index(indexName)
+            .queryMaxSize(40)
+            .resultConverter(Function.identity())
+            .build();
+        Assertions.assertEquals(40, scroller.scroll().size());
+
         client.shutdown();
         server.stop();
     }
diff --git a/oap-server/server-storage-plugin/storage-elasticsearch-plugin/src/main/java/org/apache/skywalking/oap/server/storage/plugin/elasticsearch/query/MetadataQueryEsDAO.java b/oap-server/server-storage-plugin/storage-elasticsearch-plugin/src/main/java/org/apache/skywalking/oap/server/storage/plugin/elasticsearch/query/MetadataQueryEsDAO.java
index 9a103b536c916..5bae5a6431231 100644
--- a/oap-server/server-storage-plugin/storage-elasticsearch-plugin/src/main/java/org/apache/skywalking/oap/server/storage/plugin/elasticsearch/query/MetadataQueryEsDAO.java
+++ b/oap-server/server-storage-plugin/storage-elasticsearch-plugin/src/main/java/org/apache/skywalking/oap/server/storage/plugin/elasticsearch/query/MetadataQueryEsDAO.java
@@ -238,7 +238,7 @@ public List<Endpoint> findEndpoint(String keyword, String serviceId, int limit)
             .client(getClient())
             .search(search.build())
             .index(index)
-            .queryMaxSize(queryMaxSize)
+            .queryMaxSize(limit)
             .resultConverter(searchHit -> {
                 final var sourceAsMap = searchHit.getSource();