Delay publishing the initial event until fully wired (#865)

* Delay publishing the initial event until fully wired

There's a race condition that can happen when we add the initial
request using onNext, that call will call flush while the downstream
subscriber is not yet set. To avoid this we just enqueue the initial
request, and wait till the upstream requests more input.

* Make spotless happy

* Avoid exposing the enqueueItem method to keep th Rx contract

* Update core/src/main/java/software/amazon/smithy/java/core/serde/event/EventStreamFrameEncodingProcessor.java

Co-authored-by: Adwait Kumar Singh <[email protected]>

* Avoid flushing twice if onError is called

---------

Co-authored-by: Adwait Kumar Singh <[email protected]>

Author: sugmanue

aws/aws-event-streams/build.gradle.kts

@@ -9,5 +9,6 @@ extra["moduleName"] = "software.amazon.smithy.java.aws.events"
 
 dependencies {
     api(project(":core"))
+    implementation(project(":logging"))
     api("software.amazon.eventstream:eventstream:1.0.1")
 }

aws/aws-event-streams/src/main/java/software/amazon/smithy/java/aws/events/RpcEventStreamsUtil.java

@@ -29,10 +29,7 @@ public static Flow.Publisher<ByteBuffer> bodyForEventStreaming(
             SerializableStruct input
     ) {
         Flow.Publisher<SerializableStruct> eventStream = input.getMemberValue(streamingMember(input.schema()));
-        var publisher = EventStreamFrameEncodingProcessor.create(eventStream, eventStreamEncodingFactory);
-        // Queue the input as the initial-request.
-        publisher.onNext(input);
-        return publisher;
+        return EventStreamFrameEncodingProcessor.create(eventStream, eventStreamEncodingFactory, input);
     }
 
     public static <O extends SerializableStruct> CompletableFuture<O> deserializeResponse(
@@ -62,7 +59,7 @@ public void onError(Throwable throwable) {
 
             @Override
             public void onComplete() {
-                result.completeExceptionally(new RuntimeException("Unexpected vent stream completion"));
+                result.completeExceptionally(new RuntimeException("Unexpected event stream completion"));
             }
         });
 

core/build.gradle.kts

@@ -14,6 +14,7 @@ dependencies {
     api(project(":io"))
     api(project(":retries-api"))
     api(libs.smithy.model)
+    implementation(project(":logging"))
 }
 
 jmh {}

core/src/main/java/software/amazon/smithy/java/core/serde/BufferingFlatMapProcessor.java

@@ -12,6 +12,7 @@
 import java.util.concurrent.atomic.AtomicLong;
 import java.util.concurrent.atomic.AtomicReference;
 import java.util.stream.Stream;
+import software.amazon.smithy.java.logging.InternalLogger;
 
 /**
  * A processor abstraction that maps inputs of type I from an upstream publisher to 0-n items of type O
@@ -27,12 +28,12 @@
 public abstract class BufferingFlatMapProcessor<I, O> implements
         Flow.Processor<I, O>,
         Flow.Subscription {
+    private static final InternalLogger LOG = InternalLogger.getLogger(BufferingFlatMapProcessor.class);
     private static final Throwable COMPLETE_SENTINEL = new RuntimeException();
 
-    private final AtomicReference<Throwable> terminalEvent = new AtomicReference<>();
+    private final AtomicReference<Throwable> terminalEventHolder = new AtomicReference<>();
     private final AtomicLong pendingRequests = new AtomicLong();
     private final AtomicInteger pendingFlushes = new AtomicInteger();
-    private final Flow.Publisher<I> publisher;
     private final BlockingQueue<O> queue = new LinkedBlockingQueue<>();
 
     private volatile Flow.Subscription upstreamSubscription;
@@ -42,7 +43,6 @@ public abstract class BufferingFlatMapProcessor<I, O> implements
     public BufferingFlatMapProcessor(
             Flow.Publisher<I> publisher
     ) {
-        this.publisher = publisher;
         publisher.subscribe(this);
     }
 
@@ -67,28 +67,46 @@ public final void onNext(I item) {
         try {
             map(item).forEach(this::addToQueue);
         } catch (Exception e) {
+            LOG.warn("Malformed input", e);
             onError(new SerializationException("Malformed input", e));
             return;
         }
         flush();
     }
 
+    /**
+     * Used to add the initial message to the queue. This message won't be sent until
+     * a flush happens, either by a calling {@link #request(long)} or {@link #onNext(Object)}.
+     * <p>
+     * This method will re-throw any exception caught when calling {@link #map} without calling
+     * {@link #onError(Throwable)} since it's assumed that the processor is not yet fully setup
+     * when this method is called.
+     */
+    protected final void enqueueItem(I item) {
+        try {
+            map(item).forEach(this::addToQueue);
+        } catch (RuntimeException e) {
+            LOG.warn("Malformed input", e);
+            throw e;
+        }
+    }
+
     private void addToQueue(O item) {
         queue.add(item);
     }
 
     @Override
     public final void onError(Throwable t) {
         upstreamSubscription.cancel();
-        terminalEvent.compareAndSet(null, t);
+        terminalEventHolder.compareAndSet(null, t);
         if (upstreamSubscription != null && downstream != null) {
             flush();
         }
     }
 
     @Override
     public final void onComplete() {
-        terminalEvent.compareAndSet(null, COMPLETE_SENTINEL);
+        terminalEventHolder.compareAndSet(null, COMPLETE_SENTINEL);
         if (upstreamSubscription != null && downstream != null) {
             flush();
         }
@@ -107,6 +125,10 @@ public final void request(long n) {
 
     private void flush() {
         if (upstreamSubscription == null || downstream == null) {
+            LOG.warn("flush() requested before upstream and downstream fully wired, " +
+                    "upstreamSubscription is null: {}, downstream is null: {}",
+                    upstreamSubscription == null,
+                    downstream == null);
             onError(new IllegalStateException("flush() requested before upstream and downstream fully wired."));
             return;
         }
@@ -126,8 +148,8 @@ private void flush() {
             Flow.Subscriber<? super O> subscriber = downstream;
             long delivered = sendMessages(subscriber, pending);
             boolean empty = queue.isEmpty();
-            Throwable term = terminalEvent.get();
-            if (term != null && attemptTermination(subscriber, term, empty)) {
+            Throwable terminalEvent = terminalEventHolder.get();
+            if (terminalEvent != null && attemptTermination(subscriber, terminalEvent, empty)) {
                 terminated = true;
                 return;
             }
@@ -184,12 +206,12 @@ protected void handleError(Throwable error, Flow.Subscriber<? super O> subscribe
     /**
      * @return true if this decoder is in a terminal state
      */
-    private boolean attemptTermination(Flow.Subscriber<? super O> subscriber, Throwable term, boolean done) {
+    private boolean attemptTermination(Flow.Subscriber<? super O> subscriber, Throwable terminalEvent, boolean done) {
         if (done && subscriber != null) {
-            if (term == COMPLETE_SENTINEL) {
+            if (terminalEvent == COMPLETE_SENTINEL) {
                 subscriber.onComplete();
             } else {
-                handleError(term, subscriber);
+                handleError(terminalEvent, subscriber);
             }
             return true;
         }
@@ -238,7 +260,7 @@ private static long accumulate(long current, long n) {
         }
     }
 
-    private static long accumulate(AtomicLong l, long n) {
-        return l.accumulateAndGet(n, BufferingFlatMapProcessor::accumulate);
+    private static void accumulate(AtomicLong l, long n) {
+        l.accumulateAndGet(n, BufferingFlatMapProcessor::accumulate);
     }
 }

core/src/main/java/software/amazon/smithy/java/core/serde/event/EventStreamFrameEncodingProcessor.java

@@ -10,9 +10,11 @@
 import java.util.stream.Stream;
 import software.amazon.smithy.java.core.schema.SerializableStruct;
 import software.amazon.smithy.java.core.serde.BufferingFlatMapProcessor;
+import software.amazon.smithy.java.logging.InternalLogger;
 
 public final class EventStreamFrameEncodingProcessor<F extends Frame<?>, T extends SerializableStruct>
         extends BufferingFlatMapProcessor<T, ByteBuffer> {
+    private static final InternalLogger LOG = InternalLogger.getLogger(EventStreamFrameEncodingProcessor.class);
     private final EventEncoder<F> eventEncoder;
     private final FrameEncoder<F> encoder;
 
@@ -30,10 +32,24 @@ public static <F extends Frame<?>> EventStreamFrameEncodingProcessor<F, Serializ
             Flow.Publisher<SerializableStruct> publisher,
             EventEncoderFactory<F> encoderFactory
     ) {
-        return new EventStreamFrameEncodingProcessor<>(
+        var processor = new EventStreamFrameEncodingProcessor<>(
                 publisher,
                 encoderFactory.newEventEncoder(),
                 encoderFactory.newFrameEncoder());
+        return processor;
+    }
+
+    public static <F extends Frame<?>> EventStreamFrameEncodingProcessor<F, SerializableStruct> create(
+            Flow.Publisher<SerializableStruct> publisher,
+            EventEncoderFactory<F> encoderFactory,
+            SerializableStruct firstItem
+    ) {
+        var processor = new EventStreamFrameEncodingProcessor<>(
+                publisher,
+                encoderFactory.newEventEncoder(),
+                encoderFactory.newFrameEncoder());
+        processor.enqueueItem(firstItem);
+        return processor;
     }
 
     @Override
@@ -44,6 +60,7 @@ protected Stream<ByteBuffer> map(T item) {
     @Override
     protected void handleError(Throwable error, Flow.Subscriber<? super ByteBuffer> subscriber) {
         subscriber.onNext(encoder.encode(eventEncoder.encodeFailure(error)));
+        LOG.warn("Unexpected error", error);
         subscriber.onComplete();
     }
 }

http/http-binding/src/main/java/software/amazon/smithy/java/http/binding/ResponseSerializer.java

@@ -121,6 +121,7 @@ public ResponseSerializer errorSchema(Schema errorSchema) {
      *
      * @return Returns the created response.
      */
+    @SuppressWarnings("unchecked")
     public HttpResponse serializeResponse() {
         Objects.requireNonNull(shapeValue, "shapeValue is not set");
         Objects.requireNonNull(operation, "operation is not set");
@@ -151,8 +152,7 @@ public HttpResponse serializeResponse() {
 
         var eventStream = (Flow.Publisher<SerializableStruct>) serializer.getEventStream();
         if (eventStream != null && operation instanceof OutputEventStreamingApiOperation<?, ?, ?>) {
-            builder.body(
-                    EventStreamFrameEncodingProcessor.create(eventStream, eventEncoderFactory));
+            builder.body(EventStreamFrameEncodingProcessor.create(eventStream, eventEncoderFactory));
             serializer.setContentType(eventEncoderFactory.contentType());
         } else if (serializer.hasBody()) {
             builder.body(serializer.getBody());