Use a sizehint instead of relying on a new parameter

Author: alextwoods

services/s3/src/main/java/software/amazon/awssdk/services/s3/internal/crt/S3CrtRequestBodyStreamAdapter.java

@@ -27,17 +27,15 @@
  */
 @SdkInternalApi
 public final class S3CrtRequestBodyStreamAdapter implements HttpRequestBodyStream {
-    private static final long MINIMUM_BYTES_BUFFERED = 1024 * 1024L;
-    // for 16 kb chunks, this limits to about 16 MB (2x the standard crt provided buffer size)
-    private static final int MAXIMUM_OUTSTANDING_DEMAND = 1024;
+    private static final long MINIMUM_BYTES_BUFFERED = 16 * 1024 * 1024L;
     private final SdkHttpContentPublisher bodyPublisher;
     private final ByteBufferStoringSubscriber requestBodySubscriber;
 
     private final AtomicBoolean subscribed = new AtomicBoolean(false);
 
     public S3CrtRequestBodyStreamAdapter(SdkHttpContentPublisher bodyPublisher) {
         this.bodyPublisher = bodyPublisher;
-        this.requestBodySubscriber = new ByteBufferStoringSubscriber(MINIMUM_BYTES_BUFFERED, MAXIMUM_OUTSTANDING_DEMAND);
+        this.requestBodySubscriber = new ByteBufferStoringSubscriber(MINIMUM_BYTES_BUFFERED);
     }
 
     @Override

services/s3/src/test/java/software/amazon/awssdk/services/s3/internal/crt/S3CrtRequestBodyStreamAdapterTest.java

@@ -60,32 +60,25 @@ void getRequestData_fillsInputBuffer_publisherBuffersAreSmaller() {
 
     @Test
     void getRequestData_fillsInputBuffer_limitsOutstandingDemand() {
-        int inputBufferSize = 2000;
-        int maximumOutstandingDemand = 1024;
+        int minBytesBuffered  = 16 * 1024 * 1024;
+        int inputBufferSize = 1024;
 
         RequestTrackingPublisher requestTrackingPublisher = new RequestTrackingPublisher();
-        SdkHttpContentPublisher requestBody = requestBody(requestTrackingPublisher, 42L);
+        SdkHttpContentPublisher requestBody = requestBody(requestTrackingPublisher, minBytesBuffered);
 
         S3CrtRequestBodyStreamAdapter adapter = new S3CrtRequestBodyStreamAdapter(requestBody);
 
         ByteBuffer inputBuffer = ByteBuffer.allocate(inputBufferSize);
-        for (int i = 0; i < maximumOutstandingDemand; i++) {
-            // we are under the minimum buffer size, so each request here increases outstanding demand by 1
-            adapter.sendRequestBody(inputBuffer);
-            // release 1 byte of data, calling onNext (satisfies one request, but then requests 1 more)
-            requestTrackingPublisher.release(1);
-        }
-        // we should have 2x requests
-        assertThat(requestTrackingPublisher.requests()).isEqualTo(maximumOutstandingDemand * 2);
-        // but the total released bytes is only maximumOutstandingDemand
-        assertThat(inputBuffer.remaining()).isEqualTo(inputBufferSize - maximumOutstandingDemand + 1);
-
-        // now that we have reached maximum outstanding demand, new requests won't be sent
-        adapter.sendRequestBody(inputBuffer);
-        assertThat(requestTrackingPublisher.requests()).isEqualTo(maximumOutstandingDemand * 2);
-
+        adapter.sendRequestBody(inputBuffer); // initiate the subscription, but no bytes available, makes 1 request
 
+        // release 1 request of minBytesBuffered bytes of data, calling onNext (satisfies one request, but then requests 1 more)
+        requestTrackingPublisher.release(1, minBytesBuffered-100);
+        assertThat(requestTrackingPublisher.requests()).isEqualTo(2);
 
+        // call sendRequestBody, outstandingDemand=1, sizeHint=16*1024*1024-100 + existing data buffered is > our min
+        // so no more requests will be made
+        adapter.sendRequestBody(inputBuffer);
+        assertThat(requestTrackingPublisher.requests()).isEqualTo(2);
     }
 
     private static SdkHttpContentPublisher requestBody(Publisher<ByteBuffer> delegate, long size) {
@@ -159,9 +152,9 @@ public void subscribe(Subscriber<? super ByteBuffer> subscriber) {
         }
 
         // publish up to n requests
-        public void release(int n) {
+        public void release(int n, int size) {
             for (int i = 0; i < n; i++) {
-                ByteBuffer buffer = ByteBuffer.allocate(1);
+                ByteBuffer buffer = ByteBuffer.allocate(size);
                 subscriber.onNext(buffer);
             }
         }

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

@@ -59,32 +59,21 @@ public class ByteBufferStoringSubscriber implements Subscriber<ByteBuffer> {
 
     private final AtomicInteger outstandingDemand = new AtomicInteger(0);
 
-    private final Optional<Integer> maximumOutstandingDemand;
+    private volatile long byteBufferSizeHint = 0L;
 
     /**
      * The active subscription. Set when {@link #onSubscribe(Subscription)} is invoked.
      */
     private Subscription subscription;
 
     /**
-     * Create a subscriber that stores at least {@code minimumBytesBuffered} in memory for retrieval.
+     * Create a subscriber that stores at least {@code minimumBytesBuffered} in memory for retrieval. The subscriber will
+     * only request more from the subscription when fewer bytes are buffered AND in flight requests from the subscription will
+     * likely be under minimumBytesBuffered.
      */
     public ByteBufferStoringSubscriber(long minimumBytesBuffered) {
         this.minimumBytesBuffered = Validate.isPositive(minimumBytesBuffered, "Data buffer minimum must be positive");
         this.storingSubscriber = new StoringSubscriber<>(Integer.MAX_VALUE);
-        this.maximumOutstandingDemand = Optional.empty();
-    }
-
-    /**
-     * Create a subscriber that stores at least {@code minimumBytesBuffered} in memory for retrieval and which limits the
-     * maximum outstanding demand (requests) to its subscription.
-     */
-    public ByteBufferStoringSubscriber(long minimumBytesBuffered, int maximumOutstandingDemand) {
-        this.minimumBytesBuffered = Validate.isPositive(minimumBytesBuffered, "Data buffer minimum must be positive");
-        this.maximumOutstandingDemand = Optional.of(Validate.isPositive(maximumOutstandingDemand,
-                                                                    "maximumOutstandingDemand must be positive"));
-
-        this.storingSubscriber = new StoringSubscriber<>(Integer.MAX_VALUE);
     }
 
     /**
@@ -201,6 +190,10 @@ public void onSubscribe(Subscription s) {
     public void onNext(ByteBuffer byteBuffer) {
         int remaining = byteBuffer.remaining();
         outstandingDemand.decrementAndGet();
+        // atomic update not required here, in a race it does not matter which thread sets this value since it is not being
+        // incremented, just set.
+        byteBufferSizeHint = byteBuffer.remaining();
+
         storingSubscriber.onNext(byteBuffer.duplicate());
         addBufferedDataAmount(remaining);
         phaser.arrive();
@@ -224,12 +217,8 @@ private void addBufferedDataAmount(long amountToAdd) {
     }
 
     private void maybeRequestMore(long currentDataBuffered) {
-        // if we have too many outstanding requests, no need to make more requests
-        if (maximumOutstandingDemand.isPresent() && outstandingDemand.get() >= maximumOutstandingDemand.get()) {
-            return;
-        }
-
-        if (currentDataBuffered < minimumBytesBuffered) {
+        long dataBufferedAndInFlight = currentDataBuffered + (byteBufferSizeHint * outstandingDemand.get());
+        if (dataBufferedAndInFlight < minimumBytesBuffered) {
             outstandingDemand.incrementAndGet();
             subscription.request(1);
         }

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

@@ -53,13 +53,10 @@ public class ByteBufferStoringSubscriberTest {
     public void constructorCalled_withNonPositiveSize_throwsException() {
         assertThatCode(() -> new ByteBufferStoringSubscriber(1)).doesNotThrowAnyException();
         assertThatCode(() -> new ByteBufferStoringSubscriber(Integer.MAX_VALUE)).doesNotThrowAnyException();
-        assertThatCode(() -> new ByteBufferStoringSubscriber(1, 1)).doesNotThrowAnyException();
 
         assertThatThrownBy(() -> new ByteBufferStoringSubscriber(0)).isInstanceOf(IllegalArgumentException.class);
         assertThatThrownBy(() -> new ByteBufferStoringSubscriber(-1)).isInstanceOf(IllegalArgumentException.class);
         assertThatThrownBy(() -> new ByteBufferStoringSubscriber(Integer.MIN_VALUE)).isInstanceOf(IllegalArgumentException.class);
-        assertThatThrownBy(() -> new ByteBufferStoringSubscriber(1, 0)).isInstanceOf(IllegalArgumentException.class);
-        assertThatThrownBy(() -> new ByteBufferStoringSubscriber(1, -1)).isInstanceOf(IllegalArgumentException.class);
     }
 
     @Test
@@ -80,16 +77,17 @@ public void doesNotRequestMoreThanMaxBytes() {
     }
 
     @Test
-    public void doesNotRequestMoreThanMaxDemand() {
-        ByteBufferStoringSubscriber subscriber = new ByteBufferStoringSubscriber(5, 2);
+    public void doesNotRequestMoreWhenInflightMoreThanMinBytes() {
+        ByteBufferStoringSubscriber subscriber = new ByteBufferStoringSubscriber(5);
 
         subscriber.onSubscribe(subscription); // request 1, demand = 1
-        subscriber.onNext(fullByteBufferOfSize(3)); // demand = 0
+        subscriber.onNext(fullByteBufferOfSize(3)); // demand = 0, sizeHint=3
         subscriber.transferTo(emptyByteBufferOfSize(1)); // requests more, demand = 1
         subscriber.transferTo(emptyByteBufferOfSize(1)); // requests more, demand = 2
         verify(subscription, times(3)).request(1);
 
-        subscriber.transferTo(emptyByteBufferOfSize(1)); // demand already maximum, no request
+        //sizeHint=3, demand=2, dataBufferedAndInFlight=6. 6 > 5, so no new request
+        subscriber.transferTo(emptyByteBufferOfSize(1));
         verifyNoMoreInteractions(subscription);
     }