Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Original file line number Diff line number Diff line change
Expand Up @@ -18,6 +18,7 @@

package org.apache.flink.fs.s3native.writer;

import org.apache.flink.annotation.VisibleForTesting;
import org.apache.flink.core.fs.RecoverableFsDataOutputStream;
import org.apache.flink.core.fs.RecoverableWriter;
import org.apache.flink.fs.s3native.writer.NativeS3Recoverable.PartETag;
Expand All @@ -26,7 +27,6 @@
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import javax.annotation.Nullable;
import javax.annotation.concurrent.NotThreadSafe;

import java.io.BufferedOutputStream;
Expand Down Expand Up @@ -96,7 +96,7 @@ public NativeS3RecoverableFsDataOutputStream(
long minPartSize,
List<PartETag> existingParts,
long numBytesInParts,
@Nullable File incompleteTailFile)
File incompleteTailFile)
throws IOException {
this.s3AccessHelper = s3AccessHelper;
this.key = key;
Expand Down Expand Up @@ -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);
Expand All @@ -219,17 +219,30 @@ 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;
try {
tryAbortUploadAndReleaseResources();
} catch (IOException cleanup) {
e.addSuppressed(cleanup);
}
throw e;
}

closed = true;
return new NativeS3Committer(s3AccessHelper, recoverable);
Expand Down Expand Up @@ -272,41 +285,51 @@ 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);
}
if (cleanupException != null) {
throw cleanupException;
}
tryAbortUploadAndReleaseResources();
}
} finally {
unlock();
}
}

/** Aborts the multipart upload and releases local resources on the best effort basis. */
private void tryAbortUploadAndReleaseResources() throws IOException {
IOException collected = null;
if (currentOutputStream != null) {
try {
currentOutputStream.close();
} catch (IOException e) {
collected = ExceptionUtils.firstOrSuppressed(e, collected);
Comment thread
Samrat002 marked this conversation as resolved.
}
}
if (currentTempFile != null && currentTempFile.exists()) {
try {
deleteTempFile(currentTempFile);
} 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);
}
if (collected != null) {
throw collected;
}
}

@VisibleForTesting
protected void deleteTempFile(File file) throws IOException {
Files.delete(file.toPath());
}

private void lock() throws IOException {
try {
lock.lockInterruptibly();
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Comment thread
Samrat002 marked this conversation as resolved.

/** 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();
Expand All @@ -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);
Expand Down Expand Up @@ -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);
}

Expand Down
Original file line number Diff line number Diff line change
@@ -0,0 +1,192 @@
/*
* 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.BeforeEach;
import org.junit.jupiter.api.Test;
import org.junit.jupiter.api.io.TempDir;

import java.io.File;
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;

InMemoryNativeS3Operations s3;
String uploadId;
NativeS3RecoverableFsDataOutputStream stream;

@BeforeEach
void setUp() throws IOException {
s3 = new InMemoryNativeS3Operations();
uploadId = s3.startMultiPartUpload(KEY);
stream = newStream(s3, uploadId);
stream.write(bytes('A', 5), 0, 5); // < MIN_PART_SIZE, so it is uploaded during commit
}

@Test
void closeForCommitAbortsMultipartUploadWhenPartUploadFails() throws Exception {
s3.failUploadPart = true;

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 {
s3.failUploadPart = true;
s3.failAbortMultiPartUpload = true;

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 {
s3.failAbortMultiPartUpload = true;

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 {
stream.close();

assertThat(s3.abortAttempts).isEqualTo(1);
assertThat(s3.openMultipartUploads).doesNotContainKey(uploadId);
assertThat(countLocalFilesIn(tmp)).isZero();
}

@Test
void closeSurfacesTempFileDeletionFailure() throws Exception {
NativeS3RecoverableFsDataOutputStream failingStream = newFailingDeleteStream();
failingStream.write(bytes('A', 5), 0, 5);

assertThatThrownBy(failingStream::close)
.isInstanceOf(IOException.class)
.hasMessageContaining("injected temp-file delete failure");

assertThat(s3.abortAttempts)
.as("abort is still attempted despite delete failure")
.isEqualTo(1);
}

@Test
void closeForCommitDoesNotAbortOnSuccess() throws Exception {
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 {
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 ops, String uid) throws IOException {
return new NativeS3RecoverableFsDataOutputStream(
ops, KEY, uid, tmp.toString(), MIN_PART_SIZE);
}

private NativeS3RecoverableFsDataOutputStream newFailingDeleteStream() throws IOException {
String uid = s3.startMultiPartUpload(KEY);
return new NativeS3RecoverableFsDataOutputStream(
s3, KEY, uid, tmp.toString(), MIN_PART_SIZE) {
@Override
protected void deleteTempFile(File file) throws IOException {
throw new IOException("injected temp-file delete failure");
}
};
}

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;
}
}