Limit size of shardDeleteResults

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

@@ -33,6 +33,7 @@ public class BytesStreamOutput extends BytesStream {
     @Nullable
     protected ByteArray bytes;
     protected int count;
+    protected int maximumSize;
 
     /**
      * Create a non recycling {@link BytesStreamOutput} with an initial capacity of 0.
@@ -54,10 +55,15 @@ public BytesStreamOutput(int expectedSize) {
     }
 
     protected BytesStreamOutput(int expectedSize, BigArrays bigArrays) {
+        this(expectedSize, bigArrays, Integer.MAX_VALUE);
+    }
+
+    protected BytesStreamOutput(int expectedSize, BigArrays bigArrays, int maximumSize) {
         this.bigArrays = bigArrays;
         if (expectedSize != 0) {
             this.bytes = bigArrays.newByteArray(expectedSize, false);
         }
+        this.maximumSize = maximumSize;
     }
 
     @Override
@@ -171,8 +177,8 @@ private static void copyToArray(BytesReference bytesReference, byte[] arr) {
     }
 
     protected void ensureCapacity(long offset) {
-        if (offset > Integer.MAX_VALUE) {
-            throw new IllegalArgumentException(getClass().getSimpleName() + " cannot hold more than 2GB of data");
+        if (offset > this.maximumSize) {
+            throw new IllegalArgumentException(getClass().getSimpleName() + " has exceeded it's max size of " + this.maximumSize);
         }
         if (bytes == null) {
             this.bytes = bigArrays.newByteArray(BigArrays.overSize(offset, PageCacheRecycler.PAGE_SIZE_IN_BYTES, 1), false);

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

@@ -30,6 +30,10 @@ public ReleasableBytesStreamOutput(BigArrays bigarrays) {
         this(PageCacheRecycler.PAGE_SIZE_IN_BYTES, bigarrays);
     }
 
+    public ReleasableBytesStreamOutput(BigArrays bigArrays, int maximumSize) {
+        super(PageCacheRecycler.PAGE_SIZE_IN_BYTES, bigArrays, maximumSize);
+    }
+
     public ReleasableBytesStreamOutput(int expectedSize, BigArrays bigArrays) {
         super(expectedSize, bigArrays);
     }

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

@@ -151,6 +151,9 @@
 import java.io.InputStream;
 import java.io.OutputStream;
 import java.io.UncheckedIOException;
+import java.lang.management.ManagementFactory;
+import java.lang.management.MemoryMXBean;
+import java.lang.management.MemoryUsage;
 import java.nio.file.NoSuchFileException;
 import java.util.ArrayList;
 import java.util.Collection;
@@ -1477,6 +1480,7 @@ private void cleanupUnlinkedShardLevelBlobs(ActionListener<Void> listener) {
                 listener.onResponse(null);
                 return;
             }
+
             snapshotExecutor.execute(ActionRunnable.wrap(listener, l -> {
                 try {
                     deleteFromContainer(OperationPurpose.SNAPSHOT_DATA, blobContainer(), filesToDelete);
@@ -1678,26 +1682,66 @@ 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));
+            int maxSizeOfShardDeleteResults = calculateMaximumShardDeleteResultsSize();
+
+            if (maxSizeOfShardDeleteResults > 0) {
+                this.shardDeleteResults = new ReleasableBytesStreamOutput(bigArrays, maxSizeOfShardDeleteResults);
+                this.compressed = new OutputStreamStreamOutput(
+                    new BufferedOutputStream(
+                        new DeflaterOutputStream(Streams.flushOnCloseStream(shardDeleteResults)),
+                        DeflateCompressor.BUFFER_SIZE
+                    )
+                );
+                resources.add(compressed);
+                resources.add(LeakTracker.wrap((Releasable) shardDeleteResults));
+            } else {
+                this.shardDeleteResults = null;
+                this.compressed = null;
+            }
+
+        }
+
+        MemoryMXBean getMemoryMXBean() {
+            return ManagementFactory.getMemoryMXBean();
+        }
+
+        /**
+         * 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
+         * @return The maximum number of bytes the shardDeleteResults BytesStreamOutput can consume in the heap
+         */
+        int calculateMaximumShardDeleteResultsSize() {
+            MemoryMXBean memoryBean = getMemoryMXBean();
+            MemoryUsage heapUsage = memoryBean.getHeapMemoryUsage();
+
+            long heapUsageInBytes = heapUsage.getUsed();
+            long maxHeapUsageInBytes = heapUsage.getMax();
+
+            // Undefined heap
+            if (maxHeapUsageInBytes == -1) {
+                logger.warn("The heap is undefined when attempting to instantiate shardDeleteResults");
+                return 0;
+            }
+
+            long heapSpaceLeftInBytes = maxHeapUsageInBytes - heapUsageInBytes;
+            double maxPercentageHeapSpaceAllocation = 0.25;
+            long heapSpaceLeftToBeUtilisedInBytes = (long) Math.floor(heapSpaceLeftInBytes * maxPercentageHeapSpaceAllocation);
+            if (heapSpaceLeftToBeUtilisedInBytes > Integer.MAX_VALUE) {
+                return Integer.MAX_VALUE;
+            } else {
+                return (int) heapSpaceLeftInBytes;
+            }
         }
 
         synchronized void addShardDeleteResult(
@@ -1706,6 +1750,11 @@ 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);
@@ -1714,13 +1763,18 @@ synchronized void addShardDeleteResult(
                 assert false : e; // no IO actually happens here
                 throw new UncheckedIOException(e);
             }
+            // TODO - Catch the out of memory error
         }
 
         public ShardGenerations getUpdatedShardGenerations() {
             return shardGenerationsBuilder.build();
         }
 
         public Iterator<String> getBlobPaths() {
+            if (compressed == null || shardDeleteResults == null) {
+                // No output stream: nothing to return
+                return Collections.emptyIterator();
+            }
             final StreamInput input;
             try {
                 compressed.close();
@@ -1750,6 +1804,9 @@ public Iterator<String> getBlobPaths() {
 
         @Override
         public void close() {
+            if (resources.isEmpty()) {
+                return;
+            }
             try {
                 IOUtils.close(resources);
             } catch (IOException e) {
@@ -1760,7 +1817,7 @@ public void close() {
 
         // exposed for tests
         int sizeInBytes() {
-            return shardDeleteResults.size();
+            return shardDeleteResults == null ? 0 : shardDeleteResults.size();
         }
     }
 

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

@@ -16,6 +16,7 @@
 import org.elasticsearch.common.collect.ImmutableOpenMap;
 import org.elasticsearch.common.geo.GeoPoint;
 import org.elasticsearch.common.lucene.BytesRefs;
+import org.elasticsearch.common.util.BigArrays;
 import org.elasticsearch.common.util.Maps;
 import org.elasticsearch.common.util.PageCacheRecycler;
 import org.elasticsearch.core.TimeValue;
@@ -63,7 +64,7 @@ public void testEmpty() throws Exception {
         out.close();
     }
 
-    public void testSingleByte() throws Exception {
+    public void testSingleByte() {
         TestStreamOutput out = new TestStreamOutput();
         assertEquals(0, out.size());
 
@@ -78,7 +79,7 @@ public void testSingleByte() throws Exception {
         out.close();
     }
 
-    public void testSingleShortPage() throws Exception {
+    public void testSingleShortPage() {
         TestStreamOutput out = new TestStreamOutput();
 
         int expectedSize = 10;
@@ -95,7 +96,7 @@ public void testSingleShortPage() throws Exception {
         out.close();
     }
 
-    public void testIllegalBulkWrite() throws Exception {
+    public void testIllegalBulkWrite() {
         TestStreamOutput out = new TestStreamOutput();
 
         // bulk-write with wrong args
@@ -138,7 +139,7 @@ public void testSingleFullPageBulkWrite() throws Exception {
         out.close();
     }
 
-    public void testSingleFullPageBulkWriteWithOffset() throws Exception {
+    public void testSingleFullPageBulkWriteWithOffset() {
         TestStreamOutput out = new TestStreamOutput();
 
         int initialOffset = 10;
@@ -157,7 +158,7 @@ public void testSingleFullPageBulkWriteWithOffset() throws Exception {
         out.close();
     }
 
-    public void testSingleFullPageBulkWriteWithOffsetCrossover() throws Exception {
+    public void testSingleFullPageBulkWriteWithOffsetCrossover() {
         TestStreamOutput out = new TestStreamOutput();
 
         int initialOffset = 10;
@@ -176,7 +177,7 @@ public void testSingleFullPageBulkWriteWithOffsetCrossover() throws Exception {
         out.close();
     }
 
-    public void testSingleFullPage() throws Exception {
+    public void testSingleFullPage() {
         TestStreamOutput out = new TestStreamOutput();
 
         int expectedSize = PageCacheRecycler.BYTE_PAGE_SIZE;
@@ -193,7 +194,7 @@ public void testSingleFullPage() throws Exception {
         out.close();
     }
 
-    public void testOneFullOneShortPage() throws Exception {
+    public void testOneFullOneShortPage() {
         TestStreamOutput out = new TestStreamOutput();
 
         int expectedSize = PageCacheRecycler.BYTE_PAGE_SIZE + 10;
@@ -210,7 +211,7 @@ public void testOneFullOneShortPage() throws Exception {
         out.close();
     }
 
-    public void testTwoFullOneShortPage() throws Exception {
+    public void testTwoFullOneShortPage() {
         TestStreamOutput out = new TestStreamOutput();
 
         int expectedSize = (PageCacheRecycler.BYTE_PAGE_SIZE * 2) + 1;
@@ -227,8 +228,9 @@ public void testTwoFullOneShortPage() throws Exception {
         out.close();
     }
 
-    public void testSeek() throws Exception {
-        BytesStreamOutput out = new BytesStreamOutput();
+    public void testSeek() {
+        int maximumSize = randomIntBetween(0, Integer.MAX_VALUE);
+        BytesStreamOutput out = new BytesStreamOutput(0, BigArrays.NON_RECYCLING_INSTANCE, maximumSize);
 
         int position = 0;
         assertEquals(position, out.position());
@@ -241,13 +243,14 @@ public void testSeek() throws Exception {
         assertEquals(position, BytesReference.toBytes(out.bytes()).length);
 
         IllegalArgumentException iae = expectThrows(IllegalArgumentException.class, () -> out.seek(Integer.MAX_VALUE + 1L));
-        assertEquals("BytesStreamOutput cannot hold more than 2GB of data", iae.getMessage());
+        assertEquals("BytesStreamOutput has exceeded it's max size of " + maximumSize, iae.getMessage());
 
         out.close();
     }
 
-    public void testSkip() throws Exception {
-        BytesStreamOutput out = new BytesStreamOutput();
+    public void testSkip() {
+        int maximumSize = randomIntBetween(0, Integer.MAX_VALUE);
+        BytesStreamOutput out = new BytesStreamOutput(0, BigArrays.NON_RECYCLING_INSTANCE, maximumSize);
 
         int position = 0;
         assertEquals(position, out.position());
@@ -256,8 +259,8 @@ public void testSkip() throws Exception {
         out.skip(forward);
         assertEquals(position + forward, out.position());
 
-        IllegalArgumentException iae = expectThrows(IllegalArgumentException.class, () -> out.skip(Integer.MAX_VALUE - 50));
-        assertEquals("BytesStreamOutput cannot hold more than 2GB of data", iae.getMessage());
+        IllegalArgumentException iae = expectThrows(IllegalArgumentException.class, () -> out.skip(maximumSize - 50));
+        assertEquals("BytesStreamOutput has exceeded it's max size of " + maximumSize, iae.getMessage());
 
         out.close();
     }