Tablets support

Introduces basic tablets support for version 4.x of the driver.
Metadata about tablets will be kept in TabletMap that gets continuously updated
through the tablets-routing-v1 extension. Each time the BoundStatement targets
the wrong node and shard combination the server supporting tablets should
respond with tablet metadata inside custom payload of its response.
This metadata will be transparently processed and used for future queries.

Tablets metadata will by enabled by default. Until now driver was using
TokenMaps to choose replicas and appropriate shards. Having a token was enough
information to do that. Now driver will first attempt tablet-based lookup
and only after failing to find corresponding tablet it will defer to TokenMap
lookup. Since to find a correct tablet besides the token we need the keyspace
and table names, many of the methods were extended to also accept those
as parameters.

RequestHandlerTestHarness was adjusted to mock also MetadataManager.
Before it used to mock only `session.getMetadata()` call but the same can
be obtained by `context.getMetadataManager().getMetadata()`. Using the
second method was causing test failures.

Author: Bouncheck

core/src/main/java/com/datastax/oss/driver/api/core/cql/PreparedStatement.java

@@ -21,6 +21,7 @@
  */
 package com.datastax.oss.driver.api.core.cql;
 
+import com.datastax.oss.driver.api.core.CqlIdentifier;
 import com.datastax.oss.driver.api.core.CqlSession;
 import com.datastax.oss.driver.api.core.DefaultProtocolVersion;
 import com.datastax.oss.driver.api.core.metadata.token.Partitioner;
@@ -58,6 +59,13 @@ public interface PreparedStatement {
   @NonNull
   ColumnDefinitions getVariableDefinitions();
 
+  /**
+   * Table name inferred from {@link PreparedStatement#getVariableDefinitions()} or {@code null} if
+   * not possible.
+   */
+  @Nullable
+  CqlIdentifier inferTable();
+
   /**
    * The partitioner to use for token-aware routing. If {@code null}, the cluster-wide partitioner
    * will be used.

core/src/main/java/com/datastax/oss/driver/api/core/metadata/KeyspaceTableNamePair.java

@@ -0,0 +1,47 @@
+package com.datastax.oss.driver.api.core.metadata;
+
+import java.util.Objects;
+
+/** Simple keyspace name and table name pair. */
+public class KeyspaceTableNamePair {
+  private final String keyspace;
+  private final String tableName;
+
+  public KeyspaceTableNamePair(String keyspace, String tableName) {
+    this.keyspace = keyspace;
+    this.tableName = tableName;
+  }
+
+  public String getKeyspace() {
+    return keyspace;
+  }
+
+  public String getTableName() {
+    return tableName;
+  }
+
+  @Override
+  public String toString() {
+    return "KeyspaceTableNamePair{"
+        + "keyspace='"
+        + keyspace
+        + '\''
+        + ", tableName='"
+        + tableName
+        + '\''
+        + '}';
+  }
+
+  @Override
+  public boolean equals(Object o) {
+    if (this == o) return true;
+    if (o == null || !(o instanceof KeyspaceTableNamePair)) return false;
+    KeyspaceTableNamePair that = (KeyspaceTableNamePair) o;
+    return keyspace.equals(that.keyspace) && tableName.equals(that.tableName);
+  }
+
+  @Override
+  public int hashCode() {
+    return Objects.hash(keyspace, tableName);
+  }
+}

core/src/main/java/com/datastax/oss/driver/api/core/metadata/Metadata.java

@@ -115,6 +115,14 @@ default Optional<KeyspaceMetadata> getKeyspace(@NonNull String keyspaceName) {
   @NonNull
   Optional<TokenMap> getTokenMap();
 
+  /**
+   * The tablet map for this cluster.
+   *
+   * <p>Starts as an empty map that will gradually receive updates on each query of a yet unknown
+   * tablet.
+   */
+  TabletMap getTabletMap();
+
   /**
    * The cluster name to which this session is connected. The Optional returned should contain the
    * value from the server for <b>system.local.cluster_name</b>.

core/src/main/java/com/datastax/oss/driver/api/core/metadata/NodeShardPair.java

@@ -0,0 +1,28 @@
+package com.datastax.oss.driver.api.core.metadata;
+
+/**
+ * Simple class to hold Node representing a host and a shard number. Class itself makes no checks or
+ * guarantees about existence of a shard on corresponding host.
+ */
+public class NodeShardPair {
+  private final Node node;
+  private final int shard;
+
+  public NodeShardPair(Node node, int shard) {
+    this.node = node;
+    this.shard = shard;
+  }
+
+  public Node getNode() {
+    return node;
+  }
+
+  public int getShard() {
+    return shard;
+  }
+
+  @Override
+  public String toString() {
+    return "NodeShardPair{" + "host=" + node + ", shard=" + shard + '}';
+  }
+}

core/src/main/java/com/datastax/oss/driver/api/core/metadata/Tablet.java

@@ -0,0 +1,44 @@
+package com.datastax.oss.driver.api.core.metadata;
+
+import com.datastax.oss.driver.shaded.guava.common.annotations.Beta;
+import java.util.Set;
+import java.util.UUID;
+
+/**
+ * Represents a tablet as described in tablets-routing-v1 protocol extension with some additional
+ * fields for ease of use.
+ */
+@Beta
+public interface Tablet extends Comparable<Tablet> {
+  public String getKeyspaceName();
+
+  public UUID getTableId();
+
+  public String getTableName();
+
+  /**
+   * Returns left endpoint of an interval. This interval is left-open, meaning the tablet does not
+   * own the token equal to the first token.
+   *
+   * @return {@code long} value representing first token.
+   */
+  public long getFirstToken();
+
+  /**
+   * Returns right endpoint of an interval. This interval is right-closed, which means that last
+   * token is owned by this tablet.
+   *
+   * @return {@code long} value representing last token.
+   */
+  public long getLastToken();
+
+  public Set<Node> getReplicaNodes();
+
+  /**
+   * Looks up the shard number for specific replica Node.
+   *
+   * @param node one of the replica nodes of this tablet.
+   * @return Shard number for the replica or -1 if no such Node found.
+   */
+  public int getShardForNode(Node node);
+}

core/src/main/java/com/datastax/oss/driver/api/core/metadata/TabletMap.java

@@ -0,0 +1,34 @@
+package com.datastax.oss.driver.api.core.metadata;
+
+import com.datastax.oss.driver.shaded.guava.common.annotations.Beta;
+import java.util.concurrent.ConcurrentMap;
+import java.util.concurrent.ConcurrentSkipListSet;
+
+/** Holds all currently known tablet metadata. */
+@Beta
+public interface TabletMap {
+  /**
+   * Returns mapping from tables to the sets of their tablets.
+   *
+   * @return the Map keyed by (keyspace,table) pairs with Set of tablets as value type.
+   */
+  public ConcurrentMap<KeyspaceTableNamePair, ConcurrentSkipListSet<Tablet>> getMapping();
+
+  /**
+   * Adds a single tablet to the map. Handles removal of overlapping tablets.
+   *
+   * @param tablet
+   */
+  public void addTablet(Tablet tablet);
+
+  /**
+   * Returns {@link Tablet} instance
+   *
+   * @param keyspace tablet's keyspace
+   * @param table tablet's table
+   * @param token target token
+   * @return {@link Tablet} responsible for provided token or {@code null} if no such tablet is
+   *     present.
+   */
+  public Tablet getTablet(String keyspace, String table, long token);
+}

core/src/main/java/com/datastax/oss/driver/internal/core/channel/ProtocolInitHandler.java

@@ -43,6 +43,7 @@
 import com.datastax.oss.driver.internal.core.protocol.SegmentToFrameDecoder;
 import com.datastax.oss.driver.internal.core.protocol.ShardingInfo;
 import com.datastax.oss.driver.internal.core.protocol.ShardingInfo.ConnectionShardingInfo;
+import com.datastax.oss.driver.internal.core.protocol.TabletInfo;
 import com.datastax.oss.driver.internal.core.util.ProtocolUtils;
 import com.datastax.oss.driver.internal.core.util.concurrent.UncaughtExceptions;
 import com.datastax.oss.protocol.internal.Message;
@@ -94,6 +95,7 @@ class ProtocolInitHandler extends ConnectInitHandler {
   private ChannelHandlerContext ctx;
   private final boolean querySupportedOptions;
   private LwtInfo lwtInfo;
+  private TabletInfo tabletInfo;
 
   /**
    * @param querySupportedOptions whether to send OPTIONS as the first message, to request which
@@ -191,6 +193,9 @@ Message getRequest() {
           if (lwtInfo != null) {
             lwtInfo.addOption(startupOptions);
           }
+          if (tabletInfo != null && tabletInfo.isEnabled()) {
+            TabletInfo.addOption(startupOptions);
+          }
           return request = new Startup(startupOptions);
         case GET_CLUSTER_NAME:
           return request = CLUSTER_NAME_QUERY;
@@ -230,6 +235,7 @@ void onResponse(Message response) {
           if (lwtInfo != null) {
             channel.attr(LWT_INFO_KEY).set(lwtInfo);
           }
+          tabletInfo = TabletInfo.parseTabletInfo(res.options);
           step = Step.STARTUP;
           send();
         } else if (step == Step.STARTUP && response instanceof Ready) {

core/src/main/java/com/datastax/oss/driver/internal/core/cql/CqlRequestHandler.java

@@ -31,9 +31,13 @@
 import com.datastax.oss.driver.api.core.config.DriverExecutionProfile;
 import com.datastax.oss.driver.api.core.connection.FrameTooLongException;
 import com.datastax.oss.driver.api.core.cql.AsyncResultSet;
+import com.datastax.oss.driver.api.core.cql.BoundStatement;
 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.Statement;
 import com.datastax.oss.driver.api.core.metadata.Node;
+import com.datastax.oss.driver.api.core.metadata.Tablet;
+import com.datastax.oss.driver.api.core.metadata.TabletMap;
 import com.datastax.oss.driver.api.core.metadata.TokenMap;
 import com.datastax.oss.driver.api.core.metadata.token.Partitioner;
 import com.datastax.oss.driver.api.core.metadata.token.Token;
@@ -59,8 +63,10 @@
 import com.datastax.oss.driver.internal.core.context.InternalDriverContext;
 import com.datastax.oss.driver.internal.core.metadata.DefaultNode;
 import com.datastax.oss.driver.internal.core.metadata.token.DefaultTokenMap;
+import com.datastax.oss.driver.internal.core.metadata.token.TokenLong64;
 import com.datastax.oss.driver.internal.core.metrics.NodeMetricUpdater;
 import com.datastax.oss.driver.internal.core.metrics.SessionMetricUpdater;
+import com.datastax.oss.driver.internal.core.protocol.TabletInfo;
 import com.datastax.oss.driver.internal.core.session.DefaultSession;
 import com.datastax.oss.driver.internal.core.session.RepreparePayload;
 import com.datastax.oss.driver.internal.core.tracker.NoopRequestTracker;
@@ -284,6 +290,73 @@ private Token getRoutingToken(Statement statement) {
     return tokenMap == null ? null : ((DefaultTokenMap) tokenMap).getTokenFactory().hash(key);
   }
 
+  private CqlIdentifier getTabletRoutingKeyspace(Statement statement) {
+    // Bundles keyspace checks into one function call
+    if (statement == null) {
+      return null;
+    }
+    CqlIdentifier result = statement.getKeyspace();
+    if (result == null) {
+      result = statement.getRoutingKeyspace();
+    }
+    if (result == null) {
+      result = session.getKeyspace().orElse(null);
+    }
+    return result;
+  }
+
+  private CqlIdentifier getTabletRoutingTable(Statement statement) {
+    // Tries to infer target table from column definitions of this statement
+    if (statement == null) {
+      return null;
+    }
+    if (statement instanceof BoundStatement) {
+      return ((BoundStatement) statement).getPreparedStatement().inferTable();
+    } else if (statement instanceof PreparedStatement) {
+      return ((PreparedStatement) statement).inferTable();
+    } else {
+      return null;
+    }
+  }
+
+  public Integer getShardFromTabletMap(Statement statement, Node node, Token token) {
+    TabletMap tabletMap = context.getMetadataManager().getMetadata().getTabletMap();
+    if (!(token instanceof TokenLong64)) {
+      LOG.trace(
+          "Token ({}) is not a TokenLong64. Not performing tablet shard lookup for statement {}.",
+          token,
+          statement);
+      return null;
+    }
+    if (getTabletRoutingKeyspace(statement) == null || getTabletRoutingTable(statement) == null) {
+      return null;
+    }
+    long tokenValue = ((TokenLong64) token).getValue();
+    String statementKeyspace = getTabletRoutingKeyspace(statement).asInternal();
+    String statementTable = getTabletRoutingTable(statement).asInternal();
+    Tablet targetTablet = tabletMap.getTablet(statementKeyspace, statementTable, tokenValue);
+    if (targetTablet == null) {
+      LOG.trace(
+          "Could not determine shard for token {} and table {}.{} on Node {}: Could not find corresponding tablet, returning null.",
+          token,
+          statementKeyspace,
+          statementTable,
+          node);
+      return null;
+    }
+    int shard = targetTablet.getShardForNode(node);
+    if (shard == -1) {
+      LOG.trace(
+          "Could not find shard corresponding to token {} and Node {} for table {} in keyspace {}. Returning null.",
+          token,
+          node,
+          statementTable,
+          statementKeyspace);
+      return null;
+    }
+    return shard;
+  }
+
   /**
    * Sends the request to the next available node.
    *
@@ -309,9 +382,20 @@ private void sendRequest(
     Node node = retriedNode;
     DriverChannel channel = null;
     if (node == null
-        || (channel = session.getChannel(node, logPrefix, getRoutingToken(statement))) == null) {
+        || (channel =
+                session.getChannel(
+                    node,
+                    logPrefix,
+                    getRoutingToken(statement),
+                    getShardFromTabletMap(statement, node, getRoutingToken(statement))))
+            == null) {
       while (!result.isDone() && (node = queryPlan.poll()) != null) {
-        channel = session.getChannel(node, logPrefix, getRoutingToken(statement));
+        channel =
+            session.getChannel(
+                node,
+                logPrefix,
+                getRoutingToken(statement),
+                getShardFromTabletMap(statement, node, getRoutingToken(statement)));
         if (channel != null) {
           break;
         } else {
@@ -420,6 +504,18 @@ private void setFinalResult(
               totalLatencyNanos,
               TimeUnit.NANOSECONDS);
         }
+        if (resultSet.getColumnDefinitions().size() > 0
+            && resultSet
+                .getExecutionInfo()
+                .getIncomingPayload()
+                .containsKey(TabletInfo.TABLETS_ROUTING_V1_CUSTOM_PAYLOAD_KEY)) {
+          context
+              .getMetadataManager()
+              .addTabletFromPayload(
+                  resultSet.getColumnDefinitions().get(0).getKeyspace(),
+                  resultSet.getColumnDefinitions().get(0).getTable(),
+                  resultSet.getExecutionInfo().getIncomingPayload());
+        }
       }
       // log the warnings if they have NOT been disabled
       if (!executionInfo.getWarnings().isEmpty()

core/src/main/java/com/datastax/oss/driver/internal/core/cql/DefaultPreparedStatement.java

@@ -36,6 +36,7 @@
 import com.datastax.oss.driver.internal.core.data.ValuesHelper;
 import com.datastax.oss.driver.internal.core.session.RepreparePayload;
 import edu.umd.cs.findbugs.annotations.NonNull;
+import edu.umd.cs.findbugs.annotations.Nullable;
 import java.nio.ByteBuffer;
 import java.time.Duration;
 import java.util.List;
@@ -140,6 +141,15 @@ public ColumnDefinitions getVariableDefinitions() {
     return variableDefinitions;
   }
 
+  @Nullable
+  @Override
+  public CqlIdentifier inferTable() {
+    if (variableDefinitions.size() > 0) {
+      return variableDefinitions.get(0).getTable();
+    }
+    return null;
+  }
+
   @Override
   public Partitioner getPartitioner() {
     return partitioner;