Merge branch 'Squiry-concurrent-exchange-fix'

[resolves pgjdbc#207][pgjdbc#230]

Author: mp911de

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

@@ -68,6 +68,7 @@ default Flux<BackendMessage> exchange(Publisher<FrontendMessage> requests) {
      * @param requests the publisher of outbound messages
      * @return a {@link Flux} of incoming messages that ends with the end of the frame (i.e. reception of a {@link ReadyForQuery} message.
      * @throws IllegalArgumentException if {@code requests} is {@code null}
+     * @since 0.9
      */
     Flux<BackendMessage> exchange(Predicate<BackendMessage> takeUntil, Publisher<FrontendMessage> requests);
 
@@ -76,6 +77,7 @@ default Flux<BackendMessage> exchange(Publisher<FrontendMessage> requests) {
      *
      * @param message outbound message
      * @throws IllegalArgumentException if {@code message} is {@code null}
+     * @since 0.9
      */
     void send(FrontendMessage message);
 

src/main/java/io/r2dbc/postgresql/client/ReactorNettyClient.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.
@@ -29,7 +29,6 @@
 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;
@@ -72,9 +71,9 @@
 import java.util.Queue;
 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.Predicate;
 import java.util.function.Supplier;
 
@@ -99,11 +98,11 @@ public final class ReactorNettyClient implements Client {
 
     private final Connection connection;
 
-    private final EmitterProcessor<FrontendMessage> requestProcessor = EmitterProcessor.create(false);
+    private final EmitterProcessor<Publisher<FrontendMessage>> requestProcessor = EmitterProcessor.create(false);
 
-    private final FluxSink<FrontendMessage> requests = this.requestProcessor.sink();
+    private final FluxSink<Publisher<FrontendMessage>> requests = this.requestProcessor.sink();
 
-    private final Queue<ResponseReceiver> responseReceivers = Queues.<ResponseReceiver>unbounded().get();
+    private final Queue<Conversation> conversations = Queues.<Conversation>unbounded().get();
 
     private final DirectProcessor<NotificationResponse> notificationProcessor = DirectProcessor.create();
 
@@ -140,11 +139,11 @@ private ReactorNettyClient(Connection connection) {
                 handleConnectionError(throwable);
             })
             .handle((message, sink) -> {
-                ResponseReceiver receiver = this.responseReceivers.peek();
+                Conversation receiver = this.conversations.peek();
                 if (receiver != null) {
                     if (receiver.takeUntil.test(message)) {
                         receiver.sink.complete();
-                        this.responseReceivers.poll();
+                        this.conversations.poll();
                     } else {
                         receiver.sink.next(message);
                     }
@@ -154,6 +153,7 @@ private ReactorNettyClient(Connection connection) {
             .then();
 
         Mono<Void> request = this.requestProcessor
+            .concatMap(Function.identity())
             .flatMap(message -> {
                 if (DEBUG_ENABLED) {
                     logger.debug("Request:  {}", message);
@@ -173,47 +173,61 @@ private ReactorNettyClient(Connection connection) {
     }
 
     @Override
-    public Flux<BackendMessage> exchange(Predicate<BackendMessage> takeUntil, Publisher<FrontendMessage> requests) {
-        Assert.requireNonNull(requests, "requests must not be null");
+    public Disposable addNotificationListener(Consumer<NotificationResponse> consumer) {
+        return this.notificationProcessor.subscribe(consumer);
+    }
 
-        return Flux
-            .create(sink -> {
+    @Override
+    public Mono<Void> close() {
+        return Mono.defer(() -> {
 
-                final AtomicInteger once = new AtomicInteger();
+            drainError(EXPECTED);
+            if (this.isClosed.compareAndSet(false, true)) {
 
-                Flux.from(requests)
-                    .subscribe(message -> {
+                if (!isConnected() || this.processId == null) {
+                    this.connection.dispose();
+                    return this.connection.onDispose();
+                }
 
-                        if (!isConnected()) {
-                            ReferenceCountUtil.safeRelease(message);
-                            sink.error(new PostgresConnectionClosedException("Cannot exchange messages because the connection is closed"));
-                            return;
-                        }
+                return Flux.just(Terminate.INSTANCE)
+                    .doOnNext(message -> logger.debug("Request:  {}", message))
+                    .concatMap(message -> this.connection.outbound().send(message.encode(this.connection.outbound().alloc())))
+                    .then()
+                    .doOnSuccess(v -> this.connection.dispose())
+                    .then(this.connection.onDispose());
+            }
 
-                        if (once.get() == 0 && once.compareAndSet(0, 1)) {
-                            synchronized (this) {
-                                this.responseReceivers.add(new ResponseReceiver(sink, takeUntil));
-                                this.requests.next(message);
-                            }
-                        } else {
-                            this.requests.next(message);
-                        }
+            return Mono.empty();
+        });
+    }
 
-                    }, this.requests::error, () -> {
+    @Override
+    public Flux<BackendMessage> exchange(Predicate<BackendMessage> takeUntil, Publisher<FrontendMessage> requests) {
+        Assert.requireNonNull(takeUntil, "takeUntil must not be null");
+        Assert.requireNonNull(requests, "requests must not be null");
 
+        return Flux
+            .create(sink -> {
+                if (!isConnected()) {
+                    sink.error(new PostgresConnectionClosedException("Cannot exchange messages because the connection is closed"));
+                    return;
+                }
+                synchronized (this) {
+                    this.conversations.add(new Conversation(sink, takeUntil));
+                    this.requests.next(Flux.from(requests).doOnNext(m -> {
                         if (!isConnected()) {
                             sink.error(new PostgresConnectionClosedException("Cannot exchange messages because the connection is closed"));
                         }
-                    });
-
+                    }));
+                }
             });
     }
 
     @Override
     public void send(FrontendMessage message) {
         Assert.requireNonNull(message, "requests must not be null");
 
-        this.requests.next(message);
+        this.requests.next(Mono.just(message));
     }
 
     private Mono<Void> resumeError(Throwable throwable) {
@@ -370,38 +384,6 @@ private static Mono<? extends Void> registerSslHandler(SSLConfig sslConfig, Conn
         return Mono.empty();
     }
 
-    @Override
-    public Mono<Void> close() {
-        return Mono.defer(() -> {
-
-            drainError(EXPECTED);
-            if (this.isClosed.compareAndSet(false, true)) {
-
-                if (!isConnected() || this.processId == null) {
-                    this.connection.dispose();
-                    return this.connection.onDispose();
-                }
-
-                return Flux.just(Terminate.INSTANCE)
-                    .doOnNext(message -> logger.debug("Request:  {}", message))
-                    .concatMap(message -> this.connection.outbound().send(message.encode(this.connection.outbound().alloc())))
-                    .then()
-                    .doOnSuccess(v -> this.connection.dispose())
-                    .then(this.connection.onDispose());
-            }
-
-            return Mono.empty();
-        });
-    }
-
-    private void drainError(Supplier<? extends Throwable> supplier) {
-        ResponseReceiver receiver;
-
-        while ((receiver = this.responseReceivers.poll()) != null) {
-            receiver.sink.error(supplier.get());
-        }
-    }
-
     @Override
     public ByteBufAllocator getByteBufAllocator() {
         return this.byteBufAllocator;
@@ -441,11 +423,6 @@ public boolean isConnected() {
         return channel.isOpen();
     }
 
-    @Override
-    public Disposable addNotificationListener(Consumer<NotificationResponse> consumer) {
-        return this.notificationProcessor.subscribe(consumer);
-    }
-
     private static String toString(List<Field> fields) {
 
         StringJoiner joiner = new StringJoiner(", ");
@@ -468,23 +445,37 @@ private void handleConnectionError(Throwable error) {
         drainError(() -> new PostgresConnectionException(error));
     }
 
-    private static class ResponseReceiver {
+    private void drainError(Supplier<? extends Throwable> supplier) {
+        Conversation receiver;
+
+        while ((receiver = this.conversations.poll()) != null) {
+            receiver.sink.error(supplier.get());
+        }
+    }
+
+    /**
+     * Value object representing a single conversation. The driver permits a single conversation at a time to ensure that request messages get routed to the proper response receiver and do not leak
+     * into other conversations. A conversation must be finished in the sense that the {@link Publisher} of {@link FrontendMessage} has completed before the next conversation is started.
+     * <p>
+     * A single conversation can make use of pipelining.
+     */
+    private static class Conversation {
 
         private final FluxSink<BackendMessage> sink;
 
         private final Predicate<BackendMessage> takeUntil;
 
-        private ResponseReceiver(FluxSink<BackendMessage> sink, Predicate<BackendMessage> takeUntil) {
+        private Conversation(FluxSink<BackendMessage> sink, Predicate<BackendMessage> takeUntil) {
             this.sink = sink;
             this.takeUntil = takeUntil;
         }
     }
 
     private final class EnsureSubscribersCompleteChannelHandler extends ChannelDuplexHandler {
 
-        private final EmitterProcessor<FrontendMessage> requestProcessor;
+        private final EmitterProcessor<Publisher<FrontendMessage>> requestProcessor;
 
-        private EnsureSubscribersCompleteChannelHandler(EmitterProcessor<FrontendMessage> requestProcessor) {
+        private EnsureSubscribersCompleteChannelHandler(EmitterProcessor<Publisher<FrontendMessage>> requestProcessor) {
             this.requestProcessor = requestProcessor;
         }
 

src/test/java/io/r2dbc/postgresql/client/ReactorNettyClientIntegrationTests.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.
@@ -23,13 +23,20 @@
 import io.r2dbc.postgresql.PostgresqlConnectionFactory;
 import io.r2dbc.postgresql.api.PostgresqlConnection;
 import io.r2dbc.postgresql.authentication.PasswordAuthenticationHandler;
+import io.r2dbc.postgresql.message.Format;
 import io.r2dbc.postgresql.message.backend.BackendMessage;
+import io.r2dbc.postgresql.message.backend.BindComplete;
 import io.r2dbc.postgresql.message.backend.CommandComplete;
 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.Bind;
+import io.r2dbc.postgresql.message.frontend.Describe;
+import io.r2dbc.postgresql.message.frontend.Execute;
+import io.r2dbc.postgresql.message.frontend.ExecutionType;
 import io.r2dbc.postgresql.message.frontend.FrontendMessage;
 import io.r2dbc.postgresql.message.frontend.Query;
+import io.r2dbc.postgresql.message.frontend.Sync;
 import io.r2dbc.postgresql.util.PostgresqlServerExtension;
 import io.r2dbc.spi.R2dbcNonTransientResourceException;
 import io.r2dbc.spi.R2dbcPermissionDeniedException;
@@ -62,6 +69,7 @@
 import java.util.function.Function;
 import java.util.stream.IntStream;
 
+import static io.r2dbc.postgresql.type.PostgresqlObjectId.INT4;
 import static org.assertj.core.api.Assertions.assertThat;
 import static org.assertj.core.api.Assertions.assertThatIllegalArgumentException;
 import static org.assertj.core.api.Assertions.fail;
@@ -274,6 +282,56 @@ void parallelExchange() {
             .verifyComplete();
     }
 
+    @Test
+    void parallelExchangeExtendedFlow() {
+        ExtendedQueryMessageFlow.parse(this.client, "S_1", "SELECT $1", Arrays.asList(INT4.getObjectId()))
+            .as(StepVerifier::create)
+            .verifyComplete();
+
+        this.client.exchange(Mono.just(Sync.INSTANCE))
+            .as(StepVerifier::create)
+            .verifyComplete();
+
+        FrontendMessage bind1 = new Bind(
+            "P_1",
+            Collections.singletonList(Format.FORMAT_BINARY),
+            Collections.singletonList(this.client.getByteBufAllocator().buffer(4).writeInt(42)),
+            Collections.singletonList(Format.FORMAT_BINARY),
+            "S_1"
+        );
+        Describe describe1 = new Describe("P_1", ExecutionType.PORTAL);
+        Execute execute1 = new Execute("P_1", Integer.MAX_VALUE);
+        FrontendMessage bind2 = new Bind(
+            "P_2",
+            Collections.singletonList(Format.FORMAT_BINARY),
+            Collections.singletonList(this.client.getByteBufAllocator().buffer(4).writeInt(42)),
+            Collections.singletonList(Format.FORMAT_BINARY),
+            "S_1"
+        );
+        Describe describe2 = new Describe("P_2", ExecutionType.PORTAL);
+        Execute execute2 = new Execute("P_2", Integer.MAX_VALUE);
+
+        Flux<FrontendMessage> flow1 = Flux.just(bind1, describe1, execute1, Sync.INSTANCE).delayElements(Duration.ofMillis(10));
+        Flux<FrontendMessage> flow2 = Flux.just(bind2, describe2, execute2, Sync.INSTANCE).delayElements(Duration.ofMillis(20));
+
+        // Actual response sequence is BindComplete/RowDescription/DataRow/CommandComplete for each of the extended flows
+        // Zipping and flattening makes it appear in the assertions that it would be BindComplete/BindComplete/RowDescription/RowDescription…
+
+        this.datarowCleanup(Flux.zip(this.client.exchange(flow1), this.client.exchange(flow2))
+            .flatMapIterable(t -> Arrays.asList(t.getT1(), t.getT2()))
+        )
+            .as(StepVerifier::create)
+            .assertNext(message -> assertThat(message).isInstanceOf(BindComplete.class))
+            .assertNext(message -> assertThat(message).isInstanceOf(BindComplete.class))
+            .assertNext(message -> assertThat(message).isInstanceOf(RowDescription.class))
+            .assertNext(message -> assertThat(message).isInstanceOf(RowDescription.class))
+            .assertNext(message -> assertThat(message).isInstanceOf(DataRow.class))
+            .assertNext(message -> assertThat(message).isInstanceOf(DataRow.class))
+            .expectNext(new CommandComplete("SELECT", null, 1))
+            .expectNext(new CommandComplete("SELECT", null, 1))
+            .verifyComplete();
+    }
+
     @Test
     void timeoutTest() {
         PostgresqlConnectionFactory postgresqlConnectionFactory = new PostgresqlConnectionFactory(PostgresqlConnectionConfiguration.builder()