Publish S3 dataset cache files atomically

docs/release notes/4.0.2/745.bugfix.md

@@ -0,0 +1,20 @@
+### Publish S3 Dataset Downloads Only After Completion
+
+**Description**
+
+Download S3 dataset files into private temporary files and atomically publish the
+cache path after successful completion and content-length validation. This prevents
+other benchmark processes from reading an in-progress download or losing a completed
+cache file when another download attempt fails.
+
+**How to Enable**
+
+No configuration is required. This applies to S3 downloads made by the example
+benchmark dataset loader; HTTP downloads already use atomic publication.
+
+**Notes**
+
+Normal failures remove their temporary files. A hard process termination may leave
+an unreferenced temporary file, which is not accepted as cached data. Existing cache
+files are not retroactively validated; remove any files known to come from an
+interrupted download before reusing that cache.

jvector-examples/src/main/java/io/github/jbellis/jvector/example/benchmarks/datasets/DataSetLoaderSimpleMFD.java

@@ -854,38 +854,38 @@ private void downloadFileS3(String s3Url, Path localPath) throws IOException {
 
         S3TransferManager tm = getS3TransferManager();
 
-        DownloadFileRequest request = DownloadFileRequest.builder()
-                .getObjectRequest(b -> b.bucket(bucket).key(key))
-                .addTransferListener(LoggingTransferListener.create())
-                .destination(localPath)
-                .build();
-
-        boolean downloaded = false;
+        Path targetDir = localPath.toAbsolutePath().getParent();
+        Files.createDirectories(targetDir);
         for (int i = 0; i < 3; i++) { // 3 retries
+            // Only this attempt owns the temporary path; the shared cache name stays unpublished.
+            Path tempFile = Files.createTempFile(targetDir, "download-", ".tmp");
             try {
+                DownloadFileRequest request = DownloadFileRequest.builder()
+                        .getObjectRequest(b -> b.bucket(bucket).key(key))
+                        .addTransferListener(LoggingTransferListener.create())
+                        .destination(tempFile)
+                        .build();
                 FileDownload downloadFile = tm.downloadFile(request);
                 CompletedFileDownload result = downloadFile.completionFuture().join();
-                long downloadedSize = Files.size(localPath);
+                long downloadedSize = Files.size(tempFile);
                 Long expectedSize = result.response().contentLength();
 
-                // Null check prevents NullPointerException during unboxing.
-                // If expectedSize is null, we trust the transfer manager's successful completion.
+                // If size is unavailable, trust the transfer manager's successful completion.
                 if (expectedSize != null && downloadedSize != expectedSize) {
                     logger.error("Incomplete download (got {} of {} bytes). Retrying...", downloadedSize, expectedSize);
-                    Files.deleteIfExists(localPath);
                     continue;
                 }
 
-                downloaded = true;
-                break;
+                // Readers see a complete file, and a failed concurrent attempt cannot delete it.
+                Files.move(tempFile, localPath, StandardCopyOption.ATOMIC_MOVE, StandardCopyOption.REPLACE_EXISTING);
+                return;
             } catch (Exception e) {
                 logger.error("Download attempt {} failed for {}: {}", i + 1, redact(key), redact(e.getMessage()));
-                Files.deleteIfExists(localPath);
+            } finally {
+                Files.deleteIfExists(tempFile);
             }
         }
-        if (!downloaded) {
-            throw new IOException("Failed to download " + redact(s3Url) + " after 3 attempts");
-        }
+        throw new IOException("Failed to download " + redact(s3Url) + " after 3 attempts");
     }
 
     private static S3AsyncClient s3AsyncClient() {

jvector-examples/src/test/java/io/github/jbellis/jvector/example/benchmarks/datasets/TestDatasetDownloadPublication.java

@@ -0,0 +1,158 @@
+/*
+ * Copyright DataStax, Inc.
+ *
+ * Licensed 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 io.github.jbellis.jvector.example.benchmarks.datasets;
+
+import org.junit.Rule;
+import org.junit.Test;
+import org.junit.rules.TemporaryFolder;
+import software.amazon.awssdk.services.s3.model.GetObjectResponse;
+import software.amazon.awssdk.transfer.s3.S3TransferManager;
+import software.amazon.awssdk.transfer.s3.model.CompletedFileDownload;
+import software.amazon.awssdk.transfer.s3.model.DownloadFileRequest;
+import software.amazon.awssdk.transfer.s3.model.FileDownload;
+
+import java.io.IOException;
+import java.lang.reflect.InvocationTargetException;
+import java.lang.reflect.Proxy;
+import java.nio.file.Files;
+import java.nio.file.Path;
+import java.util.concurrent.*;
+
+import static org.junit.Assert.*;
+
+public class TestDatasetDownloadPublication {
+    @Rule public TemporaryFolder folder = new TemporaryFolder();
+
+    @Test(timeout = 10000) public void incompleteS3DownloadIsNotVisibleInCache() throws Exception {
+        var cache = folder.newFolder().toPath();
+        var downloads = new Downloads(cache);
+        var executor = Executors.newSingleThreadExecutor();
+        try {
+            var result = executor.submit(() -> { downloads.ensure(); return null; });
+            var pending = downloads.next();
+            assertArrayEquals(new byte[] {1}, Files.readAllBytes(pending.path));
+            assertFalse("A cache hit must not expose a partial download", Files.exists(downloads.target));
+            pending.complete(new byte[] {1, 2, 3}, 3);
+            result.get(5, TimeUnit.SECONDS);
+            assertArrayEquals(new byte[] {1, 2, 3}, Files.readAllBytes(downloads.target));
+            assertEquals(1, fileCount(cache));
+        } finally { downloads.abort(); executor.shutdownNow(); }
+    }
+
+    @Test(timeout = 10000) public void failedConcurrentS3DownloadCannotDeletePublishedFile() throws Exception {
+        var cache = folder.newFolder().toPath();
+        var downloads = new Downloads(cache);
+        var executor = Executors.newFixedThreadPool(2);
+        try {
+            var slow = executor.submit(() -> { downloads.ensure(); return null; });
+            var first = downloads.next();
+            var fast = executor.submit(() -> { downloads.ensure(); return null; });
+            var second = downloads.next();
+            assertNotEquals(first.path, second.path);
+            second.complete(new byte[] {4, 5, 6}, 3);
+            fast.get(5, TimeUnit.SECONDS);
+            first.fail();
+            downloads.next().fail();
+            downloads.next().fail();
+            var failure = assertThrows(ExecutionException.class, () -> slow.get(5, TimeUnit.SECONDS));
+            assertTrue(failure.getCause() instanceof IOException);
+            assertArrayEquals(new byte[] {4, 5, 6}, Files.readAllBytes(downloads.target));
+            assertEquals(1, fileCount(cache));
+        } finally { downloads.abort(); executor.shutdownNow(); }
+    }
+
+    @Test(timeout = 10000) public void invalidS3LengthIsNotPublishedAndRetriesAreCleanedUp() throws Exception {
+        var cache = folder.newFolder().toPath();
+        var downloads = new Downloads(cache);
+        var executor = Executors.newSingleThreadExecutor();
+        try {
+            var result = executor.submit(() -> { downloads.ensure(); return null; });
+            for (int attempt = 0; attempt < 3; attempt++) {
+                downloads.next().complete(new byte[] {1}, 2);
+            }
+            var failure = assertThrows(ExecutionException.class, () -> result.get(5, TimeUnit.SECONDS));
+            assertTrue(failure.getCause() instanceof IOException);
+            assertFalse(Files.exists(downloads.target));
+            assertEquals(0, fileCount(cache));
+        } finally { downloads.abort(); executor.shutdownNow(); }
+    }
+
+    private static long fileCount(Path directory) throws IOException {
+        try (var files = Files.list(directory)) { return files.count(); }
+    }
+
+    private static class Pending {
+        final Path path;
+        final CompletableFuture<CompletedFileDownload> future = new CompletableFuture<>();
+        Pending(Path path) { this.path = path; }
+        void complete(byte[] content, long expectedSize) throws IOException {
+            Files.write(path, content);
+            future.complete(CompletedFileDownload.builder()
+                    .response(GetObjectResponse.builder().contentLength(expectedSize).build()).build());
+        }
+        void fail() { future.completeExceptionally(new IOException("Injected download failure")); }
+    }
+
+    private static class Downloads {
+        final Path cache, target;
+        final DataSetLoaderSimpleMFD loader;
+        final BlockingQueue<Pending> pending = new LinkedBlockingQueue<>();
+        final ConcurrentLinkedQueue<Pending> all = new ConcurrentLinkedQueue<>();
+        volatile boolean aborted;
+        Downloads(Path cache) throws Exception {
+            this.cache = cache;
+            target = cache.resolve("vectors.fvecs");
+            Path metadata = cache.getParent().resolve("metadata.yml");
+            Files.writeString(metadata, "test-ds:\n  similarity_function: DOT_PRODUCT\n  load_behavior: NO_SCRUB\n");
+            loader = new DataSetLoaderSimpleMFD(null, cache.toString(), false,
+                    DataSetMetadataReader.load(metadata.toString()));
+            // Replace only the network transport; exercise the real cache and retry logic.
+            var manager = (S3TransferManager) Proxy.newProxyInstance(S3TransferManager.class.getClassLoader(),
+                    new Class<?>[] {S3TransferManager.class}, (proxy, method, args) -> {
+                if (!method.getName().equals("downloadFile")) throw new UnsupportedOperationException(method.getName());
+                var request = (DownloadFileRequest) args[0];
+                var download = new Pending(request.destination());
+                Files.write(download.path, new byte[] {1});
+                all.add(download);
+                pending.add(download);
+                if (aborted) download.fail();
+                return Proxy.newProxyInstance(FileDownload.class.getClassLoader(),
+                        new Class<?>[] {FileDownload.class}, (file, call, arguments) -> {
+                    if (call.getName().equals("completionFuture")) return download.future;
+                    throw new UnsupportedOperationException(call.getName());
+                });
+            });
+            var field = DataSetLoaderSimpleMFD.class.getDeclaredField("s3TransferManager");
+            field.setAccessible(true);
+            field.set(loader, manager);
+        }
+        void abort() {
+            aborted = true;
+            all.forEach(Pending::fail);
+        }
+        Pending next() throws Exception {
+            var download = pending.poll(5, TimeUnit.SECONDS);
+            assertNotNull("Expected a download attempt", download);
+            return download;
+        }
+        void ensure() throws Exception {
+            var method = DataSetLoaderSimpleMFD.class.getDeclaredMethod("ensureFileAvailable", String.class, Path.class, String.class);
+            method.setAccessible(true);
+            try { method.invoke(loader, "vectors.fvecs", cache, "s3://bucket/"); }
+            catch (InvocationTargetException failure) { throw (Exception) failure.getCause(); }
+        }
+    }
+}