Add DiscardOnCancel operator

src/main/java/io/r2dbc/postgresql/ExtendedQueryPostgresqlStatement.java

@@ -1,5 +1,5 @@
 /*
- * Copyright 2017-2019 the original author or authors.
+ * Copyright 2017-2020 the original author or authors.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -16,6 +16,8 @@
 
 package io.r2dbc.postgresql;
 
+import io.netty.util.ReferenceCountUtil;
+import io.netty.util.ReferenceCounted;
 import io.r2dbc.postgresql.api.PostgresqlStatement;
 import io.r2dbc.postgresql.client.Binding;
 import io.r2dbc.postgresql.client.Client;
@@ -28,6 +30,7 @@
 import io.r2dbc.postgresql.message.backend.NoData;
 import io.r2dbc.postgresql.util.Assert;
 import io.r2dbc.postgresql.util.GeneratedValuesUtils;
+import io.r2dbc.postgresql.util.Operators;
 import reactor.core.publisher.Flux;
 
 import java.util.ArrayList;
@@ -182,11 +185,16 @@ private Flux<io.r2dbc.postgresql.api.PostgresqlResult> execute(String sql) {
 
         ExceptionFactory factory = ExceptionFactory.withSql(sql);
         return this.statementCache.getName(this.bindings.first(), sql)
-            .flatMapMany(name -> ExtendedQueryMessageFlow
-                .execute(Flux.fromIterable(this.bindings.bindings), this.client, this.portalNameSupplier, name, sql, this.forceBinary))
+            .flatMapMany(name -> {
+                return ExtendedQueryMessageFlow
+                    .execute(Flux.fromIterable(this.bindings.bindings), this.client, this.portalNameSupplier, name, sql, this.forceBinary);
+            })
             .filter(RESULT_FRAME_FILTER)
             .windowUntil(CloseComplete.class::isInstance)
-            .map(messages -> PostgresqlResult.toResult(this.codecs, messages, factory));
+            .map(messages -> PostgresqlResult.toResult(this.codecs, messages, factory))
+            .cast(io.r2dbc.postgresql.api.PostgresqlResult.class)
+            .as(Operators::discardOnCancel)
+            .doOnDiscard(ReferenceCounted.class, ReferenceCountUtil::release);
     }
 
     private int getIndex(String identifier) {

src/main/java/io/r2dbc/postgresql/PostgresqlConnection.java

@@ -1,5 +1,5 @@
 /*
- * Copyright 2017-2019 the original author or authors.
+ * Copyright 2017-2020 the original author or authors.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -25,6 +25,7 @@
 import io.r2dbc.postgresql.client.TransactionStatus;
 import io.r2dbc.postgresql.codec.Codecs;
 import io.r2dbc.postgresql.util.Assert;
+import io.r2dbc.postgresql.util.Operators;
 import io.r2dbc.spi.Connection;
 import io.r2dbc.spi.IsolationLevel;
 import io.r2dbc.spi.ValidationDepth;
@@ -334,6 +335,7 @@ private static Function<TransactionStatus, String> getTransactionIsolationLevelQ
 
     private Mono<Void> useTransactionStatus(Function<TransactionStatus, Publisher<?>> f) {
         return Flux.defer(() -> f.apply(this.client.getTransactionStatus()))
+            .as(Operators::discardOnCancel)
             .then();
     }
 
@@ -342,9 +344,10 @@ private <T> Mono<T> withTransactionStatus(Function<TransactionStatus, T> f) {
     }
 
     private Publisher<?> exchange(String sql) {
-        ExceptionFactory exceptionFactory = ExceptionFactory.withSql("BEGIN");
+        ExceptionFactory exceptionFactory = ExceptionFactory.withSql(sql);
         return SimpleQueryMessageFlow.exchange(this.client, sql)
-            .handle(exceptionFactory::handleErrorResponse);
+            .handle(exceptionFactory::handleErrorResponse)
+            .as(Operators::discardOnCancel);
     }
 
     /**

src/main/java/io/r2dbc/postgresql/PostgresqlResult.java

@@ -1,5 +1,5 @@
 /*
- * Copyright 2017-2019 the original author or authors.
+ * Copyright 2017-2020 the original author or authors.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -16,7 +16,9 @@
 
 package io.r2dbc.postgresql;
 
+import io.netty.util.AbstractReferenceCounted;
 import io.netty.util.ReferenceCountUtil;
+import io.netty.util.ReferenceCounted;
 import io.r2dbc.postgresql.codec.Codecs;
 import io.r2dbc.postgresql.message.backend.BackendMessage;
 import io.r2dbc.postgresql.message.backend.CommandComplete;
@@ -39,7 +41,7 @@
 /**
  * An implementation of {@link Result} representing the results of a query against a PostgreSQL database.
  */
-final class PostgresqlResult implements io.r2dbc.postgresql.api.PostgresqlResult {
+final class PostgresqlResult extends AbstractReferenceCounted implements io.r2dbc.postgresql.api.PostgresqlResult {
 
     private static final Predicate<BackendMessage> TAKE_UNTIL = or(CommandComplete.class::isInstance, EmptyQueryResponse.class::isInstance, PortalSuspended.class::isInstance);
 
@@ -105,6 +107,18 @@ public <T> Flux<T> map(BiFunction<Row, RowMetadata, ? extends T> f) {
             });
     }
 
+    @Override
+    protected void deallocate() {
+
+        // drain messages for cleanup
+        this.getRowsUpdated().subscribe();
+    }
+
+    @Override
+    public ReferenceCounted touch(Object hint) {
+        return this;
+    }
+
     @Override
     public String toString() {
         return "PostgresqlResult{" +

src/main/java/io/r2dbc/postgresql/SimpleQueryPostgresqlStatement.java

@@ -1,5 +1,5 @@
 /*
- * Copyright 2017-2019 the original author or authors.
+ * Copyright 2017-2020 the original author or authors.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -16,6 +16,8 @@
 
 package io.r2dbc.postgresql;
 
+import io.netty.util.ReferenceCountUtil;
+import io.netty.util.ReferenceCounted;
 import io.r2dbc.postgresql.api.PostgresqlStatement;
 import io.r2dbc.postgresql.client.Client;
 import io.r2dbc.postgresql.client.SimpleQueryMessageFlow;
@@ -26,6 +28,7 @@
 import io.r2dbc.postgresql.message.backend.ErrorResponse;
 import io.r2dbc.postgresql.util.Assert;
 import io.r2dbc.postgresql.util.GeneratedValuesUtils;
+import io.r2dbc.postgresql.util.Operators;
 import reactor.core.publisher.Flux;
 import reactor.util.annotation.Nullable;
 
@@ -126,7 +129,10 @@ private Flux<io.r2dbc.postgresql.api.PostgresqlResult> execute(String sql) {
         return SimpleQueryMessageFlow
             .exchange(this.client, sql)
             .windowUntil(WINDOW_UNTIL)
-            .map(dataRow -> PostgresqlResult.toResult(this.codecs, dataRow, factory));
+            .map(dataRow -> PostgresqlResult.toResult(this.codecs, dataRow, factory))
+            .cast(io.r2dbc.postgresql.api.PostgresqlResult.class)
+            .as(Operators::discardOnCancel)
+            .doOnDiscard(ReferenceCounted.class, ReferenceCountUtil::release);
     }
 
 }

src/main/java/io/r2dbc/postgresql/client/ExtendedQueryMessageFlow.java

@@ -1,5 +1,5 @@
 /*
- * Copyright 2017-2019 the original author or authors.
+ * Copyright 2017-2020 the original author or authors.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.

src/main/java/io/r2dbc/postgresql/client/ReactorNettyClient.java

@@ -29,6 +29,7 @@
 import io.netty.handler.codec.LengthFieldBasedFrameDecoder;
 import io.netty.handler.logging.LogLevel;
 import io.netty.handler.logging.LoggingHandler;
+import io.netty.util.ReferenceCountUtil;
 import io.netty.util.internal.logging.InternalLogger;
 import io.netty.util.internal.logging.InternalLoggerFactory;
 import io.r2dbc.postgresql.message.backend.BackendKeyData;
@@ -145,7 +146,12 @@ private ReactorNettyClient(Connection connection) {
                         receiver.sink.complete();
                         this.conversations.poll();
                     } else {
-                        receiver.sink.next(message);
+
+                        if (receiver.sink.isCancelled()) {
+                            ReferenceCountUtil.release(message);
+                        } else {
+                            receiver.sink.next(message);
+                        }
                     }
                 }
             })

src/main/java/io/r2dbc/postgresql/util/FluxDiscardOnCancel.java

@@ -0,0 +1,126 @@
+/*
+ * Copyright 2019-2020 the original author or authors.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      https://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package io.r2dbc.postgresql.util;
+
+import org.reactivestreams.Subscriber;
+import org.reactivestreams.Subscription;
+import reactor.core.CoreSubscriber;
+import reactor.core.publisher.Flux;
+import reactor.core.publisher.FluxOperator;
+import reactor.core.publisher.Operators;
+import reactor.util.Logger;
+import reactor.util.Loggers;
+import reactor.util.context.Context;
+
+import java.util.concurrent.atomic.AtomicBoolean;
+
+/**
+ * A decorating operator that replays signals from its source to a {@link Subscriber} and drains the source upon {@link Subscription#cancel() cancel} and drops data signals until termination.
+ * Draining data is required to complete a particular request/response window and clear the protocol state as client code expects to start a request/response conversation without any previous
+ * response state.
+ */
+class FluxDiscardOnCancel<T> extends FluxOperator<T, T> {
+
+    private static final Logger logger = Loggers.getLogger(FluxDiscardOnCancel.class);
+
+    private final Runnable cancelConsumer;
+
+    FluxDiscardOnCancel(Flux<? extends T> source, Runnable cancelConsumer) {
+        super(source);
+        this.cancelConsumer = cancelConsumer;
+    }
+
+    @Override
+    public void subscribe(CoreSubscriber<? super T> actual) {
+        this.source.subscribe(new FluxDiscardOnCancelSubscriber<>(actual, this.cancelConsumer));
+    }
+
+    static class FluxDiscardOnCancelSubscriber<T> extends AtomicBoolean implements CoreSubscriber<T>, Subscription {
+
+        final CoreSubscriber<T> actual;
+
+        final Context ctx;
+
+        final Runnable cancelConsumer;
+
+        Subscription s;
+
+        FluxDiscardOnCancelSubscriber(CoreSubscriber<T> actual, Runnable cancelConsumer) {
+
+            this.actual = actual;
+            this.ctx = actual.currentContext();
+            this.cancelConsumer = cancelConsumer;
+        }
+
+        @Override
+        public void onSubscribe(Subscription s) {
+
+            if (Operators.validate(this.s, s)) {
+                this.s = s;
+                this.actual.onSubscribe(this);
+            }
+        }
+
+        @Override
+        public void onNext(T t) {
+
+            if (this.get()) {
+                Operators.onDiscard(t, this.ctx);
+                return;
+            }
+
+            this.actual.onNext(t);
+        }
+
+        @Override
+        public void onError(Throwable t) {
+            if (this.get()) {
+                Operators.onErrorDropped(t, this.ctx);
+            } else {
+                this.actual.onError(t);
+            }
+        }
+
+        @Override
+        public void onComplete() {
+            if (!this.get()) {
+                this.actual.onComplete();
+            }
+        }
+
+        @Override
+        public void request(long n) {
+            this.s.request(n);
+        }
+
+        @Override
+        public void cancel() {
+
+            if (compareAndSet(false, true)) {
+                if (logger.isDebugEnabled()) {
+                    logger.debug("received cancel signal");
+                }
+                try {
+                    this.cancelConsumer.run();
+                } catch (Exception e) {
+                    Operators.onErrorDropped(e, this.ctx);
+                }
+                this.s.request(Long.MAX_VALUE);
+            }
+        }
+    }
+}

src/main/java/io/r2dbc/postgresql/util/Operators.java

@@ -0,0 +1,62 @@
+/*
+ * Copyright 2019-2020 the original author or authors.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      https://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package io.r2dbc.postgresql.util;
+
+import org.reactivestreams.Subscription;
+import reactor.core.publisher.Flux;
+
+/**
+ * Operator utility.
+ *
+ * @since 0.8.1
+ */
+public final class Operators {
+
+    private Operators() {
+    }
+
+    /**
+     * Replay signals from {@link Flux the source} until cancellation. Drains the source for data signals if the subscriber cancels the subscription.
+     * <p>
+     * Draining data is required to complete a particular request/response window and clear the protocol state as client code expects to start a request/response conversation without leaving
+     * previous frames on the stack.
+     *
+     * @param source the source to decorate.
+     * @param <T>    The type of values in both source and output sequences.
+     * @return decorated {@link Flux}.
+     */
+    public static <T> Flux<T> discardOnCancel(Flux<? extends T> source) {
+        return new FluxDiscardOnCancel<>(source, () -> {
+        });
+    }
+
+    /**
+     * Replay signals from {@link Flux the source} until cancellation. Drains the source for data signals if the subscriber cancels the subscription.
+     * <p>
+     * Draining data is required to complete a particular request/response window and clear the protocol state as client code expects to start a request/response conversation without leaving
+     * previous frames on the stack.
+     * <p>Propagates the {@link Subscription#cancel()}  signal to a {@link Runnable consumer}.
+     *
+     * @param source         the source to decorate.
+     * @param cancelConsumer {@link Runnable} notified when the resulting {@link Flux} receives a {@link Subscription#cancel() cancel} signal.
+     * @param <T>            The type of values in both source and output sequences.
+     * @return decorated {@link Flux}.
+     */
+    public static <T> Flux<T> discardOnCancel(Flux<? extends T> source, Runnable cancelConsumer) {
+        return new FluxDiscardOnCancel<>(source, cancelConsumer);
+    }
+}