Propagate connection closed and connection errors to message exchange

mp911de · mp911de · commit 60b1138aa1cf · 2019-10-11T19:38:09.000-05:00
We now fail a stream if the connection was closed in flight and if connection errors occurred. [resolves #177]
diff --git a/src/main/java/io/r2dbc/postgresql/client/ReactorNettyClient.java b/src/main/java/io/r2dbc/postgresql/client/ReactorNettyClient.java
@@ -22,6 +22,7 @@
 import io.netty.channel.ChannelHandlerContext;
 import io.netty.channel.ChannelOption;
 import io.netty.handler.codec.LengthFieldBasedFrameDecoder;
+import io.netty.util.ReferenceCountUtil;
 import io.r2dbc.postgresql.message.backend.BackendKeyData;
 import io.r2dbc.postgresql.message.backend.BackendMessage;
 import io.r2dbc.postgresql.message.backend.BackendMessageDecoder;
@@ -34,6 +35,7 @@
 import io.r2dbc.postgresql.message.frontend.FrontendMessage;
 import io.r2dbc.postgresql.message.frontend.Terminate;
 import io.r2dbc.postgresql.util.Assert;
+import io.r2dbc.spi.R2dbcNonTransientResourceException;
 import org.reactivestreams.Publisher;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
@@ -59,8 +61,10 @@
 import java.util.StringJoiner;
 import java.util.concurrent.atomic.AtomicBoolean;
 import java.util.concurrent.atomic.AtomicInteger;
+import java.util.concurrent.atomic.AtomicReference;
 import java.util.function.Consumer;
 import java.util.function.Function;
+import java.util.function.Supplier;
 
 import static io.r2dbc.postgresql.client.TransactionStatus.IDLE;
 
@@ -75,6 +79,10 @@ public final class ReactorNettyClient implements Client {
 
     private static final boolean DEBUG_ENABLED = logger.isDebugEnabled();
 
+    private static final Supplier<PostgresConnectionClosedException> UNEXPECTED = () -> new PostgresConnectionClosedException("Connection unexpectedly closed");
+
+    private static final Supplier<PostgresConnectionClosedException> EXPECTED = () -> new PostgresConnectionClosedException("Connection closed");
+
     private final ByteBufAllocator byteBufAllocator;
 
     private final Connection connection;
@@ -107,26 +115,36 @@ private ReactorNettyClient(Connection connection) {
         Assert.requireNonNull(connection, "Connection must not be null");
 
         connection.addHandler(new LengthFieldBasedFrameDecoder(Integer.MAX_VALUE - 5, 1, 4, -4, 0));
-        connection.addHandler(new EnsureSubscribersCompleteChannelHandler(this.requestProcessor, this.responseReceivers));
+        connection.addHandler(new EnsureSubscribersCompleteChannelHandler(this.requestProcessor));
         this.connection = connection;
         this.byteBufAllocator = connection.outbound().alloc();
 
+        AtomicReference<Throwable> receiveError = new AtomicReference<>();
         Mono<Void> receive = connection.inbound().receive()
             .map(BackendMessageDecoder::decode)
             .handle(this::handleResponse)
+            .doOnError(throwable -> {
+                receiveError.set(throwable);
+                handleConnectionError(throwable);
+            })
             .windowWhile(it -> it.getClass() != ReadyForQuery.class)
             .doOnNext(fluxOfMessages -> {
                 MonoSink<Flux<BackendMessage>> receiver = this.responseReceivers.poll();
                 if (receiver != null) {
-                    receiver.success(fluxOfMessages);
-                }
-            })
-            .doOnComplete(() -> {
-                MonoSink<Flux<BackendMessage>> receiver = this.responseReceivers.poll();
-                if (receiver != null) {
-                    receiver.success(Flux.empty());
+                    receiver.success(fluxOfMessages.doOnComplete(() -> {
+
+                        Throwable throwable = receiveError.get();
+                        if (throwable != null) {
+                            throw new PostgresConnectionException(throwable);
+                        }
+
+                        if (!isConnected()) {
+                            throw EXPECTED.get();
+                        }
+                    }));
                 }
             })
+            .doOnComplete(this::handleClose)
             .then();
 
         Mono<Void> request = this.requestProcessor
@@ -139,35 +157,27 @@ private ReactorNettyClient(Connection connection) {
             .then();
 
         receive
-            .onErrorResume(throwable -> {
-
-                MonoSink<Flux<BackendMessage>> receiver = this.responseReceivers.poll();
-                if (receiver != null) {
-                    receiver.error(throwable);
-                }
-                this.requestProcessor.onComplete();
-
-                if (isSslException(throwable)) {
-                    logger.debug("Connection Error", throwable);
-                } else {
-
-                    logger.error("Connection Error", throwable);
-                }
-                return close();
-            })
+            .onErrorResume(this::resumeError)
             .subscribe();
 
         request
-            .onErrorResume(throwable -> {
+            .onErrorResume(this::resumeError)
+            .doAfterTerminate(this::handleClose)
+            .subscribe();
+    }
 
-                if (isSslException(throwable)) {
-                    logger.debug("Connection Error", throwable);
-                }
+    private Mono<Void> resumeError(Throwable throwable) {
 
-                logger.error("Connection Error", throwable);
-                return close();
-            })
-            .subscribe();
+        handleConnectionError(throwable);
+        this.requestProcessor.onComplete();
+
+        if (isSslException(throwable)) {
+            logger.debug("Connection Error", throwable);
+        } else {
+            logger.error("Connection Error", throwable);
+        }
+
+        return close();
     }
 
     private static boolean isSslException(Throwable throwable) {
@@ -301,9 +311,10 @@ private static Mono<? extends Void> registerSslHandler(SSLConfig sslConfig, Conn
     public Mono<Void> close() {
         return Mono.defer(() -> {
 
+            drainError(EXPECTED);
             if (this.isClosed.compareAndSet(false, true)) {
 
-                if (!this.connection.channel().isOpen() || this.processId == null) {
+                if (!isConnected() || this.processId == null) {
                     this.connection.dispose();
                     return this.connection.onDispose();
                 }
@@ -326,24 +337,33 @@ public Flux<BackendMessage> exchange(Publisher<FrontendMessage> requests) {
 
         return Mono
             .<Flux<BackendMessage>>create(sink -> {
-                if (this.isClosed.get()) {
-                    sink.error(new IllegalStateException("Cannot exchange messages because the connection is closed"));
-                }
 
                 final AtomicInteger once = new AtomicInteger();
 
                 Flux.from(requests)
                     .subscribe(message -> {
+
+                        if (!isConnected()) {
+                            ReferenceCountUtil.safeRelease(message);
+                            sink.error(new PostgresConnectionClosedException("Cannot exchange messages because the connection is closed"));
+                            return;
+                        }
+
                         if (once.get() == 0 && once.compareAndSet(0, 1)) {
                             synchronized (this) {
                                 this.responseReceivers.add(sink);
                                 this.requests.next(message);
                             }
-                            return;
+                        } else {
+                            this.requests.next(message);
                         }
 
-                        this.requests.next(message);
-                    }, this.requests::error);
+                    }, this.requests::error, () -> {
+
+                        if (!isConnected()) {
+                            sink.error(new PostgresConnectionClosedException("Cannot exchange messages because the connection is closed"));
+                        }
+                    });
 
             })
             .flatMapMany(Function.identity());
@@ -380,6 +400,10 @@ public boolean isConnected() {
             return false;
         }
 
+        if (this.requestProcessor.isDisposed()) {
+            return false;
+        }
+
         Channel channel = this.connection.channel();
         return channel.isOpen();
     }
@@ -399,26 +423,59 @@ private static String toString(List<Field> fields) {
         return joiner.toString();
     }
 
-    private static final class EnsureSubscribersCompleteChannelHandler extends ChannelDuplexHandler {
+    private void handleClose() {
+        if (this.isClosed.compareAndSet(false, true)) {
+            drainError(UNEXPECTED);
+        } else {
+            drainError(EXPECTED);
+        }
+    }
 
-        private final EmitterProcessor<FrontendMessage> requestProcessor;
+    private void handleConnectionError(Throwable error) {
+        drainError(() -> new PostgresConnectionException(error));
+    }
 
-        private final Queue<MonoSink<Flux<BackendMessage>>> responseReceivers;
+    private void drainError(Supplier<? extends Throwable> supplier) {
+        MonoSink<Flux<BackendMessage>> receiver;
+
+        while ((receiver = this.responseReceivers.poll()) != null) {
+            receiver.error(supplier.get());
+        }
+    }
 
-        private EnsureSubscribersCompleteChannelHandler(EmitterProcessor<FrontendMessage> requestProcessor, Queue<MonoSink<Flux<BackendMessage>>> responseReceivers) {
+    private final class EnsureSubscribersCompleteChannelHandler extends ChannelDuplexHandler {
+
+        private final EmitterProcessor<FrontendMessage> requestProcessor;
+
+        private EnsureSubscribersCompleteChannelHandler(EmitterProcessor<FrontendMessage> requestProcessor) {
             this.requestProcessor = requestProcessor;
-            this.responseReceivers = responseReceivers;
+        }
+
+        @Override
+        public void channelInactive(ChannelHandlerContext ctx) throws Exception {
+            super.channelInactive(ctx);
         }
 
         @Override
         public void channelUnregistered(ChannelHandlerContext ctx) throws Exception {
             super.channelUnregistered(ctx);
 
             this.requestProcessor.onComplete();
+            handleClose();
+        }
+    }
 
-            for (MonoSink<Flux<BackendMessage>> responseReceiver = this.responseReceivers.poll(); responseReceiver != null; responseReceiver = this.responseReceivers.poll()) {
-                responseReceiver.success(Flux.empty());
-            }
+    static class PostgresConnectionClosedException extends R2dbcNonTransientResourceException {
+
+        public PostgresConnectionClosedException(String reason) {
+            super(reason);
+        }
+    }
+
+    static class PostgresConnectionException extends R2dbcNonTransientResourceException {
+
+        public PostgresConnectionException(Throwable cause) {
+            super(cause);
         }
     }
 
diff --git a/src/test/java/io/r2dbc/postgresql/client/ReactorNettyClientTest.java b/src/test/java/io/r2dbc/postgresql/client/ReactorNettyClientTest.java
@@ -26,6 +26,7 @@
 import io.r2dbc.postgresql.message.backend.DataRow;
 import io.r2dbc.postgresql.message.backend.NotificationResponse;
 import io.r2dbc.postgresql.message.backend.RowDescription;
+import io.r2dbc.postgresql.message.frontend.FrontendMessage;
 import io.r2dbc.postgresql.message.frontend.Query;
 import io.r2dbc.postgresql.util.PostgresqlServerExtension;
 import io.r2dbc.spi.R2dbcNonTransientResourceException;
@@ -36,28 +37,43 @@
 import org.junit.jupiter.api.Test;
 import org.junit.jupiter.api.TestInstance;
 import org.junit.jupiter.api.extension.RegisterExtension;
+import org.springframework.util.ReflectionUtils;
+import reactor.core.publisher.EmitterProcessor;
 import reactor.core.publisher.Flux;
 import reactor.core.publisher.Mono;
 import reactor.core.publisher.UnicastProcessor;
+import reactor.netty.Connection;
 import reactor.test.StepVerifier;
 
 import javax.net.ssl.HostnameVerifier;
 import javax.net.ssl.SSLSession;
+import java.lang.reflect.Field;
 import java.time.Duration;
 import java.util.Arrays;
 import java.util.Collections;
+import java.util.List;
+import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.ExecutionException;
+import java.util.concurrent.TimeUnit;
 import java.util.function.Consumer;
 import java.util.function.Function;
 import java.util.stream.IntStream;
 
 import static org.assertj.core.api.Assertions.assertThat;
 import static org.assertj.core.api.Assertions.assertThatIllegalArgumentException;
+import static org.assertj.core.api.Assertions.fail;
 
 final class ReactorNettyClientTest {
 
     @RegisterExtension
     static final PostgresqlServerExtension SERVER = new PostgresqlServerExtension();
 
+    static final Field CONNECTION = ReflectionUtils.findField(ReactorNettyClient.class, "connection");
+
+    static {
+        ReflectionUtils.makeAccessible(CONNECTION);
+    }
+
     private final ReactorNettyClient client = ReactorNettyClient.connect(SERVER.getHost(), SERVER.getPort())
         .delayUntil(client -> StartupMessageFlow
             .exchange(this.getClass().getName(), m -> new PasswordAuthenticationHandler(SERVER.getPassword(), SERVER.getUsername()), client, SERVER.getDatabase(), SERVER.getUsername(),
@@ -69,7 +85,65 @@ void close() {
         this.client.close()
             .thenMany(this.client.exchange(Mono.empty()))
             .as(StepVerifier::create)
-            .verifyErrorSatisfies(t -> assertThat(t).isInstanceOf(IllegalStateException.class).hasMessage("Cannot exchange messages because the connection is closed"));
+            .verifyErrorSatisfies(t -> assertThat(t).isInstanceOf(R2dbcNonTransientResourceException.class).hasMessage("Cannot exchange messages because the connection is closed"));
+    }
+
+    @Test
+    void disconnectedShouldRejectExchange() {
+
+        Connection connection = (Connection) ReflectionUtils.getField(CONNECTION, this.client);
+        connection.channel().close().awaitUninterruptibly();
+
+        this.client.close()
+            .thenMany(this.client.exchange(Mono.empty()))
+            .as(StepVerifier::create)
+            .verifyErrorSatisfies(t -> assertThat(t).isInstanceOf(R2dbcNonTransientResourceException.class).hasMessage("Cannot exchange messages because the connection is closed"));
+    }
+
+    @Test
+    void shouldCancelExchangeOnCloseFirstMessage() throws Exception {
+
+        Connection connection = (Connection) ReflectionUtils.getField(CONNECTION, this.client);
+
+        EmitterProcessor<FrontendMessage> messages = EmitterProcessor.create();
+        Flux<BackendMessage> query = this.client.exchange(messages);
+        CompletableFuture<List<BackendMessage>> future = query.collectList().toFuture();
+
+        connection.channel().eventLoop().execute(() -> {
+
+            connection.channel().close();
+            messages.onNext(new Query("SELECT value FROM test"));
+        });
+
+        try {
+            future.get(5, TimeUnit.SECONDS);
+            fail("Expected PostgresConnectionClosedException");
+        } catch (ExecutionException e) {
+            assertThat(e).hasCauseInstanceOf(ReactorNettyClient.PostgresConnectionClosedException.class).hasMessageContaining("Cannot exchange messages");
+        }
+    }
+
+    @Test
+    void shouldCancelExchangeOnCloseInFlight() throws Exception {
+
+        Connection connection = (Connection) ReflectionUtils.getField(CONNECTION, this.client);
+
+        EmitterProcessor<FrontendMessage> messages = EmitterProcessor.create();
+        Flux<BackendMessage> query = this.client.exchange(messages);
+        CompletableFuture<List<BackendMessage>> future = query.doOnNext(ignore -> {
+            connection.channel().close();
+            messages.onNext(new Query("SELECT value FROM test"));
+
+        }).collectList().toFuture();
+
+        messages.onNext(new Query("SELECT value FROM test;SELECT value FROM test;"));
+
+        try {
+            future.get(5, TimeUnit.SECONDS);
+            fail("Expected PostgresConnectionClosedException");
+        } catch (ExecutionException e) {
+            assertThat(e).hasCauseInstanceOf(ReactorNettyClient.PostgresConnectionClosedException.class).hasMessageContaining("Connection closed");
+        }
     }
 
     @AfterEach
