Fix limit doesn't work for findEndpoint API in ES storage

apache · Nov 9, 2023 · 589b7f6 · 589b7f6
1 parent 92af797
commit 589b7f6
Show file tree

Hide file tree

Showing 3 changed files with 32 additions and 5 deletions.
diff --git a/.../org/apache/skywalking/oap/server/library/client/elasticsearch/ElasticSearchScroller.java b/.../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/...lient/src/test/java/org/apache/skywalking/library/elasticsearch/bulk/ElasticSearchIT.java b/...lient/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/...g/apache/skywalking/oap/server/storage/plugin/elasticsearch/query/MetadataQueryEsDAO.java b/...g/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();