Anchor prepare cache entry via PreparedStatement back-reference

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

@@ -164,16 +164,17 @@ public CompletionStage<PreparedStatement> process(
                       mine.completeExceptionally(error);
                       cache.invalidate(request); // Make sure failure isn't cached indefinitely
                     } else {
+                      // Anchor the cache entry on the statement, so that it survives for as long as
+                      // the application holds the statement. The cache holds its values weakly, and
+                      // callers routinely keep only the statement, not the future we return here.
+                      if (preparedStatement instanceof PrepareCacheAnchor) {
+                        ((PrepareCacheAnchor) preparedStatement).setPrepareCacheAnchor(mine);
+                      }
                       mine.complete(preparedStatement);
                     }
                   });
         }
       }
-      // If the future is already completed, return it directly to maintain a strong reference
-      // in the cache and avoid premature GC with weakValues() (ScyllaDB PR #892).
-      if (result.isDone()) {
-        return result;
-      }
       // Return a defensive copy. So if a client cancels its request, the cache won't be impacted
       // nor a potential concurrent request.
       return result.thenApply(x -> x); // copy() is available only since Java 9

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

@@ -53,12 +53,14 @@
 import java.time.Duration;
 import java.util.List;
 import java.util.Map;
+import java.util.concurrent.CompletableFuture;
 import net.jcip.annotations.ThreadSafe;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
 @ThreadSafe
-public class DefaultPreparedStatement implements PreparedStatement, RequestRoutingTypeAccessor {
+public class DefaultPreparedStatement
+    implements PreparedStatement, RequestRoutingTypeAccessor, PrepareCacheAnchor {
   private static final Logger LOGGER = LoggerFactory.getLogger(DefaultPreparedStatement.class);
   private static final Splitter SPACE_SPLITTER = Splitter.onPattern("\\s+");
   private static final Splitter COMMA_SPLITTER = Splitter.onPattern(",");
@@ -87,6 +89,15 @@ public class DefaultPreparedStatement implements PreparedStatement, RequestRouti
   @Nullable private final RequestRoutingType requestRoutingType;
   private volatile boolean skipMetadata;
 
+  /**
+   * Retains the prepare cache entry that produced this statement. Never read: it exists purely so
+   * that the entry, which the cache holds weakly, stays reachable for as long as this statement is.
+   *
+   * @see PrepareCacheAnchor
+   */
+  @SuppressWarnings("unused")
+  private volatile CompletableFuture<PreparedStatement> prepareCacheAnchor;
+
   public DefaultPreparedStatement(
       ByteBuffer id,
       String query,
@@ -144,6 +155,11 @@ public DefaultPreparedStatement(
             query, resultMetadataId, resultSetDefinitions, this.executionProfileForBoundStatements);
   }
 
+  @Override
+  public void setPrepareCacheAnchor(@Nullable CompletableFuture<PreparedStatement> anchor) {
+    this.prepareCacheAnchor = anchor;
+  }
+
   @NonNull
   @Override
   public ByteBuffer getId() {

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

@@ -0,0 +1,43 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you 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.cql;
+
+import com.datastax.oss.driver.api.core.cql.PreparedStatement;
+import edu.umd.cs.findbugs.annotations.Nullable;
+import java.util.concurrent.CompletableFuture;
+
+/**
+ * Internal hook allowing a {@link PreparedStatement} to keep its prepare cache entry reachable.
+ *
+ * <p>{@link CqlPrepareAsyncProcessor} caches prepare futures with weak values, so an entry can be
+ * collected while the application still holds the resulting statement, causing a needless
+ * re-PREPARE. Storing the cached future on the statement ties the entry's lifetime to the
+ * statement's: the cache holds a weak reference to the future, the future references the statement,
+ * and the statement references the future back. The cycle stays reachable while the application
+ * holds the statement, and becomes collectible as a whole once it does not.
+ *
+ * <p>Implementations only need to retain the reference; the anchor is never read back.
+ */
+public interface PrepareCacheAnchor {
+
+  /**
+   * Retains the prepare cache entry for this statement, preventing its weak-value eviction for as
+   * long as this statement is reachable.
+   */
+  void setPrepareCacheAnchor(@Nullable CompletableFuture<PreparedStatement> anchor);
+}

core/src/test/java/com/datastax/oss/driver/internal/core/cql/CqlPrepareAsyncProcessorTest.java

@@ -17,6 +17,7 @@
  */
 package com.datastax.oss.driver.internal.core.cql;
 
+import static com.datastax.oss.driver.internal.core.cql.PreparedStatementTestHelper.newPreparedStatement;
 import static org.assertj.core.api.Assertions.assertThat;
 
 import com.datastax.oss.driver.api.core.cql.PrepareRequest;
@@ -25,6 +26,7 @@
 import com.datastax.oss.driver.api.core.type.UserDefinedType;
 import com.datastax.oss.driver.internal.core.type.UserDefinedTypeBuilder;
 import com.datastax.oss.driver.shaded.guava.common.cache.Cache;
+import java.lang.ref.WeakReference;
 import java.util.Optional;
 import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.CompletionStage;
@@ -48,23 +50,26 @@ public void setup() {
   }
 
   /**
-   * When the cached future is already completed, process() should return the exact same instance
-   * (identity). This ensures callers hold a strong reference to the cached CF, preventing
-   * weak-value eviction under GC pressure.
+   * process() always hands out a defensive copy, including for an already-completed entry: keeping
+   * the entry alive is the anchor's job, so the cached future never needs to be exposed. A caller
+   * that obtrudes on its copy must not corrupt what the cache holds.
    */
   @Test
-  public void should_return_cached_future_directly_when_already_completed() throws Exception {
+  public void should_return_defensive_copy_when_future_is_already_completed() throws Exception {
     PrepareRequest request = new DefaultPrepareRequest("SELECT 1");
     PreparedStatement ps = Mockito.mock(PreparedStatement.class);
 
     // Pre-populate cache with a completed future
     CompletableFuture<PreparedStatement> completed = CompletableFuture.completedFuture(ps);
     cache.put(request, completed);
 
-    // process() should return the exact same object
     CompletionStage<PreparedStatement> returned = processor.process(request, null, null, "test");
 
-    assertThat(returned).isSameAs(completed);
+    assertThat(returned).isNotSameAs(completed);
+    assertThat(returned.toCompletableFuture().get()).isSameAs(ps);
+
+    returned.toCompletableFuture().obtrudeValue(null);
+    assertThat(completed.get()).isSameAs(ps);
   }
 
   /**
@@ -106,4 +111,63 @@ public void should_match_udt_by_name_when_field_definitions_differ() {
     assertThat(CqlPrepareAsyncProcessor.typeMatches(oldType, DataTypes.listOf(resultType)))
         .isTrue();
   }
+
+  /**
+   * The anchor is what keeps a weakly-held cache entry alive: while the application holds the
+   * statement, the entry survives GC even though nothing else references the cached future.
+   */
+  @Test
+  public void should_keep_cache_entry_alive_via_prepared_statement_anchor() throws Exception {
+    PrepareRequest request = new DefaultPrepareRequest("SELECT 1");
+
+    // The only surviving reference is the statement; the future stays in the callee's frame.
+    DefaultPreparedStatement ps = anchorNewEntry(request);
+
+    collectGarbage();
+
+    assertThat(cache.getIfPresent(request)).isNotNull();
+    assertThat(cache.getIfPresent(request).get()).isSameAs(ps);
+  }
+
+  /**
+   * The reverse: once the statement becomes unreachable the whole cycle is collectible, so the
+   * anchor cannot turn the cache into a leak.
+   */
+  @Test
+  public void should_evict_cache_entry_when_prepared_statement_is_unreachable() throws Exception {
+    PrepareRequest request = new DefaultPrepareRequest("SELECT 1");
+
+    // Wrapping in a WeakReference lets us drop the statement without keeping it in a local.
+    WeakReference<DefaultPreparedStatement> ps = new WeakReference<>(anchorNewEntry(request));
+
+    collectGarbage();
+
+    assertThat(ps.get())
+        .as("statement was not collected, so the cache assertion below proves nothing")
+        .isNull();
+    assertThat(cache.getIfPresent(request)).isNull();
+  }
+
+  /**
+   * Reproduces what {@link CqlPrepareAsyncProcessor#process} does on a successful prepare: cache
+   * the future, anchor it on the resulting statement, then complete it. The future is deliberately
+   * a local of this method, so it becomes unreachable as soon as this frame returns.
+   */
+  private DefaultPreparedStatement anchorNewEntry(PrepareRequest request) {
+    CompletableFuture<PreparedStatement> cachedFuture = new CompletableFuture<>();
+    cache.put(request, cachedFuture);
+
+    DefaultPreparedStatement ps = newPreparedStatement();
+    ps.setPrepareCacheAnchor(cachedFuture);
+    cachedFuture.complete(ps);
+    return ps;
+  }
+
+  private void collectGarbage() throws InterruptedException {
+    for (int i = 0; i < 10; i++) {
+      System.gc();
+      Thread.sleep(50);
+      cache.cleanUp();
+    }
+  }
 }

core/src/test/java/com/datastax/oss/driver/internal/core/cql/DefaultPreparedStatementTest.java

@@ -17,17 +17,12 @@
  */
 package com.datastax.oss.driver.internal.core.cql;
 
+import static com.datastax.oss.driver.internal.core.cql.PreparedStatementTestHelper.newPreparedStatement;
 import static org.assertj.core.api.Assertions.assertThat;
 
-import com.datastax.oss.driver.api.core.ConsistencyLevel;
 import com.datastax.oss.driver.api.core.DefaultConsistencyLevel;
-import com.datastax.oss.driver.api.core.DefaultProtocolVersion;
 import com.datastax.oss.driver.api.core.RequestRoutingType;
 import com.datastax.oss.driver.api.core.cql.BoundStatement;
-import com.datastax.oss.driver.api.core.cql.ColumnDefinitions;
-import com.datastax.oss.driver.api.core.type.codec.registry.CodecRegistry;
-import com.datastax.oss.protocol.internal.util.Bytes;
-import java.util.Collections;
 import org.junit.Test;
 
 public class DefaultPreparedStatementTest {
@@ -75,38 +70,4 @@ public void should_keep_detected_lwt_routing_type_after_bound_consistency_overri
 
     assertThat(boundStatement.getRequestRoutingType()).isEqualTo(RequestRoutingType.LWT);
   }
-
-  private DefaultPreparedStatement newPreparedStatement(
-      ConsistencyLevel consistencyLevel,
-      ConsistencyLevel serialConsistencyLevel,
-      RequestRoutingType requestRoutingType) {
-    ColumnDefinitions variableDefinitions =
-        DefaultColumnDefinitions.valueOf(Collections.emptyList());
-    return new DefaultPreparedStatement(
-        Bytes.fromHexString("0x"),
-        "SELECT * FROM test.foo WHERE pk = ?",
-        variableDefinitions,
-        Collections.emptyList(),
-        null,
-        null,
-        null,
-        null,
-        Collections.emptyMap(),
-        null,
-        null,
-        null,
-        null,
-        null,
-        Collections.emptyMap(),
-        null,
-        null,
-        null,
-        Integer.MIN_VALUE,
-        consistencyLevel,
-        serialConsistencyLevel,
-        false,
-        CodecRegistry.DEFAULT,
-        DefaultProtocolVersion.DEFAULT,
-        requestRoutingType);
-  }
 }

core/src/test/java/com/datastax/oss/driver/internal/core/cql/PreparedStatementTestHelper.java

@@ -0,0 +1,71 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you 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.cql;
+
+import com.datastax.oss.driver.api.core.ConsistencyLevel;
+import com.datastax.oss.driver.api.core.DefaultProtocolVersion;
+import com.datastax.oss.driver.api.core.RequestRoutingType;
+import com.datastax.oss.driver.api.core.cql.ColumnDefinitions;
+import com.datastax.oss.driver.api.core.type.codec.registry.CodecRegistry;
+import com.datastax.oss.protocol.internal.util.Bytes;
+import java.util.Collections;
+
+/** Builds minimally-valid {@link DefaultPreparedStatement} instances for tests. */
+public class PreparedStatementTestHelper {
+
+  /** Returns a statement with no consistency levels and no routing type configured. */
+  public static DefaultPreparedStatement newPreparedStatement() {
+    return newPreparedStatement(null, null, null);
+  }
+
+  public static DefaultPreparedStatement newPreparedStatement(
+      ConsistencyLevel consistencyLevel,
+      ConsistencyLevel serialConsistencyLevel,
+      RequestRoutingType requestRoutingType) {
+    ColumnDefinitions variableDefinitions =
+        DefaultColumnDefinitions.valueOf(Collections.emptyList());
+    return new DefaultPreparedStatement(
+        Bytes.fromHexString("0x"),
+        "SELECT * FROM test.foo WHERE pk = ?",
+        variableDefinitions,
+        Collections.emptyList(),
+        null,
+        null,
+        null,
+        null,
+        Collections.emptyMap(),
+        null,
+        null,
+        null,
+        null,
+        null,
+        Collections.emptyMap(),
+        null,
+        null,
+        null,
+        Integer.MIN_VALUE,
+        consistencyLevel,
+        serialConsistencyLevel,
+        false,
+        CodecRegistry.DEFAULT,
+        DefaultProtocolVersion.DEFAULT,
+        requestRoutingType);
+  }
+
+  private PreparedStatementTestHelper() {}
+}