Fix the race condition in the FileAsyncRequestBody which causes the request to hang

Author: zoewangg

core/sdk-core/src/main/java/software/amazon/awssdk/core/internal/async/FileAsyncRequestBody.java

@@ -31,6 +31,7 @@
 import software.amazon.awssdk.core.async.AsyncRequestBody;
 import software.amazon.awssdk.core.internal.util.Mimetype;
 import software.amazon.awssdk.core.internal.util.NoopSubscription;
+import software.amazon.awssdk.utils.Logger;
 import software.amazon.awssdk.utils.builder.SdkBuilder;
 
 /**
@@ -41,6 +42,7 @@
  */
 @SdkInternalApi
 public final class FileAsyncRequestBody implements AsyncRequestBody {
+    private static final Logger log = Logger.loggerFor(FileAsyncRequestBody.class);
 
     /**
      * Default size (in bytes) of ByteBuffer chunks read from the file and delivered to the subscriber.
@@ -169,10 +171,11 @@ private static final class FileSubscription implements Subscription {
         private final Subscriber<? super ByteBuffer> subscriber;
         private final int chunkSize;
 
-        private long position = 0;
-        private AtomicLong outstandingDemand = new AtomicLong(0);
-        private boolean writeInProgress = false;
+        private final AtomicLong position = new AtomicLong(0);
+        private long outstandingDemand = 0;
+        private boolean readInProgress = false;
         private volatile boolean done = false;
+        private final Object lock = new Object();
 
         private FileSubscription(AsynchronousFileChannel inputChannel, Subscriber<? super ByteBuffer> subscriber, int chunkSize) {
             this.inputChannel = inputChannel;
@@ -189,23 +192,24 @@ public void request(long n) {
             if (n < 1) {
                 IllegalArgumentException ex =
                     new IllegalArgumentException(subscriber + " violated the Reactive Streams rule 3.9 by requesting a "
-                            + "non-positive number of elements.");
+                                                 + "non-positive number of elements.");
                 signalOnError(ex);
             } else {
                 try {
-                    // As governed by rule 3.17, when demand overflows `Long.MAX_VALUE` we treat the signalled demand as
-                    // "effectively unbounded"
-                    outstandingDemand.getAndUpdate(initialDemand -> {
-                        if (Long.MAX_VALUE - initialDemand < n) {
-                            return Long.MAX_VALUE;
+                    // We need to synchronize here because of the race condition
+                    // where readData finishes reading at the same time request
+                    // demand comes in
+                    synchronized (lock) {
+                        // As governed by rule 3.17, when demand overflows `Long.MAX_VALUE` we treat the signalled demand as
+                        // "effectively unbounded"
+                        if (Long.MAX_VALUE -  outstandingDemand < n) {
+                            outstandingDemand = Long.MAX_VALUE;
                         } else {
-                            return initialDemand + n;
+                            outstandingDemand += n;
                         }
-                    });
 
-                    synchronized (this) {
-                        if (!writeInProgress) {
-                            writeInProgress = true;
+                        if (!readInProgress) {
+                            readInProgress = true;
                             readData();
                         }
                     }
@@ -227,32 +231,33 @@ public void cancel() {
 
         private void readData() {
             // It's possible to have another request for data come in after we've closed the file.
-            if (!inputChannel.isOpen()) {
+            if (!inputChannel.isOpen() || done) {
                 return;
             }
 
             ByteBuffer buffer = ByteBuffer.allocate(chunkSize);
-            inputChannel.read(buffer, position, buffer, new CompletionHandler<Integer, ByteBuffer>() {
+            inputChannel.read(buffer, position.get(), buffer, new CompletionHandler<Integer, ByteBuffer>() {
                 @Override
                 public void completed(Integer result, ByteBuffer attachment) {
+
                     if (result > 0) {
                         attachment.flip();
-                        position += attachment.remaining();
+                        position.addAndGet(attachment.remaining());
                         signalOnNext(attachment);
-                        // If we have more permits, queue up another read.
-                        if (outstandingDemand.decrementAndGet() > 0) {
-                            readData();
-                            return;
+
+                        synchronized (lock) {
+                            // If we have more permits, queue up another read.
+                            if (--outstandingDemand > 0) {
+                                readData();
+                            } else {
+                                readInProgress = false;
+                            }
                         }
                     } else {
                         // Reached the end of the file, notify the subscriber and cleanup
                         signalOnComplete();
                         closeFile();
                     }
-
-                    synchronized (FileSubscription.this) {
-                        writeInProgress = false;
-                    }
                 }
 
                 @Override
@@ -267,32 +272,32 @@ private void closeFile() {
             try {
                 inputChannel.close();
             } catch (IOException e) {
-                signalOnError(e);
+                log.warn(() -> "Failed to close the file", e);
             }
         }
 
-        private void signalOnNext(ByteBuffer bb) {
+        private void signalOnNext(ByteBuffer attachment) {
             synchronized (this) {
                 if (!done) {
-                    subscriber.onNext(bb);
+                    subscriber.onNext(attachment);
                 }
             }
         }
 
         private void signalOnComplete() {
             synchronized (this) {
                 if (!done) {
-                    subscriber.onComplete();
                     done = true;
+                    subscriber.onComplete();
                 }
             }
         }
 
         private void signalOnError(Throwable t) {
             synchronized (this) {
                 if (!done) {
-                    subscriber.onError(t);
                     done = true;
+                    subscriber.onError(t);
                 }
             }
         }
@@ -301,4 +306,4 @@ private void signalOnError(Throwable t) {
     private static AsynchronousFileChannel openInputChannel(Path path) throws IOException {
         return AsynchronousFileChannel.open(path, StandardOpenOption.READ);
     }
-}
+}

test/stability-tests/src/it/java/software/amazon/awssdk/stability/tests/s3/S3BaseStabilityTest.java

@@ -15,6 +15,8 @@
 
 package software.amazon.awssdk.stability.tests.s3;
 
+import static org.assertj.core.api.Assertions.assertThat;
+
 import java.io.File;
 import java.io.IOException;
 import java.nio.file.Path;
@@ -33,10 +35,13 @@
 import software.amazon.awssdk.services.s3.model.DeleteBucketRequest;
 import software.amazon.awssdk.services.s3.model.NoSuchBucketException;
 import software.amazon.awssdk.services.s3.model.NoSuchKeyException;
+import software.amazon.awssdk.stability.tests.exceptions.StabilityTestsRetryableException;
+import software.amazon.awssdk.stability.tests.utils.RetryableTest;
 import software.amazon.awssdk.stability.tests.utils.StabilityTestRunner;
 import software.amazon.awssdk.testutils.RandomTempFile;
 import software.amazon.awssdk.testutils.service.AwsTestBase;
 import software.amazon.awssdk.utils.Logger;
+import software.amazon.awssdk.utils.Md5Utils;
 
 public abstract class S3BaseStabilityTest extends AwsTestBase {
     private static final Logger log = Logger.loggerFor(S3BaseStabilityTest.class);
@@ -60,6 +65,19 @@ public S3BaseStabilityTest(S3AsyncClient testClient) {
         this.testClient = testClient;
     }
 
+    @RetryableTest(maxRetries = 3, retryableException = StabilityTestsRetryableException.class)
+    public void largeObject_put_get_usingFile() {
+        String md5Upload = uploadLargeObjectFromFile();
+        String md5Download = downloadLargeObjectToFile();
+        assertThat(md5Upload).isEqualTo(md5Download);
+    }
+
+    @RetryableTest(maxRetries = 3, retryableException = StabilityTestsRetryableException.class)
+    public void putObject_getObject_highConcurrency() {
+        putObject();
+        getObject();
+    }
+
     protected String computeKeyName(int i) {
         return "key_" + i;
     }
@@ -79,25 +97,35 @@ protected void doGetBucketAcl_lowTpsLongInterval() {
     }
 
 
-    protected void downloadLargeObjectToFile() {
+    protected String downloadLargeObjectToFile() {
         File randomTempFile = RandomTempFile.randomUncreatedFile();
         StabilityTestRunner.newRunner()
                 .testName("S3AsyncStabilityTest.downloadLargeObjectToFile")
                 .futures(testClient.getObject(b -> b.bucket(getTestBucketName()).key(LARGE_KEY_NAME),
                         AsyncResponseTransformer.toFile(randomTempFile)))
                 .run();
-        randomTempFile.delete();
+
+
+        try {
+            return Md5Utils.md5AsBase64(randomTempFile);
+        } catch (IOException e) {
+            throw new RuntimeException(e);
+        } finally {
+            randomTempFile.delete();
+        }
     }
 
-    protected void uploadLargeObjectFromFile() {
+    protected String uploadLargeObjectFromFile() {
         RandomTempFile file = null;
         try {
             file = new RandomTempFile((long) 2e+9);
+            String md5 = Md5Utils.md5AsBase64(file);
             StabilityTestRunner.newRunner()
                     .testName("S3AsyncStabilityTest.uploadLargeObjectFromFile")
                     .futures(testClient.putObject(b -> b.bucket(getTestBucketName()).key(LARGE_KEY_NAME),
                             AsyncRequestBody.fromFile(file)))
                     .run();
+            return md5;
         } catch (IOException e) {
             throw new RuntimeException("fail to create test file", e);
         } finally {

test/stability-tests/src/it/java/software/amazon/awssdk/stability/tests/s3/S3CrtAsyncClientStabilityTest.java

@@ -17,12 +17,10 @@
 
 import org.junit.jupiter.api.AfterAll;
 import org.junit.jupiter.api.BeforeAll;
-import software.amazon.awssdk.transfer.s3.internal.S3CrtAsyncClient;
 import software.amazon.awssdk.http.nio.netty.NettyNioAsyncHttpClient;
 import software.amazon.awssdk.regions.Region;
 import software.amazon.awssdk.services.s3.S3AsyncClient;
-import software.amazon.awssdk.stability.tests.exceptions.StabilityTestsRetryableException;
-import software.amazon.awssdk.stability.tests.utils.RetryableTest;
+import software.amazon.awssdk.transfer.s3.internal.S3CrtAsyncClient;
 
 /**
  * Stability tests for {@link S3CrtAsyncClient}
@@ -64,16 +62,4 @@ public static void cleanup() {
     protected String getTestBucketName() {
         return BUCKET_NAME;
     }
-
-    @RetryableTest(maxRetries = 3, retryableException = StabilityTestsRetryableException.class)
-    public void largeObject_put_get_usingFile() {
-        uploadLargeObjectFromFile();
-        downloadLargeObjectToFile();
-    }
-
-    @RetryableTest(maxRetries = 3, retryableException = StabilityTestsRetryableException.class)
-    public void putObject_getObject_highConcurrency() {
-        putObject();
-        getObject();
-    }
 }

test/stability-tests/src/it/java/software/amazon/awssdk/stability/tests/s3/S3NettyAsyncStabilityTest.java

@@ -45,18 +45,6 @@ public static void cleanup() {
     @Override
     protected String getTestBucketName() { return bucketName; }
 
-    @RetryableTest(maxRetries = 3, retryableException = StabilityTestsRetryableException.class)
-    public void putObject_getObject_highConcurrency() {
-        putObject();
-        getObject();
-    }
-
-    @RetryableTest(maxRetries = 3, retryableException = StabilityTestsRetryableException.class)
-    public void largeObject_put_get_usingFile() {
-        uploadLargeObjectFromFile();
-        downloadLargeObjectToFile();
-    }
-
     @RetryableTest(maxRetries = 3, retryableException = StabilityTestsRetryableException.class)
     public void getBucketAcl_lowTpsLongInterval_Netty() {
         doGetBucketAcl_lowTpsLongInterval();

test/stability-tests/src/it/java/software/amazon/awssdk/stability/tests/s3/S3WithCrtAsyncHttpClientStabilityTest.java

@@ -51,18 +51,6 @@ public static void cleanup() {
     @Override
     protected String getTestBucketName() { return bucketName; }
 
-    @RetryableTest(maxRetries = 3, retryableException = StabilityTestsRetryableException.class)
-    public void putObject_getObject_highConcurrency() {
-        putObject();
-        getObject();
-    }
-
-    @RetryableTest(maxRetries = 3, retryableException = StabilityTestsRetryableException.class)
-    public void largeObject_put_get_usingFile() {
-        uploadLargeObjectFromFile();
-        downloadLargeObjectToFile();
-    }
-
     @RetryableTest(maxRetries = 3, retryableException = StabilityTestsRetryableException.class)
     public void getBucketAcl_lowTpsLongInterval_Crt() {
         doGetBucketAcl_lowTpsLongInterval();