SearchableSnapshotIndexInput should read all bytes (#52199)

This commit fixes a bug in SearchableSnapshotIndexInput which does 
not use the InputStream.read(byte[], int, int) method correctly.

The javadoc indicates that: "An attempt is made to read as many as 
len bytes, but a smaller number may be read." Consequently, 
SearchableSnapshotIndexInput should try to read more bytes when 
a first reading operation has not brought back all the requested bytes.

Author: tlrx

x-pack/plugin/searchable-snapshots/src/main/java/org/elasticsearch/index/store/SearchableSnapshotIndexInput.java

@@ -8,6 +8,7 @@
 import org.apache.lucene.store.BufferedIndexInput;
 import org.apache.lucene.store.IndexInput;
 import org.apache.lucene.util.BytesRef;
+import org.elasticsearch.common.CheckedRunnable;
 import org.elasticsearch.common.Nullable;
 import org.elasticsearch.common.blobstore.BlobContainer;
 import org.elasticsearch.core.internal.io.IOUtils;
@@ -122,7 +123,9 @@ private void readInternalBytes(final int part, long pos, final byte[] b, int off
         if (optimizedReadSize < length) {
             // we did not read everything in an optimized fashion, so read the remainder directly
             try (InputStream inputStream = openBlobStream(part, pos + optimizedReadSize, length - optimizedReadSize)) {
-                final int directReadSize = inputStream.read(b, offset + optimizedReadSize, length - optimizedReadSize);
+                final int directReadSize = readFully(inputStream, b, offset + optimizedReadSize, length - optimizedReadSize, () -> {
+                    throw new EOFException("Read past EOF at [" + position + "] with length [" + fileInfo.partBytes(part) + "]");
+                });
                 assert optimizedReadSize + directReadSize == length : optimizedReadSize + " and " + directReadSize + " vs " + length;
                 position += directReadSize;
             }
@@ -280,6 +283,23 @@ private InputStream openBlobStream(int part, long pos, long length) throws IOExc
         return stream;
     }
 
+    /**
+     * Fully read up to {@code length} bytes from the given {@link InputStream}
+     */
+    private static int readFully(InputStream inputStream, byte[] b, int offset, int length, CheckedRunnable<IOException> onEOF)
+        throws IOException {
+        int totalRead = 0;
+        while (totalRead < length) {
+            final int read = inputStream.read(b, offset + totalRead, length - totalRead);
+            if (read == -1) {
+                onEOF.run();
+                break;
+            }
+            totalRead += read;
+        }
+        return totalRead > 0 ? totalRead : -1;
+    }
+
     private static class StreamForSequentialReads implements Closeable {
         private final InputStream inputStream;
         private final int part;
@@ -299,7 +319,7 @@ boolean canContinueSequentialRead(int part, long pos) {
 
         int read(byte[] b, int offset, int length) throws IOException {
             assert this.pos < maxPos : "should not try and read from a fully-read stream";
-            int read = inputStream.read(b, offset, length);
+            final int read = readFully(inputStream, b, offset, length, () -> {});
             assert read <= length : read + " vs " + length;
             pos += read;
             return read;

x-pack/plugin/searchable-snapshots/src/test/java/org/elasticsearch/index/store/SearchableSnapshotIndexInputTests.java

@@ -16,7 +16,9 @@
 
 import java.io.ByteArrayInputStream;
 import java.io.EOFException;
+import java.io.FilterInputStream;
 import java.io.IOException;
+import java.io.InputStream;
 import java.nio.charset.StandardCharsets;
 import java.util.concurrent.atomic.AtomicInteger;
 
@@ -64,9 +66,10 @@ private SearchableSnapshotIndexInput createIndexInput(final byte[] input, long p
 
                 onReadBlob.run();
 
+                final InputStream stream;
                 if (fileInfo.numberOfParts() == 1L) {
                     assertThat("Unexpected blob name [" + name + "]", name, equalTo(fileInfo.name()));
-                    return new ByteArrayInputStream(input, Math.toIntExact(position), Math.toIntExact(length));
+                    stream = new ByteArrayInputStream(input, Math.toIntExact(position), Math.toIntExact(length));
 
                 } else {
                     assertThat("Unexpected blob name [" + name + "]", name, allOf(startsWith(fileInfo.name()), containsString(".part")));
@@ -75,7 +78,19 @@ private SearchableSnapshotIndexInput createIndexInput(final byte[] input, long p
                     assertThat("Unexpected part number [" + partNumber + "] for [" + name + "]", partNumber,
                         allOf(greaterThanOrEqualTo(0L), lessThan(fileInfo.numberOfParts())));
 
-                    return new ByteArrayInputStream(input, Math.toIntExact(partNumber * partSize + position), Math.toIntExact(length));
+                    stream = new ByteArrayInputStream(input, Math.toIntExact(partNumber * partSize + position), Math.toIntExact(length));
+                }
+
+                if (randomBoolean()) {
+                    return stream;
+                } else {
+                    // sometimes serve less bytes than expected, in agreement with InputStream{@link #read(byte[], int, int)} javadoc
+                    return new FilterInputStream(stream) {
+                        @Override
+                        public int read(byte[] b, int off, int len) throws IOException {
+                            return super.read(b, off, randomIntBetween(1, len));
+                        }
+                    };
                 }
             });
         return new SearchableSnapshotIndexInput(blobContainer, fileInfo, minimumReadSize,