[GOBBLIN-2266] Add rename order functionality for Iceberg table replication

gobblin-data-management/src/main/java/org/apache/gobblin/data/management/copy/publisher/CopyDataPublisher.java

@@ -342,7 +342,8 @@ private void publishFileSet(CopyEntity.DatasetAndPartition datasetAndPartition,
 
     if (statesHelper.hasAnyCopyableFile()) {
       // Targets are always absolute, so we start moving from root (will skip any existing directories).
-      HadoopUtils.renameRecursively(this.fs, datasetWriterOutputPath, new Path("/"));
+      HadoopUtils.renameRecursivelyOrdered(this.fs, datasetWriterOutputPath, new Path("/"),
+          HadoopUtils::isIcebergMetadataDir);
     } else {
       log.info("[{}] No copyable files in dataset. Proceeding to post-publish steps.", datasetAndPartition.identifier());
     }

gobblin-utility/src/main/java/org/apache/gobblin/util/HadoopUtils.java

@@ -28,6 +28,7 @@
 import java.io.OutputStream;
 import java.net.URI;
 import java.nio.file.AccessDeniedException;
+import java.util.AbstractMap;
 import java.util.Collection;
 import java.util.Iterator;
 import java.util.List;
@@ -38,6 +39,7 @@
 import java.util.concurrent.ExecutorService;
 import java.util.concurrent.Future;
 import java.util.concurrent.TimeUnit;
+import java.util.function.Predicate;
 
 import org.apache.commons.lang.StringUtils;
 import org.apache.hadoop.conf.Configuration;
@@ -551,8 +553,8 @@ public static void renameRecursively(FileSystem fileSystem, Path from, Path to)
         throw new IOException("Trying to rename a path that does not exist! " + from);
       }
 
-      futures.add(executorService
-          .submit(new RenameRecursively(throttledFS, fileSystem.getFileStatus(from), to, executorService, futures)));
+      futures.add(executorService.submit(
+          new RenameRecursively(throttledFS, fileSystem.getFileStatus(from), to, executorService, futures, null, null)));
       int futuresUsed = 0;
       while (!futures.isEmpty()) {
         try {
@@ -570,6 +572,95 @@ public static void renameRecursively(FileSystem fileSystem, Path from, Path to)
     }
   }
 
+  /**
+   * Like {@link #renameRecursively(FileSystem, Path, Path)} but renames in two phases: every directory subtree whose
+   * root matches {@code deferToLastPhase} is held back and renamed only after all other files have been fully
+   * renamed. Both phases use the same concurrent, atomic-subtree-move machinery as the unordered variant.
+   *
+   * <p>
+   * Used by Iceberg distcp publishing to move the {@code metadata/} subtree only after every other file (the
+   * {@code data/} files and anything else) has landed, since Iceberg metadata references data files by path and must
+   * never become visible ahead of the data it points to. {@code deferToLastPhase} is evaluated only on directory
+   * nodes encountered during the descent -- not on every file -- so the {@code metadata/} subtree is identified by
+   * its directory name and never enumerated file-by-file.
+   * </p>
+   *
+   * @param fileSystem on which the data needs to be moved
+   * @param from path of the data to be moved
+   * @param to path of the data to be moved
+   * @param deferToLastPhase predicate on a directory {@link Path}; matching subtrees are renamed in the last phase
+   */
+  public static void renameRecursivelyOrdered(FileSystem fileSystem, Path from, Path to,
+      Predicate<Path> deferToLastPhase) throws IOException {
+
+    log.info(String.format("Recursively renaming (ordered) %s in %s to %s.", from, fileSystem.getUri(), to));
+
+    FileSystem throttledFS = getOptionallyThrottledFileSystem(fileSystem, 10000);
+
+    ExecutorService executorService = ScalingThreadPoolExecutor.newScalingThreadPool(1, 100, 100,
+        ExecutorsUtils.newThreadFactory(Optional.of(log), Optional.of("rename-thread-%d")));
+    Queue<Future<?>> futures = Queues.newConcurrentLinkedQueue();
+    // Subtrees skipped in phase 1 (their root FileStatus and target path), to be renamed in phase 2.
+    Queue<Entry<FileStatus, Path>> deferred = Queues.newConcurrentLinkedQueue();
+
+    try {
+      if (!fileSystem.exists(from)) {
+        throw new IOException("Trying to rename a path that does not exist! " + from);
+      }
+
+      // Phase 1: rename everything except the deferred subtrees (which are collected into `deferred`).
+      log.info(String.format("[ordered-rename] Phase 1: renaming all non-deferred files under %s to %s.", from, to));
+      futures.add(executorService.submit(new RenameRecursively(
+          throttledFS, fileSystem.getFileStatus(from), to, executorService, futures, deferToLastPhase, deferred)));
+      drainFutures(futures);
+      log.info(String.format("[ordered-rename] Phase 1 complete for %s. Deferred %d subtree(s) for phase 2: %s",
+          from, deferred.size(), deferredSourcePaths(deferred)));
+
+      // Phase 2: rename the deferred subtrees, now that everything else is in place. No further deferral.
+      for (Entry<FileStatus, Path> entry : deferred) {
+        log.info(String.format("[ordered-rename] Phase 2: renaming deferred directory %s to %s.",
+            entry.getKey().getPath(), entry.getValue()));
+        futures.add(executorService.submit(new RenameRecursively(
+            throttledFS, entry.getKey(), entry.getValue(), executorService, futures, null, null)));
+      }
+      drainFutures(futures);
+
+      log.info(String.format("[ordered-rename] Recursive ordered renaming of %s to %s complete (%d deferred subtree(s)).",
+          from, to, deferred.size()));
+
+    } finally {
+      ExecutorsUtils.shutdownExecutorService(executorService, Optional.of(log), 1, TimeUnit.SECONDS);
+    }
+  }
+
+  private static void drainFutures(Queue<Future<?>> futures) throws IOException {
+    while (!futures.isEmpty()) {
+      try {
+        futures.poll().get();
+      } catch (ExecutionException | InterruptedException ee) {
+        throw new IOException(ee.getCause());
+      }
+    }
+  }
+
+  private static List<Path> deferredSourcePaths(Queue<Entry<FileStatus, Path>> deferred) {
+    List<Path> paths = Lists.newArrayListWithCapacity(deferred.size());
+    for (Entry<FileStatus, Path> entry : deferred) {
+      paths.add(entry.getKey().getPath());
+    }
+    return paths;
+  }
+
+  /**
+   * Whether {@code path} is an Iceberg {@code metadata/} directory. Used as the {@code deferToLastPhase} predicate of
+   * {@link #renameRecursivelyOrdered(FileSystem, Path, Path, Predicate)} so the {@code metadata/} subtree is renamed
+   * after the {@code data/} files it references. Matches on the directory name alone -- unlike {@code data}, there is
+   * no like-named HDFS mount to disambiguate.
+   */
+  public static boolean isIcebergMetadataDir(Path path) {
+    return "metadata".equals(path.getName());
+  }
+
   /**
    * Calls {@link #getOptionallyThrottledFileSystem(FileSystem, int)} parsing the qps from the input {@link State}
    * at key {@link #MAX_FILESYSTEM_QPS}.
@@ -619,11 +710,23 @@ private static class RenameRecursively implements Runnable {
     private final Path to;
     private final ExecutorService executorService;
     private final Queue<Future<?>> futures;
+    // When non-null, directory subtrees matching this predicate are not renamed here but recorded in `deferred`
+    // for a later phase. Evaluated only on directories, so files are never visited individually for this check.
+    private final Predicate<Path> deferPredicate;
+    private final Queue<Entry<FileStatus, Path>> deferred;
 
     @Override
     public void run() {
       try {
 
+        // Hold back deferred subtrees (e.g. Iceberg `metadata/`) for a later phase instead of renaming them now.
+        if (this.deferPredicate != null && this.from.isDirectory() && this.deferPredicate.test(this.from.getPath())) {
+          log.info(String.format("[ordered-rename] Deferring directory %s (-> %s) to the last phase.",
+              this.from.getPath(), this.to));
+          this.deferred.add(new AbstractMap.SimpleImmutableEntry<>(this.from, this.to));
+          return;
+        }
+
         // Attempt to move safely if directory, unsafely if file (for performance, files are much less likely to collide on target)
         boolean moveSucessful;
 
@@ -646,8 +749,9 @@ public void run() {
               Path relativeFilePath = new Path(StringUtils.substringAfter(fromFile.getPath().toString(),
                   this.from.getPath().toString() + Path.SEPARATOR));
               Path toFilePath = new Path(this.to, relativeFilePath);
-              this.futures.add(this.executorService.submit(
-                  new RenameRecursively(this.fileSystem, fromFile, toFilePath, this.executorService, this.futures)));
+              this.futures.add(this.executorService.submit(new RenameRecursively(
+                  this.fileSystem, fromFile, toFilePath, this.executorService, this.futures, this.deferPredicate,
+                  this.deferred)));
             }
           } else {
             log.info(String.format("File already exists %s. Will not rewrite", this.to));

gobblin-utility/src/test/java/org/apache/gobblin/util/HadoopUtilsTest.java

@@ -264,6 +264,163 @@ public void run() {
 
   }
 
+  @Test
+  public void testRenameRecursivelyOrderedMovesDataBeforeMetadata() throws Exception {
+    final Path testDir = new Path(Files.createTempDir().getAbsolutePath(), "HadoopUtilsTestDir");
+    final FileSystem fs = Mockito.spy(FileSystem.getLocal(new Configuration()));
+
+    // Record the destination of every successful rename, in the order they actually happen. Data renames are
+    // slowed down so that, WITHOUT the phase barrier, the fast metadata renames would land first -- making this a
+    // real regression guard for the ordering (not a pass-by-luck on fast local renames).
+    final List<String> renameDestinations = Collections.synchronizedList(Lists.<String>newArrayList());
+    Mockito.doAnswer(new Answer<Boolean>() {
+      @Override
+      public Boolean answer(InvocationOnMock invocation) throws Throwable {
+        Path dst = (Path) invocation.getArguments()[1];
+        if (dst.toString().endsWith(".orc")) {
+          Thread.sleep(200);
+        }
+        Boolean result = (Boolean) invocation.callRealMethod();
+        if (Boolean.TRUE.equals(result)) {
+          renameDestinations.add(dst.toString());
+        }
+        return result;
+      }
+    }).when(fs).rename(Mockito.any(Path.class), Mockito.any(Path.class));
+
+    try {
+      // Staging tree mirrors an Iceberg table: data/ holds the data files, metadata/ references them by path.
+      Path staging = new Path(testDir, "staging");
+      fs.mkdirs(new Path(staging, "table1/data"));
+      fs.mkdirs(new Path(staging, "table1/metadata"));
+      fs.create(new Path(staging, "table1/data/f1.orc")).close();
+      fs.create(new Path(staging, "table1/data/f2.orc")).close();
+      fs.create(new Path(staging, "table1/metadata/m1.avro")).close();
+      fs.create(new Path(staging, "table1/metadata/m2.avro")).close();
+
+      // Pre-create the target table dir and its data/metadata children so the rename descends to file level --
+      // the incremental/retry case where per-file ordering actually matters (no atomic whole-subtree move).
+      Path target = new Path(testDir, "target");
+      fs.mkdirs(new Path(target, "table1/data"));
+      fs.mkdirs(new Path(target, "table1/metadata"));
+
+      HadoopUtils.renameRecursivelyOrdered(fs, staging, target, HadoopUtils::isIcebergMetadataDir);
+
+      // Every file landed at the target.
+      Assert.assertTrue(fs.exists(new Path(target, "table1/data/f1.orc")));
+      Assert.assertTrue(fs.exists(new Path(target, "table1/data/f2.orc")));
+      Assert.assertTrue(fs.exists(new Path(target, "table1/metadata/m1.avro")));
+      Assert.assertTrue(fs.exists(new Path(target, "table1/metadata/m2.avro")));
+
+      // Every data file was renamed strictly before any metadata file.
+      int lastDataIdx = -1;
+      int firstMetadataIdx = Integer.MAX_VALUE;
+      for (int i = 0; i < renameDestinations.size(); i++) {
+        String dst = renameDestinations.get(i);
+        if (dst.endsWith(".orc")) {
+          lastDataIdx = Math.max(lastDataIdx, i);
+        } else if (dst.endsWith(".avro")) {
+          firstMetadataIdx = Math.min(firstMetadataIdx, i);
+        }
+      }
+      Assert.assertTrue(lastDataIdx >= 0, "Expected at least one data rename, got: " + renameDestinations);
+      Assert.assertTrue(firstMetadataIdx != Integer.MAX_VALUE,
+          "Expected at least one metadata rename, got: " + renameDestinations);
+      Assert.assertTrue(lastDataIdx < firstMetadataIdx,
+          "All data files must be renamed before any metadata file. Actual order: " + renameDestinations);
+    } finally {
+      fs.delete(testDir, true);
+    }
+  }
+
+  @Test
+  public void testRenameRecursivelyOrderedFreshTableCopiesEverything() throws Exception {
+    // When the target table dir does not yet exist, the whole subtree moves atomically (nothing is deferred);
+    // verify the ordered variant still copies every file correctly.
+    final Path testDir = new Path(Files.createTempDir().getAbsolutePath(), "HadoopUtilsTestDir");
+    final FileSystem fs = FileSystem.getLocal(new Configuration());
+    try {
+      Path staging = new Path(testDir, "staging");
+      fs.mkdirs(new Path(staging, "table1/data"));
+      fs.mkdirs(new Path(staging, "table1/metadata"));
+      fs.create(new Path(staging, "table1/data/f1.orc")).close();
+      fs.create(new Path(staging, "table1/metadata/m1.avro")).close();
+
+      Path target = new Path(testDir, "target");
+
+      HadoopUtils.renameRecursivelyOrdered(fs, staging, target, HadoopUtils::isIcebergMetadataDir);
+
+      Assert.assertTrue(fs.exists(new Path(target, "table1/data/f1.orc")));
+      Assert.assertTrue(fs.exists(new Path(target, "table1/metadata/m1.avro")));
+    } finally {
+      fs.delete(testDir, true);
+    }
+  }
+
+  @Test
+  public void testRenameRecursivelyOrderedAbortsBeforeMetadataWhenDataRenameFails() throws Exception {
+    // If a data-phase rename fails, the metadata phase must never run (Iceberg metadata must not be published
+    // when the data it references failed to land).
+    final Path testDir = new Path(Files.createTempDir().getAbsolutePath(), "HadoopUtilsTestDir");
+    final FileSystem fs = Mockito.spy(FileSystem.getLocal(new Configuration()));
+
+    // Fail every data (.orc) rename; record any metadata (.avro) rename that is attempted.
+    final List<String> metadataRenameAttempts = Collections.synchronizedList(Lists.<String>newArrayList());
+    Mockito.doAnswer(new Answer<Object>() {
+      @Override
+      public Object answer(InvocationOnMock invocation) throws Throwable {
+        Path dst = (Path) invocation.getArguments()[1];
+        if (dst.toString().endsWith(".avro")) {
+          metadataRenameAttempts.add(dst.toString());
+        }
+        if (dst.toString().endsWith(".orc")) {
+          throw new IOException("Injected failure renaming data file " + dst);
+        }
+        return invocation.callRealMethod();
+      }
+    }).when(fs).rename(Mockito.any(Path.class), Mockito.any(Path.class));
+
+    try {
+      Path staging = new Path(testDir, "staging");
+      fs.mkdirs(new Path(staging, "table1/data"));
+      fs.mkdirs(new Path(staging, "table1/metadata"));
+      fs.create(new Path(staging, "table1/data/f1.orc")).close();
+      fs.create(new Path(staging, "table1/data/f2.orc")).close();
+      fs.create(new Path(staging, "table1/metadata/m1.avro")).close();
+      fs.create(new Path(staging, "table1/metadata/m2.avro")).close();
+
+      // Pre-create the target table dir and children so the rename descends to file level.
+      Path target = new Path(testDir, "target");
+      fs.mkdirs(new Path(target, "table1/data"));
+      fs.mkdirs(new Path(target, "table1/metadata"));
+
+      try {
+        HadoopUtils.renameRecursivelyOrdered(fs, staging, target, HadoopUtils::isIcebergMetadataDir);
+        Assert.fail("Expected the rename to fail during the data phase");
+      } catch (IOException expected) {
+        // expected: the data-phase failure aborts the whole operation
+      }
+
+      Assert.assertTrue(metadataRenameAttempts.isEmpty(),
+          "Metadata must not be renamed after a data-phase failure, but these were attempted: "
+              + metadataRenameAttempts);
+      // And no metadata file should have actually landed at the target.
+      Assert.assertFalse(fs.exists(new Path(target, "table1/metadata/m1.avro")));
+      Assert.assertFalse(fs.exists(new Path(target, "table1/metadata/m2.avro")));
+    } finally {
+      fs.delete(testDir, true);
+    }
+  }
+
+  @Test
+  public void testIsIcebergMetadataDir() {
+    Assert.assertTrue(HadoopUtils.isIcebergMetadataDir(new Path("/data/openhouse/db/tbl-uuid/metadata")));
+    Assert.assertFalse(HadoopUtils.isIcebergMetadataDir(new Path("/data/openhouse/db/tbl-uuid/data")));
+    // The `/data` HDFS mount is not an Iceberg metadata dir, and the predicate matches the dir, not its files.
+    Assert.assertFalse(HadoopUtils.isIcebergMetadataDir(new Path("/data")));
+    Assert.assertFalse(HadoopUtils.isIcebergMetadataDir(new Path("/data/openhouse/db/tbl-uuid/metadata/m1.avro")));
+  }
+
   @Test
   public void testSanitizePath() throws Exception {
     Assert.assertEquals(HadoopUtils.sanitizePath("/A:B/::C:::D\\", "abc"), "/AabcB/abcabcCabcabcabcDabc");