Limit size of shardDeleteResults

server/src/main/java/org/elasticsearch/common/io/stream/StreamInput.java

@@ -72,7 +72,7 @@
  * it means that the "barrier to entry" for adding new methods to this class is relatively low even though it is a shared class with code
  * everywhere. That being said, this class deals primarily with {@code List}s rather than Arrays. For the most part calls should adapt to
  * lists, either by storing {@code List}s internally or just converting to and from a {@code List} when calling. This comment is repeated
- * on {@link StreamInput}.
+ * on {@link StreamOutput}.
  */
 public abstract class StreamInput extends InputStream {
 

server/src/main/java/org/elasticsearch/common/io/stream/TruncatedOutputStream.java

@@ -0,0 +1,60 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the "Elastic License
+ * 2.0", the "GNU Affero General Public License v3.0 only", and the "Server Side
+ * Public License v 1"; you may not use this file except in compliance with, at
+ * your election, the "Elastic License 2.0", the "GNU Affero General Public
+ * License v3.0 only", or the "Server Side Public License, v 1".
+ */
+
+package org.elasticsearch.common.io.stream;
+
+import java.io.FilterOutputStream;
+import java.io.IOException;
+import java.io.OutputStream;
+import java.util.function.IntSupplier;
+
+/**
+ * Truncates writes once the max size is exceeded.
+ * However, when writing byte arrays, the stream does not check whether there is capacity for the full
+ * array prior to writing, so there is overspill of up to b.length - 1.
+ */
+public class TruncatedOutputStream extends FilterOutputStream {
+    private final IntSupplier currentSizeSupplier;
+    private final int maxSize;
+    private boolean hasCapacity = true;
+
+    public TruncatedOutputStream(OutputStream out, IntSupplier currentSizeSupplier, int maxSize) {
+        super(out);
+        this.currentSizeSupplier = currentSizeSupplier;
+        this.maxSize = maxSize;
+    }
+
+    private boolean hasCapacity() {
+        if (hasCapacity) {
+            hasCapacity = currentSizeSupplier.getAsInt() < maxSize;
+        }
+        return hasCapacity;
+    }
+
+    @Override
+    public void write(int b) throws IOException {
+        if (hasCapacity()) {
+            out.write(b);
+        }
+    }
+
+    @Override
+    public void write(byte[] b) throws IOException {
+        if (hasCapacity()) {
+            out.write(b);
+        }
+    }
+
+    @Override
+    public void write(byte[] b, int off, int len) throws IOException {
+        if (hasCapacity()) {
+            out.write(b, off, len);
+        }
+    }
+}

server/src/main/java/org/elasticsearch/repositories/blobstore/BlobStoreRepository.java

@@ -71,6 +71,7 @@
 import org.elasticsearch.common.io.stream.ReleasableBytesStreamOutput;
 import org.elasticsearch.common.io.stream.StreamInput;
 import org.elasticsearch.common.io.stream.StreamOutput;
+import org.elasticsearch.common.io.stream.TruncatedOutputStream;
 import org.elasticsearch.common.lucene.Lucene;
 import org.elasticsearch.common.lucene.store.InputStreamIndexInput;
 import org.elasticsearch.common.settings.Setting;
@@ -1006,7 +1007,8 @@ private void createSnapshotsDeletion(
                     SnapshotsServiceUtils.minCompatibleVersion(minimumNodeVersion, originalRepositoryData, snapshotIds),
                     originalRootBlobs,
                     blobStore().blobContainer(indicesPath()).children(OperationPurpose.SNAPSHOT_DATA),
-                    originalRepositoryData
+                    originalRepositoryData,
+                    metadata.settings()
                 );
             }));
         }
@@ -1096,15 +1098,16 @@ class SnapshotsDeletion {
         /**
          * Tracks the shard-level blobs which can be deleted once all the metadata updates have completed.
          */
-        private final ShardBlobsToDelete shardBlobsToDelete = new ShardBlobsToDelete();
+        private final ShardBlobsToDelete shardBlobsToDelete;
 
         SnapshotsDeletion(
             Collection<SnapshotId> snapshotIds,
             long originalRepositoryDataGeneration,
             IndexVersion repositoryFormatIndexVersion,
             Map<String, BlobMetadata> originalRootBlobs,
             Map<String, BlobContainer> originalIndexContainers,
-            RepositoryData originalRepositoryData
+            RepositoryData originalRepositoryData,
+            Settings settings
         ) {
             this.snapshotIds = snapshotIds;
             this.originalRepositoryDataGeneration = originalRepositoryDataGeneration;
@@ -1113,6 +1116,8 @@ class SnapshotsDeletion {
             this.originalRootBlobs = originalRootBlobs;
             this.originalIndexContainers = originalIndexContainers;
             this.originalRepositoryData = originalRepositoryData;
+
+            shardBlobsToDelete = new ShardBlobsToDelete(settings);
         }
 
         // ---------------------------------------------------------------------------------------------------------------------------------
@@ -1666,6 +1671,7 @@ void writeTo(StreamOutput out) throws IOException {
             }
         }
 
+        private final int shardDeleteResultsMaxSize;
         /**
          * <p>
          *     Shard-level results, i.e. a sequence of {@link ShardSnapshotMetaDeleteResult} objects, except serialized, concatenated, and
@@ -1678,28 +1684,57 @@ void writeTo(StreamOutput out) throws IOException {
          *     need no further synchronization.
          * </p>
          */
-        // If the size of this continues to be a problem even after compression, consider either a hard limit on its size (preferring leaked
-        // blobs over an OOME on the master) or else offloading it to disk or to the repository itself.
-        private final BytesStreamOutput shardDeleteResults = new ReleasableBytesStreamOutput(bigArrays);
+        private final BytesStreamOutput shardDeleteResults;
 
         private int resultCount = 0;
 
-        private final StreamOutput compressed = new OutputStreamStreamOutput(
-            new BufferedOutputStream(
-                new DeflaterOutputStream(Streams.flushOnCloseStream(shardDeleteResults)),
-                DeflateCompressor.BUFFER_SIZE
-            )
-        );
+        private final StreamOutput compressed;
 
         private final ArrayList<Closeable> resources = new ArrayList<>();
 
         private final ShardGenerations.Builder shardGenerationsBuilder = ShardGenerations.builder();
 
-        ShardBlobsToDelete() {
+        // Gets 25% of the heap size to be allocated to the shard_delete_results stream
+        public final Setting<ByteSizeValue> MAX_SHARD_DELETE_RESULTS_SIZE_SETTING = Setting.memorySizeSetting(
+            "repositories.blobstore.max_shard_delete_results_size",
+            "25%",
+            Setting.Property.NodeScope
+        );
+
+        ShardBlobsToDelete(Settings settings) {
+            this.shardDeleteResultsMaxSize = calculateMaximumShardDeleteResultsSize(settings);
+            this.shardDeleteResults = new ReleasableBytesStreamOutput(bigArrays);
+            this.compressed = new OutputStreamStreamOutput(
+                new TruncatedOutputStream(
+                    new BufferedOutputStream(
+                        new DeflaterOutputStream(Streams.flushOnCloseStream(shardDeleteResults)),
+                        DeflateCompressor.BUFFER_SIZE
+                    ),
+                    shardDeleteResults::size,
+                    this.shardDeleteResultsMaxSize
+                )
+            );
             resources.add(compressed);
             resources.add(LeakTracker.wrap((Releasable) shardDeleteResults));
         }
 
+        /**
+         * Calculates the maximum size of the shardDeleteResults BytesStreamOutput.
+         * The size cannot exceed 2GB, without {@code BytesStreamOutput} throwing an IAE,
+         * but should also be no more than 25% of the total remaining heap space.
+         * A buffer of 1MB is maintained, so that even if the stream is of max size, there is room to flush
+         * @return The maximum number of bytes the shardDeleteResults BytesStreamOutput can consume in the heap
+         */
+        int calculateMaximumShardDeleteResultsSize(Settings settings) {
+            long maxHeapSizeInBytes = MAX_SHARD_DELETE_RESULTS_SIZE_SETTING.get(settings).getBytes();
+            int oneMBBuffer = 1024 * 1024;
+            int maxShardDeleteResultsSize = Integer.MAX_VALUE - oneMBBuffer;
+            if (maxHeapSizeInBytes > maxShardDeleteResultsSize) {
+                return maxShardDeleteResultsSize;
+            }
+            return (int) maxHeapSizeInBytes;
+        }
+
         synchronized void addShardDeleteResult(
             IndexId indexId,
             int shardId,
@@ -1708,8 +1743,22 @@ synchronized void addShardDeleteResult(
         ) {
             try {
                 shardGenerationsBuilder.put(indexId, shardId, newGeneration);
-                new ShardSnapshotMetaDeleteResult(Objects.requireNonNull(indexId.getId()), shardId, blobsToDelete).writeTo(compressed);
-                resultCount += 1;
+                // Only write if we have capacity
+                if (shardDeleteResults.size() < this.shardDeleteResultsMaxSize) {
+                    new ShardSnapshotMetaDeleteResult(Objects.requireNonNull(indexId.getId()), shardId, blobsToDelete).writeTo(compressed);
+                    // We only want to read this shard delete result if we were able to write the entire object.
+                    // Otherwise, for partial writes, an EOFException will be thrown upon reading
+                    if (shardDeleteResults.size() < this.shardDeleteResultsMaxSize) {
+                        resultCount += 1;
+                    }
+                } else {
+                    logger.warn(
+                        "Failure to clean up the following dangling blobs, {}, for index {} and shard {}",
+                        blobsToDelete,
+                        indexId,
+                        shardId
+                    );
+                }
             } catch (IOException e) {
                 assert false : e; // no IO actually happens here
                 throw new UncheckedIOException(e);
@@ -1736,16 +1785,20 @@ public Iterator<String> getBlobPaths() {
                 throw new UncheckedIOException(e);
             }
 
-            return Iterators.flatMap(Iterators.forRange(0, resultCount, i -> {
+            List<String> blobPaths = new ArrayList<>();
+            for (int i = 0; i < resultCount; i++) {
+                ShardSnapshotMetaDeleteResult shardResult;
                 try {
-                    return new ShardSnapshotMetaDeleteResult(input);
+                    shardResult = new ShardSnapshotMetaDeleteResult(input);
                 } catch (IOException e) {
                     throw new UncheckedIOException(e);
                 }
-            }), shardResult -> {
-                final var shardPath = shardPath(new IndexId("_na_", shardResult.indexId), shardResult.shardId).buildAsString();
-                return Iterators.map(shardResult.blobsToDelete.iterator(), blob -> shardPath + blob);
-            });
+                String shardPath = shardPath(new IndexId("_na_", shardResult.indexId), shardResult.shardId).buildAsString();
+                for (String blob : shardResult.blobsToDelete) {
+                    blobPaths.add(shardPath + blob);
+                }
+            }
+            return blobPaths.iterator();
         }
 
         @Override

server/src/test/java/org/elasticsearch/common/io/stream/TruncatedOutputStreamTests.java

@@ -0,0 +1,122 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the "Elastic License
+ * 2.0", the "GNU Affero General Public License v3.0 only", and the "Server Side
+ * Public License v 1"; you may not use this file except in compliance with, at
+ * your election, the "Elastic License 2.0", the "GNU Affero General Public
+ * License v3.0 only", or the "Server Side Public License, v 1".
+ */
+
+package org.elasticsearch.common.io.stream;
+
+import org.elasticsearch.test.ESTestCase;
+
+import java.io.ByteArrayOutputStream;
+import java.io.IOException;
+import java.util.ArrayList;
+import java.util.List;
+
+public class TruncatedOutputStreamTests extends ESTestCase {
+
+    public void testWriteSingleBytes() throws IOException {
+        ByteArrayOutputStream byteArrayOutputStream = new ByteArrayOutputStream();
+        int maxSize = randomIntBetween(0, 100);
+        TruncatedOutputStream truncatedOutputStream = new TruncatedOutputStream(
+            byteArrayOutputStream,
+            byteArrayOutputStream::size,
+            maxSize
+        );
+
+        byte[] values = new byte[maxSize];
+
+        // Write enough bytes within the defined maxSize
+        for (int i = 0; i < maxSize; i++) {
+            byte b = randomByte();
+            truncatedOutputStream.write(b);
+            values[i] = b;
+        }
+
+        // The stream should be truncated now that it is filled
+        for (int i = 0; i < randomIntBetween(0, 20); i++) {
+            truncatedOutputStream.write(randomByte());
+        }
+
+        assertArrayEquals(values, byteArrayOutputStream.toByteArray());
+    }
+
+    public void testWriteByteArray() throws IOException {
+        ByteArrayOutputStream byteArrayOutputStream = new ByteArrayOutputStream();
+        int maxSize = randomIntBetween(100, 200);
+        TruncatedOutputStream truncatedOutputStream = new TruncatedOutputStream(
+            byteArrayOutputStream,
+            byteArrayOutputStream::size,
+            maxSize
+        );
+
+        List<Byte> values = new ArrayList<>();
+        int bytesWritten = 0;
+        // Write beyond the streams capacity
+        while (bytesWritten <= maxSize * 2) {
+            byte[] bytes = randomByteArrayOfLength(randomIntBetween(0, 20));
+            truncatedOutputStream.write(bytes);
+
+            // If there was capacity before writing, then the stream wrote the entire array
+            // even if that meant overflowing
+            if (bytesWritten < maxSize) {
+                for (byte b : bytes) {
+                    values.add(b);
+                }
+            }
+
+            bytesWritten += bytes.length;
+        }
+
+        byte[] valuesAsByteArray = new byte[values.size()];
+        int i = 0;
+        for (byte b : values) {
+            valuesAsByteArray[i] = b;
+            i++;
+        }
+
+        assertArrayEquals(valuesAsByteArray, byteArrayOutputStream.toByteArray());
+    }
+
+    public void testWriteByteArrayWithOffsetAndLength() throws IOException {
+        ByteArrayOutputStream byteArrayOutputStream = new ByteArrayOutputStream();
+        int maxSize = randomIntBetween(100, 200);
+        TruncatedOutputStream truncatedOutputStream = new TruncatedOutputStream(
+            byteArrayOutputStream,
+            byteArrayOutputStream::size,
+            maxSize
+        );
+
+        List<Byte> values = new ArrayList<>();
+        int bytesWritten = 0;
+        // Write beyond the streams capacity
+        while (bytesWritten <= maxSize * 2) {
+            byte[] bytes = randomByteArrayOfLength(randomIntBetween(0, 20));
+            int offset = randomIntBetween(0, bytes.length);
+            int length = randomIntBetween(0, bytes.length - offset);
+            truncatedOutputStream.write(bytes, offset, length);
+
+            // If there was capacity before writing, then the stream wrote the sub array
+            // even if that meant overflowing
+            if (bytesWritten < maxSize) {
+                for (int i = offset; i < offset + length; i++) {
+                    values.add(bytes[i]);
+                }
+            }
+
+            bytesWritten += length;
+        }
+
+        byte[] valuesAsByteArray = new byte[values.size()];
+        int i = 0;
+        for (byte b : values) {
+            valuesAsByteArray[i] = b;
+            i++;
+        }
+
+        assertArrayEquals(valuesAsByteArray, byteArrayOutputStream.toByteArray());
+    }
+}