Racy cancel from transport does not always cancel the message body (#2369)

Motivation:

When `cancel` from the server's transport races with HTTP message completion, we can receive `cancel` after `Single<HttpResponseMetaData>` completes but before we subscribe to the `messageBody` publisher. In this case, we won't subscribe to the `messageBody` because we use `concat` operator. This does not allow `BeforeFinallyHttpOperator` to observe a terminal event. The race is possible when offloading is enabled. `cancel` comes from an IO thread, HTTP meta-data completes on the offloaded thread.

Modifications:

- For the server-side control flow, use the recently added `concatPropagateCancel` operator to guarantee propagation of the cancel for the entire flat HTTP stream;

Result:

`HttpLifecycleObserverTest.testClientCancelsRequestBeforeResponse()` does not race anymore.

Author: idelpivnitskiy

servicetalk-http-netty/src/main/java/io/servicetalk/http/netty/AbstractStreamingHttpConnection.java

@@ -166,12 +166,18 @@ public Single<StreamingHttpResponse> request(final StreamingHttpRequest request)
             if (canAddRequestContentLength(request)) {
                 flatRequest = setRequestContentLength(connectionContext().protocol(), request);
             } else {
-                if (emptyMessageBody(request, request.messageBody())) {
-                    flatRequest = flatEmptyMessage(connectionContext().protocol(), request, request.messageBody());
+                final Publisher<Object> messageBody = request.messageBody();
+                // Do not propagate cancel to the messageBody if cancel arrives before meta-data completes. Client-side
+                // state machine does not depend on termination of the messageBody until after transport subscribes to
+                // it. It's preferable to avoid subscribe to the messageBody in case of cancellation to allow requests
+                // with non-replayable messageBody to retry.
+                if (emptyMessageBody(request, messageBody)) {
+                    flatRequest = flatEmptyMessage(connectionContext().protocol(), request, messageBody,
+                            /* propagateCancel */ false);
                 } else {
                     // Defer subscribe to the messageBody until transport requests it to allow clients retry failed
                     // requests with non-replayable messageBody
-                    flatRequest = Single.<Object>succeeded(request).concat(request.messageBody(), true);
+                    flatRequest = Single.<Object>succeeded(request).concat(messageBody, /* deferSubscribe */ true);
                     if (shouldAppendTrailers(connectionContext().protocol(), request)) {
                         flatRequest = flatRequest.scanWith(HeaderUtils::appendTrailersMapper);
                     }

servicetalk-http-netty/src/main/java/io/servicetalk/http/netty/HeaderUtils.java

@@ -139,12 +139,13 @@ static Publisher<Object> setRequestContentLength(final HttpProtocolVersion proto
                                                      final StreamingHttpRequest request) {
         return setContentLength(request, request.messageBody(),
                 shouldAddZeroContentLength(request.method()) ? HeaderUtils::updateContentLength :
-                        HeaderUtils::updateRequestContentLengthNonZero, protocolVersion);
+                        HeaderUtils::updateRequestContentLengthNonZero, protocolVersion, /* propagateCancel */ false);
     }
 
     static Publisher<Object> setResponseContentLength(final HttpProtocolVersion protocolVersion,
                                                       final StreamingHttpResponse response) {
-        return setContentLength(response, response.messageBody(), HeaderUtils::updateContentLength, protocolVersion);
+        return setContentLength(response, response.messageBody(), HeaderUtils::updateContentLength, protocolVersion,
+                /* propagateCancel */ true);
     }
 
     private static void updateRequestContentLengthNonZero(final int contentLength, final HttpHeaders headers) {
@@ -210,16 +211,22 @@ static boolean emptyMessageBody(final HttpMetaData metadata) {
     }
 
     static Publisher<Object> flatEmptyMessage(final HttpProtocolVersion protocolVersion,
-                                              final HttpMetaData metadata, final Publisher<Object> messageBody) {
+                                              final HttpMetaData metadata,
+                                              final Publisher<Object> messageBody,
+                                              final boolean propagateCancel) {
         assert emptyMessageBody(metadata, messageBody);
         // HTTP/2 and above can write meta-data as a single frame with endStream=true flag. To check the version, use
         // HttpProtocolVersion from ConnectionInfo because HttpMetaData may have different version.
         final Publisher<Object> flatMessage =
                 protocolVersion.major() > 1 || !shouldAppendTrailers(protocolVersion, metadata) ? from(metadata) :
                         from(metadata, EmptyHttpHeaders.INSTANCE);
-        return messageBody == empty() ? flatMessage :
-                // Subscribe to the messageBody publisher to trigger any applied transformations, but ignore its
-                // content because the PayloadInfo indicated it's effectively empty and does not contain trailers
+        if (messageBody == empty()) {
+            return flatMessage;
+        }
+        // Subscribe to the messageBody publisher to trigger any applied transformations, but ignore its content because
+        // the PayloadInfo indicated it's effectively empty and does not contain trailers.
+        return propagateCancel ?
+                flatMessage.concatPropagateCancel(messageBody.ignoreElements()) :
                 flatMessage.concat(messageBody.ignoreElements());
     }
 
@@ -246,10 +253,11 @@ public boolean equals(Object o) {
     private static Publisher<Object> setContentLength(final HttpMetaData metadata,
                                                       final Publisher<Object> messageBody,
                                                       final BiIntConsumer<HttpHeaders> contentLengthUpdater,
-                                                      final HttpProtocolVersion protocolVersion) {
+                                                      final HttpProtocolVersion protocolVersion,
+                                                      final boolean propagateCancel) {
         if (emptyMessageBody(metadata, messageBody)) {
             contentLengthUpdater.apply(0, metadata.headers());
-            return flatEmptyMessage(protocolVersion, metadata, messageBody);
+            return flatEmptyMessage(protocolVersion, metadata, messageBody, propagateCancel);
         }
         return messageBody.collect(() -> null, (reduction, item) -> {
             if (reduction == null) {

servicetalk-http-netty/src/main/java/io/servicetalk/http/netty/NettyHttpServer.java

@@ -453,11 +453,14 @@ private static Publisher<Object> handleResponse(final HttpProtocolVersion protoc
                 return setResponseContentLength(protocolVersion, response);
             } else {
                 Publisher<Object> flatResponse;
-                if (emptyMessageBody(response, response.messageBody())) {
-                    flatResponse = flatEmptyMessage(protocolVersion, response, response.messageBody());
+                final Publisher<Object> messageBody = response.messageBody();
+                // Ensure cancel is propagated through the messageBody. Otherwise, if cancel from transport races with
+                // execution of this method and wins, BeforeFinallyHttpOperator won't trigger and observers won't
+                // complete the exchange.
+                if (emptyMessageBody(response, messageBody)) {
+                    flatResponse = flatEmptyMessage(protocolVersion, response, messageBody, /* propagateCancel */ true);
                 } else {
-                    // Not necessary to defer subscribe to the messageBody because server does not retry responses
-                    flatResponse = Single.<Object>succeeded(response).concat(response.messageBody());
+                    flatResponse = Single.<Object>succeeded(response).concatPropagateCancel(messageBody);
                     if (shouldAppendTrailers(protocolVersion, response)) {
                         flatResponse = flatResponse.scanWith(HeaderUtils::appendTrailersMapper);
                     }

servicetalk-http-netty/src/test/java/io/servicetalk/http/netty/HttpLifecycleObserverTest.java

@@ -290,6 +290,7 @@ public Single<StreamingHttpResponse> handle(HttpServiceContext ctx,
         verify(serverExchangeObserver, atMostOnce()).onResponse(any(StreamingHttpResponse.class));
         verify(serverResponseObserver, atMostOnce()).onResponseDataRequested(anyLong());
         verify(serverResponseObserver, atMostOnce()).onResponseComplete();
+        verify(serverResponseObserver, atMostOnce()).onResponseCancel();
         serverInOrder.verify(serverExchangeObserver).onExchangeFinally();
         verifyNoMoreInteractions(serverLifecycleObserver, serverExchangeObserver,
                 serverRequestObserver, serverResponseObserver);

servicetalk-http-netty/src/test/java/io/servicetalk/http/netty/RetryRequestWithNonRepeatablePayloadTest.java

@@ -96,7 +96,7 @@ public Single<StreamingHttpResponse> request(StreamingHttpRequest request) {
                                 try {
                                     assertThat("Unexpected exception type", t,
                                             instanceOf(RetryableException.class));
-                                    assertThat("Unexpected exception type",
+                                    assertThat("Unexpected exception cause type",
                                             t.getCause(), instanceOf(DeliberateException.class));
                                     assertThat("Unexpected subscribe to payload body",
                                             payloadBody.isSubscribed(), is(false));