Fixed an issue where ByteBufferStoringSubscriber is requesting more than specified in minimumBytesBuffered (#5014)

Author: zoewangg

.changes/next-release/bugfix-AWSSDKforJavav2-b56c8a4.json

@@ -0,0 +1,6 @@
+{
+    "type": "bugfix",
+    "category": "AWS SDK for Java v2",
+    "contributor": "",
+    "description": "Fixed an issue in `ByteBufferStoringSubscriber` where it could buffer more data than configured, resulting in out of memory error. See [#4999](https://github.com/aws/aws-sdk-java-v2/issues/4999)."
+}

utils/src/main/java/software/amazon/awssdk/utils/async/ByteBufferStoringSubscriber.java

@@ -99,7 +99,9 @@ public TransferResult transferTo(ByteBuffer out) {
             next = storingSubscriber.peek();
         }
 
-        addBufferedDataAmount(-transferred);
+        if (transferred != 0) {
+            addBufferedDataAmount(-transferred);
+        }
 
         if (!next.isPresent()) {
             return TransferResult.SUCCESS;
@@ -178,8 +180,9 @@ public void onSubscribe(Subscription s) {
 
     @Override
     public void onNext(ByteBuffer byteBuffer) {
+        int remaining = byteBuffer.remaining();
         storingSubscriber.onNext(byteBuffer.duplicate());
-        addBufferedDataAmount(byteBuffer.remaining());
+        addBufferedDataAmount(remaining);
         phaser.arrive();
     }
 

utils/src/test/java/software/amazon/awssdk/utils/async/ByteBufferStoringSubscriberTest.java

@@ -18,20 +18,13 @@
 import static org.assertj.core.api.Assertions.assertThat;
 import static org.assertj.core.api.Assertions.assertThatCode;
 import static org.assertj.core.api.Assertions.assertThatThrownBy;
-import static org.mockito.ArgumentMatchers.any;
-import static org.mockito.ArgumentMatchers.anyInt;
 import static org.mockito.ArgumentMatchers.anyLong;
 import static org.mockito.Mockito.doAnswer;
-import static org.mockito.Mockito.mock;
 import static org.mockito.Mockito.times;
 import static org.mockito.Mockito.verify;
 import static org.mockito.Mockito.verifyNoMoreInteractions;
-import static org.mockito.Mockito.when;
 
 import java.nio.ByteBuffer;
-import java.util.ArrayList;
-import java.util.Collections;
-import java.util.List;
 import java.util.concurrent.CountDownLatch;
 import java.util.concurrent.ExecutionException;
 import java.util.concurrent.ExecutorService;
@@ -43,11 +36,19 @@
 import java.util.concurrent.atomic.AtomicReference;
 import org.junit.jupiter.api.Test;
 import org.junit.jupiter.api.Timeout;
+import org.junit.jupiter.api.extension.ExtendWith;
+import org.mockito.Mock;
+import org.mockito.junit.jupiter.MockitoExtension;
 import org.reactivestreams.Subscription;
 import software.amazon.awssdk.utils.ThreadFactoryBuilder;
 import software.amazon.awssdk.utils.async.ByteBufferStoringSubscriber.TransferResult;
 
+@ExtendWith(MockitoExtension.class)
 public class ByteBufferStoringSubscriberTest {
+
+    @Mock
+    private Subscription subscription;
+
     @Test
     public void constructorCalled_withNonPositiveSize_throwsException() {
         assertThatCode(() -> new ByteBufferStoringSubscriber(1)).doesNotThrowAnyException();
@@ -61,7 +62,6 @@ public void constructorCalled_withNonPositiveSize_throwsException() {
     @Test
     public void doesNotRequestMoreThanMaxBytes() {
         ByteBufferStoringSubscriber subscriber = new ByteBufferStoringSubscriber(3);
-        Subscription subscription = mock(Subscription.class);
 
         subscriber.onSubscribe(subscription);
         verify(subscription).request(1);
@@ -79,7 +79,6 @@ public void doesNotRequestMoreThanMaxBytes() {
     @Test
     public void canStoreMoreThanMaxBytesButWontAskForMoreUntilBelowMax() {
         ByteBufferStoringSubscriber subscriber = new ByteBufferStoringSubscriber(3);
-        Subscription subscription = mock(Subscription.class);
 
         subscriber.onSubscribe(subscription);
         verify(subscription).request(1);
@@ -104,7 +103,6 @@ public void blockingTransfer_waitsForFullOutputBuffer() throws InterruptedExcept
 
         try {
             ByteBufferStoringSubscriber subscriber = new ByteBufferStoringSubscriber(1);
-            Subscription subscription = mock(Subscription.class);
 
             AtomicInteger bufferNumber = new AtomicInteger(0);
             doAnswer(i -> {
@@ -146,17 +144,18 @@ public void blockingTransfer_stopsOnComplete() throws ExecutionException, Interr
         try {
             ByteBufferStoringSubscriber subscriber = new ByteBufferStoringSubscriber(Long.MAX_VALUE);
 
-            subscriber.onSubscribe(mock(Subscription.class));
+            subscriber.onSubscribe(subscription); // request 1
             Future<ByteBuffer> blockingRead = executor.submit(() -> {
                 ByteBuffer out = ByteBuffer.allocate(1024);
+                // request 2
                 TransferResult transferResult = subscriber.blockingTransferTo(out);
                 assertThat(transferResult).isEqualTo(TransferResult.END_OF_STREAM);
                 return out;
             });
 
             ByteBuffer input = fullByteBufferOfSize(1);
 
-            subscriber.onNext(input);
+            subscriber.onNext(input); // request 3
             subscriber.onComplete();
 
             ByteBuffer output = blockingRead.get();
@@ -165,6 +164,8 @@ public void blockingTransfer_stopsOnComplete() throws ExecutionException, Interr
         } finally {
             executor.shutdownNow();
         }
+
+        verify(subscription, times(3)).request(1);
     }
 
     @Test
@@ -175,7 +176,7 @@ public void blockingTransfer_stopsOnError() throws ExecutionException, Interrupt
         try {
             ByteBufferStoringSubscriber subscriber = new ByteBufferStoringSubscriber(Long.MAX_VALUE);
 
-            subscriber.onSubscribe(mock(Subscription.class));
+            subscriber.onSubscribe(subscription);
             Future<?> blockingRead = executor.submit(() -> {
                 subscriber.blockingTransferTo(ByteBuffer.allocate(1024)); // Expected to throw
                 return null;
@@ -202,7 +203,7 @@ public void blockingTransfer_stopsOnInterrupt() throws InterruptedException {
         AtomicBoolean threadIsInterruptedInCatch = new AtomicBoolean(false);
         AtomicReference<Throwable> failureReason = new AtomicReference<>();
 
-        subscriber.onSubscribe(mock(Subscription.class));
+        subscriber.onSubscribe(subscription);
 
         CountDownLatch threadIsRunning = new CountDownLatch(1);
         Thread thread = new Thread(() -> {
@@ -230,98 +231,105 @@ public void blockingTransfer_stopsOnInterrupt() throws InterruptedException {
     public void blockingTransfer_returnsEndOfStreamWithRepeatedCalls() {
         ByteBufferStoringSubscriber subscriber = new ByteBufferStoringSubscriber(Long.MAX_VALUE);
 
-        subscriber.onSubscribe(mock(Subscription.class));
+        subscriber.onSubscribe(subscription);
         subscriber.onComplete();
 
         ByteBuffer buffer = ByteBuffer.allocate(0);
 
         assertThat(subscriber.blockingTransferTo(buffer)).isEqualTo(TransferResult.END_OF_STREAM);
         assertThat(subscriber.blockingTransferTo(buffer)).isEqualTo(TransferResult.END_OF_STREAM);
         assertThat(subscriber.blockingTransferTo(buffer)).isEqualTo(TransferResult.END_OF_STREAM);
+        verify(subscription).request(1);
     }
 
 
     @Test
     public void noDataTransferredIfNoDataBuffered() {
         ByteBufferStoringSubscriber subscriber = new ByteBufferStoringSubscriber(2);
-        subscriber.onSubscribe(mock(Subscription.class));
+        subscriber.onSubscribe(subscription);
 
         ByteBuffer out = emptyByteBufferOfSize(1);
 
         assertThat(subscriber.transferTo(out)).isEqualTo(TransferResult.SUCCESS);
         assertThat(out.remaining()).isEqualTo(1);
+        verify(subscription).request(1);
     }
 
     @Test
     public void noDataTransferredIfComplete() {
         ByteBufferStoringSubscriber subscriber = new ByteBufferStoringSubscriber(2);
-        subscriber.onSubscribe(mock(Subscription.class));
+        subscriber.onSubscribe(subscription);
         subscriber.onComplete();
 
         ByteBuffer out = emptyByteBufferOfSize(1);
 
         assertThat(subscriber.transferTo(out)).isEqualTo(TransferResult.END_OF_STREAM);
         assertThat(out.remaining()).isEqualTo(1);
+        verify(subscription).request(1);
     }
 
     @Test
     public void noDataTransferredIfError() {
         RuntimeException error = new RuntimeException();
 
         ByteBufferStoringSubscriber subscriber = new ByteBufferStoringSubscriber(2);
-        subscriber.onSubscribe(mock(Subscription.class));
+        subscriber.onSubscribe(subscription);
         subscriber.onError(error);
 
         ByteBuffer out = emptyByteBufferOfSize(1);
 
         assertThatThrownBy(() -> subscriber.transferTo(out)).isEqualTo(error);
         assertThat(out.remaining()).isEqualTo(1);
+        verify(subscription).request(1);
     }
 
     @Test
     public void checkedExceptionsAreWrapped() {
         Exception error = new Exception();
 
         ByteBufferStoringSubscriber subscriber = new ByteBufferStoringSubscriber(2);
-        subscriber.onSubscribe(mock(Subscription.class));
+        subscriber.onSubscribe(subscription);
         subscriber.onError(error);
 
         ByteBuffer out = emptyByteBufferOfSize(1);
 
         assertThatThrownBy(() -> subscriber.transferTo(out)).hasCause(error);
         assertThat(out.remaining()).isEqualTo(1);
+        verify(subscription).request(1);
     }
 
     @Test
     public void completeIsReportedEvenWithExactOutSize() {
         ByteBufferStoringSubscriber subscriber = new ByteBufferStoringSubscriber(2);
-        subscriber.onSubscribe(mock(Subscription.class));
+        subscriber.onSubscribe(subscription);
         subscriber.onNext(fullByteBufferOfSize(2));
         subscriber.onComplete();
 
         ByteBuffer out = emptyByteBufferOfSize(2);
         assertThat(subscriber.transferTo(out)).isEqualTo(TransferResult.END_OF_STREAM);
         assertThat(out.remaining()).isEqualTo(0);
+        verify(subscription, times(2)).request(1);
     }
 
     @Test
     public void completeIsReportedEvenWithExtraOutSize() {
         ByteBufferStoringSubscriber subscriber = new ByteBufferStoringSubscriber(2);
-        subscriber.onSubscribe(mock(Subscription.class));
+        subscriber.onSubscribe(subscription);
         subscriber.onNext(fullByteBufferOfSize(2));
         subscriber.onComplete();
 
         ByteBuffer out = emptyByteBufferOfSize(3);
         assertThat(subscriber.transferTo(out)).isEqualTo(TransferResult.END_OF_STREAM);
         assertThat(out.remaining()).isEqualTo(1);
+        verify(subscription, times(2)).request(1);
     }
 
     @Test
     public void errorIsReportedEvenWithExactOutSize() {
         RuntimeException error = new RuntimeException();
 
         ByteBufferStoringSubscriber subscriber = new ByteBufferStoringSubscriber(2);
-        subscriber.onSubscribe(mock(Subscription.class));
+        subscriber.onSubscribe(subscription);
         subscriber.onNext(fullByteBufferOfSize(2));
         subscriber.onError(error);
 
@@ -335,7 +343,7 @@ public void errorIsReportedEvenWithExtraOutSize() {
         RuntimeException error = new RuntimeException();
 
         ByteBufferStoringSubscriber subscriber = new ByteBufferStoringSubscriber(2);
-        subscriber.onSubscribe(mock(Subscription.class));
+        subscriber.onSubscribe(subscription);
         subscriber.onNext(fullByteBufferOfSize(2));
         subscriber.onError(error);
 
@@ -351,7 +359,7 @@ public void dataIsDeliveredInTheRightOrder() {
         ByteBuffer buffer3 = fullByteBufferOfSize(1);
 
         ByteBufferStoringSubscriber subscriber = new ByteBufferStoringSubscriber(3);
-        subscriber.onSubscribe(mock(Subscription.class));
+        subscriber.onSubscribe(subscription);
         subscriber.onNext(buffer1);
         subscriber.onNext(buffer2);
         subscriber.onNext(buffer3);