register listener first, create prepare later

Author: SiyaoIsHiding

integration-tests/src/test/java/com/datastax/oss/driver/core/cql/PreparedStatementCachingIT.java

@@ -239,20 +239,34 @@ private boolean waitForCacheRemovalWithCleanup(
     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");
+        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, forcing cleanup. Current cache size: {}",
+          "Cache removal latch not triggered yet (count: {}), forcing cleanup. Current cache size: {}",
+          latch.getCount(),
           getPreparedCacheSize(session));
 
       try {
@@ -268,9 +282,10 @@ private boolean waitForCacheRemovalWithCleanup(
     }
 
     LOG.error(
-        "Cache removal latch failed to trigger within timeout. Final cache size: {}",
+        "Cache removal latch failed to trigger within timeout. Final latch count: {}, cache size: {}",
+        latch.getCount(),
         getPreparedCacheSize(session));
-    return false;
+    return latch.getCount() == 0; // Return true if latch reached 0 even if await timed out
   }
 
   private void invalidationResultSetTest(
@@ -313,15 +328,7 @@ private void invalidationTestInner(
 
       setupTestSchema.accept(session);
 
-      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());
-
+      // Set up event handlers BEFORE creating prepared statements to avoid race conditions
       CountDownLatch preparedStmtCacheRemoveLatch = new CountDownLatch(1);
       CountDownLatch typeChangeEventLatch = new CountDownLatch(expectedChangedTypes.size());
 
@@ -333,6 +340,8 @@ 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,
@@ -367,9 +376,24 @@ private void invalidationTestInner(
                   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);
@@ -380,10 +404,6 @@ private void invalidationTestInner(
       Uninterruptibles.sleepUninterruptibly(500, TimeUnit.MILLISECONDS);
       session.checkSchemaAgreement();
 
-      // Additional delay to allow event processing to complete
-      LOG.error("Waiting for event processing to complete...");
-      Uninterruptibles.sleepUninterruptibly(1000, TimeUnit.MILLISECONDS);
-
       // 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 =