[GitHub] [pinot] jasperjiaguo commented on a diff in pull request #8753: Streamed segment download & untar with rate limiter to control disk usage

[GitBox <gi...@apache.org>]

jasperjiaguo commented on code in PR #8753:
URL: https://github.com/apache/pinot/pull/8753#discussion_r894004559


##########
pinot-common/src/main/java/org/apache/pinot/common/utils/TarGzCompressionUtils.java:
##########
@@ -198,4 +238,30 @@ public static void untarOneFile(File inputFile, String fileName, File outputFile
       throw new IOException(String.format("Failed to find file: %s in: %s", fileName, inputFile));
     }
   }
+
+  public static long copyWithRateLimiter(InputStream inputStream, FileOutputStream outputStream, long rateLimit)
+      throws IOException {
+    Objects.requireNonNull(inputStream, "inputStream is null");
+    Objects.requireNonNull(outputStream, "outputStream is null");
+    FileDescriptor fd = outputStream.getFD();
+    LOGGER.info("Using rate limiter for stream copy, target limit {} bytes/s", rateLimit);
+    byte[] buffer = new byte[DEFAULT_BUFFER_SIZE];
+    RateLimiter rateLimiter = RateLimiter.create(rateLimit);
+    long count;
+    int n;
+
+    if (rateLimit == SYNC_DISK_WRITE_WITH_UPSTREAM_RATE) {
+      for (count = 0L; -1 != (n = inputStream.read(buffer)); count += (long) n) {
+        outputStream.write(buffer, 0, n);
+        fd.sync(); // flush the buffer timely to the disk so that the disk bandwidth wouldn't get saturated
+      }
+    } else {
+      for (count = 0L; -1 != (n = inputStream.read(buffer)); count += (long) n) {
+        rateLimiter.acquire(n);
+        outputStream.write(buffer, 0, n);
+        fd.sync();
+      }
+    }
+    return count;

Review Comment:
   This function is actually called per file. So I removed all the logging here as it will be too much. We already have similar logging in `private File downloadAndStreamUntarRateLimit(String segmentName, SegmentZKMetadata zkMetadata, File tempRootDir, long maxDownloadRateInByte)` which should be enough



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: commits-unsubscribe@pinot.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@pinot.apache.org
For additional commands, e-mail: commits-help@pinot.apache.org