More precise total data set size verification in FrozenSearchableSnapshotsIntegTests (#73243)

In #70625 we added the total data set size of shards 
to the Indices Stats API and we enhanced the test 
testCreateAndRestorePartialSearchableSnapshot to 
also verify the correctness of this data set size.

Because restoring a searchable snapshot shard 
creates a new in-memory segment size, the 
verification of the data set size was implemented 
in an approximative fashion: between the 
expected size and twice the expected size. This 
approximation sometimes fails for shards that 
have no documents indexed (see #73194).

This commit changes the test so that it now 
verifies the exact data set size returned by the 
Indices Stats API, which should be the sum of 
the original expected size of the snapshotted 
size + the length of the extra segment file in 
memory.

Closes #73194

Author: tlrx

x-pack/plugin/searchable-snapshots/src/internalClusterTest/java/org/elasticsearch/xpack/searchablesnapshots/FrozenSearchableSnapshotsIntegTests.java

@@ -7,7 +7,11 @@
 
 package org.elasticsearch.xpack.searchablesnapshots;
 
+import org.apache.lucene.index.SegmentInfos;
 import org.apache.lucene.search.TotalHits;
+import org.apache.lucene.store.ByteBuffersDirectory;
+import org.apache.lucene.store.Directory;
+import org.apache.lucene.store.FilterDirectory;
 import org.elasticsearch.ResourceNotFoundException;
 import org.elasticsearch.action.admin.cluster.snapshots.restore.RestoreSnapshotResponse;
 import org.elasticsearch.action.admin.cluster.snapshots.status.SnapshotIndexShardStatus;
@@ -17,6 +21,7 @@
 import org.elasticsearch.action.admin.indices.stats.ShardStats;
 import org.elasticsearch.cluster.metadata.IndexMetadata;
 import org.elasticsearch.cluster.node.DiscoveryNode;
+import org.elasticsearch.cluster.routing.ShardRouting;
 import org.elasticsearch.cluster.routing.allocation.decider.Decision;
 import org.elasticsearch.cluster.routing.allocation.decider.DiskThresholdDecider;
 import org.elasticsearch.common.Priority;
@@ -32,6 +37,7 @@
 import org.elasticsearch.index.store.Store;
 import org.elasticsearch.index.store.StoreStats;
 import org.elasticsearch.indices.IndexClosedException;
+import org.elasticsearch.indices.IndicesService;
 import org.elasticsearch.snapshots.SnapshotInfo;
 import org.elasticsearch.xpack.cluster.routing.allocation.DataTierAllocationDecider;
 import org.elasticsearch.xpack.core.DataTier;
@@ -54,10 +60,11 @@
 import static org.elasticsearch.xpack.searchablesnapshots.SearchableSnapshots.getDataTiersPreference;
 import static org.elasticsearch.xpack.searchablesnapshots.SearchableSnapshotsConstants.SNAPSHOT_DIRECTORY_FACTORY_KEY;
 import static org.elasticsearch.xpack.searchablesnapshots.SearchableSnapshotsConstants.SNAPSHOT_RECOVERY_STATE_FACTORY_KEY;
+import static org.hamcrest.Matchers.arrayWithSize;
 import static org.hamcrest.Matchers.equalTo;
 import static org.hamcrest.Matchers.greaterThan;
-import static org.hamcrest.Matchers.greaterThanOrEqualTo;
-import static org.hamcrest.Matchers.lessThan;
+import static org.hamcrest.Matchers.instanceOf;
+import static org.hamcrest.Matchers.notNullValue;
 import static org.hamcrest.Matchers.oneOf;
 import static org.hamcrest.Matchers.sameInstance;
 
@@ -219,24 +226,40 @@ public void testCreateAndRestorePartialSearchableSnapshot() throws Exception {
         long totalExpectedSize = 0;
         for (ShardStats shardStats : indicesStatsResponse.getShards()) {
             StoreStats store = shardStats.getStats().getStore();
-            assertThat(shardStats.getShardRouting().toString(), store.getReservedSize().getBytes(), equalTo(0L));
-            assertThat(shardStats.getShardRouting().toString(), store.getSize().getBytes(), equalTo(0L));
-
-            // the extra segments_N file created for bootstrap new history and associate translog makes us unable to precisely assert this.
-            final long expectedSize = snapshotShards.get(shardStats.getShardRouting().getId()).getStats().getTotalSize();
-            assertThat(shardStats.getShardRouting().toString(), store.getTotalDataSetSize().getBytes(), greaterThanOrEqualTo(expectedSize));
-            // the extra segments_N file only has a new history UUID and translog UUID, both of which have constant size. It's size is
-            // therefore identical to the original segments_N file from the snapshot. We expect at least 1 byte of other content, making
-            // it safe to assert that the total data set size is less than 2x the size.
-            assertThat(shardStats.getShardRouting().toString(), store.getTotalDataSetSize().getBytes(), lessThan(expectedSize * 2));
-
-            totalExpectedSize += expectedSize;
+
+            final ShardRouting shardRouting = shardStats.getShardRouting();
+            assertThat(shardRouting.toString(), store.getReservedSize().getBytes(), equalTo(0L));
+            assertThat(shardRouting.toString(), store.getSize().getBytes(), equalTo(0L));
+
+            // the original shard size from the snapshot
+            final long originalSize = snapshotShards.get(shardRouting.getId()).getStats().getTotalSize();
+            totalExpectedSize += originalSize;
+
+            // an extra segments_N file is created for bootstrapping new history and associating translog. We can extract the size of this
+            // extra file but we have to unwrap the in-memory directory first.
+            final Directory unwrappedDir = FilterDirectory.unwrap(
+                internalCluster().getInstance(IndicesService.class, getDiscoveryNodes().resolveNode(shardRouting.currentNodeId()).getName())
+                    .indexServiceSafe(shardRouting.index())
+                    .getShard(shardRouting.getId())
+                    .store()
+                    .directory()
+            );
+            assertThat(shardRouting.toString(), unwrappedDir, notNullValue());
+            assertThat(shardRouting.toString(), unwrappedDir, instanceOf(ByteBuffersDirectory.class));
+
+            final ByteBuffersDirectory inMemoryDir = (ByteBuffersDirectory) unwrappedDir;
+            assertThat(inMemoryDir.listAll(), arrayWithSize(1));
+
+            final String segmentsFileName = SegmentInfos.getLastCommitSegmentsFileName(inMemoryDir);
+            assertThat("Fail to find segment file name directory for " + shardRouting.toString(), segmentsFileName, notNullValue());
+            final long extraSegmentFileSize = inMemoryDir.fileLength(segmentsFileName);
+
+            assertThat(shardRouting.toString(), store.getTotalDataSetSize().getBytes(), equalTo(originalSize + extraSegmentFileSize));
+            totalExpectedSize += extraSegmentFileSize;
         }
 
-        // the extra segments_N file created for bootstrap new history and associate translog makes us unable to precisely assert this.
         final StoreStats store = indicesStatsResponse.getTotal().getStore();
-        assertThat(store.getTotalDataSetSize().getBytes(), greaterThanOrEqualTo(totalExpectedSize));
-        assertThat(store.getTotalDataSetSize().getBytes(), lessThan(totalExpectedSize * 2));
+        assertThat(store.getTotalDataSetSize().getBytes(), equalTo(totalExpectedSize));
 
         statsWatcherRunning.set(false);
         statsWatcher.join();