http-utils: fix leak in AbstractTimeoutHttpFilter

servicetalk-http-utils/src/main/java/io/servicetalk/http/utils/AbstractTimeoutHttpFilter.java

@@ -15,9 +15,12 @@
  */
 package io.servicetalk.http.utils;
 
+import io.servicetalk.concurrent.Cancellable;
 import io.servicetalk.concurrent.Executor;
+import io.servicetalk.concurrent.SingleSource;
 import io.servicetalk.concurrent.TimeSource;
 import io.servicetalk.concurrent.api.Single;
+import io.servicetalk.concurrent.internal.CancelImmediatelySubscriber;
 import io.servicetalk.http.api.HttpExecutionStrategies;
 import io.servicetalk.http.api.HttpExecutionStrategy;
 import io.servicetalk.http.api.HttpExecutionStrategyInfluencer;
@@ -28,11 +31,13 @@
 
 import java.time.Duration;
 import java.util.concurrent.TimeoutException;
+import java.util.concurrent.atomic.AtomicReferenceFieldUpdater;
 import java.util.function.BiFunction;
 import java.util.function.Function;
 import javax.annotation.Nullable;
 
 import static io.servicetalk.concurrent.api.Publisher.defer;
+import static io.servicetalk.concurrent.api.SourceAdapters.toSource;
 import static io.servicetalk.http.api.HttpContextKeys.HTTP_EXECUTION_STRATEGY_KEY;
 import static io.servicetalk.utils.internal.DurationUtils.ensurePositive;
 import static java.time.Duration.ofNanos;
@@ -117,7 +122,9 @@ final Single<StreamingHttpResponse> withTimeout(final StreamingHttpRequest reque
             final Duration timeout = timeoutForRequest.apply(request, useForTimeout);
             Single<StreamingHttpResponse> response = responseFunction.apply(request);
             if (null != timeout) {
-                final Single<StreamingHttpResponse> timeoutResponse = response.timeout(timeout, useForTimeout);
+                final Single<StreamingHttpResponse> timeoutResponse = response
+                        .<StreamingHttpResponse>liftSync(CleanupSubscriber::new)
+                        .timeout(timeout, useForTimeout);
 
                 if (fullRequestResponse) {
                     final long deadline = useForTimeout.currentTime(NANOSECONDS) + timeout.toNanos();
@@ -151,6 +158,62 @@ private Executor contextExecutor(final HttpRequestMetaData requestMetaData,
                 context.executor() : context.ioExecutor();
     }
 
+    // This is responsible for ensuring that we don't abandon the response resources due to the use of Single.timeout.
+    // It does so by retaining a reference to the StreamingHttpResponse in case we receive a cancellation. It will
+    // then attempt to drain the resource. We're protected from multiple subscribes because the TimeoutSingle will
+    // short circuit also on upstream cancellation, although that is not obviously correct behavior.
+    private static final class CleanupSubscriber implements SingleSource.Subscriber<StreamingHttpResponse> {
+
+        private static final AtomicReferenceFieldUpdater<CleanupSubscriber, Object> stateUpdater =
+                AtomicReferenceFieldUpdater.newUpdater(CleanupSubscriber.class, Object.class, "state");
+
+        private static final String COMPLETE = "complete";
+
+        private final SingleSource.Subscriber<? super StreamingHttpResponse> delegate;
+        @Nullable
+        private volatile Object state;
+
+        CleanupSubscriber(final SingleSource.Subscriber<? super StreamingHttpResponse> delegate) {
+            this.delegate = delegate;
+        }
+
+        @Override
+        public void onSubscribe(Cancellable cancellable) {
+            delegate.onSubscribe(() -> {
+                try {
+                    Object current = stateUpdater.getAndSet(this, COMPLETE);
+                    if (current instanceof StreamingHttpResponse) {
+                        clean((StreamingHttpResponse) current);
+                    }
+                } finally {
+                    cancellable.cancel();
+                }
+            });
+        }
+
+        @Override
+        public void onSuccess(@Nullable StreamingHttpResponse result) {
+            assert result != null;
+            if (stateUpdater.compareAndSet(this, null, result)) {
+                // We win.
+                delegate.onSuccess(result);
+            } else {
+                // we lost to cancellation. No need to send it forward.
+                clean(result);
+            }
+        }
+
+        @Override
+        public void onError(Throwable t) {
+            assert !(state instanceof StreamingHttpResponse);
+            delegate.onError(t);
+        }
+
+        private void clean(StreamingHttpResponse httpResponse) {
+            toSource(httpResponse.messageBody()).subscribe(CancelImmediatelySubscriber.INSTANCE);
+        }
+    }
+
     private static final class MappedTimeoutException extends TimeoutException {
         private static final long serialVersionUID = -8230476062001221272L;
 

servicetalk-http-utils/src/test/java/io/servicetalk/http/utils/AbstractTimeoutHttpFilterTest.java

@@ -16,6 +16,7 @@
 package io.servicetalk.http.utils;
 
 import io.servicetalk.buffer.api.Buffer;
+import io.servicetalk.concurrent.Cancellable;
 import io.servicetalk.concurrent.TimeSource;
 import io.servicetalk.concurrent.api.DefaultThreadFactory;
 import io.servicetalk.concurrent.api.Executor;
@@ -65,6 +66,7 @@
 import static java.time.Duration.ofSeconds;
 import static org.hamcrest.MatcherAssert.assertThat;
 import static org.hamcrest.Matchers.is;
+import static org.hamcrest.Matchers.notNullValue;
 import static org.junit.jupiter.api.Assertions.assertThrows;
 import static org.mockito.ArgumentMatchers.any;
 import static org.mockito.Mockito.mock;
@@ -220,8 +222,64 @@ true, defaultStrategy(), responseWith(payloadBody)))
         assertThat("No subscribe for payload body", payloadBody.isSubscribed(), is(true));
     }
 
+    @ParameterizedTest(name = "{index}: fullRequestResponse={0}")
+    @ValueSource(booleans = {false, true})
+    void responseRacesWithCancellationStillHaveBodyDrained(boolean fullResponse) {
+        Duration timeout = ofMillis(100L);
+        TestSingle<StreamingHttpResponse> response = new TestSingle<>();
+        StepVerifiers.create(applyFilter(timeout, fullResponse, defaultStrategy(), response)
+
+                        .flatMapPublisher(StreamingHttpResponse::payloadBody))
+                .thenRequest(MAX_VALUE)
+                .expectErrorMatches(t -> TimeoutException.class.isInstance(t) &&
+                        (Thread.currentThread() instanceof IoThreadFactory.IoThread ^
+                                defaultStrategy().isRequestResponseOffloaded()))
+                .verify();
+
+        // We should be subscribed at this point.
+        response.awaitSubscribed();
+        TestPublisher<Buffer> payloadBody = new TestPublisher<>();
+        AtomicBoolean cancelled = new AtomicBoolean();
+        response.onSuccess(responseRawWith(payloadBody.beforeCancel(() -> cancelled.set(true))));
+        assertThat("No subscribe for payload body", payloadBody.isSubscribed(), is(true));
+        assertThat("Payload body wasn't cancelled as part of draining", cancelled.get(), is(true));
+    }
+
+    @ParameterizedTest(name = "{index}: fullRequestResponse={0}")
+    @ValueSource(booleans = {false, true})
+    void responseCompletesBeforeTimeoutWithFollowingCancel(boolean fullRequestResponse) {
+        TestSingle<StreamingHttpResponse> responseSingle = new TestSingle<>();
+        AtomicReference<Cancellable> doCancel = new AtomicReference<>();
+        AtomicBoolean cancelPropagated = new AtomicBoolean();
+        StepVerifiers.create(applyFilter(ofSeconds(DEFAULT_TIMEOUT_SECONDS / 2),
+                        fullRequestResponse, defaultStrategy(), responseSingle.beforeCancel(() ->
+                                cancelPropagated.set(true))).
+                        afterOnSubscribe(doCancel::set))
+                .then(() -> immediate().schedule(() -> {
+                            StreamingHttpResponse response = mock(StreamingHttpResponse.class);
+                            when(response.transformMessageBody(any())).thenReturn(response);
+                            responseSingle.onSuccess(response);
+                        },
+                        ofMillis(1L)))
+                .expectSuccess()
+                .verify();
+        assertThat("No subscribe for response single", responseSingle.isSubscribed(), is(true));
+        responseSingle.isSubscribed();
+        assertThat(doCancel.get(), is(notNullValue()));
+
+        // We rely on the Single.timeout(..) operator to swallow the losing cancellation so that our CleanupSubscriber
+        // doesn't end up being the second subscriber to the response body. If that behavior changes we may need to be
+        // more defensive.
+        doCancel.get().cancel();
+        assertThat(cancelPropagated.get(), is(false));
+    }
+
     private static Single<StreamingHttpResponse> responseWith(Publisher<Buffer> payloadBody) {
-        return succeeded(newResponse(OK, HTTP_1_1, EmptyHttpHeaders.INSTANCE, DEFAULT_ALLOCATOR,
-                DefaultHttpHeadersFactory.INSTANCE).payloadBody(payloadBody));
+        return succeeded(responseRawWith(payloadBody));
+    }
+
+    private static StreamingHttpResponse responseRawWith(Publisher<Buffer> payloadBody) {
+        return newResponse(OK, HTTP_1_1, EmptyHttpHeaders.INSTANCE, DEFAULT_ALLOCATOR,
+                DefaultHttpHeadersFactory.INSTANCE).payloadBody(payloadBody);
     }
 }