Add support for configuring PreparedStatementCache

src/main/java/io/r2dbc/mssql/LRUPreparedStatementCache.java

@@ -0,0 +1,102 @@
+/*
+ * Copyright 2018-2021 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.mssql;
+
+import java.util.LinkedHashMap;
+import java.util.Map;
+import java.util.function.Function;
+
+import static io.r2dbc.mssql.util.Assert.isTrue;
+import static io.r2dbc.mssql.util.Assert.requireNonNull;
+
+/**
+ * {@link PreparedStatementCache} implementation that maintains a simple "least recently used" cache.
+ * By default, this cache has a maximum size of 32.
+ *
+ * @author Suraj Vijayakumar
+ */
+class LRUPreparedStatementCache implements PreparedStatementCache {
+
+    private static final int DEFAULT_MAX_SIZE = 32;
+
+    private final Map<String, Integer> handleCache;
+
+    private final Map<String, Object> sqlCache;
+
+    public LRUPreparedStatementCache() {
+        this(DEFAULT_MAX_SIZE);
+    }
+
+    public LRUPreparedStatementCache(int maxSize) {
+        isTrue(maxSize > 0, "Max cache size must be > 0");
+
+        handleCache = new LRUCache<>(maxSize);
+        sqlCache = new LRUCache<>(maxSize);
+    }
+
+    @Override
+    public int getHandle(String sql, Binding binding) {
+        requireNonNull(sql, "SQL query must not be null");
+        requireNonNull(binding, "Binding must not be null");
+
+        String key = createKey(sql, binding);
+        return handleCache.getOrDefault(key, UNPREPARED);
+    }
+
+    @Override
+    public void putHandle(int handle, String sql, Binding binding) {
+        requireNonNull(sql, "SQL query must not be null");
+        requireNonNull(binding, "Binding must not be null");
+
+        String key = createKey(sql, binding);
+        handleCache.put(key, handle);
+    }
+
+    @SuppressWarnings("unchecked")
+    @Override
+    public <T> T getParsedSql(String sql, Function<String, T> parseFunction) {
+        requireNonNull(sql, "SQL query must not be null");
+        requireNonNull(parseFunction, "Parse function must not be null");
+
+        return (T) sqlCache.computeIfAbsent(sql, parseFunction);
+    }
+
+    @Override
+    public int size() {
+        return handleCache.size();
+    }
+
+    private static String createKey(String sql, Binding binding) {
+        return sql + "-" + binding.getFormalParameters();
+    }
+
+    private static class LRUCache<K, V> extends LinkedHashMap<K, V> {
+
+        private final int maxSize;
+
+        LRUCache(int maxSize) {
+            super(16, .75f, true);
+
+            this.maxSize = maxSize;
+        }
+
+        @Override
+        protected boolean removeEldestEntry(Map.Entry<K, V> eldest) {
+            return size() > maxSize;
+        }
+    }
+}

src/main/java/io/r2dbc/mssql/MssqlConnectionConfiguration.java

@@ -86,6 +86,8 @@ public final class MssqlConnectionConfiguration {
 
     private final Predicate<String> preferCursoredExecution;
 
+    private final PreparedStatementCache preparedStatementCache;
+
     @Nullable
     private final Duration lockWaitTimeout;
 
@@ -118,8 +120,8 @@ public final class MssqlConnectionConfiguration {
     private final String username;
 
     private MssqlConnectionConfiguration(@Nullable String applicationName, @Nullable UUID connectionId, Duration connectTimeout, @Nullable String database, String host, String hostNameInCertificate,
-                                         @Nullable Duration lockWaitTimeout, CharSequence password, Predicate<String> preferCursoredExecution, int port, boolean sendStringParametersAsUnicode,
-                                         boolean ssl,
+                                         @Nullable Duration lockWaitTimeout, CharSequence password, Predicate<String> preferCursoredExecution, PreparedStatementCache preparedStatementCache,
+                                         int port, boolean sendStringParametersAsUnicode, boolean ssl,
                                          Function<SslContextBuilder, SslContextBuilder> sslContextBuilderCustomizer,
                                          @Nullable Function<SslContextBuilder, SslContextBuilder> sslTunnelSslContextBuilderCustomizer, boolean tcpKeepAlive, boolean tcpNoDelay,
                                          boolean trustServerCertificate, @Nullable File trustStore, @Nullable String trustStoreType,
@@ -134,6 +136,7 @@ private MssqlConnectionConfiguration(@Nullable String applicationName, @Nullable
         this.lockWaitTimeout = lockWaitTimeout;
         this.password = Assert.requireNonNull(password, "password must not be null");
         this.preferCursoredExecution = Assert.requireNonNull(preferCursoredExecution, "preferCursoredExecution must not be null");
+        this.preparedStatementCache = Assert.requireNonNull(preparedStatementCache, "preparedStatementCache must not be null");
         this.port = port;
         this.sendStringParametersAsUnicode = sendStringParametersAsUnicode;
         this.ssl = ssl;
@@ -182,7 +185,7 @@ MssqlConnectionConfiguration withRedirect(Redirect redirect) {
 
         return new MssqlConnectionConfiguration(this.applicationName, this.connectionId, this.connectTimeout, this.database, redirectServerName, hostNameInCertificate, this.lockWaitTimeout,
             this.password,
-            this.preferCursoredExecution, redirect.getPort(), this.sendStringParametersAsUnicode, this.ssl, this.sslContextBuilderCustomizer,
+            this.preferCursoredExecution, this.preparedStatementCache, redirect.getPort(), this.sendStringParametersAsUnicode, this.ssl, this.sslContextBuilderCustomizer,
             this.sslTunnelSslContextBuilderCustomizer, this.tcpKeepAlive, this.tcpNoDelay, this.trustServerCertificate, this.trustStore, this.trustStoreType, this.trustStorePassword, this.username);
     }
 
@@ -192,7 +195,7 @@ ClientConfiguration toClientConfiguration() {
     }
 
     ConnectionOptions toConnectionOptions() {
-        return new ConnectionOptions(this.preferCursoredExecution, new DefaultCodecs(), new IndefinitePreparedStatementCache(), this.sendStringParametersAsUnicode);
+        return new ConnectionOptions(this.preferCursoredExecution, new DefaultCodecs(), this.preparedStatementCache, this.sendStringParametersAsUnicode);
     }
 
     @Override
@@ -355,6 +358,8 @@ public static final class Builder {
 
         private Predicate<String> preferCursoredExecution = sql -> false;
 
+        private PreparedStatementCache preparedStatementCache = new IndefinitePreparedStatementCache();
+
         private CharSequence password;
 
         private int port = DEFAULT_PORT;
@@ -551,6 +556,17 @@ public Builder preferCursoredExecution(Predicate<String> preference) {
             return this;
         }
 
+        /**
+         * Configures the {@link PreparedStatementCache}. By default, uses {@link IndefinitePreparedStatementCache}.
+         *
+         * @param cache the cache implementation to use (must not be null).
+         * @return this {@link Builder}
+         */
+        public Builder preparedStatementCache(PreparedStatementCache cache) {
+            this.preparedStatementCache = Assert.requireNonNull(cache, "Prepared statement cache must not be null");
+            return this;
+        }
+
         /**
          * Configure the port. Defaults to {@code 5432}.
          *
@@ -714,7 +730,7 @@ public MssqlConnectionConfiguration build() {
 
             return new MssqlConnectionConfiguration(this.applicationName, this.connectionId, this.connectTimeout, this.database, this.host, this.hostNameInCertificate, this.lockWaitTimeout,
                 this.password,
-                this.preferCursoredExecution, this.port, this.sendStringParametersAsUnicode, this.ssl, this.sslContextBuilderCustomizer,
+                this.preferCursoredExecution, this.preparedStatementCache, this.port, this.sendStringParametersAsUnicode, this.ssl, this.sslContextBuilderCustomizer,
                 this.sslTunnelSslContextBuilderCustomizer, this.tcpKeepAlive,
                 this.tcpNoDelay, this.trustServerCertificate, this.trustStore,
                 this.trustStoreType, this.trustStorePassword, this.username);

src/main/java/io/r2dbc/mssql/MssqlConnectionFactoryProvider.java

@@ -71,6 +71,16 @@ public final class MssqlConnectionFactoryProvider implements ConnectionFactoryPr
      */
     public static final Option<Object> PREFER_CURSORED_EXECUTION = Option.valueOf("preferCursoredExecution");
 
+    /**
+     * Configures the prepared statement cache to use.
+     * The value can be an {@link Integer}, a {@link PreparedStatementCache} or a {@link Class class name}.
+     * <p>
+     * A value of 0 disables the cache ({@link NoPreparedStatementCache}).<br/>
+     * A value of -1 (or any negative number) caches items indefinitely ({@link IndefinitePreparedStatementCache}) - this is the default value.<br/>
+     * Any other integer creates an LRU cache of that size ({@link LRUPreparedStatementCache}).<br/>
+     */
+    public static final Option<Object> PREPARED_STATEMENT_CACHE = Option.valueOf("preparedStatementCache");
+
     /**
      * Configure whether to send character data as unicode (NVARCHAR, NCHAR, NTEXT) or whether to use the database encoding. Enabled by default.
      * If disabled, {@link CharSequence} data is sent using the database-specific collation such as ASCII/MBCS instead of Unicode.
@@ -162,6 +172,7 @@ public MssqlConnectionFactory create(ConnectionFactoryOptions connectionFactoryO
         mapper.from(LOCK_WAIT_TIMEOUT).map(OptionMapper::toDuration).to(builder::lockWaitTimeout);
         mapper.from(PORT).map(OptionMapper::toInteger).to(builder::port);
         mapper.from(PREFER_CURSORED_EXECUTION).map(OptionMapper::toStringPredicate).to(builder::preferCursoredExecution);
+        mapper.from(PREPARED_STATEMENT_CACHE).map(OptionMapper::toPreparedStatementCache).to(builder::preparedStatementCache);
         mapper.from(SEND_STRING_PARAMETERS_AS_UNICODE).map(OptionMapper::toBoolean).to(builder::sendStringParametersAsUnicode);
         mapper.from(SSL).to(builder::enableSsl);
         mapper.fromTyped(SSL_CONTEXT_BUILDER_CUSTOMIZER).to(builder::sslContextBuilderCustomizer);

src/main/java/io/r2dbc/mssql/NoPreparedStatementCache.java

@@ -0,0 +1,46 @@
+/*
+ * Copyright 2018-2021 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.mssql;
+
+import java.util.function.Function;
+
+/**
+ * {@link PreparedStatementCache} implementation that does not cache anything.
+ *
+ * @author Suraj Vijayakumar
+ */
+class NoPreparedStatementCache implements PreparedStatementCache {
+
+    @Override
+    public int getHandle(String sql, Binding binding) {
+        return PreparedStatementCache.UNPREPARED;
+    }
+
+    @Override
+    public void putHandle(int handle, String sql, Binding binding) {
+    }
+
+    @Override
+    public <T> T getParsedSql(String sql, Function<String, T> parseFunction) {
+        return parseFunction.apply(sql);
+    }
+
+    @Override
+    public int size() {
+        return 0;
+    }
+}

src/main/java/io/r2dbc/mssql/OptionMapper.java

@@ -56,7 +56,6 @@ public static OptionMapper create(ConnectionFactoryOptions options) {
      * Construct a new {@link Source} for a {@link Option}. Options without a value are not bound or mapped in the later stages of {@link Source}.
      *
      * @param option the option to apply.
-     * @param <T>    inferred option type.
      * @return the source object.
      */
     public Source<Object> from(Option<?> option) {
@@ -192,6 +191,25 @@ static Predicate<String> toStringPredicate(Object value) {
         throw new IllegalArgumentException(String.format("Cannot convert value %s to Predicate", value));
     }
 
+    /**
+     * Parse an {@link Option} to a {@link PreparedStatementCache}.
+     */
+    static PreparedStatementCache toPreparedStatementCache(Object value) {
+        if (value instanceof PreparedStatementCache) {
+            return (PreparedStatementCache) value;
+        }
+
+        if (value instanceof Integer) {
+            return toPreparedStatementCache((Integer) value);
+        }
+
+        if (value instanceof String) {
+            return toPreparedStatementCache((String) value);
+        }
+
+        throw new IllegalArgumentException(String.format("Cannot convert value %s to PreparedStatementCache", value));
+    }
+
     /**
      * Parse an {@link Option} to {@link UUID}.
      */
@@ -208,6 +226,35 @@ static UUID toUuid(Object value) {
         throw new IllegalArgumentException(String.format("Cannot convert value %s to UUID", value));
     }
 
+    private static PreparedStatementCache toPreparedStatementCache(Integer value) {
+        if (value < 0) {
+            return new IndefinitePreparedStatementCache();
+        } else if (value == 0) {
+            return new NoPreparedStatementCache();
+        } else {
+            return new LRUPreparedStatementCache(value);
+        }
+    }
+
+    private static PreparedStatementCache toPreparedStatementCache(String value) {
+        try {
+            Integer number = Integer.parseInt(value);
+            return toPreparedStatementCache(number);
+        } catch (NumberFormatException ignore) {
+            // ignore - value is not a number
+        }
+
+        try {
+            Object cache = Class.forName(value).getDeclaredConstructor().newInstance();
+            if (cache instanceof PreparedStatementCache) {
+                return (PreparedStatementCache) cache;
+            }
+            throw new IllegalArgumentException("Value '" + value + "' must be an instance of PreparedStatementCache");
+        } catch (ReflectiveOperationException e) {
+            throw new IllegalArgumentException("Cannot instantiate '" + value + "'", e);
+        }
+    }
+
     public interface Source<T> {
 
         /**

src/main/java/io/r2dbc/mssql/PreparedStatementCache.java

@@ -20,10 +20,13 @@
 
 /**
  * Cache for prepared statements.
+ * <p>
+ * Implementations will need a default no-arg constructor for the {@link OptionMapper}
+ * to instantiate them from a discovery option.
  *
  * @author Mark Paluch
  */
-interface PreparedStatementCache {
+public interface PreparedStatementCache {
 
     /**
      * Marker for no prepared statement found/no prepared statement.