#6018: buffer up bytes before calling RateLimiter

Author: mikemccand

src/main/java/org/apache/lucene/store/RateLimitedFSDirectory.java

@@ -81,6 +81,7 @@ static final class RateLimitedIndexOutput extends BufferedIndexOutput {
         private final BufferedIndexOutput bufferedDelegate;
         private final RateLimiter rateLimiter;
         private final StoreRateLimiting.Listener rateListener;
+        private long bytesSinceLastRateLimit;
 
         RateLimitedIndexOutput(final RateLimiter rateLimiter, final StoreRateLimiting.Listener rateListener, final IndexOutput delegate) {
             super(delegate instanceof BufferedIndexOutput ? ((BufferedIndexOutput) delegate).getBufferSize() : BufferedIndexOutput.DEFAULT_BUFFER_SIZE);
@@ -97,7 +98,11 @@ static final class RateLimitedIndexOutput extends BufferedIndexOutput {
 
         @Override
         protected void flushBuffer(byte[] b, int offset, int len) throws IOException {
-            rateListener.onPause(rateLimiter.pause(len));
+            bytesSinceLastRateLimit += len;
+            if (bytesSinceLastRateLimit >= rateLimiter.getMinPauseCheckBytes()) {
+                bytesSinceLastRateLimit = 0;
+                rateListener.onPause(rateLimiter.pause(bytesSinceLastRateLimit));
+            }
             if (bufferedDelegate != null) {
                 bufferedDelegate.flushBuffer(b, offset, len);
             } else {

src/main/java/org/elasticsearch/index/snapshots/blobstore/RateLimitingInputStream.java

@@ -35,6 +35,8 @@ public class RateLimitingInputStream extends InputStream {
 
     private final Listener listener;
 
+    private long bytesSinceLastRateLimit;
+
     public interface Listener {
         void onPause(long nanos);
     }
@@ -45,13 +47,21 @@ public RateLimitingInputStream(InputStream delegate, RateLimiter rateLimiter, Li
         this.listener = listener;
     }
 
+    private void maybePause(int bytes) {
+        bytesSinceLastRateLimit += bytes;
+        if (bytesSinceLastRateLimit >= rateLimiter.getMinPauseCheckBytes()) {
+            long pause = rateLimiter.pause(bytesSinceLastRateLimit);
+            bytesSinceLastRateLimit = 0;
+            if (pause > 0) {
+                listener.onPause(pause);
+            }
+        }
+    }
+
     @Override
     public int read() throws IOException {
         int b = delegate.read();
-        long pause = rateLimiter.pause(1);
-        if (pause > 0) {
-            listener.onPause(pause);
-        }
+        maybePause(1);
         return b;
     }
 
@@ -64,7 +74,7 @@ public int read(byte[] b) throws IOException {
     public int read(byte[] b, int off, int len) throws IOException {
         int n = delegate.read(b, off, len);
         if (n > 0) {
-            listener.onPause(rateLimiter.pause(n));
+            maybePause(n);
         }
         return n;
     }