Replace stats API call with more efficient count call. (#24075)

* Replace stats API call with more efficient count call.

Clean unused method in StatsAPI.

Fix test.

* Removing unused StatsAPI method in older storage modules

full-backend-tests/src/test/java/org/graylog2/indexer/indices/IndicesIT.java

@@ -22,9 +22,9 @@ import com.google.common.collect.ImmutableSet;
 import com.google.common.eventbus.EventBus;
 import com.google.common.eventbus.Subscribe;
 import org.apache.commons.codec.binary.Base64;
-import org.graylog.testing.completebackend.Lifecycle;
 import org.graylog.testing.completebackend.FullBackendTest;
 import org.graylog.testing.completebackend.GraylogBackendConfiguration;
+import org.graylog.testing.completebackend.Lifecycle;
 import org.graylog.testing.elasticsearch.SearchServerBaseTest;
 import org.graylog2.audit.NullAuditEventSender;
 import org.graylog2.indexer.IgnoreIndexTemplate;
@@ -36,6 +36,7 @@ import org.graylog2.indexer.IndexTemplateNotFoundException;
 import org.graylog2.indexer.MessageIndexTemplateProvider;
 import org.graylog2.indexer.TestIndexSet;
 import org.graylog2.indexer.cluster.Node;
+import org.graylog2.indexer.counts.CountsAdapter;
 import org.graylog2.indexer.indexset.IndexSetConfig;
 import org.graylog2.indexer.indexset.profile.IndexFieldTypeProfileService;
 import org.graylog2.indexer.indices.blocks.IndicesBlockStatus;
@@ -117,7 +118,8 @@ public class IndicesIT extends SearchServerBaseTest {
                 new NullAuditEventSender(),
                 eventBus,
                 searchServer().adapters().indicesAdapter(),
-                mock(IndexFieldTypeProfileService.class)
+                mock(IndexFieldTypeProfileService.class),
+                searchServer().adapters().countsAdapter()
         );
     }
 
@@ -372,7 +374,8 @@ public class IndicesIT extends SearchServerBaseTest {
                 new NullAuditEventSender(),
                 eventBus,
                 searchServer().adapters().indicesAdapter(),
-                mock(IndexFieldTypeProfileService.class));
+                mock(IndexFieldTypeProfileService.class),
+                mock(CountsAdapter.class));
 
         assertThatCode(() -> indices.ensureIndexTemplate(indexSet)).doesNotThrowAnyException();
 
@@ -405,7 +408,8 @@ public class IndicesIT extends SearchServerBaseTest {
                 new NullAuditEventSender(),
                 eventBus,
                 searchServer().adapters().indicesAdapter(),
-                mock(IndexFieldTypeProfileService.class));
+                mock(IndexFieldTypeProfileService.class),
+                mock(CountsAdapter.class));
 
         assertThatCode(() -> indices.ensureIndexTemplate(indexSet))
                 .isExactlyInstanceOf(IndexTemplateNotFoundException.class)

graylog-storage-elasticsearch7/src/main/java/org/graylog/storage/elasticsearch7/IndicesAdapterES7.java

@@ -358,16 +358,6 @@ public class IndicesAdapterES7 implements IndicesAdapter {
                 "Unable to close index " + index);
     }
 
-    @Override
-    public long numberOfMessages(String index) {
-        final JsonNode result = statsApi.indexStats(index);
-        final JsonNode count = result.path("_all").path("primaries").path("docs").path("count");
-        if (count.isMissingNode()) {
-            throw new RuntimeException("Unable to extract count from response.");
-        }
-        return count.asLong();
-    }
-
     private GetSettingsResponse settingsFor(String indexOrAlias) {
         final GetSettingsRequest request = new GetSettingsRequest().indices(indexOrAlias)
                 .indicesOptions(IndicesOptions.LENIENT_EXPAND_OPEN_CLOSED);

graylog-storage-elasticsearch7/src/main/java/org/graylog/storage/elasticsearch7/stats/StatsApi.java

@@ -19,12 +19,11 @@ package org.graylog.storage.elasticsearch7.stats;
 import com.fasterxml.jackson.databind.JsonNode;
 import com.fasterxml.jackson.databind.ObjectMapper;
 import com.google.common.collect.ImmutableSet;
+import jakarta.inject.Inject;
 import org.graylog.shaded.elasticsearch7.org.elasticsearch.client.Request;
 import org.graylog.shaded.elasticsearch7.org.elasticsearch.client.Response;
 import org.graylog.storage.elasticsearch7.ElasticsearchClient;
 
-import jakarta.inject.Inject;
-
 import java.util.Collection;
 import java.util.Collections;
 import java.util.Optional;
@@ -41,10 +40,6 @@ public class StatsApi {
         this.client = client;
     }
 
-    public JsonNode indexStats(String index) {
-        return stats(index);
-    }
-
     public JsonNode indexStatsWithShardLevel(String index) {
         return indexStatsWithShardLevel(Collections.singleton(index)).path(index);
     }

graylog-storage-opensearch2/src/main/java/org/graylog/storage/opensearch2/IndicesAdapterOS2.java

@@ -359,16 +359,6 @@ public class IndicesAdapterOS2 implements IndicesAdapter {
                 "Unable to close index " + index);
     }
 
-    @Override
-    public long numberOfMessages(String index) {
-        final JsonNode result = statsApi.indexStats(index);
-        final JsonNode count = result.path("_all").path("primaries").path("docs").path("count");
-        if (count.isMissingNode()) {
-            throw new RuntimeException("Unable to extract count from response.");
-        }
-        return count.asLong();
-    }
-
     private GetSettingsResponse settingsFor(String indexOrAlias) {
         final GetSettingsRequest request = new GetSettingsRequest().indices(indexOrAlias)
                 .indicesOptions(IndicesOptions.LENIENT_EXPAND_OPEN_CLOSED);

graylog-storage-opensearch2/src/main/java/org/graylog/storage/opensearch2/stats/StatsApi.java

@@ -19,12 +19,11 @@ package org.graylog.storage.opensearch2.stats;
 import com.fasterxml.jackson.databind.JsonNode;
 import com.fasterxml.jackson.databind.ObjectMapper;
 import com.google.common.collect.ImmutableSet;
+import jakarta.inject.Inject;
 import org.graylog.shaded.opensearch2.org.opensearch.client.Request;
 import org.graylog.shaded.opensearch2.org.opensearch.client.Response;
 import org.graylog.storage.opensearch2.OpenSearchClient;
 
-import jakarta.inject.Inject;
-
 import java.util.Collection;
 import java.util.Collections;
 import java.util.Optional;
@@ -41,10 +40,6 @@ public class StatsApi {
         this.client = client;
     }
 
-    public JsonNode indexStats(String index) {
-        return stats(index);
-    }
-
     public JsonNode indexStatsWithShardLevel(String index) {
         return indexStatsWithShardLevel(Collections.singleton(index)).path(index);
     }

graylog-storage-opensearch3/src/main/java/org/graylog/storage/opensearch3/IndicesAdapterOS.java

@@ -369,11 +369,6 @@ public class IndicesAdapterOS implements IndicesAdapter {
                 "Unable to close index " + index);
     }
 
-    @Override
-    public long numberOfMessages(final String index) {
-        return statsApi.numberOfMessagesInIndex(index);
-    }
-
     private GetSettingsResponse settingsFor(String indexOrAlias) {
         final GetSettingsRequest request = new GetSettingsRequest().indices(indexOrAlias)
                 .indicesOptions(IndicesOptions.LENIENT_EXPAND_OPEN_CLOSED);

graylog-storage-opensearch3/src/main/java/org/graylog/storage/opensearch3/stats/StatsApi.java

@@ -18,7 +18,6 @@ package org.graylog.storage.opensearch3.stats;
 
 import jakarta.inject.Inject;
 import org.graylog.storage.opensearch3.OfficialOpensearchClient;
-import org.opensearch.client.opensearch._types.DocStats;
 import org.opensearch.client.opensearch._types.Level;
 import org.opensearch.client.opensearch._types.StoreStats;
 import org.opensearch.client.opensearch.indices.IndicesStatsRequest;
@@ -43,19 +42,6 @@ public class StatsApi {
         this.client = client;
     }
 
-    public long numberOfMessagesInIndex(final String index) {
-        return Optional.ofNullable(
-                        stats(Collections.singleton(index),
-                                List.of(),
-                                false)
-                                .all()
-                                .primaries()
-                                .docs()
-                ).map(DocStats::count)
-                .orElse(0L);
-
-    }
-
     public IndicesStats indexStatsWithShardLevel(final String index) {
         return indicesStatsWithShardLevel(Collections.singleton(index)).get(index);
     }

graylog-storage-opensearch3/src/test/java/org/graylog/storage/opensearch3/IndicesAdapterOSTest.java

@@ -78,15 +78,6 @@ class IndicesAdapterOSTest {
         );
     }
 
-    @Test
-    void testNumberOfMessages() {
-        doReturn(42L).when(statsApi).numberOfMessagesInIndex("index_1");
-        doReturn(13L).when(statsApi).numberOfMessagesInIndex("index_2");
-
-        assertEquals(42L, toTest.numberOfMessages("index_1"));
-        assertEquals(13L, toTest.numberOfMessages("index_2"));
-    }
-
     @Test
     void testIndicesStats() {
         IndicesStats stats1 = mock(IndicesStats.class);

graylog-storage-opensearch3/src/test/java/org/graylog/storage/opensearch3/stats/StatsApiTest.java

@@ -65,62 +65,6 @@ class StatsApiTest {
         statsApi = new StatsApi(client);
     }
 
-    @Test
-    void testNumberOfMessagesInIndex() throws Exception {
-        final IndicesStatsResponse singleIndexResponse = IndicesStatsResponse.builder()
-                .all(
-                        AllIndicesStats.builder()
-                                .primaries(
-                                        IndexStats.builder()
-                                                .docs(builder -> builder.count(13L))
-                                                .build()
-                                )
-                                .total(emptyIndexStats())
-                                .build()
-                )
-                .indices(
-                        Map.of(
-                                "graylog_13",
-                                buildIndicesStats("P999CCTTFOhEGdspOYh9Q", 13L))
-                )
-                .shards(mock(ShardStatistics.class))
-                .build();
-        doReturn(singleIndexResponse)
-                .when(indicesClient)
-                .stats(IndicesStatsRequest.builder().index(List.of("graylog_13")).build());
-
-        final long returned = statsApi.numberOfMessagesInIndex("graylog_13");
-        assertEquals(13, returned);
-
-        final IndicesStatsResponse multipleIndicesResponse = IndicesStatsResponse.builder()
-                .all(
-                        AllIndicesStats.builder()
-                                .primaries(
-                                        IndexStats.builder()
-                                                .docs(builder -> builder.count(30L))
-                                                .build()
-                                )
-                                .total(emptyIndexStats())
-                                .build()
-                )
-                .indices(
-                        Map.of(
-                                "graylog_13",
-                                buildIndicesStats("P999CCTTFOhEGdspOYh9Q", 13L),
-                                "graylog_12",
-                                buildIndicesStats("P999CCTTFOhEGdspOYh9V", 17L)
-                        )
-                )
-                .shards(mock(ShardStatistics.class))
-                .build();
-
-        doReturn(multipleIndicesResponse)
-                .when(indicesClient)
-                .stats(IndicesStatsRequest.builder().index(List.of("graylog_*")).build());
-
-        assertEquals(30, statsApi.numberOfMessagesInIndex("graylog_*"));
-    }
-
     @Test
     void testStoreSizes() throws Exception {
         final IndicesStatsResponse indicesStatsResponse = IndicesStatsResponse.builder()

graylog2-server/src/main/java/org/graylog2/indexer/indices/Indices.java

@@ -36,6 +36,7 @@ import org.graylog2.indexer.IndexMappingTemplate;
 import org.graylog2.indexer.IndexNotFoundException;
 import org.graylog2.indexer.IndexSet;
 import org.graylog2.indexer.IndexTemplateNotFoundException;
+import org.graylog2.indexer.counts.CountsAdapter;
 import org.graylog2.indexer.indexset.CustomFieldMappings;
 import org.graylog2.indexer.indexset.IndexSetConfig;
 import org.graylog2.indexer.indexset.IndexSetMappingTemplate;
@@ -78,6 +79,7 @@ public class Indices {
     private final EventBus eventBus;
     private final IndicesAdapter indicesAdapter;
     private final IndexFieldTypeProfileService profileService;
+    private final CountsAdapter countsAdapter;
 
     @Inject
     public Indices(IndexMappingFactory indexMappingFactory,
@@ -85,13 +87,15 @@ public class Indices {
                    AuditEventSender auditEventSender,
                    EventBus eventBus,
                    IndicesAdapter indicesAdapter,
-                   IndexFieldTypeProfileService profileService) {
+                   IndexFieldTypeProfileService profileService,
+                   CountsAdapter countsAdapter) {
         this.indexMappingFactory = indexMappingFactory;
         this.nodeId = nodeId;
         this.auditEventSender = auditEventSender;
         this.eventBus = eventBus;
         this.indicesAdapter = indicesAdapter;
         this.profileService = profileService;
+        this.countsAdapter = countsAdapter;
     }
 
     public IndicesBlockStatus getIndicesBlocksStatus(final List<String> indices) {
@@ -138,7 +142,7 @@ public class Indices {
     }
 
     public long numberOfMessages(String indexName) throws IndexNotFoundException {
-        return indicesAdapter.numberOfMessages(indexName);
+        return countsAdapter.totalCount(List.of(indexName));
     }
 
     public IndexSetStats getIndexSetStats() {

graylog2-server/src/main/java/org/graylog2/indexer/indices/IndicesAdapter.java

@@ -88,8 +88,6 @@ public interface IndicesAdapter {
 
     void close(String indexName);
 
-    long numberOfMessages(String indexName);
-
     boolean aliasExists(String alias) throws IOException;
 
     Map<String, Set<String>> aliases(String indexPattern);

graylog2-server/src/test/java/org/graylog2/indexer/fieldtypes/IndexFieldTypePollerIT.java

@@ -27,6 +27,7 @@ import org.graylog2.indexer.IndexMappingFactory;
 import org.graylog2.indexer.MessageIndexTemplateProvider;
 import org.graylog2.indexer.TestIndexSet;
 import org.graylog2.indexer.cluster.Node;
+import org.graylog2.indexer.counts.CountsAdapter;
 import org.graylog2.indexer.indexset.IndexSetConfig;
 import org.graylog2.indexer.indexset.profile.IndexFieldTypeProfileService;
 import org.graylog2.indexer.indices.Indices;
@@ -101,7 +102,8 @@ public abstract class IndexFieldTypePollerIT extends ElasticsearchBaseTest {
                 new NullAuditEventSender(),
                 mock(EventBus.class),
                 createIndicesAdapter(),
-                mock(IndexFieldTypeProfileService.class)
+                mock(IndexFieldTypeProfileService.class),
+                mock(CountsAdapter.class)
         );
         final Configuration withStreamAwarenessOff = spy(new Configuration());
         doReturn(false).when(withStreamAwarenessOff).maintainsStreamAwareFieldTypes();

graylog2-server/src/test/java/org/graylog2/indexer/indices/IndicesGetAllMessageFieldsIT.java

@@ -25,6 +25,7 @@ import org.graylog2.indexer.IndexMappingFactory;
 import org.graylog2.indexer.MessageIndexTemplateProvider;
 import org.graylog2.indexer.cluster.Node;
 import org.graylog2.indexer.cluster.NodeAdapter;
+import org.graylog2.indexer.counts.CountsAdapter;
 import org.graylog2.indexer.indexset.profile.IndexFieldTypeProfileService;
 import org.graylog2.plugin.system.SimpleNodeId;
 import org.junit.Before;
@@ -60,7 +61,8 @@ public abstract class IndicesGetAllMessageFieldsIT extends ElasticsearchBaseTest
                 new NullAuditEventSender(),
                 new EventBus(),
                 indicesAdapter(),
-                mock(IndexFieldTypeProfileService.class)
+                mock(IndexFieldTypeProfileService.class),
+                mock(CountsAdapter.class)
         );
     }
 

graylog2-server/src/test/java/org/graylog2/indexer/indices/IndicesTest.java

@@ -23,6 +23,7 @@ import org.graylog2.indexer.IndexMappingFactory;
 import org.graylog2.indexer.IndexMappingTemplate;
 import org.graylog2.indexer.IndexTemplateNotFoundException;
 import org.graylog2.indexer.TestIndexSet;
+import org.graylog2.indexer.counts.CountsAdapter;
 import org.graylog2.indexer.indexset.CustomFieldMapping;
 import org.graylog2.indexer.indexset.CustomFieldMappings;
 import org.graylog2.indexer.indexset.IndexSetConfig;
@@ -76,6 +77,8 @@ class IndicesTest {
     private IndicesAdapter indicesAdapter;
     @Mock
     private IndexFieldTypeProfileService profileService;
+    @Mock
+    private CountsAdapter countsAdapter;
 
     @BeforeEach
     public void setup() {
@@ -85,7 +88,8 @@ class IndicesTest {
                 auditEventSender,
                 eventBus,
                 indicesAdapter,
-                profileService
+                profileService,
+                countsAdapter
         );
     }