strong values of cache

SiyaoIsHiding · SiyaoIsHiding · commit 98eb7570b809 · 2025-07-11T12:18:15.000-07:00
diff --git a/core/src/main/java/com/datastax/oss/driver/internal/core/cql/CqlPrepareAsyncProcessor.java b/core/src/main/java/com/datastax/oss/driver/internal/core/cql/CqlPrepareAsyncProcessor.java
@@ -83,6 +83,19 @@ protected CqlPrepareAsyncProcessor(
         });
   }
 
+  protected CqlPrepareAsyncProcessor(
+      Optional<? extends DefaultDriverContext> context, CacheBuilder<Object, Object> cacheBuilder) {
+    this.cache = cacheBuilder.build();
+    context.ifPresent(
+        (ctx) -> {
+          LOG.info("Adding handler to invalidate cached prepared statements on type changes");
+          EventExecutor adminExecutor = ctx.getNettyOptions().adminEventExecutorGroup().next();
+          ctx.getEventBus()
+              .register(
+                  TypeChangeEvent.class, RunOrSchedule.on(adminExecutor, this::onTypeChanged));
+        });
+  }
+
   private static boolean typeMatches(UserDefinedType oldType, DataType typeToCheck) {
 
     switch (typeToCheck.getProtocolCode()) {
diff --git a/integration-tests/src/test/java/com/datastax/oss/driver/core/cql/PreparedStatementCachingIT.java b/integration-tests/src/test/java/com/datastax/oss/driver/core/cql/PreparedStatementCachingIT.java
@@ -40,6 +40,7 @@
 import com.datastax.oss.driver.internal.core.session.BuiltInRequestProcessors;
 import com.datastax.oss.driver.internal.core.session.RequestProcessor;
 import com.datastax.oss.driver.internal.core.session.RequestProcessorRegistry;
+import com.datastax.oss.driver.shaded.guava.common.cache.CacheBuilder;
 import com.datastax.oss.driver.shaded.guava.common.cache.RemovalListener;
 import com.datastax.oss.driver.shaded.guava.common.util.concurrent.Uninterruptibles;
 import com.google.common.collect.ImmutableList;
@@ -72,8 +73,6 @@
 @Category(IsolatedTests.class)
 public class PreparedStatementCachingIT {
 
-  private static final Logger LOG = LoggerFactory.getLogger(PreparedStatementCachingIT.class);
-
   private CustomCcmRule ccmRule = CustomCcmRule.builder().build();
 
   private SessionRule<CqlSession> sessionRule =
@@ -123,30 +122,11 @@ private static RemovalListener<Object, Object> buildCacheRemoveCallback(
         @NonNull Optional<DefaultDriverContext> context) {
       return (evt) -> {
         try {
-          LOG.error(
-              "Cache removal callback triggered, cause: {}, key: {}", evt.getCause(), evt.getKey());
           CompletableFuture<PreparedStatement> future =
               (CompletableFuture<PreparedStatement>) evt.getValue();
-
-          // Add more detailed logging about the future state
-          LOG.error(
-              "Future state - done: {}, cancelled: {}, completedExceptionally: {}",
-              future.isDone(),
-              future.isCancelled(),
-              future.isCompletedExceptionally());
-
-          if (future.isDone() && !future.isCompletedExceptionally() && !future.isCancelled()) {
-            ByteBuffer queryId = Uninterruptibles.getUninterruptibly(future).getId();
-            LOG.error("Firing PreparedStatementRemovalEvent for queryId: {}", queryId);
-            context.ifPresent(
-                ctx -> {
-                  LOG.error("About to fire PreparedStatementRemovalEvent on event bus");
-                  ctx.getEventBus().fire(new PreparedStatementRemovalEvent(queryId));
-                  LOG.error("PreparedStatementRemovalEvent fired successfully");
-                });
-          } else {
-            LOG.error("Skipping removal event - future not in valid state for extraction");
-          }
+          ByteBuffer queryId = Uninterruptibles.getUninterruptibly(future).getId();
+          context.ifPresent(
+              ctx -> ctx.getEventBus().fire(new PreparedStatementRemovalEvent(queryId)));
         } catch (Exception e) {
           LOG.error("Unable to register removal handler", e);
         }
@@ -156,7 +136,8 @@ private static RemovalListener<Object, Object> buildCacheRemoveCallback(
     public TestCqlPrepareAsyncProcessor(@NonNull Optional<DefaultDriverContext> context) {
       // Default CqlPrepareAsyncProcessor uses weak values here as well.  We avoid doing so
       // to prevent cache entries from unexpectedly disappearing mid-test.
-      super(context, builder -> builder.removalListener(buildCacheRemoveCallback(context)));
+      // TODO: it was still weak value cuz it's only a decorator.
+      super(context, CacheBuilder.newBuilder().removalListener(buildCacheRemoveCallback(context)));
     }
   }
 
@@ -211,83 +192,6 @@ public static SessionBuilder builder() {
     return new TestSessionBuilder();
   }
 
-  private void debugCacheInvalidation(CqlSession session, TypeChangeEvent event) {
-    try {
-      DefaultDriverContext ctx = (DefaultDriverContext) session.getContext();
-      // Get the processor to check cache state
-      RequestProcessorRegistry registry = ctx.getRequestProcessorRegistry();
-
-      LOG.error(
-          "Debug: TypeChangeEvent received for type: {} (changeType: {})",
-          event.oldType.getName(),
-          event.changeType);
-      LOG.error("Debug: Current cache size: {}", getPreparedCacheSize(session));
-
-      // Force cache cleanup to trigger any pending removals
-      if (registry != null) {
-        LOG.error("Debug: Forcing cache cleanup...");
-        // We can't directly access the cache from here, but we can log that we're trying
-      }
-    } catch (Exception e) {
-      LOG.error("Debug: Error during cache invalidation debugging", e);
-    }
-  }
-
-  private boolean waitForCacheRemovalWithCleanup(
-      CountDownLatch latch, CqlSession session, long timeout, TimeUnit unit) {
-    long timeoutNanos = unit.toNanos(timeout);
-    long startTime = System.nanoTime();
-    long remainingNanos = timeoutNanos;
-
-    LOG.error("Starting cache removal wait, initial latch count: {}", latch.getCount());
-
-    // First check if the latch is already at 0
-    if (latch.getCount() == 0) {
-      LOG.error("Cache removal latch already at 0, returning success immediately");
-      return true;
-    }
-
-    while (remainingNanos > 0 && latch.getCount() > 0) {
-      // Wait for a short period
-      long waitTime = Math.min(remainingNanos, TimeUnit.SECONDS.toNanos(5));
-      LOG.error(
-          "Waiting for cache removal latch, current count: {}, wait time: {}s",
-          latch.getCount(),
-          TimeUnit.NANOSECONDS.toSeconds(waitTime));
-
-      boolean success =
-          Uninterruptibles.awaitUninterruptibly(latch, waitTime, TimeUnit.NANOSECONDS);
-
-      if (success) {
-        LOG.error("Cache removal latch triggered successfully, final count: {}", latch.getCount());
-        return true;
-      }
-
-      // If we haven't succeeded yet, try to force cache cleanup
-      LOG.error(
-          "Cache removal latch not triggered yet (count: {}), forcing cleanup. Current cache size: {}",
-          latch.getCount(),
-          getPreparedCacheSize(session));
-
-      try {
-        // Force garbage collection to help with weak references
-        System.gc();
-        Thread.sleep(100);
-      } catch (InterruptedException e) {
-        Thread.currentThread().interrupt();
-        break;
-      }
-
-      remainingNanos = timeoutNanos - (System.nanoTime() - startTime);
-    }
-
-    LOG.error(
-        "Cache removal latch failed to trigger within timeout. Final latch count: {}, cache size: {}",
-        latch.getCount(),
-        getPreparedCacheSize(session));
-    return latch.getCount() == 0; // Return true if latch reached 0 even if await timed out
-  }
-
   private void invalidationResultSetTest(
       Consumer<CqlSession> setupTestSchema, Set<String> expectedChangedTypes) {
     invalidationTestInner(
@@ -317,18 +221,13 @@ private void invalidationTestInner(
 
     try (CqlSession session = sessionWithCacheSizeMetric()) {
 
-      // Ensure we start with a clean cache
-      long initialCacheSize = getPreparedCacheSize(session);
-      LOG.error("Starting test with cache size: {}", initialCacheSize);
-      assertThat(initialCacheSize).isEqualTo(0);
-
-      // Force garbage collection to ensure clean state
-      System.gc();
-      Uninterruptibles.sleepUninterruptibly(100, TimeUnit.MILLISECONDS);
-
+      assertThat(getPreparedCacheSize(session)).isEqualTo(0);
       setupTestSchema.accept(session);
 
-      // Set up event handlers BEFORE creating prepared statements to avoid race conditions
+      session.prepare(preparedStmtQueryType1);
+      ByteBuffer queryId2 = session.prepare(preparedStmtQueryType2).getId();
+      assertThat(getPreparedCacheSize(session)).isEqualTo(2);
+
       CountDownLatch preparedStmtCacheRemoveLatch = new CountDownLatch(1);
       CountDownLatch typeChangeEventLatch = new CountDownLatch(expectedChangedTypes.size());
 
@@ -340,21 +239,11 @@ private void invalidationTestInner(
           new AtomicReference<>(Optional.empty());
       AtomicReference<Optional<String>> removedQueryEventError =
           new AtomicReference<>(Optional.empty());
-
-      LOG.error("Registering event handlers before creating prepared statements");
       ctx.getEventBus()
           .register(
               TypeChangeEvent.class,
               (e) -> {
                 // expect one event per type changed and for every parent type that nests it
-                LOG.error(
-                    "Received TypeChangeEvent for type: {} (changeType: {})",
-                    e.oldType.getName(),
-                    e.changeType);
-
-                // Add detailed debugging for cache invalidation
-                debugCacheInvalidation(session, e);
-
                 if (Boolean.TRUE.equals(
                     changedTypes.putIfAbsent(e.oldType.getName().toString(), true))) {
                   // store an error if we see duplicate change event
@@ -367,85 +256,27 @@ private void invalidationTestInner(
           .register(
               PreparedStatementRemovalEvent.class,
               (e) -> {
-                LOG.error("Received PreparedStatementRemovalEvent for queryId: {}", e.queryId);
                 if (!removedQueryIds.compareAndSet(Optional.empty(), Optional.of(e.queryId))) {
                   // store an error if we see multiple cache invalidation events
                   // any non-empty error will fail the test so it's OK to do this multiple times
                   removedQueryEventError.set(
                       Optional.of("Unable to set reference for PS removal event"));
-                  LOG.warn(
-                      "Multiple PreparedStatementRemovalEvents received, ignoring subsequent ones");
                 }
-                LOG.error(
-                    "About to countdown preparedStmtCacheRemoveLatch, current count: {}",
-                    preparedStmtCacheRemoveLatch.getCount());
                 preparedStmtCacheRemoveLatch.countDown();
-                LOG.error(
-                    "Countdown completed, new count: {}", preparedStmtCacheRemoveLatch.getCount());
               });
 
-      // Now create the prepared statements
-      PreparedStatement stmt1 = session.prepare(preparedStmtQueryType1);
-      PreparedStatement stmt2 = session.prepare(preparedStmtQueryType2);
-      ByteBuffer queryId2 = stmt2.getId();
-      assertThat(getPreparedCacheSize(session)).isEqualTo(2);
-
-      LOG.error("Prepared statements in cache:");
-      LOG.error("  Statement 1: {} (queryId: {})", preparedStmtQueryType1, stmt1.getId());
-      LOG.error("  Statement 2: {} (queryId: {})", preparedStmtQueryType2, stmt2.getId());
-
       // alter test_type_caching_2 to trigger cache invalidation and above events
-      LOG.error("Executing ALTER TYPE test_type_caching_2 add i blob");
-      LOG.error("Expected to invalidate statement 2 (queryId: {}) due to type change", queryId2);
       session.execute("ALTER TYPE test_type_caching_2 add i blob");
 
-      // Give a longer delay to allow the schema change to propagate before checking agreement
-      LOG.error("Waiting for schema change to propagate...");
-      Uninterruptibles.sleepUninterruptibly(500, TimeUnit.MILLISECONDS);
       session.checkSchemaAgreement();
 
       // wait for latches and fail if they don't reach zero before timeout
-      // Use longer timeout for cache removal as it depends on complex event chain
-      boolean typeChangeSuccess =
-          Uninterruptibles.awaitUninterruptibly(typeChangeEventLatch, 60, TimeUnit.SECONDS);
-
-      // For cache removal, use a more robust waiting mechanism with periodic cleanup
-      boolean cacheRemovalSuccess =
-          waitForCacheRemovalWithCleanup(
-              preparedStmtCacheRemoveLatch, session, 180, TimeUnit.SECONDS);
-
-      // Provide detailed diagnostics if either latch fails
-      if (!cacheRemovalSuccess || !typeChangeSuccess) {
-        String diagnostics =
-            String.format(
-                "Test failure diagnostics:\n"
-                    + "  - Cache removal latch success: %s (count: %d)\n"
-                    + "  - Type change latch success: %s (count: %d)\n"
-                    + "  - Current cache size: %d\n"
-                    + "  - Expected changed types: %s\n"
-                    + "  - Actual changed types detected: %s\n"
-                    + "  - Expected removed query ID: %s\n"
-                    + "  - Actual removed query IDs: %s\n"
-                    + "  - Type change errors: %s\n"
-                    + "  - Removal event errors: %s",
-                cacheRemovalSuccess,
-                preparedStmtCacheRemoveLatch.getCount(),
-                typeChangeSuccess,
-                typeChangeEventLatch.getCount(),
-                getPreparedCacheSize(session),
-                expectedChangedTypes,
-                changedTypes.keySet(),
-                queryId2,
-                removedQueryIds.get(),
-                typeChangeEventError.get(),
-                removedQueryEventError.get());
-        LOG.error("Prepared statement cache invalidation test failed: {}", diagnostics);
-      }
-
-      assertThat(cacheRemovalSuccess)
+      assertThat(
+              Uninterruptibles.awaitUninterruptibly(
+                  preparedStmtCacheRemoveLatch, 120, TimeUnit.SECONDS))
           .withFailMessage("preparedStmtCacheRemoveLatch did not trigger before timeout")
           .isTrue();
-      assertThat(typeChangeSuccess)
+      assertThat(Uninterruptibles.awaitUninterruptibly(typeChangeEventLatch, 20, TimeUnit.SECONDS))
           .withFailMessage("typeChangeEventLatch did not trigger before timeout")
           .isTrue();
 
