Prepared statement cache size configuration

Added preparedStatementCacheQueries parameter, which controls cache behaviour.
When cache limit is reached last used statement is closed before parsing new one.
When cache limit is 0 caching of statement is disabled and unnamed statements are used in extended message flow.
Negative values stand for IndefiniteStatementCache (default one).

[#223][#225][#277]

Author: Squiry

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

@@ -0,0 +1,54 @@
+/*
+ * Copyright 2019 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;
+
+import io.r2dbc.postgresql.client.Binding;
+import io.r2dbc.postgresql.client.Client;
+import io.r2dbc.postgresql.client.ExtendedQueryMessageFlow;
+import io.r2dbc.postgresql.util.Assert;
+import reactor.core.publisher.Mono;
+
+class DisabledStatementCache implements StatementCache {
+
+    private static final String UNNAMED_STATEMENT_NAME = "";
+
+    private final Client client;
+
+    DisabledStatementCache(Client client) {
+        this.client = Assert.requireNonNull(client, "client must not be null");
+    }
+
+    @Override
+    public Mono<String> getName(Binding binding, String sql) {
+        Assert.requireNonNull(binding, "binding must not be null");
+        Assert.requireNonNull(sql, "sql must not be null");
+        String name = UNNAMED_STATEMENT_NAME;
+
+        ExceptionFactory factory = ExceptionFactory.withSql(name);
+        return ExtendedQueryMessageFlow
+            .parse(this.client, name, sql, binding.getParameterTypes())
+            .handle(factory::handleErrorResponse)
+            .then(Mono.just(name));
+    }
+
+    @Override
+    public String toString() {
+        return "DisabledStatementCache{" +
+            "client=" + this.client +
+            '}';
+    }
+}

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

@@ -0,0 +1,98 @@
+/*
+ * Copyright 2019 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;
+
+import io.r2dbc.postgresql.client.Binding;
+import io.r2dbc.postgresql.client.Client;
+import io.r2dbc.postgresql.client.ExtendedQueryMessageFlow;
+import io.r2dbc.postgresql.util.Assert;
+import reactor.core.publisher.Mono;
+import reactor.util.function.Tuple2;
+import reactor.util.function.Tuples;
+
+import java.util.Arrays;
+import java.util.HashMap;
+import java.util.LinkedHashMap;
+import java.util.List;
+import java.util.concurrent.atomic.AtomicInteger;
+import java.util.stream.Collectors;
+
+public class LimitedStatementCache implements StatementCache {
+
+    private final HashMap<Tuple2<String, List<Integer>>, String> cache = new LinkedHashMap<>(16, 0.75f, true);
+
+    private final Client client;
+
+    private final AtomicInteger counter = new AtomicInteger();
+
+    private final int limit;
+
+    public LimitedStatementCache(Client client, int limit) {
+        this.client = Assert.requireNonNull(client, "client must not be null");
+        if (limit <= 0) {
+            throw new IllegalArgumentException("statement cache limit must be greater than zero");
+        }
+        this.limit = limit;
+    }
+
+    @Override
+    public Mono<String> getName(Binding binding, String sql) {
+        Assert.requireNonNull(binding, "binding must not be null");
+        Assert.requireNonNull(sql, "sql must not be null");
+        Tuple2<String, List<Integer>> key = Tuples.of(sql, Arrays.stream(binding.getParameterTypes()).boxed().collect(Collectors.toList()));
+        String name = this.cache.get(key);
+        if (name != null) {
+            return Mono.just(name);
+        }
+
+        Mono<Void> closeLastStatement = Mono.defer(() -> {
+            if (this.cache.size() < this.limit) {
+                return Mono.empty();
+            }
+            String lastAccessedStatementName = this.cache.values().iterator().next();
+            ExceptionFactory factory = ExceptionFactory.withSql(lastAccessedStatementName);
+            return ExtendedQueryMessageFlow
+                .closeStatement(this.client, lastAccessedStatementName)
+                .handle(factory::handleErrorResponse)
+                .then();
+        });
+
+        return closeLastStatement.then(this.parse(sql, binding.getParameterTypes()))
+            .doOnNext(preparedName -> this.cache.put(key, preparedName));
+    }
+
+    @Override
+    public String toString() {
+        return "LimitedStatementCache{" +
+            "cache=" + this.cache +
+            ", counter=" + this.counter +
+            ", client=" + this.client +
+            ", limit=" + this.limit +
+            '}';
+    }
+
+    private Mono<String> parse(String sql, int[] types) {
+        String name = String.format("S_%d", this.counter.getAndIncrement());
+
+        ExceptionFactory factory = ExceptionFactory.withSql(name);
+        return ExtendedQueryMessageFlow
+            .parse(this.client, name, sql, types)
+            .handle(factory::handleErrorResponse)
+            .then(Mono.just(name))
+            .cache();
+    }
+}

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

@@ -83,12 +83,14 @@ public final class PostgresqlConnectionConfiguration {
 
     private final SSLConfig sslConfig;
 
+    private final int preparedStatementCacheQueries;
+
     private PostgresqlConnectionConfiguration(String applicationName, boolean autodetectExtensions,
                                               @Nullable Duration connectTimeout, @Nullable String database, List<Extension> extensions,
                                               ToIntFunction<String> fetchSize, boolean forceBinary,
                                               @Nullable String host,
                                               @Nullable Map<String, String> options, @Nullable CharSequence password, int port, @Nullable String schema, @Nullable String socket, String username,
-                                              SSLConfig sslConfig) {
+                                              SSLConfig sslConfig, int preparedStatementCacheQueries) {
         this.applicationName = Assert.requireNonNull(applicationName, "applicationName must not be null");
         this.autodetectExtensions = autodetectExtensions;
         this.connectTimeout = connectTimeout;
@@ -108,6 +110,7 @@ private PostgresqlConnectionConfiguration(String applicationName, boolean autode
         this.socket = socket;
         this.username = Assert.requireNonNull(username, "username must not be null");
         this.sslConfig = sslConfig;
+        this.preparedStatementCacheQueries = preparedStatementCacheQueries;
     }
 
     /**
@@ -239,6 +242,21 @@ SSLConfig getSslConfig() {
         return this.sslConfig;
     }
 
+    int getPreparedStatementCacheQueries() {
+        return this.preparedStatementCacheQueries;
+    }
+
+    private static String obfuscate(int length) {
+
+        StringBuilder builder = new StringBuilder();
+
+        for (int i = 0; i < length; i++) {
+            builder.append("*");
+        }
+
+        return builder.toString();
+    }
+
     /**
      * A builder for {@link PostgresqlConnectionConfiguration} instances.
      * <p>
@@ -299,6 +317,8 @@ public static final class Builder {
         @Nullable
         private String username;
 
+        private int preparedStatementCacheQueries = -1;
+
         private Builder() {
         }
 
@@ -346,7 +366,7 @@ public PostgresqlConnectionConfiguration build() {
 
             return new PostgresqlConnectionConfiguration(this.applicationName, this.autodetectExtensions, this.connectTimeout, this.database, this.extensions, this.fetchSize, this.forceBinary,
                 this.host,
-                this.options, this.password, this.port, this.schema, this.socket, this.username, this.createSslConfig());
+                this.options, this.password, this.port, this.schema, this.socket, this.username, this.createSslConfig(), this.preparedStatementCacheQueries);
         }
 
         /**
@@ -589,6 +609,18 @@ public Builder sslPassword(@Nullable CharSequence sslPassword) {
             return this;
         }
 
+        /**
+         * Configure the preparedStatementCacheQueries.
+         *
+         * @param preparedStatementCacheQueries the preparedStatementCacheQueries
+         * @return this {@link Builder}
+         * @throws IllegalArgumentException if {@code username} is {@code null}
+         */
+        public Builder preparedStatementCacheQueries(int preparedStatementCacheQueries) {
+            this.preparedStatementCacheQueries = preparedStatementCacheQueries;
+            return this;
+        }
+
         @Override
         public String toString() {
             return "Builder{" +
@@ -612,6 +644,7 @@ public String toString() {
                 ", sslCert='" + this.sslCert + '\'' +
                 ", sslKey='" + this.sslKey + '\'' +
                 ", sslHostnameVerifier='" + this.sslHostnameVerifier + '\'' +
+                ", preparedStatementCacheQueries='" + this.preparedStatementCacheQueries + '\'' +
                 '}';
         }
 

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

@@ -158,7 +158,7 @@ private Mono<PostgresqlConnection> doCreateConnection(boolean forReplication, @N
             .flatMap(client -> {
 
                 DefaultCodecs codecs = new DefaultCodecs(client.getByteBufAllocator());
-                StatementCache statementCache = new IndefiniteStatementCache(client);
+                StatementCache statementCache = StatementCache.fromPreparedStatementCacheQueries(client, this.configuration.getPreparedStatementCacheQueries());
 
                 // early connection object to retrieve initialization details
                 PostgresqlConnection earlyConnection = new PostgresqlConnection(client, codecs, DefaultPortalNameSupplier.INSTANCE, statementCache, IsolationLevel.READ_COMMITTED,

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

@@ -25,7 +25,6 @@
 
 import javax.net.ssl.HostnameVerifier;
 import java.util.LinkedHashMap;
-import java.util.Locale;
 import java.util.Map;
 import java.util.function.Function;
 
@@ -125,6 +124,12 @@ public final class PostgresqlConnectionFactoryProvider implements ConnectionFact
      */
     public static final Option<String> SSL_ROOT_CERT = Option.valueOf("sslRootCert");
 
+    /**
+     * Determine the number of queries that are cached in each connection.
+     * The default is -1, meaning there's no limit. The value of 0 disables the cache.
+     */
+    public static final Option<Integer> PREPARED_STATEMENT_CACHE_QUERIES = Option.valueOf("preparedStatementCacheQueries");
+
     /**
      * Connection options which are applied once after the connection has been created.
      */
@@ -188,6 +193,11 @@ private static PostgresqlConnectionConfiguration createConfiguration(ConnectionF
             builder.forceBinary(convertToBoolean(forceBinary));
         }
 
+        Object preparedStatementCacheQueries = connectionFactoryOptions.getValue(PREPARED_STATEMENT_CACHE_QUERIES);
+        if (preparedStatementCacheQueries != null) {
+            builder.preparedStatementCacheQueries(convertToInt(preparedStatementCacheQueries));
+        }
+
         Object options = connectionFactoryOptions.getValue(Option.valueOf("options"));
         if (options != null) {
             builder.options(convertToMap(options));
@@ -274,7 +284,7 @@ private static boolean convertToBoolean(Object value) {
     private static int convertToInt(Object value) {
         return value instanceof Integer ? (int) value : Integer.parseInt(value.toString());
     }
-    
+
     @SuppressWarnings("unchecked")
     private static Map<String, String> convertToMap(Object options) {
         if (options instanceof Map) {

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

@@ -17,10 +17,20 @@
 package io.r2dbc.postgresql;
 
 import io.r2dbc.postgresql.client.Binding;
+import io.r2dbc.postgresql.client.Client;
 import reactor.core.publisher.Mono;
 
 interface StatementCache {
 
     Mono<String> getName(Binding binding, String sql);
 
+    static StatementCache fromPreparedStatementCacheQueries(Client client, int preparedStatementCacheQueries) {
+        if (preparedStatementCacheQueries < 0) {
+            return new IndefiniteStatementCache(client);
+        }
+        if (preparedStatementCacheQueries == 0) {
+            return new DisabledStatementCache(client);
+        }
+        return new LimitedStatementCache(client, preparedStatementCacheQueries);
+    }
 }

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

@@ -193,8 +193,7 @@ public Mono<Void> close() {
             if (this.isClosed.compareAndSet(false, true)) {
 
                 if (!connected || this.processId == null) {
-                    this.connection.dispose();
-                    return this.connection.onDispose();
+                    return closeConnection();
                 }
 
                 return Flux.just(Terminate.INSTANCE)
@@ -209,6 +208,11 @@ public Mono<Void> close() {
         });
     }
 
+    private Mono<? extends Void> closeConnection() {
+        this.connection.dispose();
+        return this.connection.onDispose();
+    }
+
     @Override
     public Flux<BackendMessage> exchange(Predicate<BackendMessage> takeUntil, Publisher<FrontendMessage> requests) {
         Assert.requireNonNull(takeUntil, "takeUntil must not be null");