Fix testCreateAndRestoreSearchableSnapshot (#55147)

Fixes a couple of related failures in SearchableSnapshotsIntegTests.

Firstly, we were not correctly accounting for the case where the cache was so
small that some/all files were read directly; fixed this by only asserting that
the cache is definitely used if the corresponding node has a cache that's large
enough to hold the whole index.

Secondly, we were not permitting shards to be completely empty, which might be
the case (rarely) if there were not many documents indexed and the distribution
of IDs was a bit unlucky; fixed this by asserting that we get stats for at
least one file for the whole index, rather than for each shard separately.

Closes #55126

x-pack/plugin/core/src/main/java/org/elasticsearch/xpack/core/searchablesnapshots/SearchableSnapshotShardStats.java

@@ -77,6 +77,7 @@ public class SearchableSnapshotShardStats implements Writeable, ToXContentObject
             builder.field("index_uuid", getIndexId().getId());
             builder.startObject("shard");
             {
+                builder.field("id", shardRouting.shardId());
                 builder.field("state", shardRouting.state());
                 builder.field("primary", shardRouting.primary());
                 builder.field("node", shardRouting.currentNodeId());

x-pack/plugin/searchable-snapshots/src/test/java/org/elasticsearch/xpack/searchablesnapshots/BaseSearchableSnapshotsIntegTestCase.java

@@ -55,7 +55,9 @@ public abstract class BaseSearchableSnapshotsIntegTestCase extends ESIntegTestCa
             builder.put(
                 CacheService.SNAPSHOT_CACHE_SIZE_SETTING.getKey(),
                 rarely()
-                    ? new ByteSizeValue(randomIntBetween(0, 10), ByteSizeUnit.KB)
+                    ? randomBoolean()
+                        ? new ByteSizeValue(randomIntBetween(0, 10), ByteSizeUnit.KB)
+                        : new ByteSizeValue(randomIntBetween(0, 1000), ByteSizeUnit.BYTES)
                     : new ByteSizeValue(randomIntBetween(1, 10), ByteSizeUnit.MB)
             );
         }

x-pack/plugin/searchable-snapshots/src/test/java/org/elasticsearch/xpack/searchablesnapshots/SearchableSnapshotsIntegTests.java

@@ -5,6 +5,7 @@
  */
 package org.elasticsearch.xpack.searchablesnapshots;
 
+import com.carrotsearch.hppc.cursors.ObjectCursor;
 import org.apache.lucene.index.IndexFileNames;
 import org.apache.lucene.search.TotalHits;
 import org.elasticsearch.ResourceNotFoundException;
@@ -21,6 +22,7 @@ import org.elasticsearch.common.settings.Settings;
 import org.elasticsearch.common.unit.ByteSizeUnit;
 import org.elasticsearch.common.unit.ByteSizeValue;
 import org.elasticsearch.common.util.concurrent.AtomicArray;
+import org.elasticsearch.env.Environment;
 import org.elasticsearch.index.IndexModule;
 import org.elasticsearch.index.IndexSettings;
 import org.elasticsearch.indices.recovery.RecoveryState;
@@ -31,13 +33,16 @@ import org.elasticsearch.xpack.core.searchablesnapshots.SearchableSnapshotShardS
 import org.elasticsearch.xpack.searchablesnapshots.action.SearchableSnapshotsStatsAction;
 import org.elasticsearch.xpack.searchablesnapshots.action.SearchableSnapshotsStatsRequest;
 import org.elasticsearch.xpack.searchablesnapshots.action.SearchableSnapshotsStatsResponse;
+import org.elasticsearch.xpack.searchablesnapshots.cache.CacheService;
 
 import java.nio.file.Path;
 import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.Collections;
+import java.util.HashSet;
 import java.util.List;
 import java.util.Locale;
+import java.util.Set;
 import java.util.concurrent.CountDownLatch;
 import java.util.stream.StreamSupport;
 
@@ -51,7 +56,6 @@ import static org.hamcrest.Matchers.lessThanOrEqualTo;
 
 public class SearchableSnapshotsIntegTests extends BaseSearchableSnapshotsIntegTestCase {
 
-    @AwaitsFix(bugUrl = "https://github.com/elastic/elasticsearch/issues/55126")
     public void testCreateAndRestoreSearchableSnapshot() throws Exception {
         final String fsRepoName = randomAlphaOfLength(10);
         final String indexName = randomAlphaOfLength(10).toLowerCase(Locale.ROOT);
@@ -261,11 +265,33 @@ public class SearchableSnapshotsIntegTests extends BaseSearchableSnapshotsIntegT
         final NumShards restoredNumShards = getNumShards(indexName);
         assertThat(statsResponse.getStats(), hasSize(restoredNumShards.totalNumShards));
 
+        final long totalSize = statsResponse.getStats()
+            .stream()
+            .flatMap(s -> s.getStats().stream())
+            .mapToLong(SearchableSnapshotShardStats.CacheIndexInputStats::getFileLength)
+            .sum();
+        final Set<String> nodeIdsWithLargeEnoughCache = new HashSet<>();
+        for (ObjectCursor<DiscoveryNode> nodeCursor : client().admin()
+            .cluster()
+            .prepareState()
+            .clear()
+            .setNodes(true)
+            .get()
+            .getState()
+            .nodes()
+            .getDataNodes()
+            .values()) {
+            final Settings nodeSettings = internalCluster().getInstance(Environment.class, nodeCursor.value.getName()).settings();
+            if (totalSize <= CacheService.SNAPSHOT_CACHE_SIZE_SETTING.get(nodeSettings).getBytes()) {
+                nodeIdsWithLargeEnoughCache.add(nodeCursor.value.getId());
+            }
+        }
+        assertThat("Expecting stats to exist for at least one Lucene file", totalSize, greaterThan(0L));
+
         for (SearchableSnapshotShardStats stats : statsResponse.getStats()) {
             final ShardRouting shardRouting = stats.getShardRouting();
             assertThat(stats.getShardRouting().getIndexName(), equalTo(indexName));
             if (shardRouting.started()) {
-                assertThat("Expecting stats to exist for at least 1 Lucene file", stats.getStats().size(), greaterThan(0));
                 for (SearchableSnapshotShardStats.CacheIndexInputStats indexInputStats : stats.getStats()) {
                     final String fileName = indexInputStats.getFileName();
                     assertThat(
@@ -295,7 +321,7 @@ public class SearchableSnapshotsIntegTests extends BaseSearchableSnapshotsIntegT
                             Math.max(indexInputStats.getCachedBytesRead().getCount(), indexInputStats.getCachedBytesWritten().getCount()),
                             equalTo(0L)
                         );
-                    } else {
+                    } else if (nodeIdsWithLargeEnoughCache.contains(stats.getShardRouting().currentNodeId())) {
                         assertThat(
                             "Expected at least 1 cache read or write for " + fileName + " of shard " + shardRouting,
                             Math.max(indexInputStats.getCachedBytesRead().getCount(), indexInputStats.getCachedBytesWritten().getCount()),
@@ -306,6 +332,26 @@ public class SearchableSnapshotsIntegTests extends BaseSearchableSnapshotsIntegT
                             indexInputStats.getOptimizedBytesRead().getCount(),
                             equalTo(0L)
                         );
+                        assertThat(
+                            "Expected no direct read for " + fileName + " of shard " + shardRouting,
+                            indexInputStats.getDirectBytesRead().getCount(),
+                            equalTo(0L)
+                        );
+                    } else {
+                        assertThat(
+                            "Expected at least 1 read or write of any kind for " + fileName + " of shard " + shardRouting,
+                            Math.max(
+                                Math.max(
+                                    indexInputStats.getCachedBytesRead().getCount(),
+                                    indexInputStats.getCachedBytesWritten().getCount()
+                                ),
+                                Math.max(
+                                    indexInputStats.getOptimizedBytesRead().getCount(),
+                                    indexInputStats.getDirectBytesRead().getCount()
+                                )
+                            ),
+                            greaterThan(0L)
+                        );
                     }
                 }
             }