Limit size of shardDeleteResults

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

@@ -0,0 +1,73 @@
+/*
+ * 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;
+
+/**
+ * Prevents writes when the max size is breached
+ */
+public class BoundedOutputStream extends FilterOutputStream {
+    private final int maxSize;
+    private int size;
+
+    // As soon as a write request exceeds maxSize, permit no more writes, even if there is capacity for them
+    private boolean closed = false;
+
+    public BoundedOutputStream(OutputStream out, int maxSize) {
+        super(out);
+        this.maxSize = maxSize;
+        this.size = 0;
+    }
+
+    private boolean hasCapacity(int bytes) {
+        return size + bytes <= maxSize;
+    }
+
+    @Override
+    public void write(int b) throws IOException {
+        if (closed == false && hasCapacity(1)) {
+            super.write(b);
+
+            /*
+                We only need to increment size here as both super.write(byte[] b) and
+                super.write(byte[] b, int off, int len) write each byte individually via this
+                method, and we have already checked in each respective method whether we have
+                sufficient capacity for that entire write
+             */
+            size++;
+        } else {
+            closed = true;
+            throw new BoundedOutputStreamFailedWriteException();
+        }
+    }
+
+    @Override
+    public void write(byte[] b) throws IOException {
+        if (closed == false && hasCapacity(b.length)) {
+            super.write(b);
+        } else {
+            closed = true;
+            throw new BoundedOutputStreamFailedWriteException();
+        }
+    }
+
+    @Override
+    public void write(byte[] b, int off, int len) throws IOException {
+        if (closed == false && hasCapacity(len)) {
+            super.write(b, off, len);
+        } else {
+            closed = true;
+            throw new BoundedOutputStreamFailedWriteException();
+        }
+    }
+}

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

@@ -0,0 +1,21 @@
+/*
+ * 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.ElasticsearchException;
+
+/**
+ * An exception indicating we have tried to write to the BoundedOutputStream and have exceeded capacity
+ */
+public class BoundedOutputStreamFailedWriteException extends ElasticsearchException {
+    public BoundedOutputStreamFailedWriteException() {
+        super("The write failed because there is no more capacity inside the BoundedOutputStream");
+    }
+}

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/StreamOutput.java

@@ -220,6 +220,26 @@ public void writeIntLE(int i) throws IOException {
         writeBytes(buffer, 0, 4);
     }
 
+    /**
+     * Returns the number of bytes needed to encode the given int as a VInt.
+     */
+    public static int bytesInVInt(int i) {
+        // Single byte shortcut: fits in 7 bits (i.e., values 0..127)
+        if (Integer.numberOfLeadingZeros(i) >= 25) {
+            return 1;
+        }
+        int byteCount = 1;
+        i >>>= 7;
+        while ((i & ~0x7F) != 0) {
+            byteCount++;
+            i >>>= 7;
+        }
+        if (i != 0) {
+            byteCount++;
+        }
+        return byteCount;
+    }
+
     /**
      * Writes an int in a variable-length format.  Writes between one and
      * five bytes.  Smaller values take fewer bytes.  Negative numbers
@@ -235,7 +255,7 @@ public void writeVInt(int i) throws IOException {
          * In that case benchmarks of the method itself are faster but
          * benchmarks of methods that use this method are slower.
          * This is philosophically in line with vint in general - it biases
-         * twoards being simple and fast for smaller numbers.
+         * towards being simple and fast for smaller numbers.
          */
         if (Integer.numberOfLeadingZeros(i) >= 25) {
             writeByte((byte) i);
@@ -430,6 +450,27 @@ public void writeString(String str) throws IOException {
         writeString(str, scratch.get(), 0);
     }
 
+    /**
+     * Returns the number of bytes needed to encode the given string in UTF-8.
+     */
+    public static int bytesInString(String str) {
+        int byteCount = 0;
+        final int charCount = str.length();
+        for (int i = 0; i < charCount; i++) {
+            final int c = str.charAt(i);
+            if (c <= 0x007F) {
+                byteCount += 1;
+            } else if (c > 0x07FF) {
+                byteCount += 3;
+            } else {
+                byteCount += 2;
+            }
+        }
+        // Add bytes for the length prefix (VInt)
+        byteCount += bytesInVInt(charCount);
+        return byteCount;
+    }
+
     /**
      * Write string as well as possibly the beginning of the given {@code buffer}. The given {@code buffer} will also be used when encoding
      * the given string.
@@ -1137,6 +1178,17 @@ public <T> void writeCollection(final Collection<T> collection, final Writer<T>
         }
     }
 
+    /**
+     * Returns the number of bytes needed to encode the given string collectiom
+     */
+    public static int bytesInStringCollection(Collection<String> collection) {
+        int byteCount = bytesInVInt(collection.size());
+        for (String item : collection) {
+            byteCount += bytesInString(item);
+        }
+        return byteCount;
+    }
+
     /**
      * Writes a collection of strings which can then be read using {@link StreamInput#readStringCollectionAsList} or another {@code
      * readStringCollectionAs*} method. Make sure to read the collection back into the same type as was originally written.

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

@@ -65,6 +65,8 @@
 import org.elasticsearch.common.compress.DeflateCompressor;
 import org.elasticsearch.common.compress.NotXContentException;
 import org.elasticsearch.common.io.Streams;
+import org.elasticsearch.common.io.stream.BoundedOutputStream;
+import org.elasticsearch.common.io.stream.BoundedOutputStreamFailedWriteException;
 import org.elasticsearch.common.io.stream.BytesStreamOutput;
 import org.elasticsearch.common.io.stream.InputStreamStreamInput;
 import org.elasticsearch.common.io.stream.OutputStreamStreamOutput;
@@ -1006,7 +1008,8 @@ private void createSnapshotsDeletion(
                     SnapshotsServiceUtils.minCompatibleVersion(minimumNodeVersion, originalRepositoryData, snapshotIds),
                     originalRootBlobs,
                     blobStore().blobContainer(indicesPath()).children(OperationPurpose.SNAPSHOT_DATA),
-                    originalRepositoryData
+                    originalRepositoryData,
+                    metadata.settings()
                 );
             }));
         }
@@ -1075,6 +1078,7 @@ class SnapshotsDeletion {
          * {@link RepositoryData} blob newer than the one identified by {@link #originalRepositoryDataGeneration}.
          */
         private final RepositoryData originalRepositoryData;
+        private final Settings settings;
 
         /**
          * Executor to use for all repository interactions.
@@ -1096,15 +1100,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 +1118,9 @@ class SnapshotsDeletion {
             this.originalRootBlobs = originalRootBlobs;
             this.originalIndexContainers = originalIndexContainers;
             this.originalRepositoryData = originalRepositoryData;
+            this.settings = settings;
+
+            shardBlobsToDelete = new ShardBlobsToDelete(this.settings);
         }
 
         // ---------------------------------------------------------------------------------------------------------------------------------
@@ -1477,6 +1485,7 @@ private void cleanupUnlinkedShardLevelBlobs(ActionListener<Void> listener) {
                 listener.onResponse(null);
                 return;
             }
+
             snapshotExecutor.execute(ActionRunnable.wrap(listener, l -> {
                 try {
                     deleteFromContainer(OperationPurpose.SNAPSHOT_DATA, blobContainer(), filesToDelete);
@@ -1666,6 +1675,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,26 +1688,56 @@ 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() {
-            resources.add(compressed);
-            resources.add(LeakTracker.wrap((Releasable) shardDeleteResults));
+        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);
+            if (this.shardDeleteResultsMaxSize > 0) {
+                this.shardDeleteResults = new ReleasableBytesStreamOutput(bigArrays);
+                this.compressed = new OutputStreamStreamOutput(
+                    new BoundedOutputStream(
+                        new BufferedOutputStream(
+                            new DeflaterOutputStream(Streams.flushOnCloseStream(shardDeleteResults)),
+                            DeflateCompressor.BUFFER_SIZE
+                        ),
+                        this.shardDeleteResultsMaxSize
+                    )
+                );
+                resources.add(compressed);
+                resources.add(LeakTracker.wrap((Releasable) shardDeleteResults));
+            } else {
+                this.shardDeleteResults = null;
+                this.compressed = null;
+            }
+        }
+
+        /**
+         * Calculates the maximum size of the shardDeleteResults BytesStreamOutput.
+         * The size should at most be 2GB, but 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 maxSizeInBytes = MAX_SHARD_DELETE_RESULTS_SIZE_SETTING.get(settings).getBytes();
+            int oneMBBuffer = 1024 * 1024;
+            if (maxSizeInBytes > Integer.MAX_VALUE) {
+                return Integer.MAX_VALUE - oneMBBuffer;
+            }
+            return (int) maxSizeInBytes;
         }
 
         synchronized void addShardDeleteResult(
@@ -1706,10 +1746,24 @@ synchronized void addShardDeleteResult(
             ShardGeneration newGeneration,
             Collection<String> blobsToDelete
         ) {
+            if (compressed == null) {
+                // No output stream: skip writing, but still update generations
+                shardGenerationsBuilder.put(indexId, shardId, newGeneration);
+                return;
+            }
             try {
                 shardGenerationsBuilder.put(indexId, shardId, newGeneration);
                 new ShardSnapshotMetaDeleteResult(Objects.requireNonNull(indexId.getId()), shardId, blobsToDelete).writeTo(compressed);
+
+                // The resultCount is only incremented after a successful complete write
                 resultCount += 1;
+            } catch (BoundedOutputStreamFailedWriteException ex) {
+                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);
@@ -1721,6 +1775,10 @@ public ShardGenerations getUpdatedShardGenerations() {
         }
 
         public Iterator<String> getBlobPaths() {
+            if (compressed == null || shardDeleteResults == null) {
+                // No output stream: nothing to return
+                return Collections.emptyIterator();
+            }
             final StreamInput input;
             try {
                 compressed.close();
@@ -1736,6 +1794,8 @@ public Iterator<String> getBlobPaths() {
                 throw new UncheckedIOException(e);
             }
 
+            // Iterates through complete ShardSnapshotMetaDeleteResults written to compressed
+            // Partially written ShardSnapshotMetaDeleteResults are dropped
             return Iterators.flatMap(Iterators.forRange(0, resultCount, i -> {
                 try {
                     return new ShardSnapshotMetaDeleteResult(input);
@@ -1750,6 +1810,9 @@ public Iterator<String> getBlobPaths() {
 
         @Override
         public void close() {
+            if (resources.isEmpty()) {
+                return;
+            }
             try {
                 IOUtils.close(resources);
             } catch (IOException e) {
@@ -1760,7 +1823,7 @@ public void close() {
 
         // exposed for tests
         int sizeInBytes() {
-            return shardDeleteResults.size();
+            return shardDeleteResults == null ? 0 : shardDeleteResults.size();
         }
     }