Fix a bug where the trailers of a streaming response was dropped by RetryingClient

core/src/main/java/com/linecorp/armeria/client/retry/RetryingClient.java

@@ -394,7 +394,7 @@ private void handleStreamingResponse(RetryConfig<HttpResponse> retryConfig, Clie
 
             derivedCtx.log().whenAvailable(RequestLogProperty.RESPONSE_HEADERS).thenRun(() -> {
                 if (retryConfig.needsContentInRule() && responseCause == null) {
-                    final HttpResponse response0 = HttpResponse.of(headers, splitResponse.body());
+                    final HttpResponse response0 = splitResponse.unsplit();
                     final HttpResponseDuplicator duplicator =
                             response0.toDuplicator(derivedCtx.eventLoop().withoutContext(),
                                                    derivedCtx.maxResponseLength());
@@ -426,7 +426,7 @@ private void handleStreamingResponse(RetryConfig<HttpResponse> retryConfig, Clie
                         splitResponse.body().abort(responseCause);
                         response0 = HttpResponse.ofFailure(responseCause);
                     } else {
-                        response0 = HttpResponse.of(headers, splitResponse.body());
+                        response0 = splitResponse.unsplit();
                     }
                     handleResponseWithoutContent(retryConfig, ctx, rootReqDuplicator, originalReq, returnedRes,
                                                  future, derivedCtx, response0, responseCause);

core/src/main/java/com/linecorp/armeria/common/HttpRequest.java

@@ -279,7 +279,6 @@ static HttpRequest of(RequestHeaders headers, Publisher<? extends HttpObject> pu
         if (publisher instanceof HttpRequest) {
             return ((HttpRequest) publisher).withHeaders(headers);
         } else if (publisher instanceof StreamMessage) {
-            //noinspection unchecked
             return new StreamMessageBasedHttpRequest(headers, (StreamMessage<? extends HttpObject>) publisher);
         } else {
             return new PublisherBasedHttpRequest(headers, publisher);
@@ -303,7 +302,24 @@ static HttpRequest of(RequestHeaders headers,
         if (trailers.isEmpty()) {
             return of(headers, publisher);
         }
-        return of(headers, new SurroundingPublisher<>(null, publisher, unused -> trailers));
+        return of(headers, SurroundingPublisher.of(null, publisher, trailers));
+    }
+
+    /**
+     * Creates a new instance from an existing {@link RequestHeaders}, {@link Publisher} and trailers.
+     *
+     * <p>Note that the {@link HttpData}s in the {@link Publisher} are not released when
+     * {@link Subscription#cancel()} or {@link #abort()} is called. You should add a hook in order to
+     * release the elements. See {@link PublisherBasedStreamMessage} for more information.
+     */
+    @UnstableApi
+    static HttpRequest of(RequestHeaders headers,
+                          Publisher<? extends HttpData> publisher,
+                          CompletableFuture<HttpHeaders> trailers) {
+        requireNonNull(headers, "headers");
+        requireNonNull(publisher, "publisher");
+        requireNonNull(trailers, "trailers");
+        return of(headers, SurroundingPublisher.of(null, publisher, trailers));
     }
 
     /**

core/src/main/java/com/linecorp/armeria/common/HttpResponse.java

@@ -504,7 +504,7 @@ static HttpResponse of(ResponseHeaders headers, Publisher<? extends HttpObject>
         requireNonNull(headers, "headers");
         requireNonNull(publisher, "publisher");
         return new StreamMessageBasedHttpResponse(
-                new SurroundingPublisher<>(headers, publisher, unused -> null));
+                SurroundingPublisher.of(headers, publisher, unused -> null));
     }
 
     /**
@@ -521,24 +521,49 @@ static HttpResponse of(ResponseHeaders headers,
         requireNonNull(headers, "headers");
         requireNonNull(publisher, "publisher");
         requireNonNull(trailers, "trailers");
-        return of(headers, publisher, ignored -> trailers);
+        return new StreamMessageBasedHttpResponse(
+                SurroundingPublisher.of(headers, publisher, trailers));
+    }
+
+    /**
+     * Creates a new HTTP response with the specified headers and trailers
+     * whose stream is produced from an existing {@link Publisher}.
+     *
+     * <p>Note that the {@link HttpData}s in the {@link Publisher} are not released when
+     * {@link Subscription#cancel()} or {@link #abort()} is called. You should add a hook in order to
+     * release the elements. See {@link PublisherBasedStreamMessage} for more information.
+     */
+    @UnstableApi
+    static HttpResponse of(ResponseHeaders headers,
+                           Publisher<? extends HttpData> publisher,
+                           CompletableFuture<HttpHeaders> trailers) {
+        requireNonNull(headers, "headers");
+        requireNonNull(publisher, "publisher");
+        requireNonNull(trailers, "trailers");
+        return new StreamMessageBasedHttpResponse(
+                SurroundingPublisher.of(headers, publisher, trailers));
     }
 
     /**
      * Creates a new HTTP response with the specified headers and trailers function
      * whose stream is produced from an existing {@link Publisher}.
      *
+     * <p>If the trailers function returns null when the cause is not null, the returned {@link HttpResponse}
+     * will be completed with the given exception. If the trailers function returns an non-null value, the cause
+     * will be ignored and the {@link HttpResponse} will end with the returned trailers.
+     *
      * <p>Note that the {@link HttpData}s in the {@link Publisher} are not released when
      * {@link Subscription#cancel()} or {@link #abort()} is called. You should add a hook in order to
      * release the elements. See {@link PublisherBasedStreamMessage} for more information.
      */
     static HttpResponse of(ResponseHeaders headers,
                            Publisher<? extends HttpData> publisher,
-                           Function<@Nullable Throwable, HttpHeaders> trailersFunction) {
+                           Function<@Nullable Throwable, @Nullable HttpHeaders> trailersFunction) {
         requireNonNull(headers, "headers");
         requireNonNull(publisher, "publisher");
         requireNonNull(trailersFunction, "trailersFunction");
-        return PublisherBasedHttpResponse.from(headers, publisher, trailersFunction);
+        return new StreamMessageBasedHttpResponse(
+                SurroundingPublisher.of(headers, publisher, trailersFunction));
     }
 
     /**

core/src/main/java/com/linecorp/armeria/common/PublisherBasedHttpResponse.java

@@ -27,14 +27,10 @@
 
 final class PublisherBasedHttpResponse extends PublisherBasedStreamMessage<HttpObject> implements HttpResponse {
 
-    static PublisherBasedHttpResponse from(ResponseHeaders headers, Publisher<? extends HttpObject> publisher) {
-        return new PublisherBasedHttpResponse(new SurroundingPublisher<>(headers, publisher, unused -> null));
-    }
-
     static PublisherBasedHttpResponse from(ResponseHeaders headers,
                                            Publisher<? extends HttpData> publisher,
                                            Function<@Nullable Throwable, HttpHeaders> trailersFunction) {
-        return new PublisherBasedHttpResponse(new SurroundingPublisher<>(headers, publisher, trailersFunction));
+        return new PublisherBasedHttpResponse(SurroundingPublisher.of(headers, publisher, trailersFunction));
     }
 
     PublisherBasedHttpResponse(Publisher<? extends HttpObject> publisher) {

core/src/main/java/com/linecorp/armeria/common/SplitHttpMessage.java

@@ -20,6 +20,7 @@
 
 import com.google.errorprone.annotations.CheckReturnValue;
 
+import com.linecorp.armeria.common.annotation.UnstableApi;
 import com.linecorp.armeria.common.stream.ByteStreamMessage;
 import com.linecorp.armeria.common.stream.StreamMessage;
 
@@ -41,4 +42,12 @@ public interface SplitHttpMessage {
      * returned {@link CompletableFuture} will be completed with an {@linkplain HttpHeaders#of() empty headers}.
      */
     CompletableFuture<HttpHeaders> trailers();
+
+    /**
+     * Combines the split {@link #body()}, {@link #trailers()} into a single {@link HttpMessage}.
+     *
+     * <p>Note that this method can only be used before subscribing to {@link #body()}.
+     */
+    @UnstableApi
+    HttpMessage unsplit();
 }

core/src/main/java/com/linecorp/armeria/common/SplitHttpRequest.java

@@ -16,6 +16,10 @@
 
 package com.linecorp.armeria.common;
 
+import java.util.concurrent.CompletableFuture;
+
+import com.linecorp.armeria.common.annotation.UnstableApi;
+
 /**
  * An {@link HttpRequest} which splits a stream of {@link HttpObject}s into HTTP headers and payloads.
  */
@@ -25,4 +29,24 @@ public interface SplitHttpRequest extends SplitHttpMessage {
      * Returns the {@link RequestHeaders}.
      */
     RequestHeaders headers();
+
+    /**
+     * Combines the split {@link #headers()}, {@link #body()}, {@link #trailers()} into
+     * a single {@link HttpRequest}.
+     *
+     * <p>Note that this method can only be used before subscribing to {@link #body()}.
+     */
+    @UnstableApi
+    @Override
+    default HttpRequest unsplit() {
+        final CompletableFuture<HttpHeaders> trailersFuture = trailers().thenApply(trailers -> {
+            if (trailers.isEmpty()) {
+                // An empty trailers means that the request does not have trailers.
+                return null;
+            } else {
+                return trailers;
+            }
+        });
+        return HttpRequest.of(headers(), body(), trailersFuture);
+    }
 }

core/src/main/java/com/linecorp/armeria/common/SplitHttpResponse.java

@@ -20,6 +20,8 @@
 
 import org.reactivestreams.Subscriber;
 
+import com.linecorp.armeria.common.annotation.UnstableApi;
+
 /**
  * An {@link HttpResponse} which splits a stream of {@link HttpObject}s into HTTP headers and payloads.
  * {@link #headers()} will be completed before publishing the first {@link HttpData}.
@@ -35,4 +37,32 @@ public interface SplitHttpResponse extends SplitHttpMessage {
      * Returns a {@link CompletableFuture} completed with a non-informational {@link ResponseHeaders}.
      */
     CompletableFuture<ResponseHeaders> headers();
+
+    /**
+     * Combines the split {@link #headers()}, {@link #body()} and {@link #trailers()} into a single
+     * {@link HttpResponse}.
+     *
+     * <p>Note that this method can only be used before subscribing to {@link #body()}.
+     */
+    @UnstableApi
+    @Override
+    default HttpResponse unsplit() {
+        // TODO(ikhoon): Provide an optimized implementation for unsplit().
+        final CompletableFuture<HttpResponse> future = headers().handle((headers, cause) -> {
+            if (cause != null) {
+                return HttpResponse.ofFailure(cause);
+            }
+
+            final CompletableFuture<HttpHeaders> trailersFuture = trailers().thenApply(trailers -> {
+                if (trailers.isEmpty()) {
+                    // An empty trailers means that the response does not have trailers.
+                    return null;
+                } else {
+                    return trailers;
+                }
+            });
+            return HttpResponse.of(headers, body(), trailersFuture);
+        });
+        return HttpResponse.of(future);
+    }
 }

core/src/main/java/com/linecorp/armeria/common/stream/StreamMessage.java

@@ -1192,7 +1192,7 @@ default InputStream toInputStream(Function<? super T, ? extends HttpData> httpDa
      */
     @UnstableApi
     default StreamMessage<T> endWith(Function<@Nullable Throwable, ? extends @Nullable T> finalizer) {
-        return new SurroundingPublisher<>(null, this, finalizer);
+        return SurroundingPublisher.of(null, this, finalizer);
     }
 
     /**

core/src/main/java/com/linecorp/armeria/internal/common/stream/SurroundingPublisher.java

@@ -39,6 +39,7 @@
 import com.linecorp.armeria.common.stream.StreamMessage;
 import com.linecorp.armeria.common.stream.SubscriptionOption;
 import com.linecorp.armeria.common.util.EventLoopCheckingFuture;
+import com.linecorp.armeria.common.util.UnmodifiableFuture;
 
 import io.netty.util.concurrent.EventExecutor;
 
@@ -50,19 +51,46 @@ public final class SurroundingPublisher<T> implements StreamMessage<T> {
     private static final AtomicIntegerFieldUpdater<SurroundingPublisher> subscribedUpdater =
             AtomicIntegerFieldUpdater.newUpdater(SurroundingPublisher.class, "subscribed");
 
+    public static <T> SurroundingPublisher<T> of(@Nullable T head, Publisher<? extends T> publisher, T tail) {
+        // The tail is ignored when the cause is not null.
+        return of(head, publisher, cause -> cause != null ? null : tail);
+    }
+
+    public static <T> SurroundingPublisher<T> of(
+            @Nullable T head, Publisher<? extends T> publisher,
+            Function<@Nullable Throwable, ? extends @Nullable T> finalizer) {
+        requireNonNull(finalizer, "finalizer");
+        final Function<@Nullable Throwable, @Nullable CompletableFuture<T>> asyncFinalizer = cause -> {
+            final T tail = finalizer.apply(cause);
+            if (tail == null) {
+                return null;
+            } else {
+                return UnmodifiableFuture.completedFuture(tail);
+            }
+        };
+        return new SurroundingPublisher<>(head, publisher, asyncFinalizer);
+    }
+
+    public static <T> SurroundingPublisher<T> of(@Nullable T head, Publisher<? extends T> publisher,
+                                                 CompletableFuture<? extends @Nullable T> tail) {
+        // The tail is ignored when the cause is not null.
+        return new SurroundingPublisher<>(head, publisher, cause -> cause != null ? null : tail);
+    }
+
     @Nullable
     private final T head;
     private final StreamMessage<T> publisher;
-    private final Function<@Nullable Throwable, ? extends @Nullable T> finalizer;
+    private final Function<@Nullable Throwable, ? extends @Nullable CompletableFuture<T>> finalizer;
 
     private volatile int subscribed;
     private final CompletableFuture<Void> completionFuture = new EventLoopCheckingFuture<>();
 
     @Nullable
     private volatile SurroundingSubscriber<T> surroundingSubscriber;
 
-    public SurroundingPublisher(@Nullable T head, Publisher<? extends T> publisher,
-                                Function<@Nullable Throwable, ? extends @Nullable T> finalizer) {
+    SurroundingPublisher(@Nullable T head, Publisher<? extends T> publisher,
+                         Function<@Nullable Throwable,
+                                 ? extends @Nullable CompletableFuture<? extends T>> finalizer) {
         requireNonNull(publisher, "publisher");
         requireNonNull(finalizer, "finalizer");
         this.head = head;
@@ -72,7 +100,8 @@ public SurroundingPublisher(@Nullable T head, Publisher<? extends T> publisher,
         } else {
             this.publisher = new PublisherBasedStreamMessage<>(publisher);
         }
-        this.finalizer = finalizer;
+        //noinspection unchecked
+        this.finalizer = (Function<Throwable, ? extends CompletableFuture<T>>) finalizer;
     }
 
     @Override
@@ -190,7 +219,7 @@ enum State {
         @Nullable
         private T head;
         private final StreamMessage<T> publisher;
-        private final Function<@Nullable Throwable, ? extends @Nullable T> finalizer;
+        private final Function<@Nullable Throwable, ? extends @Nullable CompletableFuture<T>> finalizer;
 
         private Subscriber<? super T> downstream;
         private final EventExecutor executor;
@@ -206,7 +235,7 @@ enum State {
         private final SubscriptionOption[] options;
 
         SurroundingSubscriber(@Nullable T head, StreamMessage<T> publisher,
-                              Function<@Nullable Throwable, ? extends @Nullable T> finalizer,
+                              Function<@Nullable Throwable, ? extends @Nullable CompletableFuture<T>> finalizer,
                               Subscriber<? super T> downstream, EventExecutor executor,
                               CompletableFuture<Void> completionFuture, SubscriptionOption... options) {
             requireNonNull(publisher, "publisher");
@@ -305,9 +334,9 @@ private void sendTail() {
         }
 
         private void finalize(@Nullable Throwable cause) {
-            final T tail;
+            final CompletableFuture<T> tailFuture;
             try {
-                tail = finalizer.apply(cause);
+                tailFuture = finalizer.apply(cause);
             } catch (Throwable ex) {
                 if (cause != null) {
                     logger.warn("Unexpected exception from finalizer:", ex);
@@ -318,11 +347,28 @@ private void finalize(@Nullable Throwable cause) {
                 return;
             }
 
-            if (tail == null) {
+            if (tailFuture == null) {
                 // Immediately close the stream if the finalizer returns null.
                 close0(cause);
             } else {
-                downstream.onNext(tail);
+                tailFuture.handle((tail, cause0) -> {
+                    if (executor.inEventLoop()) {
+                        handleTail(tail, cause0);
+                    } else {
+                        executor.execute(() -> handleTail(tail, cause0));
+                    }
+                    return null;
+                });
+            }
+        }
+
+        private void handleTail(@Nullable T tail, @Nullable Throwable cause) {
+            if (cause != null) {
+                close0(cause);
+            } else {
+                if (tail != null) {
+                    downstream.onNext(tail);
+                }
                 close0(null);
             }
         }