Add FilteredStreamMessage.onCancellation()

core/src/main/java/com/linecorp/armeria/client/encoding/HttpDecodedResponse.java

@@ -21,6 +21,7 @@
 import javax.annotation.Nullable;
 
 import org.reactivestreams.Subscriber;
+import org.reactivestreams.Subscription;
 
 import com.google.common.base.Ascii;
 
@@ -113,4 +114,11 @@ protected Throwable beforeError(Subscriber<? super HttpObject> subscriber, Throw
         }
         return cause;
     }
+
+    @Override
+    protected void beforeCancel(Subscription subscription) {
+        if (responseDecoder != null) {
+            responseDecoder.finish();
+        }
+    }
 }

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

@@ -103,6 +103,13 @@ protected Throwable beforeError(Subscriber<? super U> subscriber, Throwable caus
         return cause;
     }
 
+    /**
+     * A callback executed just before calling the upstream {@link Subscription#cancel()}.
+     * Override this method to execute any cleanup logic that may be needed before completing or failing the
+     * subscription.
+     */
+    protected void beforeCancel(Subscription subscription) {}
+
     @Override
     public final boolean isOpen() {
         return upstream.isOpen();
@@ -184,7 +191,7 @@ private final class FilteringSubscriber implements Subscriber<T> {
         @Override
         public void onSubscribe(Subscription s) {
             beforeSubscribe(delegate, s);
-            delegate.onSubscribe(s);
+            delegate.onSubscribe(new SubscriptionWrapper(s));
         }
 
         @Override
@@ -216,4 +223,24 @@ public void onComplete() {
             delegate.onComplete();
         }
     }
+
+    private final class SubscriptionWrapper implements Subscription {
+
+        private final Subscription subscription;
+
+        SubscriptionWrapper(Subscription s) {
+            subscription = s;
+        }
+
+        @Override
+        public void request(long n) {
+            subscription.request(n);
+        }
+
+        @Override
+        public void cancel() {
+            beforeCancel(subscription);
+            subscription.cancel();
+        }
+    }
 }

core/src/main/java/com/linecorp/armeria/internal/logging/ContentPreviewingUtil.java

@@ -21,6 +21,9 @@
 import javax.annotation.Nullable;
 
 import org.reactivestreams.Subscriber;
+import org.reactivestreams.Subscription;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
 
 import com.linecorp.armeria.common.FilteredHttpRequest;
 import com.linecorp.armeria.common.FilteredHttpResponse;
@@ -39,6 +42,8 @@
 
 public final class ContentPreviewingUtil {
 
+    private static final Logger logger = LoggerFactory.getLogger(ContentPreviewingUtil.class);
+
     /**
      * Sets up the request {@link ContentPreviewer} to set
      * {@link RequestLogBuilder#requestContentPreview(String)} when the preview is available.
@@ -71,19 +76,40 @@ protected HttpObject filter(HttpObject obj) {
 
             @Override
             protected void beforeComplete(Subscriber<? super HttpObject> subscriber) {
-                logBuilder.requestContentPreview(requestContentPreviewer.produce());
+                produceRequestContentPreview(requestContentPreviewer);
             }
 
             @Override
             protected Throwable beforeError(Subscriber<? super HttpObject> subscriber,
                                             Throwable cause) {
-                // Call produce() to release the resources in the previewer. Consider adding close() method.
-                requestContentPreviewer.produce();
+                try {
+                    // Call produce() to release the resources in the previewer. Consider adding close() method.
+                    requestContentPreviewer.produce();
+                } catch (Exception e) {
+                    logger.warn("Unexpected exception while producing the request content preview. " +
+                                "previewer: {}", requestContentPreviewer, e);
+                }
 
                 // Set null to make it sure the log is complete.
                 logBuilder.requestContentPreview(null);
                 return cause;
             }
+
+            @Override
+            protected void beforeCancel(Subscription subscription) {
+                produceRequestContentPreview(requestContentPreviewer);
+            }
+
+            private void produceRequestContentPreview(ContentPreviewer requestContentPreviewer) {
+                String produced = null;
+                try {
+                    produced = requestContentPreviewer.produce();
+                } catch (Exception e) {
+                    logger.warn("Unexpected exception while producing the request content preview. " +
+                                "previewer: {}", requestContentPreviewer, e);
+                }
+                logBuilder.requestContentPreview(produced);
+            }
         };
     }
 
@@ -125,12 +151,7 @@ protected HttpObject filter(HttpObject obj) {
 
             @Override
             protected void beforeComplete(Subscriber<? super HttpObject> subscriber) {
-                if (responseContentPreviewer != null) {
-                    ctx.logBuilder().responseContentPreview(responseContentPreviewer.produce());
-                } else {
-                    // Call requestContentPreview(null) to make sure that the log is complete.
-                    ctx.logBuilder().responseContentPreview(null);
-                }
+                produceResponseContentPreview(responseContentPreviewer);
             }
 
             @Override
@@ -143,6 +164,27 @@ protected Throwable beforeError(Subscriber<? super HttpObject> subscriber, Throw
                 ctx.logBuilder().responseContentPreview(null);
                 return cause;
             }
+
+            @Override
+            protected void beforeCancel(Subscription subscription) {
+                produceResponseContentPreview(responseContentPreviewer);
+            }
+
+            private void produceResponseContentPreview(@Nullable ContentPreviewer responseContentPreviewer) {
+                if (responseContentPreviewer != null) {
+                    String produced = null;
+                    try {
+                        produced = responseContentPreviewer.produce();
+                    } catch (Exception e) {
+                        logger.warn("Unexpected exception while producing the response content preview. " +
+                                    "previewer: {}", responseContentPreviewer, e);
+                    }
+                    ctx.logBuilder().responseContentPreview(produced);
+                } else {
+                    // Call requestContentPreview(null) to make sure that the log is complete.
+                    ctx.logBuilder().responseContentPreview(null);
+                }
+            }
         };
     }
 

core/src/main/java/com/linecorp/armeria/server/encoding/HttpDecodedRequest.java

@@ -17,6 +17,7 @@
 package com.linecorp.armeria.server.encoding;
 
 import org.reactivestreams.Subscriber;
+import org.reactivestreams.Subscription;
 
 import com.linecorp.armeria.common.FilteredHttpRequest;
 import com.linecorp.armeria.common.HttpData;
@@ -62,4 +63,9 @@ protected Throwable beforeError(Subscriber<? super HttpObject> subscriber, Throw
         responseDecoder.finish();
         return cause;
     }
+
+    @Override
+    protected void beforeCancel(Subscription subscription) {
+        responseDecoder.finish();
+    }
 }

core/src/main/java/com/linecorp/armeria/server/encoding/HttpEncodedResponse.java

@@ -26,6 +26,7 @@
 import javax.annotation.Nullable;
 
 import org.reactivestreams.Subscriber;
+import org.reactivestreams.Subscription;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
@@ -155,6 +156,11 @@ protected Throwable beforeError(Subscriber<? super HttpObject> subscriber, Throw
         return cause;
     }
 
+    @Override
+    protected void beforeCancel(Subscription subscription) {
+        closeEncoder();
+    }
+
     private void closeEncoder() {
         if (encodingStream == null) {
             return;

core/src/test/java/com/linecorp/armeria/client/encoding/HttpDecodedResponseTest.java

@@ -17,6 +17,12 @@
 package com.linecorp.armeria.client.encoding;
 
 import static org.assertj.core.api.Assertions.assertThat;
+import static org.awaitility.Awaitility.await;
+import static org.mockito.ArgumentMatchers.any;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.times;
+import static org.mockito.Mockito.verify;
+import static org.mockito.Mockito.when;
 
 import java.io.ByteArrayOutputStream;
 import java.io.IOException;
@@ -36,8 +42,10 @@
 import com.linecorp.armeria.common.HttpHeaderNames;
 import com.linecorp.armeria.common.HttpObject;
 import com.linecorp.armeria.common.HttpResponse;
+import com.linecorp.armeria.common.HttpResponseWriter;
 import com.linecorp.armeria.common.HttpStatus;
 import com.linecorp.armeria.common.ResponseHeaders;
+import com.linecorp.armeria.common.encoding.StreamDecoder;
 import com.linecorp.armeria.common.stream.SubscriptionOption;
 
 import io.netty.buffer.ByteBuf;
@@ -164,6 +172,26 @@ void pooledPayload_pooledDrain_withOldDecoder() {
         assertThat(decodedPayloadBuf.refCnt()).isZero();
     }
 
+    @Test
+    void streamDecoderFinishedIsCalledWhenRequestCanceled() throws InterruptedException {
+        final HttpResponseWriter response = HttpResponse.streaming();
+        response.write(ResponseHeaders.of(HttpStatus.OK, HttpHeaderNames.CONTENT_ENCODING, "foo"));
+        final HttpData data = HttpData.ofUtf8("bar");
+        response.write(data);
+
+        final com.linecorp.armeria.common.encoding.StreamDecoderFactory factory = mock(
+                com.linecorp.armeria.common.encoding.StreamDecoderFactory.class);
+        final com.linecorp.armeria.common.encoding.StreamDecoder streamDecoder = mock(StreamDecoder.class);
+        when(factory.newDecoder(any())).thenReturn(streamDecoder);
+        when(streamDecoder.decode(any())).thenReturn(data);
+
+        final HttpResponse decoded = new HttpDecodedResponse(response, ImmutableMap.of("foo", factory),
+                                                             ByteBufAllocator.DEFAULT);
+        decoded.subscribe(new CancelSubscriber());
+
+        await().untilAsserted(() -> verify(streamDecoder, times(1)).finish());
+    }
+
     private static HttpData responseData(HttpResponse decoded, boolean withPooledObjects) {
         final CompletableFuture<HttpData> future = new CompletableFuture<>();
         final Subscriber<HttpObject> subscriber = new Subscriber<HttpObject>() {
@@ -194,4 +222,28 @@ public void onComplete() {}
 
         return future.join();
     }
+
+    private static class CancelSubscriber implements Subscriber<HttpObject> {
+
+        private Subscription subscription;
+
+        @Override
+        public void onSubscribe(Subscription s) {
+            subscription = s;
+            s.request(Long.MAX_VALUE);
+        }
+
+        @Override
+        public void onNext(HttpObject httpObject) {
+            if (httpObject instanceof HttpData) {
+                subscription.cancel();
+            }
+        }
+
+        @Override
+        public void onError(Throwable t) {}
+
+        @Override
+        public void onComplete() {}
+    }
 }