Add per statement tests for tablet

Goal is to ensure that server provides tablet info and it reaches
session tablet map for every request possible.

Author: dkropachev

integration-tests/src/test/java/com/datastax/oss/driver/core/metadata/DefaultMetadataTabletMapIT.java

@@ -6,21 +6,31 @@
 import com.datastax.oss.driver.api.core.cql.PreparedStatement;
 import com.datastax.oss.driver.api.core.cql.ResultSet;
 import com.datastax.oss.driver.api.core.cql.SimpleStatement;
+import com.datastax.oss.driver.api.core.cql.Statement;
 import com.datastax.oss.driver.api.core.metadata.KeyspaceTableNamePair;
+import com.datastax.oss.driver.api.core.metadata.Node;
 import com.datastax.oss.driver.api.core.metadata.Tablet;
 import com.datastax.oss.driver.api.testinfra.CassandraSkip;
 import com.datastax.oss.driver.api.testinfra.ScyllaRequirement;
 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.internal.core.protocol.TabletInfo;
+
 import java.nio.ByteBuffer;
 import java.time.Duration;
+import java.util.ArrayList;
+import java.util.HashMap;
+import java.util.List;
 import java.util.Map;
 import java.util.Set;
 import java.util.concurrent.ConcurrentMap;
 import java.util.concurrent.ConcurrentSkipListSet;
+import java.util.function.Function;
+import java.util.function.Supplier;
+
 import org.junit.Assert;
+import org.junit.BeforeClass;
 import org.junit.ClassRule;
 import org.junit.Test;
 import org.junit.rules.RuleChain;
@@ -55,9 +65,9 @@ public class DefaultMetadataTabletMapIT {
   private static final int INITIAL_TABLETS = 32;
   private static final int QUERIES = 1600;
   private static final int REPLICATION_FACTOR = 2;
-  private static String KEYSPACE_NAME = "tabletsTest";
-  private static String TABLE_NAME = "tabletsTable";
-  private static String CREATE_KEYSPACE_QUERY =
+  private static final String KEYSPACE_NAME = "tabletsTest";
+  private static final String TABLE_NAME = "tabletsTable";
+  private static final String CREATE_KEYSPACE_QUERY =
       "CREATE KEYSPACE IF NOT EXISTS "
           + KEYSPACE_NAME
           + " WITH replication = {'class': "
@@ -68,49 +78,208 @@ public class DefaultMetadataTabletMapIT {
           + "{'initial': "
           + INITIAL_TABLETS
           + "};";
-  private static String CREATE_TABLE_QUERY =
+  private static final String CREATE_TABLE_QUERY =
       "CREATE TABLE IF NOT EXISTS "
           + KEYSPACE_NAME
           + "."
           + TABLE_NAME
-          + " (pk int, ck int, PRIMARY KEY(pk, ck));";
+          + " (pk int, ck int, val int, PRIMARY KEY(pk, ck));";
 
-  @Test
-  public void should_receive_each_tablet_exactly_once() {
-    CqlSession session = SESSION_RULE.session();
+  private static final SimpleStatement STMT_INSERT = SimpleStatement.builder(
+          "INSERT INTO "
+              + KEYSPACE_NAME
+              + "."
+              + TABLE_NAME
+              + " (pk,ck) VALUES (?,?);")
+      .setTracing(true)
+      .build();
 
-    session.execute(CREATE_KEYSPACE_QUERY);
-    session.execute(CREATE_TABLE_QUERY);
+  private static final SimpleStatement STMT_INSERT_NO_KS = SimpleStatement.builder(
+          "INSERT INTO "
+              + TABLE_NAME
+              + " (pk,ck) VALUES (?,?);")
+      .setTracing(true)
+      .build();
 
-    for (int i = 1; i <= QUERIES; i++) {
-      session.execute(
+  private static final SimpleStatement STMT_INSERT_CONCRETE = SimpleStatement.builder(
           "INSERT INTO "
               + KEYSPACE_NAME
               + "."
               + TABLE_NAME
-              + " (pk,ck) VALUES ("
-              + i
-              + ","
-              + i
-              + ");");
+              + " (pk,ck) VALUES (1,1);")
+      .setTracing(true)
+      .build();
+
+  private static final SimpleStatement STMT_SELECT = SimpleStatement.builder(
+          "SELECT pk,ck FROM "
+              + KEYSPACE_NAME
+              + "."
+              + TABLE_NAME
+              + " WHERE pk = ? AND ck = ?")
+      .setTracing(true)
+      .build();
+
+  private static final SimpleStatement STMT_SELECT_NO_KS = SimpleStatement.builder(
+          "SELECT pk,ck FROM "
+              + TABLE_NAME
+              + " WHERE pk = ? AND ck = ?")
+      .setTracing(true)
+      .build();
+
+  private static final SimpleStatement STMT_SELECT_CONCRETE = SimpleStatement.builder(
+          "SELECT pk,ck FROM "
+              + KEYSPACE_NAME
+              + "."
+              + TABLE_NAME
+              + " WHERE pk = 1 AND ck = 1")
+      .setTracing(true)
+      .build();
+
+  private static final SimpleStatement STMT_UPDATE = SimpleStatement.builder(
+          "UPDATE "
+              + KEYSPACE_NAME
+              + "."
+              + TABLE_NAME
+              + " SET val = 1"
+              + " WHERE pk = ? AND ck = ?")
+      .setTracing(true)
+      .build();
+
+  private static final SimpleStatement STMT_UPDATE_NO_KS = SimpleStatement.builder(
+          "UPDATE "
+              + TABLE_NAME
+              + " SET val = 1"
+              + " WHERE pk = ? AND ck = ?")
+      .setTracing(true)
+      .build();
+
+  private static final SimpleStatement STMT_UPDATE_CONCRETE = SimpleStatement.builder(
+          "UPDATE "
+              + KEYSPACE_NAME
+              + "."
+              + TABLE_NAME
+              + " SET val = 1"
+              + " WHERE pk = 1 AND ck = 1")
+      .setTracing(true)
+      .build();
+
+  private static final SimpleStatement STMT_DELETE = SimpleStatement.builder(
+          "DELETE FROM "
+              + KEYSPACE_NAME
+              + "."
+              + TABLE_NAME
+              + " WHERE pk = ? AND ck = ?")
+      .setTracing(true)
+      .build();
+
+  private static final SimpleStatement STMT_DELETE_NO_KS = SimpleStatement.builder(
+          "DELETE FROM "
+              + TABLE_NAME
+              + " WHERE pk = ? AND ck = ?")
+      .setTracing(true)
+      .build();
+
+  private static final SimpleStatement STMT_DELETE_CONCRETE = SimpleStatement.builder(
+          "DELETE FROM "
+              + KEYSPACE_NAME
+              + "."
+              + TABLE_NAME
+              + " WHERE pk = 1 AND ck = 1")
+      .setTracing(true)
+      .build();
+
+  @BeforeClass
+  public static void setup() {
+    CqlSession session = SESSION_RULE.session();
+    session.execute(CREATE_KEYSPACE_QUERY);
+    session.execute(CREATE_TABLE_QUERY);
+  }
+
+  @Test
+  public void every_statement_should_get_tablet_info() {
+    Map<String, Supplier<CqlSession>> sessions = new HashMap<>();
+    sessions.put("REGULAR", () -> CqlSession.builder().addContactEndPoints(CCM_RULE.getContactPoints()).build());
+    sessions.put("WITH_KEYSPACE", () -> CqlSession.builder().addContactEndPoints(CCM_RULE.getContactPoints()).withKeyspace(KEYSPACE_NAME).build());
+    sessions.put("USE_KEYSPACE", () -> {
+      CqlSession s = CqlSession.builder().addContactEndPoints(CCM_RULE.getContactPoints()).build();
+      s.execute("USE " + KEYSPACE_NAME);
+      return s;
+    });
+
+    Map<String, Function<CqlSession, Statement>> statements = new HashMap<>();
+    statements.put("SELECT_CONCRETE", s -> STMT_SELECT_CONCRETE);
+    statements.put("SELECT_PREPARED", s -> s.prepare(STMT_SELECT).bind(2, 2));
+    statements.put("SELECT_NO_KS_PREPARED", s -> s.prepare(STMT_SELECT_NO_KS).bind(2, 2));
+    statements.put("SELECT_CONCRETE_PREPARED", s -> s.prepare(STMT_SELECT_CONCRETE).bind());
+    statements.put("INSERT_CONCRETE", s -> STMT_INSERT_CONCRETE);
+    statements.put("INSERT_PREPARED", s -> s.prepare(STMT_INSERT).bind(2, 2));
+    statements.put("INSERT_NO_KS_PREPARED", s -> s.prepare(STMT_INSERT_NO_KS).bind(2, 2));
+    statements.put("INSERT_CONCRETE_PREPARED", s -> s.prepare(STMT_INSERT_CONCRETE).bind());
+    statements.put("UPDATE_CONCRETE", s -> STMT_UPDATE_CONCRETE);
+    statements.put("UPDATE_PREPARED", s -> s.prepare(STMT_UPDATE).bind(2, 2));
+    statements.put("UPDATE_NO_KS_PREPARED", s -> s.prepare(STMT_UPDATE_NO_KS).bind(2, 2));
+    statements.put("UPDATE_CONCRETE_PREPARED", s -> s.prepare(STMT_UPDATE_CONCRETE).bind());
+    statements.put("DELETE_CONCRETE", s -> STMT_DELETE_CONCRETE);
+    statements.put("DELETE_PREPARED", s -> s.prepare(STMT_DELETE).bind(2, 2));
+    statements.put("DELETE_NO_KS_PREPARED", s -> s.prepare(STMT_DELETE_NO_KS).bind(2, 2));
+    statements.put("DELETE_CONCRETE_PREPARED", s -> s.prepare(STMT_DELETE_CONCRETE).bind());
+
+    List<String> noTabletQueries = new ArrayList<>();
+    for (Map.Entry<String, Supplier<CqlSession>> sessionEntry : sessions.entrySet()) {
+      for (Map.Entry<String, Function<CqlSession, Statement>> stmtEntry : statements.entrySet()) {
+        if (stmtEntry.getKey().contains("CONCRETE")) {
+          // Broken on server side
+          continue;
+        }
+        if (sessionEntry.getKey().equals("REGULAR") && stmtEntry.getKey().contains("NO_KS")) {
+          // Statements without KS will fail on the session with no ks specified
+          continue;
+        }
+        CqlSession session = sessionEntry.getValue().get();
+        Statement stmt;
+        try {
+          stmt = stmtEntry.getValue().apply(session);
+        } catch (Exception e) {
+          RuntimeException ex = new RuntimeException(String.format("Failed to build statement %s on session %s", stmtEntry.getKey(), sessionEntry.getKey()));
+          ex.addSuppressed(e);
+          throw ex;
+        }
+        if (!executeOnAllHostsAndReturnIfResultHasTabletsInfo(session, stmt)) {
+          noTabletQueries.add(String.format("Statement %s on session %s got no tablet into", stmtEntry.getKey(), sessionEntry.getKey()));
+          continue;
+        }
+        if (!isSessionLearnedTabletInfo(session)) {
+          noTabletQueries.add(String.format("Statement %s on session %s did not trigger session tablets update", stmtEntry.getKey(), sessionEntry.getKey()));
+        }
+      }
     }
 
-    PreparedStatement preparedStatement =
-        session.prepare(
-            SimpleStatement.builder(
-                    "select pk,ck from "
-                        + KEYSPACE_NAME
-                        + "."
-                        + TABLE_NAME
-                        + " WHERE pk = ? AND ck = ?")
-                .setTracing(true)
-                .build());
-    // preparedStatement.enableTracing();
+    if (!noTabletQueries.isEmpty()) {
+      throw new AssertionError(String.format("Found queries that got no tablet info: \n%s", String.join("\n", noTabletQueries)));
+    }
+  }
+
+  @Test
+  public void should_receive_each_tablet_exactly_once() {
+    CqlSession session = CqlSession.builder().addContactEndPoints(CCM_RULE.getContactPoints()).build();
     int counter = 0;
+    PreparedStatement preparedStatement =
+        session.prepare(STMT_INSERT);
     for (int i = 1; i <= QUERIES; i++) {
-      ResultSet rs = session.execute(preparedStatement.bind(i, i).setTracing(true));
-      Map<String, ByteBuffer> payload = rs.getExecutionInfo().getIncomingPayload();
-      if (payload.containsKey(TabletInfo.TABLETS_ROUTING_V1_CUSTOM_PAYLOAD_KEY)) {
+      if (executeAndReturnIfResultHasTabletsInfo(session, preparedStatement.bind(i, i))) {
+        counter++;
+      }
+    }
+    Assert.assertEquals(INITIAL_TABLETS, counter);
+    assertSessionTabletMapIsFilled(session);
+    session.close();
+
+    session = CqlSession.builder().addContactEndPoints(CCM_RULE.getContactPoints()).build();
+    counter = 0;
+    preparedStatement =
+        session.prepare(STMT_SELECT);
+    for (int i = 1; i <= QUERIES; i++) {
+      if (executeAndReturnIfResultHasTabletsInfo(session, preparedStatement.bind(i, i))) {
         counter++;
       }
     }
@@ -119,7 +288,44 @@ public void should_receive_each_tablet_exactly_once() {
 
     // With enough queries we should hit a wrong node for each tablet exactly once.
     Assert.assertEquals(INITIAL_TABLETS, counter);
+    assertSessionTabletMapIsFilled(session);
+
+    // All tablet information should be available by now (unless for some reason cluster did sth on
+    // its own)
+    // We should not receive any tablet payloads now, since they are sent only on mismatch.
+    for (int i = 1; i <= QUERIES; i++) {
+
+      ResultSet rs = session.execute(preparedStatement.bind(i, i));
+      Map<String, ByteBuffer> payload = rs.getExecutionInfo().getIncomingPayload();
+
+      if (payload.containsKey(TabletInfo.TABLETS_ROUTING_V1_CUSTOM_PAYLOAD_KEY)) {
+        throw new RuntimeException(
+            "Received non empty payload with tablets routing information: " + payload);
+      }
+    }
+  }
 
+  private static boolean isSessionLearnedTabletInfo(CqlSession session) {
+    ConcurrentMap<KeyspaceTableNamePair, ConcurrentSkipListSet<Tablet>> tabletMapping =
+        session.getMetadata().getTabletMap().getMapping();
+    KeyspaceTableNamePair ktPair =
+        new KeyspaceTableNamePair(
+            CqlIdentifier.fromCql(KEYSPACE_NAME), CqlIdentifier.fromCql(TABLE_NAME));
+
+    Set<Tablet> tablets = tabletMapping.get(ktPair);
+    if (tablets == null || tablets.isEmpty()) {
+      return false;
+    }
+
+    for (Tablet tab : tablets) {
+      if (tab.getReplicaNodes().size() >= REPLICATION_FACTOR) {
+        return true;
+      };
+    }
+    return false;
+  }
+
+  private static void assertSessionTabletMapIsFilled(CqlSession session) {
     ConcurrentMap<KeyspaceTableNamePair, ConcurrentSkipListSet<Tablet>> tabletMapping =
         session.getMetadata().getTabletMap().getMapping();
     KeyspaceTableNamePair ktPair =
@@ -133,19 +339,20 @@ public void should_receive_each_tablet_exactly_once() {
     for (Tablet tab : tablets) {
       Assert.assertEquals(REPLICATION_FACTOR, tab.getReplicaNodes().size());
     }
+  }
 
-    // All tablet information should be available by now (unless for some reason cluster did sth on
-    // its own)
-    // We should not receive any tablet payloads now, since they are sent only on mismatch.
-    for (int i = 1; i <= QUERIES; i++) {
-
-      ResultSet rs = session.execute(preparedStatement.bind(i, i));
-      Map<String, ByteBuffer> payload = rs.getExecutionInfo().getIncomingPayload();
-
-      if (payload.containsKey(TabletInfo.TABLETS_ROUTING_V1_CUSTOM_PAYLOAD_KEY)) {
-        throw new RuntimeException(
-            "Received non empty payload with tablets routing information: " + payload);
+  private static boolean executeOnAllHostsAndReturnIfResultHasTabletsInfo(CqlSession session, Statement stmt) {
+    session.refreshSchema();
+    for (Node node: session.getMetadata().getNodes().values()) {
+      if(executeAndReturnIfResultHasTabletsInfo(session, stmt.setNode(node))) {
+        return true;
       }
     }
+    return false;
+  }
+
+  private static boolean executeAndReturnIfResultHasTabletsInfo(CqlSession session, Statement statement) {
+    ResultSet rs = session.execute(statement);
+    return rs.getExecutionInfo().getIncomingPayload().containsKey(TabletInfo.TABLETS_ROUTING_V1_CUSTOM_PAYLOAD_KEY);
   }
 }