JAVA-2910: Add a configuration option to support strong values for prepared statements cache

Co-authored-by: Alexandre Dutra <[email protected]>

Author: hnlzq163

changelog/README.md

@@ -2,8 +2,9 @@
 
 <!-- Note: contrary to 3.x, insert new entries *first* in their section -->
 
-### 4.11.1
+### 4.11.1 (in progress)
 
+- [bug] JAVA-2910: Add a configuration option to support strong values for prepared statements cache
 - [bug] JAVA-2936: Support Protocol V6
 - [bug] JAVA-2934: Handle empty non-final pages in ReactiveResultSetSubscription
 

core/src/main/java/com/datastax/oss/driver/api/core/config/DefaultDriverOption.java

@@ -902,6 +902,13 @@ public enum DefaultDriverOption implements DriverOption {
    * <p>Value-type: List of {@link java.time.Duration Duration}
    */
   METRICS_NODE_CQL_MESSAGES_SLO("advanced.metrics.node.cql-messages.slo"),
+
+  /**
+   * Whether the prepared statements cache use weak values.
+   *
+   * <p>Value-type: boolean
+   */
+  PREPARED_CACHE_WEAK_VALUES("advanced.prepared-statements.prepared-cache.weak-values"),
   ;
 
   private final String path;

core/src/main/java/com/datastax/oss/driver/api/core/config/TypedDriverOption.java

@@ -521,6 +521,9 @@ public String toString() {
   /** The request timeout when repreparing. */
   public static final TypedDriverOption<Duration> REPREPARE_TIMEOUT =
       new TypedDriverOption<>(DefaultDriverOption.REPREPARE_TIMEOUT, GenericType.DURATION);
+  /** Whether the prepared statements cache use weak values. */
+  public static final TypedDriverOption<Boolean> PREPARED_CACHE_WEAK_VALUES =
+      new TypedDriverOption<>(DefaultDriverOption.PREPARED_CACHE_WEAK_VALUES, GenericType.BOOLEAN);
   /** The number of threads in the I/O group. */
   public static final TypedDriverOption<Integer> NETTY_IO_SIZE =
       new TypedDriverOption<>(DefaultDriverOption.NETTY_IO_SIZE, GenericType.INTEGER);

core/src/main/java/com/datastax/oss/driver/internal/core/session/PoolManager.java

@@ -82,8 +82,7 @@ public class PoolManager implements AsyncAutoCloseable {
   // (e.g. DefaultPreparedStatement) which are handled at the protocol level (e.g.
   // CqlPrepareAsyncProcessor). We keep the two separate to avoid introducing a dependency from the
   // session to a particular processor implementation.
-  private final ConcurrentMap<ByteBuffer, RepreparePayload> repreparePayloads =
-      new MapMaker().weakValues().makeMap();
+  private final ConcurrentMap<ByteBuffer, RepreparePayload> repreparePayloads;
 
   private final String logPrefix;
   private final EventExecutor adminExecutor;
@@ -95,6 +94,14 @@ public PoolManager(InternalDriverContext context) {
     this.adminExecutor = context.getNettyOptions().adminEventExecutorGroup().next();
     this.config = context.getConfig().getDefaultProfile();
     this.singleThreaded = new SingleThreaded(context);
+
+    if (config.getBoolean(DefaultDriverOption.PREPARED_CACHE_WEAK_VALUES, true)) {
+      LOG.debug("[{}] Prepared statements cache configured to use weak values", logPrefix);
+      this.repreparePayloads = new MapMaker().weakValues().makeMap();
+    } else {
+      LOG.debug("[{}] Prepared statements cache configured to use strong values", logPrefix);
+      this.repreparePayloads = new MapMaker().makeMap();
+    }
   }
 
   public CompletionStage<Void> init(CqlIdentifier keyspace) {

core/src/main/resources/reference.conf

@@ -2161,6 +2161,18 @@ datastax-java-driver {
       # Overridable in a profile: no
       timeout = ${datastax-java-driver.advanced.connection.init-query-timeout}
     }
+
+    # How to build the cache of prepared statements.
+    prepared-cache {
+      # Whether to use weak references for the prepared statements cache values.
+      #
+      # If this option is absent, weak references will be used.
+      #
+      # Required: no
+      # Modifiable at runtime: no
+      # Overridable in a profile: no
+      // weak-values = true
+    }
   }
 
   # Options related to the Netty event loop groups used internally by the driver.

core/src/test/java/com/datastax/oss/driver/internal/core/session/PoolManagerTest.java

@@ -0,0 +1,68 @@
+/*
+ * Copyright DataStax, Inc.
+ *
+ * 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
+ *
+ * http://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 com.datastax.oss.driver.internal.core.session;
+
+import static org.assertj.core.api.Assertions.assertThat;
+import static org.mockito.Mockito.when;
+
+import com.datastax.oss.driver.api.core.config.DefaultDriverOption;
+import com.datastax.oss.driver.api.core.config.DriverConfig;
+import com.datastax.oss.driver.api.core.config.DriverExecutionProfile;
+import com.datastax.oss.driver.internal.core.context.EventBus;
+import com.datastax.oss.driver.internal.core.context.InternalDriverContext;
+import com.datastax.oss.driver.internal.core.context.NettyOptions;
+import io.netty.channel.DefaultEventLoopGroup;
+import java.util.concurrent.ConcurrentHashMap;
+import org.junit.Before;
+import org.junit.Test;
+import org.mockito.Mock;
+import org.mockito.MockitoAnnotations;
+
+public class PoolManagerTest {
+  @Mock private InternalDriverContext context;
+  @Mock private NettyOptions nettyOptions;
+  @Mock private DriverConfig config;
+  @Mock private DriverExecutionProfile defaultProfile;
+
+  @Before
+  public void setup() {
+    MockitoAnnotations.initMocks(this);
+
+    DefaultEventLoopGroup adminEventLoopGroup = new DefaultEventLoopGroup(1);
+    when(nettyOptions.adminEventExecutorGroup()).thenReturn(adminEventLoopGroup);
+    when(context.getNettyOptions()).thenReturn(nettyOptions);
+    when(context.getEventBus()).thenReturn(new EventBus("test"));
+    when(config.getDefaultProfile()).thenReturn(defaultProfile);
+    when(context.getConfig()).thenReturn(config);
+  }
+
+  @Test
+  public void should_use_weak_values_if_config_is_true_or_undefined() {
+    when(defaultProfile.getBoolean(DefaultDriverOption.PREPARED_CACHE_WEAK_VALUES, true))
+        .thenReturn(true);
+    // As weak values map class is MapMakerInternalMap
+    assertThat(new PoolManager(context).getRepreparePayloads())
+        .isNotInstanceOf(ConcurrentHashMap.class);
+  }
+
+  @Test
+  public void should_not_use_weak_values_if_config_is_false() {
+    when(defaultProfile.getBoolean(DefaultDriverOption.PREPARED_CACHE_WEAK_VALUES, true))
+        .thenReturn(false);
+    assertThat(new PoolManager(context).getRepreparePayloads())
+        .isInstanceOf(ConcurrentHashMap.class);
+  }
+}