Route FilterDirectory#copyFrom through createOutput (10x backport of #16530)

lucene/CHANGES.txt

@@ -185,7 +185,9 @@ Bug Fixes
   exceptional exits, so vector inputs switch back from sequential read advice.
   (Simon Cooper, John Maurice)
 
-* GITHUB#16173: FilterDirectory now delegates copyFrom() to the wrapped directory. (Prithvi S)
+* GITHUB#16530: FilterDirectory#copyFrom routes through createOutput so that per-file bookkeeping
+  in subclasses applies to copied files, reverting the delegation from GITHUB#16173. Also fixes
+  temp file cleanup when TrackingTmpOutputDirectoryWrapper#copyFrom fails. (Tim Allison, Prithvi S)
 
 * GITHUB#16452: Fix over-allocation in MemoryAccountingBitsetCollectorManager.Result#bitSet, which
   is now sized to highestMatchedDoc + 1 instead of the full searched range. (Sasilekha R)

lucene/core/src/java/org/apache/lucene/index/TrackingTmpOutputDirectoryWrapper.java

@@ -25,6 +25,7 @@
 import org.apache.lucene.store.IOContext;
 import org.apache.lucene.store.IndexInput;
 import org.apache.lucene.store.IndexOutput;
+import org.apache.lucene.util.IOUtils;
 
 final class TrackingTmpOutputDirectoryWrapper extends FilterDirectory {
   private final Map<String, String> fileNames = new HashMap<>();
@@ -47,6 +48,23 @@ public IndexInput openInput(String name, IOContext context) throws IOException {
     return super.openInput(tmpName, context);
   }
 
+  @Override
+  public void copyFrom(Directory from, String src, String dest, IOContext context)
+      throws IOException {
+    // the inherited failure cleanup would delete dest, but createOutput() redirects dest to a
+    // temp file; on failure remove the mapping and delete the temp file instead
+    try (IndexInput is = from.openInput(src, IOContext.READONCE);
+        IndexOutput os = createOutput(dest, context)) {
+      os.copyBytes(is, is.length());
+    } catch (Throwable t) {
+      String tmpName = fileNames.remove(dest);
+      if (tmpName != null) {
+        IOUtils.deleteFilesIgnoringExceptions(in, tmpName);
+      }
+      throw t;
+    }
+  }
+
   public Map<String, String> getTemporaryFiles() {
     return fileNames;
   }

lucene/core/src/java/org/apache/lucene/store/FilterDirectory.java

@@ -116,10 +116,16 @@ public String toString() {
     return getClass().getSimpleName() + "(" + in.toString() + ")";
   }
 
+  /**
+   * Not delegated to {@link #in}: the default implementation routes through {@link #createOutput}
+   * so per-file bookkeeping there covers copied files too. Subclasses that don't override {@link
+   * #createOutput} and want an optimized copy (e.g. {@code HardlinkCopyDirectoryWrapper}) may
+   * override to delegate: {@code in.copyFrom(from, src, dest, context)}.
+   */
   @Override
   public void copyFrom(Directory from, String src, String dest, IOContext context)
       throws IOException {
-    in.copyFrom(from, src, dest, context);
+    super.copyFrom(from, src, dest, context);
   }
 
   @Override

lucene/core/src/test/org/apache/lucene/index/TestTrackingTmpOutputDirectoryWrapper.java

@@ -0,0 +1,88 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.lucene.index;
+
+import java.io.IOException;
+import java.util.Arrays;
+import java.util.List;
+import org.apache.lucene.store.Directory;
+import org.apache.lucene.store.FilterDirectory;
+import org.apache.lucene.store.FilterIndexInput;
+import org.apache.lucene.store.IOContext;
+import org.apache.lucene.store.IndexInput;
+import org.apache.lucene.store.IndexOutput;
+import org.apache.lucene.tests.mockfile.ExtrasFS;
+import org.apache.lucene.tests.util.LuceneTestCase;
+
+public class TestTrackingTmpOutputDirectoryWrapper extends LuceneTestCase {
+
+  public void testCopyFromTracksLogicalNameInTemporaryFiles() throws IOException {
+    try (Directory source = newDirectory();
+        Directory backing = newDirectory()) {
+      try (IndexOutput out = source.createOutput("src", IOContext.DEFAULT)) {
+        out.writeBytes(new byte[] {1, 2, 3}, 3);
+      }
+      TrackingTmpOutputDirectoryWrapper wrapper = new TrackingTmpOutputDirectoryWrapper(backing);
+      wrapper.copyFrom(source, "src", "dest", IOContext.DEFAULT);
+
+      assertTrue(
+          "copyFrom must register dest in getTemporaryFiles() via createOutput()",
+          wrapper.getTemporaryFiles().containsKey("dest"));
+      String tmpName = wrapper.getTemporaryFiles().get("dest");
+      assertNotEquals("temp file name must differ from logical name", "dest", tmpName);
+      try (IndexInput in = wrapper.openInput("dest", IOContext.DEFAULT)) {
+        assertEquals(3L, in.length());
+      }
+    }
+  }
+
+  public void testCopyFromCleansUpOnFailure() throws IOException {
+    try (Directory source = newDirectory();
+        Directory backing = newDirectory()) {
+      try (IndexOutput out = source.createOutput("src", IOContext.DEFAULT)) {
+        out.writeBytes(new byte[] {1, 2, 3}, 3);
+      }
+      // Source whose reads always throw, so copyBytes fails after createOutput succeeds.
+      Directory failingSource =
+          new FilterDirectory(source) {
+            @Override
+            public IndexInput openInput(String name, IOContext context) throws IOException {
+              return new FilterIndexInput("failing:" + name, super.openInput(name, context)) {
+                @Override
+                public void readBytes(byte[] b, int offset, int len) throws IOException {
+                  throw new IOException("simulated read failure");
+                }
+              };
+            }
+          };
+
+      TrackingTmpOutputDirectoryWrapper wrapper = new TrackingTmpOutputDirectoryWrapper(backing);
+      expectThrows(
+          IOException.class,
+          () -> wrapper.copyFrom(failingSource, "src", "dest", IOContext.DEFAULT));
+
+      assertFalse(
+          "dest must not remain in getTemporaryFiles() after failed copyFrom",
+          wrapper.getTemporaryFiles().containsKey("dest"));
+      assertEquals(
+          "temp file must be deleted from backing dir after failed copyFrom",
+          List.of(),
+          Arrays.stream(backing.listAll()).filter(f -> ExtrasFS.isExtra(f) == false).toList());
+    }
+  }
+}

lucene/core/src/test/org/apache/lucene/store/TestFilterDirectory.java

@@ -44,25 +44,25 @@ public void testOverrides() throws Exception {
     }
   }
 
-  public void testCopyFromDelegates() throws IOException {
+  public void testCopyFromRoutesThroughCreateOutput() throws IOException {
     try (Directory srcDir = new ByteBuffersDirectory();
         Directory destDir = new ByteBuffersDirectory()) {
       try (IndexOutput out = srcDir.createOutput("test.txt", IOContext.DEFAULT)) {
         out.writeString("hello");
       }
-      boolean[] delegated = {false};
-      FilterDirectory wrappedDestDir =
+      Set<String> created = new HashSet<>();
+      FilterDirectory filterDestDir =
           new FilterDirectory(destDir) {
             @Override
-            public void copyFrom(Directory from, String src, String dest, IOContext context)
-                throws IOException {
-              delegated[0] = true;
-              in.copyFrom(from, src, dest, context);
+            public IndexOutput createOutput(String name, IOContext context) throws IOException {
+              created.add(name);
+              return in.createOutput(name, context);
             }
           };
-      FilterDirectory filterDestDir = new FilterDirectory(wrappedDestDir) {};
       filterDestDir.copyFrom(srcDir, "test.txt", "copied.txt", IOContext.DEFAULT);
-      assertTrue("copyFrom should delegate to wrapped directory", delegated[0]);
+      assertTrue(
+          "copyFrom should route through the wrapper's createOutput",
+          created.contains("copied.txt"));
       assertTrue(slowFileExists(destDir, "copied.txt"));
       try (IndexInput in = destDir.openInput("copied.txt", IOContext.DEFAULT)) {
         assertEquals("hello", in.readString());

lucene/misc/src/test/org/apache/lucene/misc/store/TestByteWritesTrackingDirectoryWrapper.java

@@ -105,6 +105,20 @@ public void testRandomTempOutput() throws Exception {
     assertEquals(expectedMergeBytes, dir.getMergedBytes());
   }
 
+  public void testCopyFromTracksBytes() throws Exception {
+    ByteWritesTrackingDirectoryWrapper dir =
+        new ByteWritesTrackingDirectoryWrapper(new ByteBuffersDirectory());
+    int expectedMergeBytes = 1 + random().nextInt(100);
+    try (Directory srcDir = new ByteBuffersDirectory()) {
+      try (IndexOutput out = srcDir.createOutput("src", IOContext.DEFAULT)) {
+        out.writeBytes(new byte[expectedMergeBytes], expectedMergeBytes);
+      }
+      dir.copyFrom(
+          srcDir, "src", "dest", IOContext.merge(new MergeInfo(10, expectedMergeBytes, false, 2)));
+      assertEquals(expectedMergeBytes, dir.getMergedBytes());
+    }
+  }
+
   @Override
   protected Directory getDirectory(Path path) throws IOException {
     return new ByteWritesTrackingDirectoryWrapper(new ByteBuffersDirectory());

lucene/misc/src/test/org/apache/lucene/misc/store/TestDirectIODirectoryCopyFrom.java

@@ -0,0 +1,52 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.lucene.misc.store;
+
+import java.util.HashSet;
+import java.util.OptionalLong;
+import java.util.Set;
+import org.apache.lucene.store.Directory;
+import org.apache.lucene.store.FSDirectory;
+import org.apache.lucene.store.IOContext;
+import org.apache.lucene.store.IndexOutput;
+import org.apache.lucene.store.MergeInfo;
+import org.apache.lucene.tests.util.LuceneTestCase;
+
+/** Tests {@link DirectIODirectory#copyFrom} without requiring actual O_DIRECT support. */
+public class TestDirectIODirectoryCopyFrom extends LuceneTestCase {
+
+  public void testCopyFromConsultsUseDirectIO() throws Exception {
+    Set<String> consulted = new HashSet<>();
+    try (Directory srcDir = FSDirectory.open(createTempDir("src"));
+        DirectIODirectory dir =
+            new DirectIODirectory(FSDirectory.open(createTempDir("dest"))) {
+              @Override
+              protected boolean useDirectIO(
+                  String name, IOContext context, OptionalLong fileLength) {
+                consulted.add(name);
+                return false;
+              }
+            }) {
+      try (IndexOutput out = srcDir.createOutput("src", IOContext.DEFAULT)) {
+        out.writeBytes(new byte[8], 8);
+      }
+      dir.copyFrom(srcDir, "src", "dest", IOContext.merge(new MergeInfo(10, 8, false, 2)));
+      assertTrue(
+          "copyFrom must route through createOutput/useDirectIO", consulted.contains("dest"));
+    }
+  }
+}