Parallelize per-file rewrites in GcsUtilV2 copy and move

sdks/java/extensions/google-cloud-platform-core/src/main/java/org/apache/beam/sdk/extensions/gcp/options/GcsOptions.java

@@ -171,11 +171,40 @@ public GoogleCloudStorageReadOptions create(PipelineOptions options) {
 
   void setGcsHttpRequestWriteTimeout(@Nullable Integer timeoutMs);
 
-  @Description("Batching limit for rewrite ops which will copy data.")
+  /**
+   * Batching limit for rewrite ops which will copy data.
+   *
+   * @deprecated Only applicable to GcsUtil V1. Not applicable to GcsUtil V2 (enabled with the
+   *     {@code use_gcsutil_v2} experiment), which does not batch rewrites; use {@link
+   *     #getGcsMaxConcurrentRewrites} to tune copy/move concurrency there instead.
+   */
+  @Deprecated
+  @Description(
+      "Batching limit for rewrite ops which will copy data. Deprecated: NOT applicable to GcsUtil"
+          + " V2 (enabled with the use_gcsutil_v2 experiment); use gcsMaxConcurrentRewrites"
+          + " instead.")
   @Nullable Integer getGcsRewriteDataOpBatchLimit();
 
+  /**
+   * @deprecated Only applicable to GcsUtil V1. See {@link #getGcsRewriteDataOpBatchLimit}.
+   */
+  @Deprecated
   void setGcsRewriteDataOpBatchLimit(@Nullable Integer timeoutMs);
 
+  /**
+   * Maximum number of files rewritten concurrently when copying or moving files. If unset, a
+   * default of 32 is used.
+   *
+   * <p>Only applicable to GcsUtil V2, which is enabled with the {@code use_gcsutil_v2} experiment.
+   */
+  @Description(
+      "Maximum number of files rewritten concurrently when copying or moving files in GCS."
+          + " If unset, defaults to 32. Only applicable to GcsUtil V2 (enabled with the"
+          + " use_gcsutil_v2 experiment).")
+  @Nullable Integer getGcsMaxConcurrentRewrites();
+
+  void setGcsMaxConcurrentRewrites(@Nullable Integer maxConcurrentRewrites);
+
   /** If true, reports number of bytes written to each gcs bucket. */
   @Description("Whether to report number of bytes written per GCS bucket.")
   @Default.Boolean(false)

sdks/java/extensions/google-cloud-platform-core/src/main/java/org/apache/beam/sdk/extensions/gcp/util/GcsUtilV2.java

@@ -52,6 +52,7 @@
 import com.google.cloud.storage.StorageChannelUtils;
 import com.google.cloud.storage.StorageException;
 import com.google.cloud.storage.StorageOptions;
+import java.io.Closeable;
 import java.io.FileNotFoundException;
 import java.io.IOException;
 import java.net.MalformedURLException;
@@ -67,6 +68,12 @@
 import java.util.List;
 import java.util.Objects;
 import java.util.Optional;
+import java.util.concurrent.CancellationException;
+import java.util.concurrent.ExecutionException;
+import java.util.concurrent.ExecutorService;
+import java.util.concurrent.Executors;
+import java.util.concurrent.Future;
+import java.util.concurrent.TimeUnit;
 import java.util.function.Consumer;
 import java.util.regex.Pattern;
 import org.apache.beam.runners.core.metrics.GcpResourceIdentifiers;
@@ -87,6 +94,7 @@
 import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.annotations.VisibleForTesting;
 import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableList;
 import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.Lists;
+import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.util.concurrent.ThreadFactoryBuilder;
 import org.checkerframework.checker.nullness.qual.Nullable;
 
 class GcsUtilV2 {
@@ -141,6 +149,21 @@ Storage storage() {
    */
   private static final long MEGABYTES_COPIED_PER_CHUNK = 2048L;
 
+  /**
+   * Default maximum number of files rewritten concurrently by {@link #copy} and {@link #move}, used
+   * when {@link GcsOptions#getGcsMaxConcurrentRewrites} is unset.
+   */
+  private static final int DEFAULT_MAX_CONCURRENT_REWRITES = 32;
+
+  /**
+   * How long {@link #copy} and {@link #move} wait for in-flight rewrites to finish after a failure
+   * before interrupting them.
+   */
+  private static final long REWRITE_TERMINATION_TIMEOUT_MINUTES = 5;
+
+  /** Maximum number of files rewritten concurrently by {@link #copy} and {@link #move}. */
+  private final int maxConcurrentRewrites;
+
   GcsUtilV2(PipelineOptions options) {
     GcsOptions gcsOptions = options.as(GcsOptions.class);
     this.projectId = options.as(GcpOptions.class).getProject();
@@ -173,6 +196,16 @@ Storage storage() {
                 ? gcsOptions.getGcsWriteCounterPrefix()
                 : null);
     this.gcsPerformanceMetrics = Boolean.TRUE.equals(gcsOptions.getGcsPerformanceMetrics());
+
+    Integer configuredMaxConcurrentRewrites = gcsOptions.getGcsMaxConcurrentRewrites();
+    checkArgument(
+        configuredMaxConcurrentRewrites == null || configuredMaxConcurrentRewrites > 0,
+        "gcsMaxConcurrentRewrites must be positive, but was %s",
+        configuredMaxConcurrentRewrites);
+    this.maxConcurrentRewrites =
+        configuredMaxConcurrentRewrites != null
+            ? configuredMaxConcurrentRewrites
+            : DEFAULT_MAX_CONCURRENT_REWRITES;
   }
 
   /**
@@ -594,73 +627,174 @@ private void rewriteHelper(
         srcList.size(),
         dstList.size());
 
-    for (int i = 0; i < srcList.size(); i++) {
-      GcsPath srcPath = srcList.get(i);
-      GcsPath dstPath = dstList.get(i);
-      BlobId srcId = BlobId.of(srcPath.getBucket(), srcPath.getObject());
-      BlobId dstId = BlobId.of(dstPath.getBucket(), dstPath.getObject());
+    if (srcList.isEmpty()) {
+      return;
+    }
+    if (srcList.size() == 1) {
+      // Nothing to overlap, so skip the thread pool.
+      rewriteOne(srcList.get(0), dstList.get(0), deleteSrc, srcMissing, dstOverwrite);
+      return;
+    }
 
-      CopyRequest.Builder copyRequestBuilder =
-          CopyRequest.newBuilder()
-              .setSource(srcId)
-              .setMegabytesCopiedPerChunk(MEGABYTES_COPIED_PER_CHUNK);
+    // Each file costs up to three dependent round trips (target lookup, rewrite, source delete),
+    // and java-storage cannot batch rewrites, so issue the files concurrently instead.
+    int numThreads = Math.min(srcList.size(), maxConcurrentRewrites);
+    ExecutorService executor =
+        Executors.newFixedThreadPool(
+            numThreads,
+            new ThreadFactoryBuilder()
+                .setDaemon(true)
+                .setNameFormat("gcsutil-v2-rewrite-%d")
+                .build());
+    // Requests run on the pool's threads, so carry the caller's container over to them.
+    MetricsContainer container = MetricsEnvironment.getCurrentContainer();
+    List<Future<Void>> futures = new ArrayList<>(srcList.size());
+    try {
+      for (int i = 0; i < srcList.size(); i++) {
+        GcsPath srcPath = srcList.get(i);
+        GcsPath dstPath = dstList.get(i);
+        futures.add(
+            executor.submit(
+                () -> {
+                  if (container != null) {
+                    try (Closeable scope = MetricsEnvironment.scopedMetricsContainer(container)) {
+                      rewriteOne(srcPath, dstPath, deleteSrc, srcMissing, dstOverwrite);
+                    }
+                  } else {
+                    rewriteOne(srcPath, dstPath, deleteSrc, srcMissing, dstOverwrite);
+                  }
+                  return null;
+                }));
+      }
 
-      if (dstOverwrite == OverwriteStrategy.ALWAYS_OVERWRITE) {
-        copyRequestBuilder.setTarget(dstId);
-      } else {
-        // FAIL_IF_EXISTS, SKIP_IF_EXISTS and SAFE_OVERWRITE require checking the target blob
-        BlobInfo existingTarget;
+      // Like the sequential loop this replaces, stop at the first failure: files not yet started
+      // are cancelled, while the ones already in flight are allowed to finish (awaited below).
+      @Nullable Throwable failure = null;
+      for (Future<Void> future : futures) {
         try {
-          existingTarget = storage().get(dstId);
-        } catch (StorageException e) {
-          throw translateStorageException(dstPath, e);
-        }
-
-        if (existingTarget == null) {
-          copyRequestBuilder.setTarget(dstId, Storage.BlobTargetOption.doesNotExist());
-        } else {
-          switch (dstOverwrite) {
-            case SKIP_IF_EXISTS:
-              LOG.warn("Ignoring rewriting from {} to {} because target exists.", srcPath, dstPath);
-              continue; // Skip to next file in for-loop
-
-            case SAFE_OVERWRITE:
-              copyRequestBuilder.setTarget(
-                  dstId, Storage.BlobTargetOption.generationMatch(existingTarget.getGeneration()));
-              break;
-
-            case FAIL_IF_EXISTS:
-              throw new FileAlreadyExistsException(
-                  srcPath.toString(),
-                  dstPath.toString(),
-                  "Target object already exists and strategy is FAIL_IF_EXISTS");
-            default:
-              throw new IllegalStateException("Unknown OverwriteStrategy: " + dstOverwrite);
+          future.get();
+        } catch (CancellationException e) {
+          // Cancelled below after an earlier failure.
+        } catch (ExecutionException e) {
+          Throwable cause = e.getCause() != null ? e.getCause() : e;
+          if (failure == null) {
+            failure = cause;
+            for (Future<Void> other : futures) {
+              other.cancel(false);
+            }
+          } else {
+            failure.addSuppressed(cause);
           }
         }
       }
-
+      if (failure != null) {
+        // Rethrow as is, so that callers still see the specific type (e.g. FileNotFoundException,
+        // FileAlreadyExistsException, AccessDeniedException) the sequential version threw.
+        if (failure instanceof IOException) {
+          throw (IOException) failure;
+        }
+        if (failure instanceof RuntimeException) {
+          throw (RuntimeException) failure;
+        }
+        if (failure instanceof Error) {
+          throw (Error) failure;
+        }
+        throw new IOException("Error rewriting objects", failure);
+      }
+    } catch (InterruptedException e) {
+      Thread.currentThread().interrupt();
+      throw new IOException("Interrupted while rewriting objects", e);
+    } finally {
+      // A cancelled future reports done even if its task is still running, so wait for the pool
+      // itself: no file is left mid-rewrite when this method returns. If the wait times out or the
+      // caller is interrupted (including above), stop waiting and interrupt the in-flight rewrites.
+      executor.shutdown();
       try {
-        CopyWriter copyWriter = storage().copy(copyRequestBuilder.build());
-        copyWriter.getResult();
-
-        if (deleteSrc) {
-          if (!storage().delete(srcId)) {
-            // This may happen if the source file is deleted by another process after copy.
-            LOG.warn(
-                "Source file {} could not be deleted after move to {}. It may not have existed.",
-                srcPath,
-                dstPath);
-          }
+        if (!executor.awaitTermination(REWRITE_TERMINATION_TIMEOUT_MINUTES, TimeUnit.MINUTES)) {
+          LOG.warn(
+              "Timed out after {} minutes waiting for in-flight GCS rewrites to finish; "
+                  + "interrupting them. Some objects may be left partially copied or moved.",
+              REWRITE_TERMINATION_TIMEOUT_MINUTES);
+          executor.shutdownNow();
         }
+      } catch (InterruptedException e) {
+        executor.shutdownNow();
+        Thread.currentThread().interrupt();
+      }
+    }
+  }
+
+  /** Rewrites a single object, then deletes the source if {@code deleteSrc}. */
+  private void rewriteOne(
+      GcsPath srcPath,
+      GcsPath dstPath,
+      boolean deleteSrc,
+      MissingStrategy srcMissing,
+      OverwriteStrategy dstOverwrite)
+      throws IOException {
+    BlobId srcId = BlobId.of(srcPath.getBucket(), srcPath.getObject());
+    BlobId dstId = BlobId.of(dstPath.getBucket(), dstPath.getObject());
+
+    CopyRequest.Builder copyRequestBuilder =
+        CopyRequest.newBuilder()
+            .setSource(srcId)
+            .setMegabytesCopiedPerChunk(MEGABYTES_COPIED_PER_CHUNK);
+
+    if (dstOverwrite == OverwriteStrategy.ALWAYS_OVERWRITE) {
+      copyRequestBuilder.setTarget(dstId);
+    } else {
+      // FAIL_IF_EXISTS, SKIP_IF_EXISTS and SAFE_OVERWRITE require checking the target blob
+      BlobInfo existingTarget;
+      try {
+        existingTarget = storage().get(dstId);
       } catch (StorageException e) {
-        if (e.getCode() == 404 && srcMissing == MissingStrategy.SKIP_IF_MISSING) {
+        throw translateStorageException(dstPath, e);
+      }
+
+      if (existingTarget == null) {
+        copyRequestBuilder.setTarget(dstId, Storage.BlobTargetOption.doesNotExist());
+      } else {
+        switch (dstOverwrite) {
+          case SKIP_IF_EXISTS:
+            LOG.warn("Ignoring rewriting from {} to {} because target exists.", srcPath, dstPath);
+            return; // Skip this file
+
+          case SAFE_OVERWRITE:
+            copyRequestBuilder.setTarget(
+                dstId, Storage.BlobTargetOption.generationMatch(existingTarget.getGeneration()));
+            break;
+
+          case FAIL_IF_EXISTS:
+            throw new FileAlreadyExistsException(
+                srcPath.toString(),
+                dstPath.toString(),
+                "Target object already exists and strategy is FAIL_IF_EXISTS");
+          default:
+            throw new IllegalStateException("Unknown OverwriteStrategy: " + dstOverwrite);
+        }
+      }
+    }
+
+    try {
+      CopyWriter copyWriter = storage().copy(copyRequestBuilder.build());
+      copyWriter.getResult();
+
+      if (deleteSrc) {
+        if (!storage().delete(srcId)) {
+          // This may happen if the source file is deleted by another process after copy.
           LOG.warn(
-              "Ignoring rewriting from {} to {} because source does not exist.", srcPath, dstPath);
-          continue;
+              "Source file {} could not be deleted after move to {}. It may not have existed.",
+              srcPath,
+              dstPath);
         }
-        throw translateStorageException(srcPath, e);
       }
+    } catch (StorageException e) {
+      if (e.getCode() == 404 && srcMissing == MissingStrategy.SKIP_IF_MISSING) {
+        LOG.warn(
+            "Ignoring rewriting from {} to {} because source does not exist.", srcPath, dstPath);
+        return;
+      }
+      throw translateStorageException(srcPath, e);
     }
   }