Wip of leak detection working as intended

servicetalk-grpc-netty/src/test/java/io/servicetalk/grpc/LeakRepro.java

@@ -0,0 +1,100 @@
+package io.servicetalk.grpc;
+
+import com.apple.servicetalkleak.Message;
+import com.apple.servicetalkleak.ServiceTalkLeak;
+import io.netty.buffer.ByteBufUtil;
+import io.servicetalk.concurrent.api.Publisher;
+import io.servicetalk.concurrent.api.Single;
+import io.servicetalk.grpc.api.GrpcServiceContext;
+import io.servicetalk.grpc.api.GrpcStatusCode;
+import io.servicetalk.grpc.api.GrpcStatusException;
+import io.servicetalk.grpc.netty.GrpcClients;
+import io.servicetalk.grpc.netty.GrpcServers;
+import io.servicetalk.http.netty.HttpProtocolConfigs;
+import io.servicetalk.http.netty.SpliceFlatStreamToMetaSingle;
+import io.servicetalk.logging.api.LogLevel;
+import io.servicetalk.transport.api.HostAndPort;
+import io.servicetalk.transport.api.IoExecutor;
+import io.servicetalk.transport.netty.internal.NettyIoExecutors;
+import org.junit.jupiter.api.Test;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+import static io.servicetalk.concurrent.api.internal.BlockingUtils.blockingInvocation;
+import static org.junit.jupiter.api.Assertions.assertFalse;
+
+public class LeakRepro {
+
+    private static final Logger LOGGER = LoggerFactory.getLogger(LeakRepro.class);
+
+    static boolean leakDetected = false;
+
+    static {
+        System.setProperty("io.servicetalk.http.netty.leakdetection", "strict");
+        System.setProperty("io.netty.leakDetection.level", "paranoid");
+        ByteBufUtil.setLeakListener((type, records) -> {
+            leakDetected = true;
+            LOGGER.error("ByteBuf leak detected!");
+        });
+    }
+
+    IoExecutor serverExecutor = NettyIoExecutors.createIoExecutor(1, "server");
+    IoExecutor clientExecutor = NettyIoExecutors.createIoExecutor(1, "client");
+
+    @SuppressWarnings("resource")
+    @Test
+    public void testLeak() throws Exception {
+        GrpcServers.forPort(8888)
+                .initializeHttp(b -> b
+                        .ioExecutor(serverExecutor)
+                        .executor(serverExecutor))
+                .listenAndAwait(new ServiceTalkLeak.ServiceTalkLeakService() {
+                    @Override
+                    public Publisher<Message> rpc(GrpcServiceContext ctx, Publisher<Message> request) {
+                        Publisher<Message> response = splice(request)
+                                .flatMapPublisher(pair -> {
+                                    LOGGER.info("Initial message: " + pair.head);
+                                    return Publisher.failed(new GrpcStatusException(GrpcStatusCode.INVALID_ARGUMENT.status()));
+                                });
+                        return response;
+                    }
+                });
+
+        ServiceTalkLeak.ServiceTalkLeakClient client = GrpcClients.forAddress(HostAndPort.of("127.0.0.1", 8888))
+                .initializeHttp(b -> b
+                        .protocols(HttpProtocolConfigs.h2().enableFrameLogging("CLIENT", LogLevel.INFO, () -> true).build())
+                        .ioExecutor(clientExecutor)
+                        .executor(clientExecutor))
+                .build(new ServiceTalkLeak.ClientFactory());
+
+        for (int i = 0; i < 10; i++) {
+            LOGGER.info("Iteration {}", i);
+            blockingInvocation(
+                    client.rpc(
+                            Publisher.from(
+                                    Message.newBuilder().setValue("first message").build(),
+                                    Message.newBuilder().setValue("second message (which leaks)").build()))
+                    .ignoreElements()
+                    .onErrorComplete());
+
+            System.gc();
+            System.runFinalization();
+        }
+
+        assertFalse(leakDetected);
+    }
+
+    private static Single<Pair> splice(Publisher<Message> request) {
+        return request.liftSyncToSingle(new SpliceFlatStreamToMetaSingle<>(Pair::new));
+    }
+
+    private static final class Pair {
+        final Message head;
+        final Publisher<Message> stream;
+
+        public Pair(Message head, Publisher<Message> stream) {
+            this.head = head;
+            this.stream = stream;
+        }
+    }
+}

servicetalk-grpc-netty/src/test/proto/servicetalkloeak.proto

@@ -0,0 +1,12 @@
+syntax = "proto3";
+
+option java_multiple_files = true;
+option java_package = "com.apple.servicetalkleak";
+
+message Message {
+  string value = 1;
+}
+
+service ServiceTalkLeak {
+  rpc Rpc(stream Message) returns (stream Message);
+}

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

@@ -36,6 +36,8 @@
 import java.util.concurrent.atomic.AtomicReference;
 
 import static io.servicetalk.http.netty.HttpMessageDiscardWatchdogServiceFilter.generifyAtomicReference;
+import static io.servicetalk.http.netty.WatchdogLeakDetector.REQUEST_LEAK_MESSAGE;
+import static io.servicetalk.http.netty.WatchdogLeakDetector.RESPONSE_LEAK_MESSAGE;
 
 /**
  * Filter which tracks message bodies and warns if they are not discarded properly.
@@ -67,6 +69,17 @@ public StreamingHttpConnectionFilter create(final FilterableStreamingHttpConnect
         return new StreamingHttpConnectionFilter(connection) {
             @Override
             public Single<StreamingHttpResponse> request(final StreamingHttpRequest request) {
+                return WatchdogLeakDetector.strictDetection() ? requestStrict(request) : requestSimple(request);
+            }
+
+            private Single<StreamingHttpResponse> requestStrict(final StreamingHttpRequest request) {
+                return delegate().request(request.transformMessageBody(publisher ->
+                                WatchdogLeakDetector.instrument(publisher, REQUEST_LEAK_MESSAGE)))
+                        .map(response -> response.transformMessageBody(publisher ->
+                                WatchdogLeakDetector.instrument(publisher, RESPONSE_LEAK_MESSAGE)));
+            }
+
+            private Single<StreamingHttpResponse> requestSimple(final StreamingHttpRequest request) {
                 return delegate().request(request).map(response -> {
                     // always write the buffer publisher into the request context. When a downstream subscriber
                     // arrives, mark the message as subscribed explicitly (having a message present and no
@@ -78,10 +91,7 @@ public Single<StreamingHttpResponse> request(final StreamingHttpRequest request)
                         // If a previous message exists, the Single<StreamingHttpResponse> got resubscribed to
                         // (i.e. during a retry) and so previous message body needs to be cleaned up by the
                         // user.
-                        LOGGER.warn("Discovered un-drained HTTP response message body which has " +
-                                "been dropped by user code - this is a strong indication of a bug " +
-                                "in a user-defined filter. Response payload (message) body must " +
-                                "be fully consumed before retrying.");
+                        LOGGER.warn(RESPONSE_LEAK_MESSAGE);
                     }
 
                     return response.transformMessageBody(msgPublisher -> msgPublisher.beforeSubscriber(() -> {
@@ -112,10 +122,7 @@ protected Single<StreamingHttpResponse> request(final StreamingHttpRequester del
                                 if (maybePublisher != null && maybePublisher.getAndSet(null) != null) {
                                     // No-one subscribed to the message (or there is none), so if there is a message
                                     // tell the user to clean it up.
-                                    LOGGER.warn("Discovered un-drained HTTP response message body which has " +
-                                            "been dropped by user code - this is a strong indication of a bug " +
-                                            "in a user-defined filter. Response payload (message) body must " +
-                                            "be fully consumed before discarding.", cause);
+                                    LOGGER.warn(RESPONSE_LEAK_MESSAGE, cause);
                                 }
                             });
                 }

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

@@ -39,13 +39,17 @@
 import java.util.concurrent.atomic.AtomicReference;
 import javax.annotation.Nullable;
 
+import static io.servicetalk.http.netty.WatchdogLeakDetector.REQUEST_LEAK_MESSAGE;
+import static io.servicetalk.http.netty.WatchdogLeakDetector.RESPONSE_LEAK_MESSAGE;
+
 /**
  * Filter which tracks message bodies and warns if they are not discarded properly.
  */
 final class HttpMessageDiscardWatchdogServiceFilter implements StreamingHttpServiceFilterFactory {
 
     private static final Logger LOGGER = LoggerFactory.getLogger(HttpMessageDiscardWatchdogServiceFilter.class);
 
+
     /**
      * Instance of {@link HttpMessageDiscardWatchdogServiceFilter}.
      */
@@ -69,8 +73,26 @@ private HttpMessageDiscardWatchdogServiceFilter() {
     public StreamingHttpServiceFilter create(final StreamingHttpService service) {
 
         return new StreamingHttpServiceFilter(service) {
+
             @Override
             public Single<StreamingHttpResponse> handle(final HttpServiceContext ctx,
+                                                              final StreamingHttpRequest request,
+                                                              final StreamingHttpResponseFactory responseFactory) {
+                return WatchdogLeakDetector.strictDetection() ?
+                        handleStrict(ctx, request, responseFactory) : handleSimple(ctx, request, responseFactory);
+            }
+
+            private Single<StreamingHttpResponse> handleStrict(final HttpServiceContext ctx,
+                                                               final StreamingHttpRequest request,
+                                                               final StreamingHttpResponseFactory responseFactory) {
+                return delegate()
+                        .handle(ctx, request.transformMessageBody(publisher ->
+                                WatchdogLeakDetector.instrument(publisher, REQUEST_LEAK_MESSAGE)), responseFactory)
+                        .map(response -> response.transformMessageBody(publisher ->
+                                        WatchdogLeakDetector.instrument(publisher, RESPONSE_LEAK_MESSAGE)));
+            }
+
+            private Single<StreamingHttpResponse> handleSimple(final HttpServiceContext ctx,
                                                         final StreamingHttpRequest request,
                                                         final StreamingHttpResponseFactory responseFactory) {
                 return delegate()
@@ -86,10 +108,7 @@ public Single<StreamingHttpResponse> handle(final HttpServiceContext ctx,
                                 // If a previous message exists, the Single<StreamingHttpResponse> got resubscribed to
                                 // (i.e. during a retry) and so previous message body needs to be cleaned up by the
                                 // user.
-                                LOGGER.warn("Discovered un-drained HTTP response message body which has " +
-                                        "been dropped by user code - this is a strong indication of a bug " +
-                                        "in a user-defined filter. Responses (or their message body) must " +
-                                        "be fully consumed before retrying.");
+                                LOGGER.warn(RESPONSE_LEAK_MESSAGE);
                             }
 
                             return response.transformMessageBody(msgPublisher -> msgPublisher.beforeSubscriber(() -> {
@@ -173,10 +192,7 @@ public void onExchangeFinally() {
                         if (maybePublisher != null && maybePublisher.get() != null) {
                             // No-one subscribed to the message (or there is none), so if there is a message
                             // tell the user to clean it up.
-                            LOGGER.warn("Discovered un-drained HTTP response message body which has " +
-                                    "been dropped by user code - this is a strong indication of a bug " +
-                                    "in a user-defined filter. Responses (or their message body) must " +
-                                    "be fully consumed before discarding.");
+                            LOGGER.warn(RESPONSE_LEAK_MESSAGE);
                         }
                     }
                 }

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

@@ -52,7 +52,8 @@
  * @param <MetaData> type of meta-data in front of the stream of {@link Payload}, eg. {@link HttpResponseMetaData}
  * @param <Payload> type of payload inside the {@link Data}, eg. {@link Buffer}
  */
-final class SpliceFlatStreamToMetaSingle<Data, MetaData, Payload> implements PublisherToSingleOperator<Object, Data> {
+// TODO: revert: this shouldn't be public.
+public final class SpliceFlatStreamToMetaSingle<Data, MetaData, Payload> implements PublisherToSingleOperator<Object, Data> {
     private static final Logger LOGGER = LoggerFactory.getLogger(SpliceFlatStreamToMetaSingle.class);
     private final BiFunction<MetaData, Publisher<Payload>, Data> packer;
 
@@ -64,7 +65,7 @@ final class SpliceFlatStreamToMetaSingle<Data, MetaData, Payload> implements Pub
      * @param packer function to pack the {@link Publisher}&lt;{@link Payload}&gt; and {@link MetaData} into a
      * {@link Data}
      */
-    SpliceFlatStreamToMetaSingle(BiFunction<MetaData, Publisher<Payload>, Data> packer) {
+    public SpliceFlatStreamToMetaSingle(BiFunction<MetaData, Publisher<Payload>, Data> packer) {
         this.packer = requireNonNull(packer);
     }