JAVA-2053: Cache results of session.prepare()

Author: olim7t

changelog/README.md

@@ -4,6 +4,7 @@
 
 ### 4.0.0-beta3 (in progress)
 
+- [improvement] JAVA-2053: Cache results of session.prepare()
 - [improvement] JAVA-2058: Make programmatic config reloading part of the public API
 - [improvement] JAVA-1943: Fail fast in execute() when the session is closed
 - [improvement] JAVA-2056: Reduce HashedWheelTimer tick duration

core/src/main/java/com/datastax/oss/driver/api/core/CqlSession.java

@@ -123,6 +123,21 @@ default CompletionStage<? extends AsyncResultSet> executeAsync(@NonNull String q
    *
    * If you want to customize this behavior, you can write your own implementation of {@link
    * PrepareRequest} and pass it to {@link #prepare(PrepareRequest)}.
+   *
+   * <p>The result of this method is cached: if you call it twice with the same {@link
+   * SimpleStatement}, you will get the same {@link PreparedStatement} instance. We still recommend
+   * keeping a reference to it (for example by caching it as a field in a DAO); if that's not
+   * possible (e.g. if query strings are generated dynamically), it's OK to call this method every
+   * time: there will just be a small performance overhead to check the internal cache. Note that
+   * caching is based on:
+   *
+   * <ul>
+   *   <li>the query string exactly as you provided it: the driver does not perform any kind of
+   *       trimming or sanitizing.
+   *   <li>all other execution parameters: for example, preparing two statements with identical
+   *       query strings but different {@linkplain SimpleStatement#getConsistencyLevel() consistency
+   *       levels} will yield distinct prepared statements.
+   * </ul>
    */
   @NonNull
   default PreparedStatement prepare(@NonNull SimpleStatement statement) {
@@ -134,6 +149,9 @@ default PreparedStatement prepare(@NonNull SimpleStatement statement) {
   /**
    * Prepares a CQL statement synchronously (the calling thread blocks until the statement is
    * prepared).
+   *
+   * <p>The result of this method is cached (see {@link #prepare(SimpleStatement)} for more
+   * explanations).
    */
   @NonNull
   default PreparedStatement prepare(@NonNull String query) {
@@ -150,6 +168,9 @@ default PreparedStatement prepare(@NonNull String query) {
    * customize how attributes are propagated when you prepare a {@link SimpleStatement} (see {@link
    * #prepare(SimpleStatement)} for more explanations). Otherwise, you should rarely have to deal
    * with {@link PrepareRequest} directly.
+   *
+   * <p>The result of this method is cached (see {@link #prepare(SimpleStatement)} for more
+   * explanations).
    */
   @NonNull
   default PreparedStatement prepare(@NonNull PrepareRequest request) {
@@ -165,6 +186,9 @@ default PreparedStatement prepare(@NonNull PrepareRequest request) {
    * <p>Note that the bound statements created from the resulting prepared statement will inherit
    * some of the attributes of {@code query}; see {@link #prepare(SimpleStatement)} for more
    * details.
+   *
+   * <p>The result of this method is cached (see {@link #prepare(SimpleStatement)} for more
+   * explanations).
    */
   @NonNull
   default CompletionStage<? extends PreparedStatement> prepareAsync(
@@ -177,6 +201,9 @@ default CompletionStage<? extends PreparedStatement> prepareAsync(
   /**
    * Prepares a CQL statement asynchronously (the call returns as soon as the prepare query was
    * sent, generally before the statement is prepared).
+   *
+   * <p>The result of this method is cached (see {@link #prepare(SimpleStatement)} for more
+   * explanations).
    */
   @NonNull
   default CompletionStage<? extends PreparedStatement> prepareAsync(@NonNull String query) {
@@ -193,6 +220,9 @@ default CompletionStage<? extends PreparedStatement> prepareAsync(@NonNull Strin
    * customize how attributes are propagated when you prepare a {@link SimpleStatement} (see {@link
    * #prepare(SimpleStatement)} for more explanations). Otherwise, you should rarely have to deal
    * with {@link PrepareRequest} directly.
+   *
+   * <p>The result of this method is cached (see {@link #prepare(SimpleStatement)} for more
+   * explanations).
    */
   @NonNull
   default CompletionStage<? extends PreparedStatement> prepareAsync(PrepareRequest request) {

core/src/main/java/com/datastax/oss/driver/api/core/metrics/DefaultSessionMetric.java

@@ -29,6 +29,7 @@ public enum DefaultSessionMetric implements SessionMetric {
   THROTTLING_DELAY("throttling.delay"),
   THROTTLING_QUEUE_SIZE("throttling.queue-size"),
   THROTTLING_ERRORS("throttling.errors"),
+  CQL_PREPARED_CACHE_SIZE("cql-prepared-cache-size"),
   ;
 
   private static final Map<String, DefaultSessionMetric> BY_PATH = sortByPath();

core/src/main/java/com/datastax/oss/driver/internal/core/cql/CqlPrepareAsyncHandler.java

@@ -23,20 +23,26 @@
 import com.datastax.oss.driver.internal.core.session.DefaultSession;
 import com.datastax.oss.driver.internal.core.session.RequestProcessor;
 import com.datastax.oss.driver.internal.core.util.concurrent.CompletableFutures;
-import java.nio.ByteBuffer;
+import com.datastax.oss.driver.shaded.guava.common.cache.CacheBuilder;
+import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.CompletionStage;
-import java.util.concurrent.ConcurrentMap;
 import net.jcip.annotations.ThreadSafe;
 
 @ThreadSafe
 public class CqlPrepareAsyncProcessor
     implements RequestProcessor<PrepareRequest, CompletionStage<PreparedStatement>> {
 
-  private final ConcurrentMap<ByteBuffer, DefaultPreparedStatement> preparedStatementsCache;
 
-  public CqlPrepareAsyncProcessor(
-      ConcurrentMap<ByteBuffer, DefaultPreparedStatement> preparedStatementsCache) {
-    this.preparedStatementsCache = preparedStatementsCache;
+
+  protected CqlPrepareAsyncProcessor(
+      Cache<PrepareRequest, CompletableFuture<PreparedStatement>> cache) {
+    this.cache = cache;
   }
 
   @Override
@@ -50,13 +56,38 @@ public CompletionStage<PreparedStatement> process(
       DefaultSession session,
       InternalDriverContext context,
       String sessionLogPrefix) {
-    return new CqlPrepareAsyncHandler(
-            request, preparedStatementsCache, session, context, sessionLogPrefix)
-        .handle();
+      if (result == null) {
+        CompletableFuture<PreparedStatement> mine = new CompletableFuture<>();
+        result = cache.get(request, () -> mine);
+        if (result == mine) {
+          new CqlPrepareHandler(request, session, context, sessionLogPrefix)
+              .handle()
+              .whenComplete(
+                  (preparedStatement, error) -> {
+                    if (error != null) {
+                      mine.completeExceptionally(error);
+                      cache.invalidate(request); // Make sure failure isn't cached indefinitely
+                    } else {
+                      mine.complete(preparedStatement);
+                    }
+                  });
+        }
+      }
+      return result;
+    } catch (ExecutionException e) {
+      return CompletableFutures.failedFuture(e.getCause());
+    }
   }
 
   @Override
   public CompletionStage<PreparedStatement> newFailure(RuntimeException error) {
     return CompletableFutures.failedFuture(error);
   }
+
+  public Cache<PrepareRequest, CompletableFuture<PreparedStatement>> getCache() {
+    return cache;
+  }
 }

core/src/main/java/com/datastax/oss/driver/internal/core/cql/CqlPrepareAsyncProcessor.java

@@ -55,7 +55,6 @@
 import io.netty.util.Timer;
 import io.netty.util.concurrent.Future;
 import io.netty.util.concurrent.GenericFutureListener;
-import java.nio.ByteBuffer;
 import java.time.Duration;
 import java.util.AbstractMap;
 import java.util.ArrayList;
@@ -65,7 +64,6 @@
 import java.util.concurrent.CancellationException;
 import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.CompletionStage;
-import java.util.concurrent.ConcurrentMap;
 import java.util.concurrent.CopyOnWriteArrayList;
 import java.util.concurrent.TimeUnit;
 import net.jcip.annotations.ThreadSafe;
@@ -74,14 +72,13 @@
 
 /** Handles the lifecycle of the preparation of a CQL statement. */
 @ThreadSafe
-public abstract class CqlPrepareHandlerBase implements Throttled {
 
-  private static final Logger LOG = LoggerFactory.getLogger(CqlPrepareHandlerBase.class);
 
   private final long startTimeNanos;
   private final String logPrefix;
   private final PrepareRequest request;
-  private final ConcurrentMap<ByteBuffer, DefaultPreparedStatement> preparedStatementsCache;
   private final DefaultSession session;
   private final InternalDriverContext context;
   private final DriverExecutionProfile executionProfile;
@@ -100,9 +97,8 @@ public abstract class CqlPrepareHandlerBase implements Throttled {
   // We don't use a map because nodes can appear multiple times.
   private volatile List<Map.Entry<Node, Throwable>> errors;
 
-  protected CqlPrepareHandlerBase(
       PrepareRequest request,
-      ConcurrentMap<ByteBuffer, DefaultPreparedStatement> preparedStatementsCache,
       DefaultSession session,
       InternalDriverContext context,
       String sessionLogPrefix) {
@@ -112,7 +108,6 @@ protected CqlPrepareHandlerBase(
     LOG.trace("[{}] Creating new handler for prepare request {}", logPrefix, request);
 
     this.request = request;
-    this.preparedStatementsCache = preparedStatementsCache;
     this.session = session;
     this.context = context;
     this.executionProfile = Conversions.resolveExecutionProfile(request, context);
@@ -171,6 +166,10 @@ public void onThrottleReady(boolean wasDelayed) {
     sendRequest(null, 0);
   }
 
+  public CompletableFuture<PreparedStatement> handle() {
+    return result;
+  }
+
   private Timeout scheduleTimeout(Duration timeoutDuration) {
     if (timeoutDuration.toNanos() > 0) {
       return this.timer.newTimeout(
@@ -221,7 +220,7 @@ private void recordError(Node node, Throwable error) {
     // Use a local variable to do only a single single volatile read in the nominal case
     List<Map.Entry<Node, Throwable>> errorsSnapshot = this.errors;
     if (errorsSnapshot == null) {
-      synchronized (CqlPrepareHandlerBase.this) {
         errorsSnapshot = this.errors;
         if (errorsSnapshot == null) {
           this.errors = errorsSnapshot = new CopyOnWriteArrayList<>();
@@ -236,58 +235,28 @@ private void setFinalResult(Prepared prepared) {
     // Whatever happens below, we're done with this stream id
     throttler.signalSuccess(this);
 
-    DefaultPreparedStatement newStatement =
         Conversions.toPreparedStatement(prepared, request, context);
 
-    DefaultPreparedStatement cachedStatement = cache(newStatement);
-
-    if (cachedStatement != newStatement) {
-      // The statement already existed in the cache, assume it's because the client called
-      // prepare() twice, and therefore it's already been prepared on other nodes.
-      result.complete(cachedStatement);
-    } else {
-      session
-          .getRepreparePayloads()
-          .put(cachedStatement.getId(), cachedStatement.getRepreparePayload());
-      if (prepareOnAllNodes) {
-        prepareOnOtherNodes()
-            .thenRun(
-                () -> {
-                  LOG.trace(
-                      "[{}] Done repreparing on other nodes, completing the request", logPrefix);
-                  result.complete(cachedStatement);
-                })
-            .exceptionally(
-                error -> {
-                  result.completeExceptionally(error);
-                  return null;
-                });
-      } else {
-        LOG.trace("[{}] Prepare on all nodes is disabled, completing the request", logPrefix);
-        result.complete(cachedStatement);
-      }
-    }
-  }
-
-  private DefaultPreparedStatement cache(DefaultPreparedStatement preparedStatement) {
-    DefaultPreparedStatement previous =
-        preparedStatementsCache.putIfAbsent(preparedStatement.getId(), preparedStatement);
-    if (previous != null) {
-      LOG.warn(
-          "Re-preparing already prepared query. "
-              + "This is generally an anti-pattern and will likely affect performance. "
-              + "The cached version of the PreparedStatement will be returned, which may use "
-              + "different bound statement execution parameters (CL, timeout, etc.) from the "
-              + "current session.prepare call. Consider preparing the statement only once. "
-              + "Query='{}'",
-          preparedStatement.getQuery());
-
-      // The one object in the cache will get GCed once it's not referenced by the client anymore
-      // since we use a weak reference. So we need to make sure that the instance we do return to
-      // the user is the one that is in the cache.
-      return previous;
     } else {
-      return preparedStatement;
+      result.complete(preparedStatement);
     }
   }