[FLINK-39786][s3]closeForCommit() leaks orphan multipart uploads on upload failure

flink-filesystems/flink-s3-fs-native/src/main/java/org/apache/flink/fs/s3native/writer/NativeS3RecoverableFsDataOutputStream.java

@@ -198,8 +198,8 @@ private void uploadCurrentPart() throws IOException {
         currentOutputStream.close();
 
         // Do not delete the temp file if uploadPart fails: propagate the original exception
-        // unmasked and let close() perform cleanup. nextPartNumber is only advanced on success so a
-        // failed attempt does not leave a gap in the part sequence.
+        // unmasked and let the cleanup path (close() or the closeForCommit() failure handler)
+        // delete it and abort the upload. nextPartNumber is only advanced on success.
         NativeS3ObjectOperations.UploadPartResult result =
                 s3AccessHelper.uploadPart(
                         key, uploadId, nextPartNumber, currentTempFile, currentPartSize);
@@ -219,17 +219,25 @@ public Committer closeForCommit() throws IOException {
                 throw new IOException("Stream is already closed");
             }
 
-            currentOutputStream.close();
+            final NativeS3Recoverable recoverable;
+            try {
+                currentOutputStream.close();
 
-            if (currentPartSize > 0) {
-                uploadCurrentPart();
-            } else {
-                Files.delete(currentTempFile.toPath());
-            }
+                if (currentPartSize > 0) {
+                    uploadCurrentPart();
+                } else {
+                    Files.delete(currentTempFile.toPath());
+                }
 
-            NativeS3Recoverable recoverable =
-                    new NativeS3Recoverable(
-                            key, uploadId, new ArrayList<>(completedParts), numBytesInParts);
+                recoverable =
+                        new NativeS3Recoverable(
+                                key, uploadId, new ArrayList<>(completedParts), numBytesInParts);
+            } catch (IOException e) {
+                // The commit failed after the multipart upload had been created and parts may
+                // already have been uploaded. Abort it so it does not leak as an orphan upload.
+                closed = true;
+                throw abortUploadAndReleaseResources(e);
+            }
 
             closed = true;
             return new NativeS3Committer(s3AccessHelper, recoverable);
@@ -272,32 +280,7 @@ public void close() throws IOException {
         try {
             if (!closed) {
                 closed = true;
-                IOException cleanupException = null;
-                if (currentOutputStream != null) {
-                    try {
-                        currentOutputStream.close();
-                    } catch (IOException e) {
-                        cleanupException = ExceptionUtils.firstOrSuppressed(e, cleanupException);
-                    }
-                }
-                if (currentTempFile != null && currentTempFile.exists()) {
-                    try {
-                        Files.delete(currentTempFile.toPath());
-                    } catch (IOException e) {
-                        cleanupException = ExceptionUtils.firstOrSuppressed(e, cleanupException);
-                    }
-                }
-
-                try {
-                    s3AccessHelper.abortMultiPartUpload(key, uploadId);
-                } catch (IOException e) {
-                    LOG.warn(
-                            "Multipart upload failed (key={}, uploadId={}). "
-                                    + "S3 lifecycle rules should eventually clean up the incomplete upload.",
-                            key,
-                            uploadId,
-                            e);
-                }
+                IOException cleanupException = abortUploadAndReleaseResources(null);
                 if (cleanupException != null) {
                     throw cleanupException;
                 }
@@ -307,6 +290,38 @@ public void close() throws IOException {
         }
     }
 
+    /** Aborts the multipart upload and releases local resources on the best effort basis. */
+    @Nullable
+    private IOException abortUploadAndReleaseResources(@Nullable IOException primary) {
+        IOException collected = primary;
+        if (currentOutputStream != null) {
+            try {
+                currentOutputStream.close();
+            } catch (IOException e) {
+                collected = ExceptionUtils.firstOrSuppressed(e, collected);
+            }
+        }
+        if (currentTempFile != null && currentTempFile.exists()) {
+            try {
+                Files.delete(currentTempFile.toPath());
+            } catch (IOException e) {
+                collected = ExceptionUtils.firstOrSuppressed(e, collected);
+            }
+        }
+        try {
+            s3AccessHelper.abortMultiPartUpload(key, uploadId);
+        } catch (IOException e) {
+            LOG.warn(
+                    "Failed to abort multipart upload (key={}, uploadId={}); it may be left as an "
+                            + "orphan upload in S3. Propagating the failure to the caller.",
+                    key,
+                    uploadId,
+                    e);
+            collected = ExceptionUtils.firstOrSuppressed(e, collected);
+        }
+        return collected;
+    }
+
     private void lock() throws IOException {
         try {
             lock.lockInterruptibly();

flink-filesystems/flink-s3-fs-native/src/test/java/org/apache/flink/fs/s3native/writer/InMemoryNativeS3Operations.java

@@ -67,6 +67,15 @@ public final class InMemoryNativeS3Operations extends NativeS3ObjectOperations {
     /** uploadId → partNumber → uploaded bytes for in-flight MPUs. */
     public final Map<String, Map<Integer, byte[]>> openMultipartUploads = new HashMap<>();
 
+    /** When {@code true}, {@link #uploadPart} throws to simulate a part-upload failure. */
+    public boolean failUploadPart = false;
+
+    /** When {@code true}, {@link #abortMultiPartUpload} throws to simulate an abort failure. */
+    public boolean failAbortMultiPartUpload = false;
+
+    /** Number of times {@link #abortMultiPartUpload} was invoked, including failed attempts. */
+    public int abortAttempts = 0;
+
     private final String bucketName;
     private final AtomicInteger uploadIdSeq = new AtomicInteger();
     private final AtomicInteger putObjectSeq = new AtomicInteger();
@@ -91,6 +100,9 @@ public String startMultiPartUpload(String key) {
     public UploadPartResult uploadPart(
             String key, String uploadId, int partNumber, File file, long length)
             throws IOException {
+        if (failUploadPart) {
+            throw new IOException("injected uploadPart failure for uploadId: " + uploadId);
+        }
         Map<Integer, byte[]> parts = openMultipartUploads.get(uploadId);
         if (parts == null) {
             throw new IOException("unknown uploadId: " + uploadId);
@@ -154,7 +166,11 @@ public CompleteMultipartUploadResult commitMultiPartUpload(
     }
 
     @Override
-    public void abortMultiPartUpload(String key, String uploadId) {
+    public void abortMultiPartUpload(String key, String uploadId) throws IOException {
+        abortAttempts++;
+        if (failAbortMultiPartUpload) {
+            throw new IOException("injected abort failure for uploadId: " + uploadId);
+        }
         openMultipartUploads.remove(uploadId);
     }
 

flink-filesystems/flink-s3-fs-native/src/test/java/org/apache/flink/fs/s3native/writer/NativeS3RecoverableFsDataOutputStreamTest.java

@@ -0,0 +1,188 @@
+/*
+ * 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.flink.fs.s3native.writer;
+
+import org.apache.flink.core.fs.RecoverableFsDataOutputStream;
+
+import org.junit.jupiter.api.Test;
+import org.junit.jupiter.api.io.TempDir;
+
+import java.io.IOException;
+import java.nio.file.Files;
+import java.nio.file.Path;
+import java.util.Arrays;
+
+import static org.assertj.core.api.Assertions.assertThat;
+import static org.assertj.core.api.Assertions.assertThatThrownBy;
+
+/** Test {@link NativeS3RecoverableFsDataOutputStream}. */
+class NativeS3RecoverableFsDataOutputStreamTest {
+
+    private static final String KEY = "out.txt";
+    private static final long MIN_PART_SIZE = 10L;
+
+    @TempDir Path tmp;
+
+    @Test
+    void closeForCommitAbortsMultipartUploadWhenPartUploadFails() throws Exception {
+        InMemoryNativeS3Operations s3 = new InMemoryNativeS3Operations();
+        s3.failUploadPart = true;
+
+        String uploadId = s3.startMultiPartUpload(KEY);
+        assertThat(s3.openMultipartUploads).containsKey(uploadId);
+
+        NativeS3RecoverableFsDataOutputStream stream = newStream(s3, uploadId);
+        stream.write(bytes('A', 5), 0, 5); // < MIN_PART_SIZE, so it is uploaded during commit
+
+        assertThatThrownBy(stream::closeForCommit)
+                .isInstanceOf(IOException.class)
+                .hasMessageContaining("injected uploadPart failure");
+
+        assertThat(s3.abortAttempts)
+                .as("closeForCommit must abort the upload on failure")
+                .isEqualTo(1);
+        assertThat(s3.openMultipartUploads)
+                .as("the multipart upload must not leak after a failed commit")
+                .doesNotContainKey(uploadId);
+        assertThat(countLocalFilesIn(tmp)).as("the local temp file must be cleaned up").isZero();
+    }
+
+    @Test
+    void closeForCommitSurfacesAbortFailureWhenBothUploadAndAbortFail() throws Exception {
+        InMemoryNativeS3Operations s3 = new InMemoryNativeS3Operations();
+        s3.failUploadPart = true;
+        s3.failAbortMultiPartUpload = true;
+
+        String uploadId = s3.startMultiPartUpload(KEY);
+        NativeS3RecoverableFsDataOutputStream stream = newStream(s3, uploadId);
+        stream.write(bytes('A', 5), 0, 5);
+
+        assertThatThrownBy(stream::closeForCommit)
+                .isInstanceOf(IOException.class)
+                .hasMessageContaining("injected uploadPart failure")
+                .satisfies(
+                        t ->
+                                assertThat(t.getSuppressed())
+                                        .as("the abort failure must be surfaced, not swallowed")
+                                        .anySatisfy(
+                                                s ->
+                                                        assertThat(s)
+                                                                .hasMessageContaining(
+                                                                        "injected abort failure")));
+
+        assertThat(s3.abortAttempts).isEqualTo(1);
+    }
+
+    @Test
+    void closeSurfacesAbortFailureInsteadOfSwallowingIt() throws Exception {
+        InMemoryNativeS3Operations s3 = new InMemoryNativeS3Operations();
+        s3.failAbortMultiPartUpload = true;
+
+        String uploadId = s3.startMultiPartUpload(KEY);
+        NativeS3RecoverableFsDataOutputStream stream = newStream(s3, uploadId);
+        stream.write(bytes('A', 5), 0, 5);
+
+        assertThatThrownBy(stream::close)
+                .isInstanceOf(IOException.class)
+                .hasMessageContaining("injected abort failure");
+
+        assertThat(s3.abortAttempts).isEqualTo(1);
+        assertThat(countLocalFilesIn(tmp))
+                .as("local resources are still released even when the abort fails")
+                .isZero();
+    }
+
+    /** An abnormal {@code close()} aborts the upload and releases local state. */
+    @Test
+    void closeAbortsMultipartUploadOnAbnormalClose() throws Exception {
+        InMemoryNativeS3Operations s3 = new InMemoryNativeS3Operations();
+
+        String uploadId = s3.startMultiPartUpload(KEY);
+        NativeS3RecoverableFsDataOutputStream stream = newStream(s3, uploadId);
+        stream.write(bytes('A', 5), 0, 5);
+
+        stream.close();
+
+        assertThat(s3.abortAttempts).isEqualTo(1);
+        assertThat(s3.openMultipartUploads).doesNotContainKey(uploadId);
+        assertThat(countLocalFilesIn(tmp)).isZero();
+    }
+
+    @Test
+    void closeForCommitDoesNotAbortOnSuccess() throws Exception {
+        InMemoryNativeS3Operations s3 = new InMemoryNativeS3Operations();
+
+        String uploadId = s3.startMultiPartUpload(KEY);
+        NativeS3RecoverableFsDataOutputStream stream = newStream(s3, uploadId);
+        stream.write(bytes('A', 5), 0, 5);
+
+        RecoverableFsDataOutputStream.Committer committer = stream.closeForCommit();
+
+        assertThat(s3.abortAttempts).as("a successful commit must not abort the upload").isZero();
+        assertThat(s3.openMultipartUploads)
+                .as("the upload stays open until the committer commits it")
+                .containsKey(uploadId);
+
+        committer.commit();
+
+        assertThat(s3.committedObjects.get(KEY)).containsExactly(bytes('A', 5));
+        assertThat(s3.openMultipartUploads).doesNotContainKey(uploadId);
+    }
+
+    @Test
+    void closeAfterSuccessfulCloseForCommitIsNoOp() throws Exception {
+        InMemoryNativeS3Operations s3 = new InMemoryNativeS3Operations();
+
+        String uploadId = s3.startMultiPartUpload(KEY);
+        NativeS3RecoverableFsDataOutputStream stream = newStream(s3, uploadId);
+        stream.write(bytes('A', 5), 0, 5);
+
+        RecoverableFsDataOutputStream.Committer committer = stream.closeForCommit();
+        stream.close();
+
+        assertThat(s3.abortAttempts)
+                .as("close() after a successful commit must not abort the pending upload")
+                .isZero();
+        assertThat(s3.openMultipartUploads).containsKey(uploadId);
+
+        committer.commit();
+        assertThat(s3.committedObjects.get(KEY)).containsExactly(bytes('A', 5));
+    }
+
+    private NativeS3RecoverableFsDataOutputStream newStream(
+            InMemoryNativeS3Operations s3, String uploadId) throws IOException {
+        return new NativeS3RecoverableFsDataOutputStream(
+                s3, KEY, uploadId, tmp.toString(), MIN_PART_SIZE);
+    }
+
+    private static long countLocalFilesIn(Path dir) throws IOException {
+        if (!Files.isDirectory(dir)) {
+            return 0;
+        }
+        try (java.util.stream.Stream<Path> s = Files.list(dir)) {
+            return s.count();
+        }
+    }
+
+    private static byte[] bytes(char c, int n) {
+        byte[] b = new byte[n];
+        Arrays.fill(b, (byte) c);
+        return b;
+    }
+}