From 02b72d6680ce51fa14c254cf2426b5cdab69822c Mon Sep 17 00:00:00 2001
From: Bouncheck <36934780+Bouncheck@users.noreply.github.com>
Date: Tue, 8 Apr 2025 10:58:04 +0200
Subject: [PATCH] Occasionally force cache cleanup in
 PreparedStatementCachingIT

This integration test periodically fails on one of its methods.
Evidence points to the cache itself being the culprit, as there is
no guarantee that elements will be removed right away.
This leads to removal event not being fired and timing out on wait for it.
Calling `cache.cleanUp()` seems to greatly improve stability of this test.

Fixes #476 and #477
---
 .../core/cql/PreparedStatementCachingIT.java  | 30 +++++++++++++++----
 1 file changed, 25 insertions(+), 5 deletions(-)

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
index 617d489fb95..f5dc49cc03e 100644
--- 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
@@ -267,11 +267,17 @@ private void invalidationTestInner(
       session.execute("ALTER TYPE test_type_2 add i blob");
 
       // wait for latches and fail if they don't reach zero before timeout
-      assertThat(
-              Uninterruptibles.awaitUninterruptibly(
-                  preparedStmtCacheRemoveLatch, 10, TimeUnit.SECONDS))
-          .withFailMessage("preparedStmtCacheRemoveLatch did not trigger before timeout")
-          .isTrue();
+      if (!Uninterruptibles.awaitUninterruptibly(
+          preparedStmtCacheRemoveLatch, 10, TimeUnit.SECONDS)) {
+        // On rare occasions cache does not trigger removal shortly after alter.
+        forceCacheCleanUp((DefaultDriverContext) session.getContext());
+        assertThat(
+                Uninterruptibles.awaitUninterruptibly(
+                    preparedStmtCacheRemoveLatch, 10, TimeUnit.SECONDS))
+            .withFailMessage("preparedStmtCacheRemoveLatch did not trigger before timeout")
+            .isTrue();
+      }
+
       assertThat(Uninterruptibles.awaitUninterruptibly(typeChangeEventLatch, 10, TimeUnit.SECONDS))
           .withFailMessage("typeChangeEventLatch did not trigger before timeout")
           .isTrue();
@@ -426,4 +432,18 @@ private static long getPreparedCacheSize(CqlSession session) {
                     "Could not access metric "
                         + DefaultSessionMetric.CQL_PREPARED_CACHE_SIZE.getPath()));
   }
+
+  private void forceCacheCleanUp(DefaultDriverContext context) {
+    context
+        .getRequestProcessorRegistry()
+        .getProcessors()
+        .forEach(
+            requestProcessor -> {
+              if (requestProcessor instanceof CqlPrepareAsyncProcessor) {
+                ((CqlPrepareAsyncProcessor) requestProcessor).getCache().cleanUp();
+              } else if (requestProcessor instanceof CqlPrepareSyncProcessor) {
+                ((CqlPrepareSyncProcessor) requestProcessor).getCache().cleanUp();
+              }
+            });
+  }
 }