Allow ByteSizeDirectory to expose their data set sizes

In elastic#97822 we introduced the ByteSizeDirectory class that can be
implemented by plugins to allow custom directory implementations to
expose their sizes.

With this change, we allow to also expose the "data set" size.

Relates ES-5995

server/src/main/java/org/elasticsearch/index/store/ByteSizeCachingDirectory.java

@@ -35,20 +35,6 @@ private static class SizeAndModCount {
         }
     }
 
-    private static long estimateSizeInBytes(Directory directory) throws IOException {
-        long estimatedSize = 0;
-        String[] files = directory.listAll();
-        for (String file : files) {
-            try {
-                estimatedSize += directory.fileLength(file);
-            } catch (NoSuchFileException | FileNotFoundException | AccessDeniedException e) {
-                // ignore, the file is not there no more; on Windows, if one thread concurrently deletes a file while
-                // calling Files.size, you can also sometimes hit AccessDeniedException
-            }
-        }
-        return estimatedSize;
-    }
-
     private final SingleObjectCache<SizeAndModCount> size;
     // Both these variables need to be accessed under `this` lock.
     private long modCount = 0;
@@ -112,6 +98,11 @@ public long estimateSizeInBytes() throws IOException {
         }
     }
 
+    @Override
+    public long estimateDataSetSizeInBytes() throws IOException {
+           return estimateSizeInBytes(); // data set size is equal to directory size for most implementations
+    }
+
     @Override
     public IndexOutput createOutput(String name, IOContext context) throws IOException {
         return wrapIndexOutput(super.createOutput(name, context));

server/src/main/java/org/elasticsearch/index/store/ByteSizeDirectory.java

@@ -11,14 +11,43 @@
 import org.apache.lucene.store.Directory;
 import org.apache.lucene.store.FilterDirectory;
 
+import java.io.FileNotFoundException;
 import java.io.IOException;
+import java.nio.file.AccessDeniedException;
+import java.nio.file.NoSuchFileException;
 
 public abstract class ByteSizeDirectory extends FilterDirectory {
 
+    protected static long estimateSizeInBytes(Directory directory) throws IOException {
+        long estimatedSize = 0;
+        String[] files = directory.listAll();
+        for (String file : files) {
+            try {
+                estimatedSize += directory.fileLength(file);
+            } catch (NoSuchFileException | FileNotFoundException | AccessDeniedException e) {
+                // ignore, the file is not there no more; on Windows, if one thread concurrently deletes a file while
+                // calling Files.size, you can also sometimes hit AccessDeniedException
+            }
+        }
+        return estimatedSize;
+    }
+
     protected ByteSizeDirectory(Directory in) {
         super(in);
     }
 
-    /** Return the cumulative size of all files in this directory. */
+    /**
+     * @return the size of the directory
+     *
+     * @throws IOException if an I/O error occurs
+     */
     public abstract long estimateSizeInBytes() throws IOException;
+
+    /**
+     * @return the size of the total data set of the directory (which can differ from {{@link #estimateSizeInBytes()}})
+     *
+     * @throws IOException if an I/O error occurs
+     */
+    public abstract long estimateDataSetSizeInBytes() throws IOException;
+
 }

server/src/main/java/org/elasticsearch/index/store/Store.java

@@ -355,7 +355,8 @@ public CheckIndex.Status checkIndex(PrintStream out) throws IOException {
     public StoreStats stats(long reservedBytes, LongUnaryOperator localSizeFunction) throws IOException {
         ensureOpen();
         long sizeInBytes = directory.estimateSizeInBytes();
-        return new StoreStats(localSizeFunction.applyAsLong(sizeInBytes), sizeInBytes, reservedBytes);
+        long dataSetSizeInBytes = directory.estimateDataSetSizeInBytes();
+        return new StoreStats(localSizeFunction.applyAsLong(sizeInBytes), dataSetSizeInBytes, reservedBytes);
     }
 
     /**
@@ -743,6 +744,11 @@ public long estimateSizeInBytes() throws IOException {
             return ((ByteSizeDirectory) getDelegate()).estimateSizeInBytes();
         }
 
+        @Override
+        public long estimateDataSetSizeInBytes() throws IOException {
+            return ((ByteSizeDirectory) getDelegate()).estimateDataSetSizeInBytes();
+        }
+
         @Override
         public void close() {
             assert false : "Nobody should close this directory except of the Store itself";

server/src/test/java/org/elasticsearch/index/store/ByteSizeCachingDirectoryTests.java

@@ -46,6 +46,7 @@ public void testBasics() throws IOException {
             ByteSizeCachingDirectory cachingDir = new ByteSizeCachingDirectory(countingDir, new TimeValue(0));
             assertEquals(11, cachingDir.estimateSizeInBytes());
             assertEquals(11, cachingDir.estimateSizeInBytes());
+            assertEquals(11, cachingDir.estimateDataSetSizeInBytes());
             assertEquals(1, countingDir.numFileLengthCalls);
 
             try (IndexOutput out = cachingDir.createOutput("foo", IOContext.DEFAULT)) {
@@ -63,6 +64,7 @@ public void testBasics() throws IOException {
             assertEquals(7, countingDir.numFileLengthCalls);
             assertEquals(16, cachingDir.estimateSizeInBytes());
             assertEquals(7, countingDir.numFileLengthCalls);
+            assertEquals(16, cachingDir.estimateDataSetSizeInBytes());
 
             try (IndexOutput out = cachingDir.createTempOutput("bar", "baz", IOContext.DEFAULT)) {
                 out.writeBytes(new byte[4], 4);
@@ -79,6 +81,7 @@ public void testBasics() throws IOException {
             assertEquals(16, countingDir.numFileLengthCalls);
             assertEquals(20, cachingDir.estimateSizeInBytes());
             assertEquals(16, countingDir.numFileLengthCalls);
+            assertEquals(20, cachingDir.estimateDataSetSizeInBytes());
 
             cachingDir.deleteFile("foo");
 
@@ -87,6 +90,7 @@ public void testBasics() throws IOException {
             assertEquals(18, countingDir.numFileLengthCalls);
             assertEquals(15, cachingDir.estimateSizeInBytes());
             assertEquals(18, countingDir.numFileLengthCalls);
+            assertEquals(15, cachingDir.estimateDataSetSizeInBytes());
 
             // Close more than once
             IndexOutput out = cachingDir.createOutput("foo", IOContext.DEFAULT);
@@ -106,6 +110,7 @@ public void testBasics() throws IOException {
             }
             out.close();
             assertEquals(20, cachingDir.estimateSizeInBytes());
+            assertEquals(20, cachingDir.estimateDataSetSizeInBytes());
             assertEquals(27, countingDir.numFileLengthCalls);
         }
     }

server/src/test/java/org/elasticsearch/index/store/StoreTests.java

@@ -49,6 +49,7 @@
 import org.elasticsearch.common.io.stream.InputStreamStreamInput;
 import org.elasticsearch.common.io.stream.OutputStreamStreamOutput;
 import org.elasticsearch.common.lucene.Lucene;
+import org.elasticsearch.common.lucene.store.FilterIndexOutput;
 import org.elasticsearch.common.settings.Settings;
 import org.elasticsearch.common.util.Maps;
 import org.elasticsearch.core.IOUtils;
@@ -82,6 +83,7 @@
 import java.util.Map;
 import java.util.Random;
 import java.util.concurrent.atomic.AtomicInteger;
+import java.util.concurrent.atomic.AtomicLong;
 import java.util.function.LongUnaryOperator;
 
 import static java.util.Collections.emptyList;
@@ -809,6 +811,105 @@ public void testStoreStats() throws IOException {
         IOUtils.close(store);
     }
 
+    public void testStoreSizes() throws IOException {
+        // directory that returns total written bytes as the data set size
+        final var directory = new ByteSizeDirectory(StoreTests.newDirectory(random())) {
+
+            final AtomicLong dataSetBytes = new AtomicLong(0L);
+
+            @Override
+            public long estimateSizeInBytes() throws IOException {
+                return estimateSizeInBytes(getDelegate());
+            }
+
+            @Override
+            public long estimateDataSetSizeInBytes() {
+                return dataSetBytes.get();
+            }
+
+            @Override
+            public IndexOutput createOutput(String name, IOContext context) throws IOException {
+                return wrap(super.createOutput(name, context));
+            }
+
+            @Override
+            public IndexOutput createTempOutput(String prefix, String suffix, IOContext context) throws IOException {
+                return wrap(super.createTempOutput(prefix, suffix, context));
+            }
+
+            private IndexOutput wrap(IndexOutput output) {
+                return new FilterIndexOutput("wrapper", output) {
+                    @Override
+                    public void writeByte(byte b) throws IOException {
+                        super.writeByte(b);
+                        dataSetBytes.incrementAndGet();
+                    }
+
+                    @Override
+                    public void writeBytes(byte[] b, int offset, int length) throws IOException {
+                        super.writeBytes(b, offset, length);
+                        dataSetBytes.addAndGet(length);
+                    }
+                };
+            }
+        };
+
+        final ShardId shardId = new ShardId("index", "_na_", 1);
+        final Store store = new Store(
+            shardId,
+            IndexSettingsModule.newIndexSettings("index", Settings.builder()
+                .put(IndexMetadata.SETTING_VERSION_CREATED, org.elasticsearch.Version.CURRENT)
+                .put(Store.INDEX_STORE_STATS_REFRESH_INTERVAL_SETTING.getKey(), TimeValue.timeValueMinutes(0))
+                .build()),
+            directory,
+            new DummyShardLock(shardId)
+        );
+        long initialStoreSize = 0L;
+        for (String extraFiles : store.directory().listAll()) {
+            assertTrue("expected extraFS file but got: " + extraFiles, extraFiles.startsWith("extra"));
+            initialStoreSize += store.directory().fileLength(extraFiles);
+        }
+
+        StoreStats stats = store.stats(0L, LongUnaryOperator.identity());
+        assertThat(stats.sizeInBytes(), equalTo(initialStoreSize));
+        assertThat(stats.totalDataSetSizeInBytes(), equalTo(initialStoreSize));
+
+        long additionalStoreSize = 0L;
+
+        int iters = randomIntBetween(1, 10);
+        for (int i = 0; i < iters; i++) {
+            try (IndexOutput output = directory.createOutput(i + ".bar", IOContext.DEFAULT)) {
+                BytesRef bytesRef = new BytesRef(TestUtil.randomRealisticUnicodeString(random(), 10, 1024));
+                output.writeBytes(bytesRef.bytes, bytesRef.offset, bytesRef.length);
+                additionalStoreSize += output.getFilePointer();
+            }
+        }
+
+        stats = store.stats(0L, LongUnaryOperator.identity());
+        assertThat(stats.sizeInBytes(), equalTo(initialStoreSize + additionalStoreSize));
+        assertThat(stats.totalDataSetSizeInBytes(), equalTo(initialStoreSize + additionalStoreSize));
+
+        long deletionsStoreSize = 0L;
+
+        var randomFiles = randomSubsetOf(Arrays.asList(directory.listAll()));
+        for (String randomFile : randomFiles) {
+            try {
+                long length = directory.fileLength(randomFile);
+                directory.deleteFile(randomFile);
+                deletionsStoreSize += length;
+            } catch (NoSuchFileException | FileNotFoundException e) {
+                // ignore
+            }
+        }
+
+        stats = store.stats(0L, LongUnaryOperator.identity());
+        assertThat(stats.sizeInBytes(), equalTo(initialStoreSize + additionalStoreSize - deletionsStoreSize));
+        assertThat(stats.totalDataSetSizeInBytes(), equalTo(initialStoreSize + additionalStoreSize));
+
+        deleteContent(store.directory());
+        IOUtils.close(store);
+    }
+
     public static void deleteContent(Directory directory) throws IOException {
         final String[] files = directory.listAll();
         final List<IOException> exceptions = new ArrayList<>();