From 55aca3e86a7fcc3b76365b809b96d697d6607db4 Mon Sep 17 00:00:00 2001
From: David Kocher <dkocher@iterate.ch>
Date: Mon, 29 Mar 2021 15:19:35 +0200
Subject: [PATCH] Add parameter to limit read ahead to maximum length. Allows
 to use multiple concurrent threads reading from the same file with an offset
 without reading too much ahead for a single segment.

---
 .../java/net/schmizz/sshj/sftp/RemoteFile.java  | 17 +++++++++++++----
 1 file changed, 13 insertions(+), 4 deletions(-)

diff --git a/src/main/java/net/schmizz/sshj/sftp/RemoteFile.java b/src/main/java/net/schmizz/sshj/sftp/RemoteFile.java
index 25463bfb7..7a0febcd8 100644
--- a/src/main/java/net/schmizz/sshj/sftp/RemoteFile.java
+++ b/src/main/java/net/schmizz/sshj/sftp/RemoteFile.java
@@ -224,6 +224,7 @@ public class ReadAheadRemoteFileInputStream
         private final byte[] b = new byte[1];
 
         private final int maxUnconfirmedReads;
+        private final long maxOffset;
         private final Queue<Promise<Response, SFTPException>> unconfirmedReads = new LinkedList<Promise<Response, SFTPException>>();
         private final Queue<Long> unconfirmedReadOffsets = new LinkedList<Long>();
 
@@ -232,17 +233,22 @@ public class ReadAheadRemoteFileInputStream
         private boolean eof;
 
         public ReadAheadRemoteFileInputStream(int maxUnconfirmedReads) {
-            assert 0 <= maxUnconfirmedReads;
-
-            this.maxUnconfirmedReads = maxUnconfirmedReads;
+            this(maxUnconfirmedReads, 0L, -1L);
         }
 
-        public ReadAheadRemoteFileInputStream(int maxUnconfirmedReads, long fileOffset) {
+        /**
+         *
+         * @param maxUnconfirmedReads Maximum number of unconfirmed requests to send
+         * @param fileOffset Initial offset in file to read from
+         * @param maxLength Maximum length to read
+         */
+        public ReadAheadRemoteFileInputStream(int maxUnconfirmedReads, long fileOffset, long maxLength) {
             assert 0 <= maxUnconfirmedReads;
             assert 0 <= fileOffset;
 
             this.maxUnconfirmedReads = maxUnconfirmedReads;
             this.requestOffset = this.responseOffset = fileOffset;
+            this.maxOffset = maxLength > 0 ? fileOffset + maxLength : Long.MAX_VALUE;
         }
 
         private ByteArrayInputStream pending = new ByteArrayInputStream(new byte[0]);
@@ -296,6 +302,9 @@ public int read(byte[] into, int off, int len) throws IOException {
                     unconfirmedReads.add(RemoteFile.this.asyncRead(requestOffset, reqLen));
                     unconfirmedReadOffsets.add(requestOffset);
                     requestOffset += reqLen;
+                    if (requestOffset >= maxOffset) {
+                        break;
+                    }
                 }
 
                 long nextOffset = unconfirmedReadOffsets.peek();