From 03fa2705e7bbf38e886cc095a0e1723e6a524585 Mon Sep 17 00:00:00 2001
From: Iraklis Psaroudakis <kingherc@gmail.com>
Date: Wed, 11 Dec 2024 17:49:12 +0200
Subject: [PATCH] Specialize skip for InputStreamIndexInput (#118436)

Skip would previously defer to the default implementation that
reads bytes unnecessarily and may be slow. We now specialize it
so that it seeks quickly.

Closes ES-10234
---
 .../lucene/store/InputStreamIndexInput.java   | 11 ++++++
 .../store/InputStreamIndexInputTests.java     | 37 ++++++++++++++++++-
 2 files changed, 47 insertions(+), 1 deletion(-)

diff --git a/server/src/main/java/org/elasticsearch/common/lucene/store/InputStreamIndexInput.java b/server/src/main/java/org/elasticsearch/common/lucene/store/InputStreamIndexInput.java
index 5603a1d4f1ab0..f3a3ec91ee931 100644
--- a/server/src/main/java/org/elasticsearch/common/lucene/store/InputStreamIndexInput.java
+++ b/server/src/main/java/org/elasticsearch/common/lucene/store/InputStreamIndexInput.java
@@ -88,4 +88,15 @@ public synchronized void reset() throws IOException {
         indexInput.seek(markPointer);
         counter = markCounter;
     }
+
+    @Override
+    public long skip(long n) throws IOException {
+        long skipBytes = Math.min(n, Math.min(indexInput.length() - indexInput.getFilePointer(), limit - counter));
+        if (skipBytes <= 0) {
+            return 0;
+        }
+        indexInput.skipBytes(skipBytes);
+        counter += skipBytes;
+        return skipBytes;
+    }
 }
diff --git a/server/src/test/java/org/elasticsearch/common/lucene/store/InputStreamIndexInputTests.java b/server/src/test/java/org/elasticsearch/common/lucene/store/InputStreamIndexInputTests.java
index a1bcf1b91fa4d..4bea6f50c7c4b 100644
--- a/server/src/test/java/org/elasticsearch/common/lucene/store/InputStreamIndexInputTests.java
+++ b/server/src/test/java/org/elasticsearch/common/lucene/store/InputStreamIndexInputTests.java
@@ -218,7 +218,7 @@ public void testReadMultiFourBytesLimit() throws IOException {
         assertThat(is.read(read), equalTo(-1));
     }
 
-    public void testMarkRest() throws Exception {
+    public void testMarkReset() throws Exception {
         Directory dir = new ByteBuffersDirectory();
         IndexOutput output = dir.createOutput("test", IOContext.DEFAULT);
         for (int i = 0; i < 3; i++) {
@@ -243,6 +243,41 @@ public void testMarkRest() throws Exception {
         assertThat(is.read(), equalTo(2));
     }
 
+    public void testSkipBytes() throws Exception {
+        Directory dir = new ByteBuffersDirectory();
+        IndexOutput output = dir.createOutput("test", IOContext.DEFAULT);
+        int bytes = randomIntBetween(10, 100);
+        for (int i = 0; i < bytes; i++) {
+            output.writeByte((byte) i);
+        }
+        output.close();
+
+        int limit = randomIntBetween(0, bytes * 2);
+        int initialReadBytes = randomIntBetween(0, limit);
+        int skipBytes = randomIntBetween(0, limit);
+        int seekExpected = Math.min(Math.min(initialReadBytes + skipBytes, limit), bytes);
+        int skipBytesExpected = Math.max(seekExpected - initialReadBytes, 0);
+        logger.debug(
+            "bytes: {}, limit: {}, initialReadBytes: {}, skipBytes: {}, seekExpected: {}, skipBytesExpected: {}",
+            bytes,
+            limit,
+            initialReadBytes,
+            skipBytes,
+            seekExpected,
+            skipBytesExpected
+        );
+
+        IndexInput input = dir.openInput("test", IOContext.DEFAULT);
+        InputStreamIndexInput is = new InputStreamIndexInput(input, limit);
+        is.readNBytes(initialReadBytes);
+        assertThat(is.skip(skipBytes), equalTo((long) skipBytesExpected));
+
+        int remainingBytes = Math.min(bytes, limit) - seekExpected;
+        for (int i = seekExpected; i < seekExpected + remainingBytes; i++) {
+            assertThat(is.read(), equalTo(i));
+        }
+    }
+
     public void testReadZeroShouldReturnZero() throws IOException {
         try (Directory dir = new ByteBuffersDirectory()) {
             try (IndexOutput output = dir.createOutput("test", IOContext.DEFAULT)) {