Use StoringSubscriber in AsyncBufferingSubscriber (#3636)

Author: zoewangg

services-custom/s3-transfer-manager/src/main/java/software/amazon/awssdk/transfer/s3/internal/AsyncBufferingSubscriber.java

@@ -15,9 +15,8 @@
 
 package software.amazon.awssdk.transfer.s3.internal;
 
-import java.util.Queue;
+import java.util.Optional;
 import java.util.concurrent.CompletableFuture;
-import java.util.concurrent.ConcurrentLinkedQueue;
 import java.util.concurrent.atomic.AtomicBoolean;
 import java.util.concurrent.atomic.AtomicInteger;
 import java.util.function.Function;
@@ -26,6 +25,8 @@
 import software.amazon.awssdk.annotations.SdkInternalApi;
 import software.amazon.awssdk.utils.Logger;
 import software.amazon.awssdk.utils.Validate;
+import software.amazon.awssdk.utils.async.DemandIgnoringSubscription;
+import software.amazon.awssdk.utils.async.StoringSubscriber;
 
 /**
  * An implementation of {@link Subscriber} that execute the provided function for every event and limits the number of concurrent
@@ -36,24 +37,24 @@
 @SdkInternalApi
 public class AsyncBufferingSubscriber<T> implements Subscriber<T> {
     private static final Logger log = Logger.loggerFor(AsyncBufferingSubscriber.class);
-    private static final Object COMPLETE_EVENT = new Object();
-    private final Queue<Object> buffer;
     private final CompletableFuture<?> returnFuture;
     private final Function<T, CompletableFuture<?>> consumer;
     private final int maxConcurrentExecutions;
     private final AtomicInteger numRequestsInFlight;
     private final AtomicBoolean isDelivering = new AtomicBoolean(false);
     private volatile boolean isStreamingDone;
-    private volatile Subscription subscription;
+    private Subscription subscription;
+
+    private final StoringSubscriber<T> storingSubscriber;
 
     public AsyncBufferingSubscriber(Function<T, CompletableFuture<?>> consumer,
                                     CompletableFuture<Void> returnFuture,
                                     int maxConcurrentExecutions) {
-        this.buffer = new ConcurrentLinkedQueue<>();
         this.returnFuture = returnFuture;
         this.consumer = consumer;
         this.maxConcurrentExecutions = maxConcurrentExecutions;
         this.numRequestsInFlight = new AtomicInteger(0);
+        this.storingSubscriber = new StoringSubscriber<>(Integer.MAX_VALUE);
     }
 
     @Override
@@ -64,68 +65,53 @@ public void onSubscribe(Subscription subscription) {
             subscription.cancel();
             return;
         }
+        storingSubscriber.onSubscribe(new DemandIgnoringSubscription(subscription));
         this.subscription = subscription;
         subscription.request(maxConcurrentExecutions);
     }
 
     @Override
     public void onNext(T item) {
-        if (item == null) {
-            subscription.cancel();
-            NullPointerException exception = new NullPointerException("Item must not be null");
-            returnFuture.completeExceptionally(exception);
-            throw exception;
-        }
-
-        try {
-            buffer.add(item);
-            flushBufferIfNeeded();
-        } catch (Exception e) {
-            isStreamingDone = true;
-            subscription.cancel();
-            returnFuture.completeExceptionally(e);
-        }
+        storingSubscriber.onNext(item);
+        flushBufferIfNeeded();
     }
 
     private void flushBufferIfNeeded() {
-        if (buffer.isEmpty()) {
-            if (isStreamingDone && numRequestsInFlight.get() == 0) {
-                returnFuture.complete(null);
-            } else {
-                subscription.request(1);
-            }
-            return;
-        }
-
         if (isDelivering.compareAndSet(false, true)) {
             try {
-                Object firstEvent = buffer.peek();
-                if (isCompleteEvent(firstEvent)) {
-                    Object event = buffer.poll();
-                    handleCompleteEvent(event);
-                    return;
-                }
-
-                while (!buffer.isEmpty() && numRequestsInFlight.get() < maxConcurrentExecutions) {
-                    Object item = buffer.poll();
-                    if (item == null) {
+                Optional<StoringSubscriber.Event<T>> next = storingSubscriber.peek();
+                while (numRequestsInFlight.get() < maxConcurrentExecutions) {
+                    if (!next.isPresent()) {
+                        subscription.request(1);
                         break;
                     }
 
-                    if (isCompleteEvent(item)) {
-                        handleCompleteEvent(item);
-                        return;
+                    switch (next.get().type()) {
+                        case ON_COMPLETE:
+                            handleCompleteEvent();
+                            break;
+                        case ON_ERROR:
+                            handleError(next.get().runtimeError());
+                            break;
+                        case ON_NEXT:
+                            handleOnNext(next.get().value());
+                            break;
+                        default:
+                            handleError(new IllegalStateException("Unknown stored type: " + next.get().type()));
+                            break;
                     }
 
-                    deliverItem((T) item);
+                    next = storingSubscriber.peek();
                 }
             } finally {
                 isDelivering.set(false);
             }
         }
     }
 
-    private void deliverItem(T item) {
+    private void handleOnNext(T item) {
+        storingSubscriber.poll();
+
         int numberOfRequestInFlight = numRequestsInFlight.incrementAndGet();
         log.debug(() -> "Delivering next item, numRequestInFlight=" + numberOfRequestInFlight);
 
@@ -139,26 +125,28 @@ private void deliverItem(T item) {
         });
     }
 
-    private void handleCompleteEvent(Object event) {
-        isStreamingDone = true;
+    private void handleCompleteEvent() {
         if (numRequestsInFlight.get() == 0) {
             returnFuture.complete(null);
+            storingSubscriber.poll();
         }
     }
 
     @Override
     public void onError(Throwable t) {
         handleError(t);
+        storingSubscriber.onError(t);
     }
 
     private void handleError(Throwable t) {
         returnFuture.completeExceptionally(t);
-        buffer.clear();
+        storingSubscriber.poll();
     }
 
     @Override
     public void onComplete() {
-        buffer.add(COMPLETE_EVENT);
+        isStreamingDone = true;
+        storingSubscriber.onComplete();
         flushBufferIfNeeded();
     }
 
@@ -168,8 +156,4 @@ public void onComplete() {
     public int numRequestsInFlight() {
         return numRequestsInFlight.get();
     }
-
-    private static boolean isCompleteEvent(Object event) {
-        return COMPLETE_EVENT.equals(event);
-    }
 }

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

@@ -222,23 +222,4 @@ public enum TransferResult {
          */
         SUCCESS
     }
-
-    private static final class DemandIgnoringSubscription implements Subscription {
-        private final Subscription delegate;
-
-        private DemandIgnoringSubscription(Subscription delegate) {
-            this.delegate = delegate;
-        }
-
-        @Override
-        public void request(long n) {
-            // Ignore demand requests from downstream, they want too much.
-            // We feed them the amount that we want.
-        }
-
-        @Override
-        public void cancel() {
-            delegate.cancel();
-        }
-    }
 }

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

@@ -0,0 +1,40 @@
+/*
+ * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License").
+ * You may not use this file except in compliance with the License.
+ * A copy of the License is located at
+ *
+ *  http://aws.amazon.com/apache2.0
+ *
+ * or in the "license" file accompanying this file. This file 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 software.amazon.awssdk.utils.async;
+
+import org.reactivestreams.Subscription;
+import software.amazon.awssdk.annotations.SdkProtectedApi;
+
+@SdkProtectedApi
+public final class DemandIgnoringSubscription implements Subscription {
+
+    private final Subscription delegate;
+
+    public DemandIgnoringSubscription(Subscription delegate) {
+        this.delegate = delegate;
+    }
+
+    @Override
+    public void request(long n) {
+        // Ignore demand requests from downstream, they want too much.
+        // We feed them the amount that we want.
+    }
+
+    @Override
+    public void cancel() {
+        delegate.cancel();
+    }
+}

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

@@ -28,7 +28,7 @@
 /**
  * An implementation of {@link Subscriber} that stores the events it receives for retrieval.
  *
- * <p>Events can be observed via {@link #peek()} and {@link #drop()}. The number of events stored is limited by the
+ * <p>Events can be observed via {@link #peek()} and {@link #poll()}. The number of events stored is limited by the
  * {@code maxElements} configured at construction.
  */
 @SdkProtectedApi