Terminate Notification stream on disconnect

We now terminate the notification stream (PostgresqlConnection.getNotifications()) when the connection gets disconnected. If the connection is closed normally, the stream terminates successfully (onComplete). Unintended disconnects result in an error (onError).

[#212]

Author: mp911de

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

@@ -24,12 +24,14 @@
 import io.r2dbc.postgresql.client.SimpleQueryMessageFlow;
 import io.r2dbc.postgresql.client.TransactionStatus;
 import io.r2dbc.postgresql.codec.Codecs;
+import io.r2dbc.postgresql.message.backend.NotificationResponse;
 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;
 import org.reactivestreams.Publisher;
+import org.reactivestreams.Subscriber;
 import org.reactivestreams.Subscription;
 import reactor.core.CoreSubscriber;
 import reactor.core.Disposable;
@@ -246,17 +248,17 @@ public Mono<Void> setAutoCommit(boolean autoCommit) {
 
         return useTransactionStatus(transactionStatus -> {
 
-            logger.debug(String.format("Setting auto-commit mode to [%s]", autoCommit));
+            this.logger.debug(String.format("Setting auto-commit mode to [%s]", autoCommit));
 
             if (isAutoCommit()) {
                 if (!autoCommit) {
-                    logger.debug("Beginning transaction");
+                    this.logger.debug("Beginning transaction");
                     return beginTransaction();
                 }
             } else {
 
                 if (autoCommit) {
-                    logger.debug("Committing pending transactions");
+                    this.logger.debug("Committing pending transactions");
                     return commitTransaction();
                 }
             }
@@ -314,7 +316,7 @@ public void onNext(Integer integer) {
 
                 @Override
                 public void onError(Throwable t) {
-                    logger.debug("Validation failed", t);
+                    PostgresqlConnection.this.logger.debug("Validation failed", t);
                     sink.success(false);
                 }
 
@@ -360,7 +362,7 @@ static class NotificationAdapter {
 
         private final DirectProcessor<Notification> processor = DirectProcessor.create();
 
-        private final FluxSink<Notification> sink = processor.sink();
+        private final FluxSink<Notification> sink = this.processor.sink();
 
         @Nullable
         private volatile Disposable subscription = null;
@@ -374,13 +376,32 @@ void dispose() {
 
         void register(Client client) {
 
-            this.subscription = client.addNotificationListener(notificationResponse -> {
-                sink.next(new NotificationResponseWrapper(notificationResponse));
+            this.subscription = client.addNotificationListener(new Subscriber<NotificationResponse>() {
+
+                @Override
+                public void onSubscribe(Subscription subscription) {
+                    subscription.request(Long.MAX_VALUE);
+                }
+
+                @Override
+                public void onNext(NotificationResponse notificationResponse) {
+                    NotificationAdapter.this.sink.next(new NotificationResponseWrapper(notificationResponse));
+                }
+
+                @Override
+                public void onError(Throwable throwable) {
+                    NotificationAdapter.this.sink.error(throwable);
+                }
+
+                @Override
+                public void onComplete() {
+                    NotificationAdapter.this.sink.complete();
+                }
             });
         }
 
         Flux<Notification> getEvents() {
-            return processor;
+            return this.processor;
         }
     }
 

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

@@ -18,7 +18,9 @@
 
 import io.r2dbc.spi.Connection;
 import io.r2dbc.spi.IsolationLevel;
+import io.r2dbc.spi.R2dbcNonTransientResourceException;
 import io.r2dbc.spi.ValidationDepth;
+import org.reactivestreams.Subscriber;
 import reactor.core.publisher.Flux;
 import reactor.core.publisher.Mono;
 
@@ -64,8 +66,9 @@ public interface PostgresqlConnection extends Connection {
     PostgresqlStatement createStatement(String sql);
 
     /**
-     * Return a {@link Flux} of {@link Notification} received from {@code LISTEN} registrations.
-     * The stream is a hot stream producing messages as they are received.
+     * Return a {@link Flux} of {@link Notification} received from {@code LISTEN} registrations. The stream is a hot stream producing messages as they are received. Notifications received by this
+     * connection are published as they are received. When the client gets {@link #close() closed}, the subscription {@link Subscriber#onComplete() completes normally}. Otherwise (transport
+     * connection disconnected unintentionally) with an {@link R2dbcNonTransientResourceException error}.
      *
      * @return a hot {@link Flux} of {@link Notification Notifications}.
      */

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

@@ -21,7 +21,9 @@
 import io.r2dbc.postgresql.message.backend.NotificationResponse;
 import io.r2dbc.postgresql.message.backend.ReadyForQuery;
 import io.r2dbc.postgresql.message.frontend.FrontendMessage;
+import io.r2dbc.spi.R2dbcNonTransientResourceException;
 import org.reactivestreams.Publisher;
+import org.reactivestreams.Subscriber;
 import reactor.core.Disposable;
 import reactor.core.publisher.Flux;
 import reactor.core.publisher.Mono;
@@ -36,14 +38,26 @@
 public interface Client {
 
     /**
-     * Add a consumer of notification messages.
+     * Add a consumer of notification messages. Notifications received by this connection are sent to the {@link Consumer notification consumer}. Note that connection errors and events such as
+     * disconnects are not visible to the {@link Consumer notification consumer}.
      *
      * @param consumer the consumer of notification messages
      * @return a new {@link Disposable} that can be used to cancel the underlying subscription.
      * @throws IllegalArgumentException if {@code consumer} is {@code null}
      */
     Disposable addNotificationListener(Consumer<NotificationResponse> consumer);
 
+    /**
+     * Add a consumer of notification messages. Notifications received by this connection are sent to the {@link Subscriber notification listener}. When the client gets {@link #close() closed}, the
+     * subscription {@link Subscriber#onComplete() completes normally}. Otherwise (transport connection disconnected unintentionally) with an {@link R2dbcNonTransientResourceException error}.
+     *
+     * @param consumer the consumer of notification messages
+     * @return a new {@link Disposable} that can be used to cancel the underlying subscription.
+     * @throws IllegalArgumentException if {@code consumer} is {@code null}
+     * @since 0.8.1
+     */
+    Disposable addNotificationListener(Subscriber<NotificationResponse> consumer);
+
     /**
      * Release any resources held by the {@link Client}.
      *

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

@@ -46,6 +46,7 @@
 import io.r2dbc.postgresql.util.Assert;
 import io.r2dbc.spi.R2dbcNonTransientResourceException;
 import org.reactivestreams.Publisher;
+import org.reactivestreams.Subscriber;
 import reactor.core.Disposable;
 import reactor.core.publisher.DirectProcessor;
 import reactor.core.publisher.EmitterProcessor;
@@ -178,16 +179,16 @@ private ReactorNettyClient(Connection connection) {
             .subscribe();
     }
 
-    @Override
-    public Disposable addNotificationListener(Consumer<NotificationResponse> consumer) {
-        return this.notificationProcessor.subscribe(consumer);
-    }
-
     @Override
     public Mono<Void> close() {
         return Mono.defer(() -> {
 
+            if (!this.notificationProcessor.isTerminated()) {
+                this.notificationProcessor.onComplete();
+            }
+
             drainError(EXPECTED);
+
             if (this.isClosed.compareAndSet(false, true)) {
 
                 if (!isConnected() || this.processId == null) {
@@ -390,6 +391,16 @@ private static Mono<? extends Void> registerSslHandler(SSLConfig sslConfig, Conn
         return Mono.empty();
     }
 
+    @Override
+    public Disposable addNotificationListener(Consumer<NotificationResponse> consumer) {
+        return this.notificationProcessor.subscribe(consumer);
+    }
+
+    @Override
+    public Disposable addNotificationListener(Subscriber<NotificationResponse> consumer) {
+        return this.notificationProcessor.subscribe(consumer::onNext, consumer::onError, consumer::onComplete, consumer::onSubscribe);
+    }
+
     @Override
     public ByteBufAllocator getByteBufAllocator() {
         return this.byteBufAllocator;
@@ -457,6 +468,10 @@ private void drainError(Supplier<? extends Throwable> supplier) {
         while ((receiver = this.conversations.poll()) != null) {
             receiver.sink.error(supplier.get());
         }
+
+        if (!this.notificationProcessor.isTerminated()) {
+            this.notificationProcessor.onError(supplier.get());
+        }
     }
 
     /**

src/test/java/io/r2dbc/postgresql/PostgresNotificationIntegrationTests.java

@@ -16,7 +16,11 @@
 
 package io.r2dbc.postgresql;
 
+import io.netty.channel.Channel;
 import io.r2dbc.postgresql.api.Notification;
+import io.r2dbc.postgresql.api.PostgresqlConnection;
+import io.r2dbc.postgresql.util.ConnectionIntrospector;
+import io.r2dbc.spi.R2dbcNonTransientResourceException;
 import org.junit.jupiter.api.Test;
 import reactor.core.Disposable;
 import reactor.test.StepVerifier;
@@ -55,4 +59,28 @@ void shouldReceivePubSubNotifications() throws Exception {
 
         listener.dispose();
     }
+
+    @Test
+    void listenShouldCompleteOnConnectionClose() {
+
+        PostgresqlConnection connection = this.connectionFactory.create().block();
+
+        connection.getNotifications().as(StepVerifier::create).expectSubscription()
+            .then(() -> connection.close().subscribe())
+            .verifyComplete();
+    }
+
+    @Test
+    void listenShouldFailOnConnectionDisconnected() {
+
+        PostgresqlConnection connection = this.connectionFactory.create().block();
+
+        connection.getNotifications().as(StepVerifier::create).expectSubscription()
+            .then(() -> {
+
+                Channel channel = ConnectionIntrospector.of(connection).getChannel();
+                channel.close();
+            })
+            .verifyError(R2dbcNonTransientResourceException.class);
+    }
 }

src/test/java/io/r2dbc/postgresql/client/TestClient.java

@@ -23,6 +23,7 @@
 import io.r2dbc.postgresql.util.Assert;
 import io.r2dbc.postgresql.util.TestByteBufAllocator;
 import org.reactivestreams.Publisher;
+import org.reactivestreams.Subscriber;
 import reactor.core.Disposable;
 import reactor.core.publisher.EmitterProcessor;
 import reactor.core.publisher.Flux;
@@ -163,6 +164,11 @@ public Disposable addNotificationListener(Consumer<NotificationResponse> consume
         return this.notificationProcessor.subscribe(consumer);
     }
 
+    @Override
+    public Disposable addNotificationListener(Subscriber<NotificationResponse> consumer) {
+        return this.notificationProcessor.subscribe(consumer::onNext, consumer::onError, consumer::onComplete, consumer::onSubscribe);
+    }
+
     public void notify(NotificationResponse notification) {
         this.notificationProcessor.onNext(notification);
     }

src/test/java/io/r2dbc/postgresql/util/ConnectionIntrospector.java

@@ -0,0 +1,73 @@
+/*
+ * Copyright 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 io.netty.channel.Channel;
+import io.r2dbc.postgresql.api.PostgresqlConnection;
+import io.r2dbc.postgresql.client.ReactorNettyClient;
+import org.springframework.beans.DirectFieldAccessor;
+import reactor.netty.Connection;
+
+/**
+ * Utility to introspect a {@link PostgresqlConnection}.
+ */
+public final class ConnectionIntrospector {
+
+    private final PostgresqlConnection connection;
+
+    private ConnectionIntrospector(PostgresqlConnection connection) {
+        this.connection = connection;
+    }
+
+    /**
+     * Create a new instance.
+     *
+     * @param connection
+     * @return
+     */
+    public static ConnectionIntrospector of(PostgresqlConnection connection) {
+        return new ConnectionIntrospector(connection);
+    }
+
+    /**
+     * Return the transport {@link Channel}.
+     *
+     * @return the transport {@link Channel}.
+     */
+    public Channel getChannel() {
+
+        DirectFieldAccessor accessor = new DirectFieldAccessor(getClient());
+        Connection connection = (Connection) accessor.getPropertyValue("connection");
+
+        return connection.channel();
+    }
+
+    /**
+     * Return the underlying {@link ReactorNettyClient}.
+     *
+     * @return the underlying {@link ReactorNettyClient}.
+     */
+    public ReactorNettyClient getClient() {
+
+        DirectFieldAccessor accessor = new DirectFieldAccessor(this.connection);
+        Object value = accessor.getPropertyValue("client");
+
+        Assert.requireType(value, ReactorNettyClient.class, "Client must be of type ReactorNettyClient. Was: " + value.getClass().getName());
+
+        return (ReactorNettyClient) value;
+    }
+}