Anchor prepare cache entry via PreparedStatement back-reference

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

@@ -157,13 +157,25 @@ public CompletionStage<PreparedStatement> process(
                       mine.completeExceptionally(error);
                       cache.invalidate(request); // Make sure failure isn't cached indefinitely
                     } else {
+                      // Anchor the cached future inside the PreparedStatement so that as long as
+                      // the PS is reachable, the weak-value cache entry won't be GC'd.
+                      if (preparedStatement instanceof DefaultPreparedStatement) {
+                        ((DefaultPreparedStatement) preparedStatement).setCacheRetainer(mine);
+                      }
                       mine.complete(preparedStatement);
                     }
                   });
         }
       }
-      // Return a defensive copy. So if a client cancels its request, the cache won't be impacted
-      // nor a potential concurrent request.
+      if (result.isDone()) {
+        // Completed futures are immutable (cancel/complete/completeExceptionally are no-ops),
+        // so returning the cached instance directly is safe. This also keeps the cache entry
+        // alive via the caller's strong reference, preventing premature weak-value eviction
+        // under GC pressure.
+        return result;
+      }
+      // Defensive copy for in-flight preparations only: protects the shared cached future
+      // from cancellation by one of multiple concurrent waiters.
       return result.thenApply(x -> x); // copy() is available only since Java 9
     } catch (ExecutionException e) {
       return CompletableFutures.failedFuture(e.getCause());

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

@@ -87,6 +87,16 @@ public class DefaultPreparedStatement implements PreparedStatement, RequestRouti
   @Nullable private final RequestRoutingType requestRoutingType;
   private volatile boolean skipMetadata;
 
+  /**
+   * Holds a strong reference to the {@link java.util.concurrent.CompletableFuture} stored in the
+   * prepare cache (which uses weak values). As long as this PreparedStatement is reachable, the
+   * cache entry won't be garbage-collected, preventing unnecessary re-PREPARE requests.
+   *
+   * @see CqlPrepareAsyncProcessor
+   */
+  @SuppressWarnings("unused")
+  private volatile Object cacheRetainer;
+
   public DefaultPreparedStatement(
       ByteBuffer id,
       String query,
@@ -144,6 +154,14 @@ public DefaultPreparedStatement(
             query, resultMetadataId, resultSetDefinitions, this.executionProfileForBoundStatements);
   }
 
+  /**
+   * Attaches a strong reference to the prepare cache entry, preventing its weak-value eviction as
+   * long as this PreparedStatement is reachable.
+   */
+  public void setCacheRetainer(Object retainer) {
+    this.cacheRetainer = retainer;
+  }
+
   @NonNull
   @Override
   public ByteBuffer getId() {

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

@@ -0,0 +1,279 @@
+/*
+ * 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 static org.assertj.core.api.Assertions.assertThat;
+
+import com.datastax.oss.driver.api.core.cql.PrepareRequest;
+import com.datastax.oss.driver.api.core.cql.PreparedStatement;
+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;
+import org.junit.Before;
+import org.junit.Test;
+import org.mockito.Mockito;
+
+/**
+ * Unit tests for {@link CqlPrepareAsyncProcessor} focusing on the caching behavior of {@link
+ * CqlPrepareAsyncProcessor#process} with respect to defensive copies and weak-value retention.
+ */
+public class CqlPrepareAsyncProcessorTest {
+
+  private CqlPrepareAsyncProcessor processor;
+  private Cache<PrepareRequest, CompletableFuture<PreparedStatement>> cache;
+
+  @Before
+  public void setup() {
+    processor = new CqlPrepareAsyncProcessor(Optional.empty());
+    cache = processor.getCache();
+  }
+
+  /**
+   * 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.
+   */
+  @Test
+  public void should_return_cached_future_directly_when_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);
+  }
+
+  /**
+   * When the cached future is still in-flight (not yet done), process() should return a defensive
+   * copy to protect the cache from cancellation by the caller.
+   */
+  @Test
+  public void should_return_defensive_copy_when_future_is_in_flight() throws Exception {
+    PrepareRequest request = new DefaultPrepareRequest("SELECT 1");
+
+    // Pre-populate cache with an incomplete future
+    CompletableFuture<PreparedStatement> inFlight = new CompletableFuture<>();
+    cache.put(request, inFlight);
+
+    CompletionStage<PreparedStatement> returned = processor.process(request, null, null, "test");
+
+    // Should NOT be the same instance
+    assertThat(returned).isNotSameAs(inFlight);
+
+    // Cancelling the returned copy should NOT affect the cached future
+    returned.toCompletableFuture().cancel(false);
+    assertThat(inFlight.isCancelled()).isFalse();
+  }
+
+  /**
+   * Verifies that returning the cached future directly (for completed entries) keeps the weak-value
+   * cache entry alive as long as the caller holds a reference to the returned stage.
+   */
+  @Test
+  public void should_keep_cache_entry_alive_when_caller_holds_completed_future() throws Exception {
+    PrepareRequest request = new DefaultPrepareRequest("SELECT 1");
+    PreparedStatement ps = Mockito.mock(PreparedStatement.class);
+
+    CompletableFuture<PreparedStatement> completed = CompletableFuture.completedFuture(ps);
+    cache.put(request, completed);
+
+    // Simulate what a caller does: hold the returned stage
+    CompletionStage<PreparedStatement> held = processor.process(request, null, null, "test");
+
+    // Create a weak reference to detect if cache entry would be collected
+    WeakReference<CompletableFuture<PreparedStatement>> weakRef = new WeakReference<>(completed);
+    // Drop our local strong reference
+    //noinspection UnusedAssignment
+    completed = null;
+
+    // Force GC
+    System.gc();
+    Thread.sleep(100);
+
+    // The cache entry should still be alive because 'held' IS the cached CF
+    cache.cleanUp();
+    assertThat(cache.getIfPresent(request)).isNotNull();
+    assertThat(weakRef.get()).isNotNull();
+
+    // Verify held is usable
+    assertThat(held.toCompletableFuture().get()).isSameAs(ps);
+  }
+
+  /**
+   * Demonstrates the problem this fix addresses: without the fix, a defensive copy would be the
+   * only reference returned, and if the caller doesn't hold it long enough, GC can evict the cache
+   * entry. This test shows that with the fix, even after the caller's reference goes out of scope,
+   * the behavior is correct for the next caller who retrieves it promptly.
+   */
+  @Test
+  public void should_allow_gc_eviction_when_no_strong_references_remain() throws Exception {
+    PrepareRequest request = new DefaultPrepareRequest("SELECT 1");
+    PreparedStatement ps = Mockito.mock(PreparedStatement.class);
+
+    CompletableFuture<PreparedStatement> completed = CompletableFuture.completedFuture(ps);
+    cache.put(request, completed);
+
+    // Drop all strong references
+    //noinspection UnusedAssignment
+    completed = null;
+
+    // Force GC - weak value should be collected
+    for (int i = 0; i < 10; i++) {
+      System.gc();
+      Thread.sleep(50);
+      cache.cleanUp();
+      if (cache.getIfPresent(request) == null) {
+        break;
+      }
+    }
+
+    // Cache entry may have been evicted (weak values)
+    // This is expected behavior - the fix ensures callers who DO hold a reference keep it alive
+    // We just verify the cache doesn't throw
+    assertThat(cache.size()).isGreaterThanOrEqualTo(0);
+  }
+
+  /**
+   * Verifies that when a DefaultPreparedStatement holds the cache retainer (strong back-reference
+   * to the cached CF), the cache entry survives GC even when no other strong references exist.
+   */
+  @Test
+  public void should_keep_cache_entry_alive_via_prepared_statement_retainer() throws Exception {
+    PrepareRequest request = new DefaultPrepareRequest("SELECT 1");
+
+    CompletableFuture<PreparedStatement> cachedFuture = new CompletableFuture<>();
+    cache.put(request, cachedFuture);
+
+    // Simulate what the processor does: create a real DefaultPreparedStatement and set retainer
+    DefaultPreparedStatement ps =
+        new DefaultPreparedStatement(
+            java.nio.ByteBuffer.wrap(new byte[] {1, 2, 3, 4}),
+            "SELECT 1",
+            com.datastax.oss.driver.internal.core.cql.EmptyColumnDefinitions.INSTANCE,
+            java.util.Collections.emptyList(),
+            null,
+            com.datastax.oss.driver.internal.core.cql.EmptyColumnDefinitions.INSTANCE,
+            null,
+            null,
+            java.util.Collections.emptyMap(),
+            null,
+            null,
+            null,
+            null,
+            null,
+            java.util.Collections.emptyMap(),
+            null,
+            null,
+            null,
+            -1,
+            null,
+            null,
+            false,
+            com.datastax.oss.driver.api.core.type.codec.registry.CodecRegistry.DEFAULT,
+            com.datastax.oss.driver.api.core.ProtocolVersion.V4,
+            null);
+
+    // Set the retainer — this is the key part of PR #2
+    ps.setCacheRetainer(cachedFuture);
+    cachedFuture.complete(ps);
+
+    // Drop all references except PS itself
+    //noinspection UnusedAssignment
+    cachedFuture = null;
+
+    // Force GC
+    for (int i = 0; i < 10; i++) {
+      System.gc();
+      Thread.sleep(50);
+      cache.cleanUp();
+    }
+
+    // Cache entry should survive because PS holds the retainer
+    assertThat(cache.getIfPresent(request)).isNotNull();
+    assertThat(cache.getIfPresent(request).get()).isSameAs(ps);
+  }
+
+  /**
+   * Verifies that when the PreparedStatement (and its retainer) is no longer reachable, the cache
+   * entry CAN be evicted — confirming we don't leak memory.
+   */
+  @Test
+  public void should_evict_cache_entry_when_prepared_statement_is_unreachable() throws Exception {
+    PrepareRequest request = new DefaultPrepareRequest("SELECT 1");
+
+    CompletableFuture<PreparedStatement> cachedFuture = new CompletableFuture<>();
+    cache.put(request, cachedFuture);
+
+    DefaultPreparedStatement ps =
+        new DefaultPreparedStatement(
+            java.nio.ByteBuffer.wrap(new byte[] {1, 2, 3, 4}),
+            "SELECT 1",
+            com.datastax.oss.driver.internal.core.cql.EmptyColumnDefinitions.INSTANCE,
+            java.util.Collections.emptyList(),
+            null,
+            com.datastax.oss.driver.internal.core.cql.EmptyColumnDefinitions.INSTANCE,
+            null,
+            null,
+            java.util.Collections.emptyMap(),
+            null,
+            null,
+            null,
+            null,
+            null,
+            java.util.Collections.emptyMap(),
+            null,
+            null,
+            null,
+            -1,
+            null,
+            null,
+            false,
+            com.datastax.oss.driver.api.core.type.codec.registry.CodecRegistry.DEFAULT,
+            com.datastax.oss.driver.api.core.ProtocolVersion.V4,
+            null);
+
+    ps.setCacheRetainer(cachedFuture);
+    cachedFuture.complete(ps);
+
+    // Drop ALL strong references
+    //noinspection UnusedAssignment
+    cachedFuture = null;
+    //noinspection UnusedAssignment
+    ps = null;
+
+    // Force GC — weak value should be collected since nothing holds the CF
+    for (int i = 0; i < 10; i++) {
+      System.gc();
+      Thread.sleep(50);
+      cache.cleanUp();
+      if (cache.getIfPresent(request) == null) {
+        break;
+      }
+    }
+
+    // Cache entry should have been evicted
+    assertThat(cache.getIfPresent(request)).isNull();
+  }
+}