xiangfu0 commented on code in PR #19082:
URL: https://github.com/apache/pinot/pull/19082#discussion_r4225320652


##########
pinot-plugins/pinot-batch-ingestion/pinot-batch-ingestion-common/src/main/java/org/apache/pinot/plugin/ingestion/batch/common/SegmentGenerationJobUtils.java:
##########
@@ -104,26 +114,173 @@ public static void moveLocalTarFileToRemote(File 
localMetadataTarFile, URI outpu
   /// If <overwrite> is true, and the source file exists in the destination 
directory, then replace it, otherwise
   /// log a warning and continue. We assume that source and destination 
directories are on the same filesystem,
   /// so that move() can be used.
+  /// Uses [#DEFAULT_STAGING_COPY_PARALLELISM] worker threads.
   ///
-  /// @param fs
-  /// @param sourceDir
-  /// @param destDir
-  /// @param overwrite
-  /// @throws IOException
-  /// @throws URISyntaxException
+  /// The shared [PinotFS] instance must support concurrent `move` of distinct 
paths (true for
+  /// `LocalPinotFS` and typical remote implementations).
+  ///
+  /// @param fs filesystem used for both source and destination
+  /// @param sourceDir source directory URI
+  /// @param destDir destination directory URI
+  /// @param overwrite whether to overwrite existing destination files
+  /// @throws IOException on listing or move failure
+  /// @throws URISyntaxException on URI construction failure
   public static void moveFiles(PinotFS fs, URI sourceDir, URI destDir, boolean 
overwrite)
-          throws IOException, URISyntaxException {
-    for (String sourcePath : fs.listFiles(sourceDir, true)) {
-      URI sourceFileUri = SegmentGenerationUtils.getFileURI(sourcePath, 
sourceDir);
-      String sourceFilename = 
SegmentGenerationUtils.getFileName(sourceFileUri);
-      URI destFileUri =
-              SegmentGenerationUtils.getRelativeOutputPath(sourceDir, 
sourceFileUri, destDir).resolve(sourceFilename);
+      throws IOException, URISyntaxException {
+    moveFiles(fs, sourceDir, destDir, overwrite, 
DEFAULT_STAGING_COPY_PARALLELISM);
+  }
+
+  /// Move all files from the <sourceDir> to the <destDir> using up to 
`parallelism` threads.
+  /// Directories in the source listing are skipped; parent directories on the 
destination are created by
+  /// [PinotFS#move]. Relative path layout under `sourceDir` is preserved.
+  ///
+  /// `parallelism` is clamped to at most [#MAX_STAGING_COPY_PARALLELISM] and 
to the number of files to
+  /// move, so direct callers of this overload cannot exceed the cap that
+  /// [#getStagingCopyParallelism] enforces.
+  ///
+  /// On partial failure, remaining in-flight moves are allowed to finish, 
then an [IOException] is thrown
+  /// with any additional failures attached as suppressed exceptions. On 
interruption, all outstanding
+  /// moves are cancelled, the interrupt status is restored, and an 
[IOException] is thrown.
+  ///
+  /// @param fs filesystem used for both source and destination (must be safe 
for concurrent move of
+  ///     distinct paths)
+  /// @param sourceDir source directory URI
+  /// @param destDir destination directory URI
+  /// @param overwrite whether to overwrite existing destination files
+  /// @param parallelism number of concurrent move workers; values <= 1 run 
serially
+  /// @throws IOException on listing or move failure
+  /// @throws URISyntaxException on URI construction failure
+  public static void moveFiles(PinotFS fs, URI sourceDir, URI destDir, boolean 
overwrite, int parallelism)
+      throws IOException, URISyntaxException {
+    List<URI> sourceFileUris = listSourceFiles(fs, sourceDir);
+    if (sourceFileUris.isEmpty()) {
+      return;
+    }
+    int effectiveParallelism =
+        Math.max(1, Math.min(Math.min(parallelism, 
MAX_STAGING_COPY_PARALLELISM), sourceFileUris.size()));
+    LOGGER.info("Moving {} files from [{}] to [{}] with parallelism {}", 
sourceFileUris.size(), sourceDir, destDir,
+        effectiveParallelism);
+    if (effectiveParallelism == 1) {
+      for (URI sourceFileUri : sourceFileUris) {
+        moveOneFile(fs, sourceDir, sourceFileUri, destDir, overwrite);
+      }
+      return;
+    }
 
-      if (!overwrite && fs.exists(destFileUri)) {
-        LOGGER.warn("Can't overwrite existing output segment tar file: {}", 
destFileUri);
-      } else {
-        fs.move(sourceFileUri, destFileUri, true);
+    ExecutorService executor = 
Executors.newFixedThreadPool(effectiveParallelism, r -> {
+      Thread t = new Thread(r, "pinot-staging-copy");
+      t.setDaemon(true);
+      return t;
+    });
+    try {
+      List<Future<Void>> futures = new ArrayList<>(sourceFileUris.size());
+      for (URI sourceFileUri : sourceFileUris) {
+        futures.add(executor.submit(() -> {
+          moveOneFile(fs, sourceDir, sourceFileUri, destDir, overwrite);
+          return null;
+        }));
+      }
+      IOException firstFailure = null;
+      for (Future<Void> future : futures) {
+        try {
+          future.get();
+        } catch (InterruptedException e) {
+          // Cancel every outstanding move and stop waiting instead of 
blocking on the remaining futures.
+          futures.forEach(f -> f.cancel(true));

Review Comment:
   [P1] Please wait for copy workers to stop before cleaning up staging.
   
   If the driver is interrupted after an overwrite worker checks that its 
source exists, `cancel(true)` and `shutdownNow()` can return while that worker 
is still running. Both runners then delete staging in `finally`, so a 
`LocalPinotFS` worker can delete the existing destination and fail to copy its 
replacement because the source is gone. Please wait for actual worker 
termination before returning to the caller's staging cleanup.



-- 
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: [email protected]

For queries about this service, please contact Infrastructure at:
[email protected]


---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]

Reply via email to