Remove tablet mappings on select events

Registers an event listener dedicated to updating the `TabletMap`
held by `Metadata`. It will trigger removal of relevant mappings whenever it's
notified of any change to the table schema or keyspace schema including
dropping them.

Introduces an integration test that verifies this listener's behaviour.

Addresses scylladb#378.

driver-core/src/main/java/com/datastax/driver/core/Metadata.java

@@ -150,6 +150,22 @@ public class Metadata {
   Metadata(Cluster.Manager cluster) {
     this.cluster = cluster;
     this.tabletMap = TabletMap.emptyMap(cluster);
+    if (cluster.configuration.getQueryOptions().isMetadataEnabled()) {
+      cluster
+          .getCluster()
+          .register(
+              new TabletMapListener(tabletMap) {
+                @Override
+                public void onRegister(Cluster cluster) {
+                  logger.info("Registered event listener for tablet map.");
+                }
+
+                @Override
+                public void onUnregister(Cluster cluster) {
+                  logger.info("Unregistered tablet map's event listener.");
+                }
+              });
+    }
   }
 
   // rebuilds the token map with the current hosts, typically when refreshing schema metadata

driver-core/src/main/java/com/datastax/driver/core/TabletMap.java

@@ -198,6 +198,24 @@ public TypeCodec<TupleValue> getTabletPayloadCodec() {
     return tabletPayloadCodec;
   }
 
+  public void removeTableMappings(KeyspaceTableNamePair key) {
+    this.mapping.remove(key);
+  }
+
+  public void removeTableMappings(String keyspace, String table) {
+    removeTableMappings(new KeyspaceTableNamePair(keyspace, table));
+  }
+
+  public void removeTableMappings(String keyspace) {
+    Iterator<TabletMap.KeyspaceTableNamePair> it = getMapping().keySet().iterator();
+    while (it.hasNext()) {
+      KeyspaceTableNamePair key = it.next();
+      if (key.getKeyspace().equals(keyspace)) {
+        it.remove();
+      }
+    }
+  }
+
   /**
    * Simple class to hold UUID of a host and a shard number. Class itself makes no checks or
    * guarantees about existence of a shard on corresponding host.

driver-core/src/main/java/com/datastax/driver/core/TabletMapListener.java

@@ -0,0 +1,29 @@
+package com.datastax.driver.core;
+
+public class TabletMapListener extends SchemaChangeListenerBase {
+  private final TabletMap tabletMap;
+
+  public TabletMapListener(TabletMap tabletMap) {
+    this.tabletMap = tabletMap;
+  }
+
+  @Override
+  public void onTableChanged(TableMetadata current, TableMetadata previous) {
+    tabletMap.removeTableMappings(previous.getKeyspace().getName(), previous.getName());
+  }
+
+  @Override
+  public void onTableRemoved(TableMetadata table) {
+    tabletMap.removeTableMappings(table.getKeyspace().getName(), table.getName());
+  }
+
+  @Override
+  public void onKeyspaceRemoved(KeyspaceMetadata keyspace) {
+    tabletMap.removeTableMappings(keyspace.getName());
+  }
+
+  @Override
+  public void onKeyspaceChanged(KeyspaceMetadata current, KeyspaceMetadata previous) {
+    tabletMap.removeTableMappings(previous.getName());
+  }
+}

driver-core/src/test/java/com/datastax/driver/core/TabletMapListenerTest.java

@@ -0,0 +1,208 @@
+package com.datastax.driver.core;
+
+import static com.datastax.driver.core.Assertions.assertThat;
+import static com.datastax.driver.core.Metadata.handleId;
+import static org.mockito.Matchers.anyObject;
+import static org.mockito.Mockito.after;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.timeout;
+import static org.mockito.Mockito.verify;
+
+import com.datastax.driver.core.policies.RoundRobinPolicy;
+import com.datastax.driver.core.utils.ScyllaVersion;
+import com.google.common.collect.Lists;
+import java.util.List;
+import java.util.concurrent.TimeUnit;
+import org.mockito.ArgumentCaptor;
+import org.testng.annotations.AfterMethod;
+import org.testng.annotations.BeforeMethod;
+import org.testng.annotations.Test;
+
+@CCMConfig(
+    createCluster = false,
+    numberOfNodes = 2,
+    jvmArgs = {
+      "--experimental-features=consistent-topology-changes",
+      "--experimental-features=tablets"
+    })
+@ScyllaVersion(minOSS = "6.0.0", minEnterprise = "2024.2", description = "Needs to support tablets")
+public class TabletMapListenerTest extends CCMTestsSupport {
+  private static final int INITIAL_TABLETS = 32;
+  private static final String KEYSPACE_NAME = "listenerTest";
+  private static final String TABLE_NAME = "testTable";
+  private static final String CREATE_TABLETS_KEYSPACE_QUERY =
+      "CREATE KEYSPACE "
+          + KEYSPACE_NAME
+          + " WITH replication = {'class': "
+          + "'NetworkTopologyStrategy', "
+          + "'replication_factor': '1'}  AND durable_writes = true AND tablets = "
+          + "{'initial': "
+          + INITIAL_TABLETS
+          + "};";
+
+  private static final String CREATE_KEYSPACE = CREATE_TABLETS_KEYSPACE_QUERY;
+  private static final String ALTER_KEYSPACE =
+      "ALTER KEYSPACE " + KEYSPACE_NAME + " WITH durable_writes = false";
+  private static final String DROP_KEYSPACE = "DROP KEYSPACE " + KEYSPACE_NAME;
+
+  private static final String CREATE_TABLE =
+      "CREATE TABLE " + KEYSPACE_NAME + "." + TABLE_NAME + "(i int primary key)";
+  private static final String INSERT_QUERY_TEMPLATE =
+      "INSERT INTO " + KEYSPACE_NAME + "." + TABLE_NAME + " (i) VALUES (%s)";
+  private static final String INSERT_ALTERED_TEMPLATE =
+      "INSERT INTO " + KEYSPACE_NAME + "." + TABLE_NAME + " (i,j) VALUES (%s,%s)";
+  private static final String SELECT_PK_WHERE =
+      "SELECT i FROM " + KEYSPACE_NAME + "." + TABLE_NAME + " WHERE  i = ?";
+  private static final String ALTER_TABLE =
+      "ALTER TABLE " + KEYSPACE_NAME + "." + TABLE_NAME + " ADD j int";
+  private static final String DROP_TABLE = "DROP TABLE " + KEYSPACE_NAME + "." + TABLE_NAME;
+
+  /** The maximum time that the test will wait to check that listeners have been notified. */
+  private static final long NOTIF_TIMEOUT_MS = TimeUnit.MINUTES.toMillis(1);
+  /** Shorter timeout for less important checks that listeners did not react to specific actions */
+  private static final long SHORT_TIMEOUT_MS = TimeUnit.SECONDS.toMillis(4);
+
+  Cluster cluster;
+  Session session;
+  SchemaChangeListener listener;
+  List<SchemaChangeListener> listeners;
+
+  @BeforeMethod(groups = "short")
+  public void setup() throws InterruptedException {
+    cluster =
+        createClusterBuilderNoDebouncing().withLoadBalancingPolicy(new RoundRobinPolicy()).build();
+
+    session = cluster.connect();
+
+    cluster.register(listener = mock(TabletMapListener.class));
+    listeners = Lists.newArrayList(listener);
+  }
+
+  // Checks for tablet removal both on table update and removal
+  @Test(groups = "short")
+  public void should_remove_tablets_on_table_alterations() throws InterruptedException {
+    session.execute(CREATE_KEYSPACE);
+    ArgumentCaptor<KeyspaceMetadata> added = null;
+    for (SchemaChangeListener listener : listeners) {
+      added = ArgumentCaptor.forClass(KeyspaceMetadata.class);
+      verify(listener, timeout(NOTIF_TIMEOUT_MS).times(1)).onKeyspaceAdded(added.capture());
+      assertThat(added.getValue()).hasName(handleId(KEYSPACE_NAME));
+    }
+    assert added != null;
+
+    TabletMap tabletMap;
+    tabletMap = cluster.getMetadata().getTabletMap();
+
+    session.execute(CREATE_TABLE);
+    assertThat(tabletMap.getMapping())
+        .doesNotContainKey(
+            new TabletMap.KeyspaceTableNamePair(handleId(KEYSPACE_NAME), handleId(TABLE_NAME)));
+
+    session.execute(String.format(INSERT_QUERY_TEMPLATE, "42"));
+    session.execute(session.prepare(SELECT_PK_WHERE).bind(42));
+    session.execute(session.prepare(SELECT_PK_WHERE).bind(42));
+    assertThat(tabletMap.getMapping())
+        .containsKey(
+            new TabletMap.KeyspaceTableNamePair(handleId(KEYSPACE_NAME), handleId(TABLE_NAME)));
+
+    session.execute(ALTER_TABLE);
+    for (SchemaChangeListener listener : listeners) {
+      ArgumentCaptor<TableMetadata> current = ArgumentCaptor.forClass(TableMetadata.class);
+      ArgumentCaptor<TableMetadata> previous = ArgumentCaptor.forClass(TableMetadata.class);
+      verify(listener, timeout(NOTIF_TIMEOUT_MS).times(1))
+          .onTableChanged(current.capture(), previous.capture());
+      assertThat(previous.getValue().getKeyspace()).hasName(handleId(KEYSPACE_NAME));
+      assertThat(previous.getValue()).hasName(handleId(TABLE_NAME));
+    }
+    assertThat(tabletMap.getMapping())
+        .doesNotContainKey(
+            new TabletMap.KeyspaceTableNamePair(handleId(KEYSPACE_NAME), handleId(TABLE_NAME)));
+
+    session.execute(String.format(INSERT_ALTERED_TEMPLATE, "42", "42"));
+    session.execute(session.prepare(SELECT_PK_WHERE).bind(42));
+    session.execute(session.prepare(SELECT_PK_WHERE).bind(42));
+    assertThat(tabletMap.getMapping())
+        .containsKey(
+            new TabletMap.KeyspaceTableNamePair(handleId(KEYSPACE_NAME), handleId(TABLE_NAME)));
+
+    session.execute(DROP_TABLE);
+    ArgumentCaptor<TableMetadata> removed = null;
+    for (SchemaChangeListener listener : listeners) {
+      removed = ArgumentCaptor.forClass(TableMetadata.class);
+      verify(listener, timeout(NOTIF_TIMEOUT_MS).times(1)).onTableRemoved(removed.capture());
+      assertThat(removed.getValue().getKeyspace()).hasName(handleId(KEYSPACE_NAME));
+      assertThat(removed.getValue()).hasName(handleId(TABLE_NAME));
+    }
+    assert removed != null;
+    assertThat(tabletMap.getMapping())
+        .doesNotContainKey(
+            new TabletMap.KeyspaceTableNamePair(handleId(KEYSPACE_NAME), handleId(TABLE_NAME)));
+
+    session.execute(DROP_KEYSPACE);
+  }
+
+  @Test(groups = "short")
+  public void should_remove_tablets_on_keyspace_alterations() {
+    session.execute(CREATE_KEYSPACE);
+    ArgumentCaptor<KeyspaceMetadata> added = null;
+    for (SchemaChangeListener listener : listeners) {
+      added = ArgumentCaptor.forClass(KeyspaceMetadata.class);
+      verify(listener, timeout(NOTIF_TIMEOUT_MS).times(1)).onKeyspaceAdded(added.capture());
+      assertThat(added.getValue()).hasName(handleId(KEYSPACE_NAME));
+    }
+    assert added != null;
+
+    TabletMap tabletMap;
+    tabletMap = cluster.getMetadata().getTabletMap();
+
+    session.execute(CREATE_TABLE);
+    session.execute(String.format(INSERT_QUERY_TEMPLATE, "42"));
+    session.execute(session.prepare(SELECT_PK_WHERE).bind(42));
+    session.execute(session.prepare(SELECT_PK_WHERE).bind(42));
+    assertThat(tabletMap.getMapping())
+        .containsKey(
+            new TabletMap.KeyspaceTableNamePair(handleId(KEYSPACE_NAME), handleId(TABLE_NAME)));
+
+    assertThat(cluster.getMetadata().getKeyspace(KEYSPACE_NAME).isDurableWrites()).isTrue();
+
+    session.execute(ALTER_KEYSPACE);
+    assertThat(cluster.getMetadata().getKeyspace(KEYSPACE_NAME)).isNotDurableWrites();
+    for (SchemaChangeListener listener : listeners) {
+      ArgumentCaptor<KeyspaceMetadata> current = ArgumentCaptor.forClass(KeyspaceMetadata.class);
+      ArgumentCaptor<KeyspaceMetadata> previous = ArgumentCaptor.forClass(KeyspaceMetadata.class);
+      verify(listener, timeout(NOTIF_TIMEOUT_MS).times(1))
+          .onKeyspaceChanged(current.capture(), previous.capture());
+      assertThat(previous.getValue()).hasName(handleId(KEYSPACE_NAME));
+    }
+    for (SchemaChangeListener listener : listeners) {
+      verify(listener, after((int) SHORT_TIMEOUT_MS).never())
+          .onTableChanged(anyObject(), anyObject());
+    }
+    assertThat(tabletMap.getMapping())
+        .doesNotContainKey(
+            new TabletMap.KeyspaceTableNamePair(handleId(KEYSPACE_NAME), handleId(TABLE_NAME)));
+
+    session.execute(session.prepare(SELECT_PK_WHERE).bind(42));
+    session.execute(session.prepare(SELECT_PK_WHERE).bind(42));
+
+    assertThat(tabletMap.getMapping())
+        .containsKey(
+            new TabletMap.KeyspaceTableNamePair(handleId(KEYSPACE_NAME), handleId(TABLE_NAME)));
+
+    session.execute(DROP_KEYSPACE);
+    for (SchemaChangeListener listener : listeners) {
+      ArgumentCaptor<KeyspaceMetadata> removed = ArgumentCaptor.forClass(KeyspaceMetadata.class);
+      verify(listener, timeout(NOTIF_TIMEOUT_MS).times(1)).onKeyspaceRemoved(removed.capture());
+      assertThat(removed.getValue()).hasName(handleId(KEYSPACE_NAME));
+    }
+    assertThat(tabletMap.getMapping())
+        .doesNotContainKey(
+            new TabletMap.KeyspaceTableNamePair(handleId(KEYSPACE_NAME), handleId(TABLE_NAME)));
+  }
+
+  @AfterMethod(groups = "short", alwaysRun = true)
+  public void teardown() {
+    if (session != null) session.close();
+    if (cluster != null) cluster.close();
+  }
+}