CASSJAVA-104: Fix Flaky tests

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()) {

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;
@@ -135,7 +136,7 @@ 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)));
+      super(context, CacheBuilder.newBuilder().removalListener(buildCacheRemoveCallback(context)));
     }
   }
 
@@ -194,8 +195,8 @@ private void invalidationResultSetTest(
       Consumer<CqlSession> setupTestSchema, Set<String> expectedChangedTypes) {
     invalidationTestInner(
         setupTestSchema,
-        "select f from test_table_1 where e = ?",
-        "select h from test_table_2 where g = ?",
+        "select f from test_table_caching_1 where e = ?",
+        "select h from test_table_caching_2 where g = ?",
         expectedChangedTypes);
   }
 
@@ -206,8 +207,8 @@ private void invalidationVariableDefsTest(
     String condition = isCollection ? "contains ?" : "= ?";
     invalidationTestInner(
         setupTestSchema,
-        String.format("select e from test_table_1 where f %s allow filtering", condition),
-        String.format("select g from test_table_2 where h %s allow filtering", condition),
+        String.format("select e from test_table_caching_1 where f %s allow filtering", condition),
+        String.format("select g from test_table_caching_2 where h %s allow filtering", condition),
         expectedChangedTypes);
   }
 
@@ -263,16 +264,18 @@ private void invalidationTestInner(
                 preparedStmtCacheRemoveLatch.countDown();
               });
 
-      // alter test_type_2 to trigger cache invalidation and above events
-      session.execute("ALTER TYPE test_type_2 add i blob");
+      // alter test_type_caching_2 to trigger cache invalidation and above events
+      session.execute("ALTER TYPE test_type_caching_2 add i blob");
+
+      session.checkSchemaAgreement();
 
       // wait for latches and fail if they don't reach zero before timeout
       assertThat(
               Uninterruptibles.awaitUninterruptibly(
-                  preparedStmtCacheRemoveLatch, 10, TimeUnit.SECONDS))
+                  preparedStmtCacheRemoveLatch, 120, TimeUnit.SECONDS))
           .withFailMessage("preparedStmtCacheRemoveLatch did not trigger before timeout")
           .isTrue();
-      assertThat(Uninterruptibles.awaitUninterruptibly(typeChangeEventLatch, 10, TimeUnit.SECONDS))
+      assertThat(Uninterruptibles.awaitUninterruptibly(typeChangeEventLatch, 20, TimeUnit.SECONDS))
           .withFailMessage("typeChangeEventLatch did not trigger before timeout")
           .isTrue();
 
@@ -295,17 +298,20 @@ private void invalidationTestInner(
 
   Consumer<CqlSession> setupCacheEntryTestBasic =
       (session) -> {
-        session.execute("CREATE TYPE test_type_1 (a text, b int)");
-        session.execute("CREATE TYPE test_type_2 (c int, d text)");
-        session.execute("CREATE TABLE test_table_1 (e int primary key, f frozen<test_type_1>)");
-        session.execute("CREATE TABLE test_table_2 (g int primary key, h frozen<test_type_2>)");
+        session.execute("CREATE TYPE test_type_caching_1 (a text, b int)");
+        session.execute("CREATE TYPE test_type_caching_2 (c int, d text)");
+        session.execute(
+            "CREATE TABLE test_table_caching_1 (e int primary key, f frozen<test_type_caching_1>)");
+        session.execute(
+            "CREATE TABLE test_table_caching_2 (g int primary key, h frozen<test_type_caching_2>)");
       };
 
   @Test
   public void should_invalidate_cache_entry_on_basic_udt_change_result_set() {
     SchemaChangeSynchronizer.withLock(
         () -> {
-          invalidationResultSetTest(setupCacheEntryTestBasic, ImmutableSet.of("test_type_2"));
+          invalidationResultSetTest(
+              setupCacheEntryTestBasic, ImmutableSet.of("test_type_caching_2"));
         });
   }
 
@@ -314,25 +320,26 @@ public void should_invalidate_cache_entry_on_basic_udt_change_variable_defs() {
     SchemaChangeSynchronizer.withLock(
         () -> {
           invalidationVariableDefsTest(
-              setupCacheEntryTestBasic, false, ImmutableSet.of("test_type_2"));
+              setupCacheEntryTestBasic, false, ImmutableSet.of("test_type_caching_2"));
         });
   }
 
   Consumer<CqlSession> setupCacheEntryTestCollection =
       (session) -> {
-        session.execute("CREATE TYPE test_type_1 (a text, b int)");
-        session.execute("CREATE TYPE test_type_2 (c int, d text)");
+        session.execute("CREATE TYPE test_type_caching_1 (a text, b int)");
+        session.execute("CREATE TYPE test_type_caching_2 (c int, d text)");
         session.execute(
-            "CREATE TABLE test_table_1 (e int primary key, f list<frozen<test_type_1>>)");
+            "CREATE TABLE test_table_caching_1 (e int primary key, f list<frozen<test_type_caching_1>>)");
         session.execute(
-            "CREATE TABLE test_table_2 (g int primary key, h list<frozen<test_type_2>>)");
+            "CREATE TABLE test_table_caching_2 (g int primary key, h list<frozen<test_type_caching_2>>)");
       };
 
   @Test
   public void should_invalidate_cache_entry_on_collection_udt_change_result_set() {
     SchemaChangeSynchronizer.withLock(
         () -> {
-          invalidationResultSetTest(setupCacheEntryTestCollection, ImmutableSet.of("test_type_2"));
+          invalidationResultSetTest(
+              setupCacheEntryTestCollection, ImmutableSet.of("test_type_caching_2"));
         });
   }
 
@@ -341,25 +348,26 @@ public void should_invalidate_cache_entry_on_collection_udt_change_variable_defs
     SchemaChangeSynchronizer.withLock(
         () -> {
           invalidationVariableDefsTest(
-              setupCacheEntryTestCollection, true, ImmutableSet.of("test_type_2"));
+              setupCacheEntryTestCollection, true, ImmutableSet.of("test_type_caching_2"));
         });
   }
 
   Consumer<CqlSession> setupCacheEntryTestTuple =
       (session) -> {
-        session.execute("CREATE TYPE test_type_1 (a text, b int)");
-        session.execute("CREATE TYPE test_type_2 (c int, d text)");
+        session.execute("CREATE TYPE test_type_caching_1 (a text, b int)");
+        session.execute("CREATE TYPE test_type_caching_2 (c int, d text)");
         session.execute(
-            "CREATE TABLE test_table_1 (e int primary key, f tuple<int, test_type_1, text>)");
+            "CREATE TABLE test_table_caching_1 (e int primary key, f tuple<int, test_type_caching_1, text>)");
         session.execute(
-            "CREATE TABLE test_table_2 (g int primary key, h tuple<text, test_type_2, int>)");
+            "CREATE TABLE test_table_caching_2 (g int primary key, h tuple<text, test_type_caching_2, int>)");
       };
 
   @Test
   public void should_invalidate_cache_entry_on_tuple_udt_change_result_set() {
     SchemaChangeSynchronizer.withLock(
         () -> {
-          invalidationResultSetTest(setupCacheEntryTestTuple, ImmutableSet.of("test_type_2"));
+          invalidationResultSetTest(
+              setupCacheEntryTestTuple, ImmutableSet.of("test_type_caching_2"));
         });
   }
 
@@ -368,26 +376,29 @@ public void should_invalidate_cache_entry_on_tuple_udt_change_variable_defs() {
     SchemaChangeSynchronizer.withLock(
         () -> {
           invalidationVariableDefsTest(
-              setupCacheEntryTestTuple, false, ImmutableSet.of("test_type_2"));
+              setupCacheEntryTestTuple, false, ImmutableSet.of("test_type_caching_2"));
         });
   }
 
   Consumer<CqlSession> setupCacheEntryTestNested =
       (session) -> {
-        session.execute("CREATE TYPE test_type_1 (a text, b int)");
-        session.execute("CREATE TYPE test_type_2 (c int, d text)");
-        session.execute("CREATE TYPE test_type_3 (e frozen<test_type_1>, f int)");
-        session.execute("CREATE TYPE test_type_4 (g int, h frozen<test_type_2>)");
-        session.execute("CREATE TABLE test_table_1 (e int primary key, f frozen<test_type_3>)");
-        session.execute("CREATE TABLE test_table_2 (g int primary key, h frozen<test_type_4>)");
+        session.execute("CREATE TYPE test_type_caching_1 (a text, b int)");
+        session.execute("CREATE TYPE test_type_caching_2 (c int, d text)");
+        session.execute("CREATE TYPE test_type_caching_3 (e frozen<test_type_caching_1>, f int)");
+        session.execute("CREATE TYPE test_type_caching_4 (g int, h frozen<test_type_caching_2>)");
+        session.execute(
+            "CREATE TABLE test_table_caching_1 (e int primary key, f frozen<test_type_caching_3>)");
+        session.execute(
+            "CREATE TABLE test_table_caching_2 (g int primary key, h frozen<test_type_caching_4>)");
       };
 
   @Test
   public void should_invalidate_cache_entry_on_nested_udt_change_result_set() {
     SchemaChangeSynchronizer.withLock(
         () -> {
           invalidationResultSetTest(
-              setupCacheEntryTestNested, ImmutableSet.of("test_type_2", "test_type_4"));
+              setupCacheEntryTestNested,
+              ImmutableSet.of("test_type_caching_2", "test_type_caching_4"));
         });
   }
 
@@ -396,7 +407,9 @@ public void should_invalidate_cache_entry_on_nested_udt_change_variable_defs() {
     SchemaChangeSynchronizer.withLock(
         () -> {
           invalidationVariableDefsTest(
-              setupCacheEntryTestNested, false, ImmutableSet.of("test_type_2", "test_type_4"));
+              setupCacheEntryTestNested,
+              false,
+              ImmutableSet.of("test_type_caching_2", "test_type_caching_4"));
         });
   }
 

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

@@ -21,20 +21,34 @@
 import static org.junit.Assert.fail;
 
 import com.datastax.oss.driver.api.core.CqlSession;
+import com.datastax.oss.driver.api.core.config.DriverConfigLoader;
+import com.datastax.oss.driver.api.core.context.DriverContext;
 import com.datastax.oss.driver.api.core.cql.PrepareRequest;
 import com.datastax.oss.driver.api.core.cql.PreparedStatement;
+import com.datastax.oss.driver.api.core.session.ProgrammaticArguments;
+import com.datastax.oss.driver.api.core.session.SessionBuilder;
 import com.datastax.oss.driver.api.testinfra.ccm.CustomCcmRule;
 import com.datastax.oss.driver.api.testinfra.session.SessionRule;
 import com.datastax.oss.driver.api.testinfra.session.SessionUtils;
 import com.datastax.oss.driver.categories.IsolatedTests;
 import com.datastax.oss.driver.internal.core.context.DefaultDriverContext;
 import com.datastax.oss.driver.internal.core.cql.CqlPrepareAsyncProcessor;
+import com.datastax.oss.driver.internal.core.cql.CqlPrepareSyncProcessor;
+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.base.Predicates;
 import com.datastax.oss.driver.shaded.guava.common.cache.Cache;
+import com.datastax.oss.driver.shaded.guava.common.cache.CacheBuilder;
 import com.datastax.oss.driver.shaded.guava.common.collect.Iterables;
+import edu.umd.cs.findbugs.annotations.NonNull;
+import java.util.List;
+import java.util.Optional;
 import java.util.concurrent.CompletableFuture;
 import org.junit.After;
+import org.junit.AfterClass;
 import org.junit.Before;
+import org.junit.BeforeClass;
 import org.junit.Rule;
 import org.junit.Test;
 import org.junit.experimental.categories.Category;
@@ -50,6 +64,69 @@ public class PreparedStatementCancellationIT {
 
   @Rule public TestRule chain = RuleChain.outerRule(ccmRule).around(sessionRule);
 
+  private static class TestCqlPrepareAsyncProcessor extends CqlPrepareAsyncProcessor {
+
+    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, CacheBuilder.newBuilder());
+    }
+  }
+
+  private static class TestDefaultDriverContext extends DefaultDriverContext {
+    public TestDefaultDriverContext(
+        DriverConfigLoader configLoader, ProgrammaticArguments programmaticArguments) {
+      super(configLoader, programmaticArguments);
+    }
+
+    @Override
+    protected RequestProcessorRegistry buildRequestProcessorRegistry() {
+      // Re-create the processor cache to insert the TestCqlPrepareAsyncProcessor with it's strong
+      // prepared statement cache, see JAVA-3062
+      List<RequestProcessor<?, ?>> processors =
+          BuiltInRequestProcessors.createDefaultProcessors(this);
+      processors.removeIf((processor) -> processor instanceof CqlPrepareAsyncProcessor);
+      processors.removeIf((processor) -> processor instanceof CqlPrepareSyncProcessor);
+      CqlPrepareAsyncProcessor asyncProcessor =
+          new PreparedStatementCancellationIT.TestCqlPrepareAsyncProcessor(Optional.of(this));
+      processors.add(2, asyncProcessor);
+      processors.add(3, new CqlPrepareSyncProcessor(asyncProcessor));
+      return new RequestProcessorRegistry(
+          getSessionName(), processors.toArray(new RequestProcessor[0]));
+    }
+  }
+
+  private static class TestSessionBuilder extends SessionBuilder {
+
+    @Override
+    protected Object wrap(@NonNull CqlSession defaultSession) {
+      return defaultSession;
+    }
+
+    @Override
+    protected DriverContext buildContext(
+        DriverConfigLoader configLoader, ProgrammaticArguments programmaticArguments) {
+      return new PreparedStatementCancellationIT.TestDefaultDriverContext(
+          configLoader, programmaticArguments);
+    }
+  }
+
+  @BeforeClass
+  public static void setupBeforeClass() {
+    System.setProperty(
+        SessionUtils.SESSION_BUILDER_CLASS_PROPERTY,
+        PreparedStatementCancellationIT.class.getName());
+  }
+
+  @AfterClass
+  public static void teardownAfterClass() {
+    System.clearProperty(SessionUtils.SESSION_BUILDER_CLASS_PROPERTY);
+  }
+
+  public static SessionBuilder builder() {
+    return new PreparedStatementCancellationIT.TestSessionBuilder();
+  }
+
   @Before
   public void setup() {
 

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

@@ -31,10 +31,10 @@
 import com.datastax.oss.driver.api.core.cql.ExecutionInfo;
 import com.datastax.oss.driver.api.core.cql.PreparedStatement;
 import com.datastax.oss.driver.api.core.cql.SimpleStatement;
+import com.datastax.oss.driver.api.core.cql.Statement;
 import com.datastax.oss.driver.api.testinfra.ccm.CcmRule;
 import com.datastax.oss.driver.api.testinfra.ccm.SchemaChangeSynchronizer;
 import com.datastax.oss.driver.api.testinfra.session.SessionRule;
-import com.datastax.oss.driver.categories.ParallelizableTests;
 import com.datastax.oss.driver.internal.core.cql.EmptyColumnDefinitions;
 import com.tngtech.java.junit.dataprovider.DataProvider;
 import com.tngtech.java.junit.dataprovider.DataProviderRunner;
@@ -47,13 +47,11 @@
 import org.junit.BeforeClass;
 import org.junit.ClassRule;
 import org.junit.Test;
-import org.junit.experimental.categories.Category;
 import org.junit.rules.RuleChain;
 import org.junit.rules.TestRule;
 import org.junit.runner.RunWith;
 
 @RunWith(DataProviderRunner.class)
-@Category(ParallelizableTests.class)
 public class DefaultReactiveResultSetIT {
 
   private static CcmRule ccmRule = CcmRule.getInstance();
@@ -67,19 +65,15 @@ public static void initialize() {
     CqlSession session = sessionRule.session();
     SchemaChangeSynchronizer.withLock(
         () -> {
-          session.execute("DROP TABLE IF EXISTS test_reactive_read");
-          session.execute("DROP TABLE IF EXISTS test_reactive_write");
+          session.execute(createSlowStatement("DROP TABLE IF EXISTS test_reactive_read"));
+          session.execute(createSlowStatement("DROP TABLE IF EXISTS test_reactive_write"));
           session.checkSchemaAgreement();
           session.execute(
-              SimpleStatement.builder(
-                      "CREATE TABLE test_reactive_read (pk int, cc int, v int, PRIMARY KEY ((pk), cc))")
-                  .setExecutionProfile(sessionRule.slowProfile())
-                  .build());
+              createSlowStatement(
+                  "CREATE TABLE test_reactive_read (pk int, cc int, v int, PRIMARY KEY ((pk), cc))"));
           session.execute(
-              SimpleStatement.builder(
-                      "CREATE TABLE test_reactive_write (pk int, cc int, v int, PRIMARY KEY ((pk), cc))")
-                  .setExecutionProfile(sessionRule.slowProfile())
-                  .build());
+              createSlowStatement(
+                  "CREATE TABLE test_reactive_write (pk int, cc int, v int, PRIMARY KEY ((pk), cc))"));
           session.checkSchemaAgreement();
         });
     for (int i = 0; i < 1000; i++) {
@@ -92,6 +86,12 @@ public static void initialize() {
     }
   }
 
+  static Statement<?> createSlowStatement(String statement) {
+    return SimpleStatement.builder(statement)
+        .setExecutionProfile(sessionRule.slowProfile())
+        .build();
+  }
+
   @Before
   public void truncateTables() throws Exception {
     CqlSession session = sessionRule.session();