Return sum of all primary shards' number of segments

Author: nielsbauman

x-pack/plugin/ilm/qa/multi-node/src/javaRestTest/java/org/elasticsearch/xpack/TimeSeriesRestDriver.java

@@ -24,6 +24,8 @@
 import org.elasticsearch.core.Nullable;
 import org.elasticsearch.core.TimeValue;
 import org.elasticsearch.test.rest.ESRestTestCase;
+import org.elasticsearch.test.rest.ObjectPath;
+import org.elasticsearch.test.rest.Stash;
 import org.elasticsearch.xcontent.ToXContent;
 import org.elasticsearch.xcontent.XContentBuilder;
 import org.elasticsearch.xcontent.XContentType;
@@ -46,9 +48,7 @@
 import java.util.List;
 import java.util.Locale;
 import java.util.Map;
-import java.util.Optional;
 import java.util.concurrent.TimeUnit;
-import java.util.stream.Collectors;
 
 import static org.elasticsearch.cluster.metadata.IndexMetadata.SETTING_NUMBER_OF_SHARDS;
 import static org.elasticsearch.test.ESTestCase.randomAlphaOfLengthBetween;
@@ -373,35 +373,22 @@ private static void ensureGreen(String index) throws IOException {
     @SuppressWarnings("unchecked")
     public static Integer getNumberOfPrimarySegments(RestClient client, String index) throws IOException {
         Response response = client.performRequest(new Request("GET", index + "/_segments"));
-        XContentType entityContentType = XContentType.fromMediaType(response.getEntity().getContentType().getValue());
-        final Map<String, Object> originalResponseEntity = XContentHelper.convertToMap(
-            entityContentType.xContent(),
-            response.getEntity().getContent(),
-            false
-        );
-        if (logger.isTraceEnabled()) {
-            logger.trace(
-                "segments response for {}: {}",
-                index,
-                originalResponseEntity.keySet()
-                    .stream()
-                    .map(key -> key + "=" + originalResponseEntity.get(key))
-                    .collect(Collectors.joining(", ", "{", "}"))
-            );
-        }
-        Map<String, Object> responseEntity = (Map<String, Object>) originalResponseEntity.get("indices");
-        responseEntity = (Map<String, Object>) responseEntity.get(index);
-        responseEntity = (Map<String, Object>) responseEntity.get("shards");
-        List<Map<String, Object>> shards = (List<Map<String, Object>>) responseEntity.get("0");
-        // We want to mamke sure to get the primary shard because there is a chance the replica doesn't have data yet:
-        Optional<Map<String, Object>> shardOptional = shards.stream()
-            .filter(shard -> ((Map<String, Object>) shard.get("routing")).get("primary").equals(true))
-            .findAny();
-        if (shardOptional.isPresent()) {
-            return (Integer) shardOptional.get().get("num_search_segments");
-        } else {
-            throw new RuntimeException("No primary shard found for index " + index);
-        }
+        final Map<String, Object> originalResponseEntity = ESRestTestCase.entityAsMap(response);
+        logger.trace("segments response for {}: {}", index, originalResponseEntity);
+        // We need to use a stash here because the index name is likely a data stream index with a dot.
+        Stash stash = new Stash();
+        stash.stashValue("index", index);
+        Map<String, Object> responseEntity = new ObjectPath(originalResponseEntity).evaluate("indices.${index}.shards", stash);
+        return responseEntity.values()
+            .stream()
+            .mapToInt(
+                shardList -> ((List<Map<String, Object>>) shardList).stream()
+                    .filter(shard -> ((Map<String, Object>) shard.get("routing")).get("primary").equals(true))
+                    .findFirst()
+                    .map(shard -> (Integer) shard.get("num_search_segments"))
+                    .orElseThrow(() -> new IllegalStateException("no primary shard found in " + shardList))
+            )
+            .sum();
     }
 
     public static void updatePolicy(RestClient client, String indexName, String policy) throws IOException {

x-pack/plugin/ilm/qa/multi-node/src/javaRestTest/java/org/elasticsearch/xpack/ilm/actions/SearchableSnapshotActionIT.java

@@ -156,7 +156,7 @@ public void testSearchableSnapshotForceMergesClonedIndex() throws Exception {
                 configureClusterAllocation(true);
             }
 
-            assertForceMergedSnapshotDone(phase, backingIndexName, numberOfPrimaries > 1, true);
+            assertForceMergedSnapshotDone(phase, backingIndexName, numberOfPrimaries, true);
         } catch (Exception | AssertionError e) {
             // Make sure we re-enable allocation in case of failure so that the remaining tests in the suite are not affected.
             configureClusterAllocation(true);
@@ -172,12 +172,13 @@ public void testSearchableSnapshotForceMergesSourceIndex() throws Exception {
         // Data streams have 1 primary shard by default.
         // The test suite runs with 4 nodes, so we can have up to 3 (allocated) replicas.
         final String phase = randomBoolean() ? "cold" : "frozen";
-        final String backingIndexName = prepareDataStreamWithDocs(phase, 1, 0);
+        final int numberOfPrimaries = 1;
+        final String backingIndexName = prepareDataStreamWithDocs(phase, numberOfPrimaries, 0);
 
         // Enable/start ILM on the data stream.
         updateIndexSettings(dataStream, Settings.builder().put(LifecycleSettings.LIFECYCLE_NAME, policy));
 
-        assertForceMergedSnapshotDone(phase, backingIndexName, false, false);
+        assertForceMergedSnapshotDone(phase, backingIndexName, numberOfPrimaries, false);
     }
 
     /**
@@ -209,7 +210,7 @@ public void testSearchableSnapshotForceMergesClonedIndexAfterRetry() throws Exce
 
             configureClusterAllocation(true);
 
-            assertForceMergedSnapshotDone(phase, backingIndexName, numberOfPrimaries > 1, true);
+            assertForceMergedSnapshotDone(phase, backingIndexName, numberOfPrimaries, true);
         } catch (Exception | AssertionError e) {
             // Make sure we re-enable allocation in case of failure so that the remaining tests in the suite are not affected.
             configureClusterAllocation(true);
@@ -1156,6 +1157,13 @@ private Step.StepKey getKeyForIndex(Response response, String indexName) throws
      * first generation backing index.
      */
     private String prepareDataStreamWithDocs(String phase, int numberOfPrimaries, int numberOfReplicas) throws Exception {
+        logger.info(
+            "--> running [{}] with [{}] primaries, [{}] replicas, in phase [{}}",
+            getTestName(),
+            numberOfPrimaries,
+            numberOfReplicas,
+            phase
+        );
         createSnapshotRepo(client(), snapshotRepo, randomBoolean());
         createNewSingletonPolicy(client(), policy, phase, new SearchableSnapshotAction(snapshotRepo, true));
 
@@ -1223,10 +1231,10 @@ private static void assertForceMergeCloneIndexSettings(String backingIndexName,
      *
      * @param phase The phase of the ILM policy that the searchable snapshot action runs in.
      * @param backingIndexName The original backing index name.
-     * @param multiplePrimaries True if the original backing index had multiple primaries, affecting segment assertions.
+     * @param numberOfPrimaries The number of primaries that the original backing index had, affecting segment count assertions.
      * @param withReplicas True if the original backing index had one or more replicas, affecting snapshot index naming assertions.
      */
-    private void assertForceMergedSnapshotDone(String phase, String backingIndexName, boolean multiplePrimaries, boolean withReplicas)
+    private void assertForceMergedSnapshotDone(String phase, String backingIndexName, int numberOfPrimaries, boolean withReplicas)
         throws Exception {
         final String prefix = phase.equals("cold")
             ? SearchableSnapshotAction.FULL_RESTORED_INDEX_PREFIX
@@ -1240,11 +1248,11 @@ private void assertForceMergedSnapshotDone(String phase, String backingIndexName
         // Regardless of whether we force merged the backing index or a clone, the cloned index should not exist (anymore).
         awaitIndexDoesNotExist(FORCE_MERGE_CLONE_INDEX_PREFIX + "-*-" + backingIndexName);
 
-        // Retrieve the number of segments in the first (random) shard of the backing index.
+        // Retrieve the total number of segments across all primary shards of the restored index.
         final Integer numberOfPrimarySegments = getNumberOfPrimarySegments(client(), restoredIndexName);
-        // If the backing index had multiple primaries, some primaries might be empty, but others should have no more than 1 segment
-        if (multiplePrimaries || phase.equals("frozen")) {
-            assertThat(numberOfPrimarySegments, lessThanOrEqualTo(1));
+        // If the backing index had multiple primaries, some primaries might be empty, but others should have no more than 1 segment.
+        if (numberOfPrimaries > 1 || phase.equals("frozen")) {
+            assertThat(numberOfPrimarySegments, lessThanOrEqualTo(numberOfPrimaries));
         } else {
             // If the backing index had only one primary, we expect exactly 1 segment after force merging.
             assertThat(numberOfPrimarySegments, equalTo(1));