BlockingStreamingHttpService: drop trailers if users didn't create any

servicetalk-http-api/src/main/java/io/servicetalk/http/api/BlockingStreamingToStreamingService.java

@@ -21,12 +21,15 @@
 import io.servicetalk.concurrent.SingleSource.Subscriber;
 import io.servicetalk.concurrent.api.Completable;
 import io.servicetalk.concurrent.api.Publisher;
+import io.servicetalk.concurrent.api.ScanMapper;
 import io.servicetalk.concurrent.api.Single;
 import io.servicetalk.concurrent.api.internal.ConnectablePayloadWriter;
 import io.servicetalk.concurrent.internal.ThreadInterruptingCancellable;
 
 import java.io.IOException;
 import java.util.function.Consumer;
+import java.util.function.Supplier;
+import javax.annotation.Nullable;
 
 import static io.servicetalk.concurrent.api.Processors.newCompletableProcessor;
 import static io.servicetalk.concurrent.api.SourceAdapters.fromSource;
@@ -76,7 +79,7 @@ protected void handleSubscribe(final Subscriber<? super StreamingHttpResponse> s
                 // still be propagated.
                 final Processor exceptionProcessor = newCompletableProcessor();
                 final BufferHttpPayloadWriter payloadWriter = new BufferHttpPayloadWriter(
-                        ctx.headersFactory().newTrailers());
+                        () -> ctx.headersFactory().newTrailers());
                 DefaultBlockingStreamingHttpServerResponse response = null;
                 try {
                     final Consumer<DefaultHttpResponseMetaData> sendMeta = (metaData) -> {
@@ -102,7 +105,7 @@ protected void handleSubscribe(final Subscriber<? super StreamingHttpResponse> s
                             Publisher<Object> messageBody = fromSource(exceptionProcessor)
                                     .merge(payloadWriter.connect());
                             if (addTrailers) {
-                                messageBody = messageBody.concat(succeeded(payloadWriter.trailers()));
+                                messageBody = messageBody.scanWithMapper(() -> new TrailersMapper(payloadWriter));
                             }
                             messageBody = messageBody.beforeSubscription(() -> new Subscription() {
                                 @Override
@@ -169,10 +172,12 @@ public Completable closeAsyncGracefully() {
 
     private static final class BufferHttpPayloadWriter implements HttpPayloadWriter<Buffer> {
         private final ConnectablePayloadWriter<Buffer> payloadWriter = new ConnectablePayloadWriter<>();
-        private final HttpHeaders trailers;
+        private final Supplier<HttpHeaders> trailersFactory;
+        @Nullable
+        private HttpHeaders trailers;
 
-        BufferHttpPayloadWriter(final HttpHeaders trailers) {
-            this.trailers = trailers;
+        BufferHttpPayloadWriter(final Supplier<HttpHeaders> trailersFactory) {
+            this.trailersFactory = trailersFactory;
         }
 
         @Override
@@ -197,11 +202,61 @@ public void close(final Throwable cause) throws IOException {
 
         @Override
         public HttpHeaders trailers() {
+            if (trailers == null) {
+                trailers = trailersFactory.get();
+            }
+            return trailers;
+        }
+
+        @Nullable
+        HttpHeaders trailers0() {
             return trailers;
         }
 
         Publisher<Buffer> connect() {
             return payloadWriter.connect();
         }
     }
+
+    private static final class TrailersMapper implements ScanMapper<Object, Object>, ScanMapper.MappedTerminal<Object> {
+        private final BufferHttpPayloadWriter payloadWriter;
+
+        TrailersMapper(final BufferHttpPayloadWriter payloadWriter) {
+            this.payloadWriter = payloadWriter;
+        }
+
+        @Nullable
+        @Override
+        public Object mapOnNext(@Nullable final Object next) {
+            return next;
+        }
+
+        @Nullable
+        @Override
+        public MappedTerminal<Object> mapOnError(final Throwable cause) {
+            return null;
+        }
+
+        @Override
+        public MappedTerminal<Object> mapOnComplete() {
+            return this;
+        }
+
+        @Nullable
+        @Override
+        public Object onNext() {
+            return payloadWriter.trailers0();
+        }
+
+        @Override
+        public boolean onNextValid() {
+            return payloadWriter.trailers0() != null;
+        }
+
+        @Nullable
+        @Override
+        public Throwable terminal() {
+            return null;
+        }
+    }
 }

servicetalk-http-api/src/main/java/io/servicetalk/http/api/DefaultHttpRequest.java

@@ -245,7 +245,7 @@ public <T> HttpRequest payloadBody(final T pojo, final HttpSerializer2<T> serial
     public HttpHeaders trailers() {
         if (trailers == null) {
             trailers = original.payloadHolder().headersFactory().newTrailers();
-            original.transform(this);
+            original.transform(this);   // Invoke "transform" to set PayloadInfo.mayHaveTrailers() flag
         }
         return trailers;
     }
@@ -277,6 +277,9 @@ public StreamingHttpRequest toStreamingRequest() {
         @Nullable
         final Publisher<Object> payload;
         if (trailers != null) {
+            // We can not drop empty Trailers here bcz users could do type conversion intermediately, while still
+            // referencing the original HttpHeaders object from an aggregated type and keep using it to add trailers
+            // before sending the message or converting it back to an aggregated one.
             payload = emptyPayloadBody ? from(trailers) : from(payloadBody, trailers);
         } else {
             payload = emptyPayloadBody ? null : from(payloadBody);

servicetalk-http-api/src/main/java/io/servicetalk/http/api/DefaultHttpResponse.java

@@ -107,6 +107,9 @@ public StreamingHttpResponse toStreamingResponse() {
         @Nullable
         final Publisher<Object> payload;
         if (trailers != null) {
+            // We can not drop empty Trailers here bcz users could do type conversion intermediately, while still
+            // referencing the original HttpHeaders object from an aggregated type and keep using it to add trailers
+            // before sending the message or converting it back to an aggregated one.
             payload = emptyPayloadBody ? from(trailers) : from(payloadBody, trailers);
         } else {
             payload = emptyPayloadBody ? null : from(payloadBody);
@@ -127,7 +130,7 @@ public BlockingStreamingHttpResponse toBlockingStreamingResponse() {
     public HttpHeaders trailers() {
         if (trailers == null) {
             trailers = original.payloadHolder().headersFactory().newTrailers();
-            original.transform(this);
+            original.transform(this);   // Invoke "transform" to set PayloadInfo.mayHaveTrailers() flag
         }
         return trailers;
     }

servicetalk-http-api/src/main/java/io/servicetalk/http/api/HttpDataSourceTransformations.java

@@ -341,6 +341,9 @@ static Single<PayloadAndTrailers> aggregatePayloadAndTrailers(final DefaultPaylo
             if (isAlwaysEmpty(pair.payload)) {
                 payloadInfo.setEmpty(true);
             }
+            // We can not drop empty Trailers here bcz users could do type conversion intermediately multiple times,
+            // while still referencing the original HttpHeaders object from an aggregated type and keep using it to add
+            // trailers before sending the message or converting it back to an aggregated one.
             if (pair.trailers == null) {
                 payloadInfo.setMayHaveTrailersAndGenericTypeBuffer(false);
             }

servicetalk-http-api/src/main/java/io/servicetalk/http/api/HttpPayloadWriter.java

@@ -27,24 +27,42 @@
  */
 public interface HttpPayloadWriter<T> extends PayloadWriter<T>, TrailersHolder {
 
+    /**
+     * <b>Note:</b> modifying trailers after the payload writer is {@link #close() closed} is not allowed.
+     */
+    @Override
+    HttpHeaders trailers();
+
+    /**
+     * <b>Note:</b> modifying trailers after the payload writer is {@link #close() closed} is not allowed.
+     */
     @Override
     default HttpPayloadWriter<T> addTrailer(final CharSequence name, final CharSequence value) {
         TrailersHolder.super.addTrailer(name, value);
         return this;
     }
 
+    /**
+     * <b>Note:</b> modifying trailers after the payload writer is {@link #close() closed} is not allowed.
+     */
     @Override
     default HttpPayloadWriter<T> addTrailers(final HttpHeaders trailers) {
         TrailersHolder.super.addTrailers(trailers);
         return this;
     }
 
+    /**
+     * <b>Note:</b> modifying trailers after the payload writer is {@link #close() closed} is not allowed.
+     */
     @Override
     default HttpPayloadWriter<T> setTrailer(final CharSequence name, final CharSequence value) {
         TrailersHolder.super.setTrailer(name, value);
         return this;
     }
 
+    /**
+     * <b>Note:</b> modifying trailers after the payload writer is {@link #close() closed} is not allowed.
+     */
     @Override
     default HttpPayloadWriter<T> setTrailers(final HttpHeaders trailers) {
         TrailersHolder.super.setTrailers(trailers);

servicetalk-http-api/src/test/java/io/servicetalk/http/api/BlockingStreamingToStreamingServiceTest.java

@@ -16,19 +16,21 @@
 package io.servicetalk.http.api;
 
 import io.servicetalk.buffer.api.Buffer;
+import io.servicetalk.buffer.api.BufferAllocator;
 import io.servicetalk.concurrent.Cancellable;
 import io.servicetalk.concurrent.PublisherSource.Subscriber;
 import io.servicetalk.concurrent.PublisherSource.Subscription;
 import io.servicetalk.concurrent.SingleSource;
 import io.servicetalk.concurrent.api.Executor;
 import io.servicetalk.concurrent.api.ExecutorExtension;
 import io.servicetalk.concurrent.api.Publisher;
-import io.servicetalk.oio.api.PayloadWriter;
 
 import org.junit.jupiter.api.BeforeEach;
 import org.junit.jupiter.api.Test;
 import org.junit.jupiter.api.extension.ExtendWith;
 import org.junit.jupiter.api.extension.RegisterExtension;
+import org.junit.jupiter.params.ParameterizedTest;
+import org.junit.jupiter.params.provider.ValueSource;
 import org.mockito.Mock;
 import org.mockito.junit.jupiter.MockitoExtension;
 
@@ -63,8 +65,10 @@
 import static java.nio.charset.StandardCharsets.UTF_8;
 import static java.util.function.Function.identity;
 import static org.hamcrest.MatcherAssert.assertThat;
+import static org.hamcrest.Matchers.containsInAnyOrder;
 import static org.hamcrest.Matchers.instanceOf;
 import static org.hamcrest.Matchers.is;
+import static org.hamcrest.Matchers.not;
 import static org.hamcrest.Matchers.notNullValue;
 import static org.junit.jupiter.api.Assertions.assertThrows;
 import static org.junit.jupiter.api.Assertions.fail;
@@ -78,7 +82,8 @@ class BlockingStreamingToStreamingServiceTest {
     private static final String HELLO_WORLD = "Hello\nWorld\n";
 
     @RegisterExtension
-    final ExecutorExtension<Executor> executorExtension = ExecutorExtension.withCachedExecutor();
+    static final ExecutorExtension<Executor> executorExtension = ExecutorExtension.withCachedExecutor()
+            .setClassLevel(true);
 
     @Mock
     private HttpExecutionContext mockExecutionCtx;
@@ -93,14 +98,26 @@ void setup() {
         mockCtx = new TestHttpServiceContext(DefaultHttpHeadersFactory.INSTANCE, reqRespFactory, mockExecutionCtx);
     }
 
-    @Test
-    void defaultResponseStatusNoPayload() throws Exception {
-        BlockingStreamingHttpService syncService = (ctx, request, response) -> response.sendMetaData().close();
+    @ParameterizedTest(name = "{displayName} [{index}] withEmptyTrailers={0}")
+    @ValueSource(booleans = {false, true})
+    void defaultResponseStatusNoPayload(boolean withEmptyTrailers) throws Exception {
+        BlockingStreamingHttpService syncService = (ctx, request, response) -> {
+            HttpPayloadWriter<Buffer> writer = response.sendMetaData();
+            if (withEmptyTrailers) {
+                writer.trailers();  // accessing trailers before close should preserve trailers in message body
+            }
+            writer.close();
+            writer.trailers();  // accessing trailers after close should not modify output
+        };
 
         List<Object> response = invokeService(syncService, reqRespFactory.get("/"));
         assertMetaData(OK, response);
         assertPayloadBody("", response, false);
-        assertEmptyTrailers(response);
+        if (withEmptyTrailers) {
+            assertEmptyTrailers(response);
+        } else {
+            assertNoTrailers(response);
+        }
     }
 
     @Test
@@ -111,7 +128,7 @@ void customResponseStatusNoPayload() throws Exception {
         List<Object> response = invokeService(syncService, reqRespFactory.get("/"));
         assertMetaData(NO_CONTENT, response);
         assertPayloadBody("", response, false);
-        assertEmptyTrailers(response);
+        assertNoTrailers(response);
     }
 
     @Test
@@ -127,24 +144,34 @@ void receivePayloadBody() throws Exception {
                 .payloadBody(from("Hello\n", "World\n"), appSerializerUtf8FixLen()));
         assertMetaData(OK, response);
         assertPayloadBody("", response, true);
-        assertEmptyTrailers(response);
+        assertNoTrailers(response);
 
         assertThat(receivedPayload.toString(), is(HELLO_WORLD));
     }
 
-    @Test
-    void respondWithPayloadBody() throws Exception {
+    @ParameterizedTest(name = "{displayName} [{index}] withEmptyTrailers={0}")
+    @ValueSource(booleans = {false, true})
+    void respondWithPayloadBody(boolean withEmptyTrailers) throws Exception {
         BlockingStreamingHttpService syncService = (ctx, request, response) -> {
-            try (PayloadWriter<Buffer> pw = response.sendMetaData()) {
-                pw.write(ctx.executionContext().bufferAllocator().fromAscii("Hello\n"));
-                pw.write(ctx.executionContext().bufferAllocator().fromAscii("World\n"));
+            BufferAllocator alloc = ctx.executionContext().bufferAllocator();
+            HttpPayloadWriter<Buffer> writer = response.sendMetaData();
+            writer.write(alloc.fromAscii("Hello\n"));
+            if (withEmptyTrailers) {
+                writer.trailers();  // accessing trailers before close should preserve trailers in message body
             }
+            writer.write(alloc.fromAscii("World\n"));
+            writer.close();
+            writer.trailers();  // accessing trailers after close should not modify output
         };
 
         List<Object> response = invokeService(syncService, reqRespFactory.get("/"));
         assertMetaData(OK, response);
         assertPayloadBody(HELLO_WORLD, response, false);
-        assertEmptyTrailers(response);
+        if (withEmptyTrailers) {
+            assertEmptyTrailers(response);
+        } else {
+            assertNoTrailers(response);
+        }
     }
 
     @Test
@@ -533,6 +560,10 @@ private static void assertPayloadBody(String expectedPayloadBody, List<Object> r
         assertThat(payloadBody, is(expectedPayloadBody));
     }
 
+    private static void assertNoTrailers(List<Object> response) {
+        assertThat(response, not(containsInAnyOrder(instanceOf(HttpHeaders.class))));
+    }
+
     private static void assertEmptyTrailers(List<Object> response) {
         HttpHeaders trailers = (HttpHeaders) response.get(response.size() - 1);
         assertThat(trailers, is(notNullValue()));