extend cat api to get shards information for a specific index

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

@@ -70,6 +70,7 @@
 import org.graylog2.indexer.indices.IndexSettings;
 import org.graylog2.indexer.indices.Indices;
 import org.graylog2.indexer.indices.IndicesAdapter;
+import org.graylog2.indexer.indices.ShardsInfo;
 import org.graylog2.indexer.indices.Template;
 import org.graylog2.indexer.indices.blocks.IndicesBlockStatus;
 import org.graylog2.indexer.indices.stats.IndexStatistics;
@@ -384,6 +385,11 @@ public JsonNode getIndexStats(Collection<String> indices) {
         return statsApi.indexStatsWithDocsAndStore(indices);
     }
 
+    @Override
+    public List<ShardsInfo> getShardsInfo(String indexName) {
+        return catApi.getShardsInfo(indexName);
+    }
+
     @Override
     public IndicesBlockStatus getIndicesBlocksStatus(final List<String> indices) {
         if (indices == null || indices.isEmpty()) {

graylog-storage-elasticsearch7/src/main/java/org/graylog/storage/elasticsearch7/cat/CatApi.java

@@ -26,8 +26,10 @@
 import org.graylog.storage.elasticsearch7.ElasticsearchClient;
 
 import jakarta.inject.Inject;
+import org.graylog2.indexer.indices.ShardsInfo;
 
 import java.io.IOException;
+import java.net.UnknownHostException;
 import java.util.Collection;
 import java.util.Collections;
 import java.util.List;
@@ -94,6 +96,21 @@ public Optional<String> indexState(String indexName, String errorMessage) {
                 .findFirst();
     }
 
+    public List<ShardsInfo> getShardsInfo(String indexName) {
+        return requestShardsInfo(indexName).stream().map(jsonNode -> {
+            try {
+                return ShardsInfo.create(jsonNode);
+            } catch (UnknownHostException e) {
+                throw new RuntimeException(e);
+            }
+        }).toList();
+    }
+
+    private List<JsonNode> requestShardsInfo(String indexName) {
+        final Request request = request("GET", "shards/" + indexName);
+        return perform(request, new TypeReference<List<JsonNode>>() {}, "Unable to retrieve index shards");
+    }
+
     private JsonNode requestIndices(String indexName, String errorMessage) {
         final Request request = request("GET", "indices/" + indexName);
         request.addParameter("h", "index,status");

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

@@ -70,6 +70,7 @@
 import org.graylog2.indexer.indices.IndexSettings;
 import org.graylog2.indexer.indices.Indices;
 import org.graylog2.indexer.indices.IndicesAdapter;
+import org.graylog2.indexer.indices.ShardsInfo;
 import org.graylog2.indexer.indices.Template;
 import org.graylog2.indexer.indices.blocks.IndicesBlockStatus;
 import org.graylog2.indexer.indices.stats.IndexStatistics;
@@ -383,6 +384,11 @@ public JsonNode getIndexStats(Collection<String> indices) {
         return statsApi.indexStatsWithDocsAndStore(indices);
     }
 
+    @Override
+    public List<ShardsInfo> getShardsInfo(String indexName) {
+        return catApi.getShardsInfo(indexName);
+    }
+
     @Override
     public IndicesBlockStatus getIndicesBlocksStatus(final List<String> indices) {
         if (indices == null || indices.isEmpty()) {

graylog-storage-opensearch2/src/main/java/org/graylog/storage/opensearch2/cat/CatApi.java

@@ -25,8 +25,10 @@
 import org.graylog.shaded.opensearch2.org.opensearch.client.Response;
 
 import jakarta.inject.Inject;
+import org.graylog2.indexer.indices.ShardsInfo;
 
 import java.io.IOException;
+import java.net.UnknownHostException;
 import java.util.Collection;
 import java.util.Collections;
 import java.util.List;
@@ -93,6 +95,21 @@ public Optional<String> indexState(String indexName, String errorMessage) {
                 .findFirst();
     }
 
+    public List<ShardsInfo> getShardsInfo(String indexName) {
+        return requestShardsInfo(indexName).stream().map(jsonNode -> {
+            try {
+                return ShardsInfo.create(jsonNode);
+            } catch (UnknownHostException e) {
+                throw new RuntimeException(e);
+            }
+        }).toList();
+    }
+
+    private List<JsonNode> requestShardsInfo(String indexName) {
+        final Request request = request("GET", "shards/" + indexName);
+        return perform(request, new TypeReference<List<JsonNode>>() {}, "Unable to retrieve index shards");
+    }
+
     private JsonNode requestIndices(String indexName, String errorMessage) {
         final Request request = request("GET", "indices/" + indexName);
         request.addParameter("h", "index,status");

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

@@ -357,6 +357,10 @@ public Set<IndexStatistics> getIndicesStats(final Collection<String> indices) {
         return indicesAdapter.indicesStats(indices);
     }
 
+    public List<ShardsInfo> getShardsInfo(String indexName) {
+        return indicesAdapter.getShardsInfo(indexName);
+    }
+
     public void cycleAlias(String aliasName, String targetIndex) {
         indicesAdapter.cycleAlias(aliasName, targetIndex);
     }

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

@@ -97,6 +97,8 @@ public interface IndicesAdapter {
 
     JsonNode getIndexStats(Collection<String> index);
 
+    List<ShardsInfo> getShardsInfo(String indexName);
+
     IndicesBlockStatus getIndicesBlocksStatus(List<String> indices);
 
     boolean exists(String indexName) throws IOException;

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

@@ -0,0 +1,71 @@
+/*
+ * Copyright (C) 2020 Graylog, Inc.
+ *
+ * This program is free software: you can redistribute it and/or modify
+ * it under the terms of the Server Side Public License, version 1,
+ * as published by MongoDB, Inc.
+ *
+ * This program is distributed in the hope that it will be useful,
+ * but WITHOUT ANY WARRANTY; without even the implied warranty of
+ * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
+ * Server Side Public License for more details.
+ *
+ * You should have received a copy of the Server Side Public License
+ * along with this program. If not, see
+ * <http://www.mongodb.com/licensing/server-side-public-license>.
+ */
+package org.graylog2.indexer.indices;
+
+import com.fasterxml.jackson.databind.JsonNode;
+import org.apache.commons.lang3.EnumUtils;
+
+import java.net.InetAddress;
+import java.net.UnknownHostException;
+import java.util.Locale;
+import java.util.function.Function;
+
+public record ShardsInfo(String index, int shard, ShardType shardType, State state, long docs, String store, InetAddress ip, String node ) {
+
+    public static ShardsInfo create(JsonNode jsonNode) throws UnknownHostException {
+
+        String index = jsonNode.get("index").asText();
+        int shard =jsonNode.get("shard").asInt();
+
+        String ipString = getValueOrDefault(jsonNode, "ip", JsonNode::asText, null);
+        InetAddress ip = ipString != null ? InetAddress.getByName(ipString) : null;
+
+        String store = getValueOrDefault(jsonNode, "store", JsonNode::asText, null);
+        String node = getValueOrDefault(jsonNode, "node", JsonNode::asText, null);
+        long docs = getValueOrDefault(jsonNode, "docs", JsonNode::asLong, 0L);
+
+        State state = EnumUtils.getEnumIgnoreCase(State.class, jsonNode.get("state").asText(), State.UNKNOWN);
+        ShardType shardType = ShardType.fromString(jsonNode.get("prirep").asText());
+
+        return new ShardsInfo(index, shard, shardType, state, docs, store, ip, node);
+    }
+
+    private static <T> T getValueOrDefault(JsonNode jsonNode, String nodeName, Function<JsonNode, T> valueConverter, T defaultValue) {
+        return jsonNode.hasNonNull(nodeName) ? valueConverter.apply(jsonNode) :  defaultValue;
+    }
+
+    public enum State {
+        INITIALIZING,
+        RELOCATING,
+        UNASSIGNED,
+        STARTED,
+        UNKNOWN
+    }
+
+    public enum ShardType {
+        PRIMARY,
+        REPLICA,
+        UNKNOWN;
+        public static ShardType fromString(String value) {
+            return switch (value.toLowerCase(Locale.ENGLISH)) {
+                case "r" -> REPLICA;
+                case "p" -> PRIMARY;
+                default -> UNKNOWN;
+            };
+        }
+    }
+}

graylog2-server/src/test/java/org/graylog2/indexer/cluster/ClusterIT.java

@@ -17,14 +17,21 @@
 package org.graylog2.indexer.cluster;
 
 import com.github.joschi.jadconfig.util.Duration;
+import com.google.common.eventbus.EventBus;
 import com.google.common.util.concurrent.ThreadFactoryBuilder;
 import org.graylog.testing.elasticsearch.ElasticsearchBaseTest;
+import org.graylog2.audit.NullAuditEventSender;
+import org.graylog2.indexer.IndexMappingFactory;
 import org.graylog2.indexer.IndexSetRegistry;
 import org.graylog2.indexer.cluster.health.ClusterAllocationDiskSettings;
 import org.graylog2.indexer.cluster.health.NodeDiskUsageStats;
 import org.graylog2.indexer.cluster.health.NodeFileDescriptorStats;
 import org.graylog2.indexer.cluster.health.WatermarkSettings;
+import org.graylog2.indexer.indexset.profile.IndexFieldTypeProfileService;
 import org.graylog2.indexer.indices.HealthStatus;
+import org.graylog2.indexer.indices.Indices;
+import org.graylog2.indexer.indices.ShardsInfo;
+import org.graylog2.plugin.system.SimpleNodeId;
 import org.graylog2.rest.models.system.indexer.responses.ClusterHealth;
 import org.junit.Before;
 import org.junit.Rule;
@@ -33,6 +40,7 @@
 import org.mockito.junit.MockitoJUnit;
 import org.mockito.junit.MockitoRule;
 
+import java.util.List;
 import java.util.Optional;
 import java.util.Set;
 import java.util.concurrent.Executors;
@@ -50,8 +58,15 @@ public abstract class ClusterIT extends ElasticsearchBaseTest {
 
     @Mock
     private IndexSetRegistry indexSetRegistry;
+    @Mock
+    private IndexMappingFactory indexMappingFactory;
+    @Mock
+    private EventBus eventBus;
+    @Mock
+    private IndexFieldTypeProfileService indexFieldTypeProfileService;
 
     protected Cluster cluster;
+    protected Indices indices;
 
     protected abstract ClusterAdapter clusterAdapter(Duration timeout);
 
@@ -70,6 +85,12 @@ public void setUp() throws Exception {
         );
         final Duration requestTimeout = Duration.seconds(1L);
         cluster = new Cluster(indexSetRegistry, scheduler, requestTimeout, clusterAdapter(requestTimeout));
+        indices = new Indices(indexMappingFactory,
+                new SimpleNodeId("123456"),
+                new NullAuditEventSender(),
+                eventBus,
+                searchServer().adapters().indicesAdapter(),
+                indexFieldTypeProfileService);
     }
 
     @Test
@@ -226,4 +247,19 @@ public void getDefaultClusterAllocationDiskSettings() {
         assertThat(clusterAllocationDiskSettings.watermarkSettings().high()).isEqualTo(90D);
         assertThat(clusterAllocationDiskSettings.watermarkSettings().floodStage()).isEqualTo(95D);
     }
+
+    @Test
+    public void getIndexShardsInfo() {
+        client().createIndex("1shard1replica", 1, 1);
+        List<ShardsInfo> shardsInfo = indices.getShardsInfo("1shard1replica");
+        assertThat(shardsInfo.size()).isEqualTo(2);
+        assertThat(shardsInfo.stream()
+                .filter(info -> info.shardType() == ShardsInfo.ShardType.PRIMARY)
+                .findFirst())
+                .isNotEmpty();
+        assertThat(shardsInfo.stream()
+                .filter(info -> info.shardType() == ShardsInfo.ShardType.REPLICA)
+                .findFirst())
+                .isNotEmpty();
+    }
 }