http-utils: fix leak in AbstractTimeoutHttpFilter (#3038)

Motivation:

The Single.timeout(..) family of operations are part of a class of operations
will short-circuit the response based on the cancellation pathway, specifically
they will return a TimeoutException down the error path. Because cancellation
and the result are intrinsically racy there is currently the possibility that
the result will be generated but we cannot return it to the subscriber having
already given them an Exception in the error pathway. For values that are
stateful this can result in a resource leak. One such leak occurs in the
AbstractTimeoutHttpFilter.

Modifications:

Modify the AbstractTimeoutHttpFilter to keep a handle on the resource and make
sure we close it if we receive a cancellation. This is a localized fix but
appears to be effective.

Result:

One less leak.

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,66 @@ 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) {
+                        // Because of the nature of the Single.timeout(..) operator we know that if we get a cancel
+                        // call the message will never be escape the `.timeout(..)` operator because it will have
+                        // 'completed'. Therefore, it is our responsibility to clean up the message body. This behavior
+                        // is verified by the `responseCompletesBeforeTimeoutWithFollowingCancel` test in the suite.
+                        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,63 @@ 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));
+
+        // 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.
+        assertThat(doCancel.get(), is(notNullValue()));
+        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);
     }
 }