Introduce support for tablets [3.x]

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

@@ -474,6 +474,8 @@ public ListenableFuture<Void> apply(Message.Response response) throws Exception
             if (lwt != null) {
               getHost().setLwtInfo(lwt);
             }
+            TabletInfo tabletInfo = TabletInfo.parseTabletInfo(msg.supported);
+            getHost().setTabletInfo(tabletInfo);
             return MoreFutures.VOID_SUCCESS;
           case ERROR:
             Responses.Error error = (Responses.Error) response;
@@ -507,6 +509,13 @@ public ListenableFuture<Void> apply(Void input) throws Exception {
         if (lwtInfo != null) {
           lwtInfo.addOption(extraOptions);
         }
+        TabletInfo tabletInfo = getHost().getTabletInfo();
+        if (tabletInfo != null
+            && tabletInfo.isEnabled()
+            && ProtocolFeature.CUSTOM_PAYLOADS.isSupportedBy(protocolVersion)) {
+          logger.debug("Enabling tablet support in OPTIONS message");
+          TabletInfo.addOption(extraOptions);
+        }
         Future startupResponseFuture =
             write(
                 new Requests.Startup(

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

@@ -70,6 +70,23 @@ public void onSet(
       switch (response.type) {
         case RESULT:
           Responses.Result rm = (Responses.Result) response;
+
+          if (rm.getCustomPayload() != null
+              && rm.getCustomPayload().containsKey(TabletInfo.TABLETS_ROUTING_V1_CUSTOM_PAYLOAD_KEY)
+              && (statement instanceof BoundStatement)) {
+            BoundStatement st = (BoundStatement) statement;
+            String keyspace = statement.getKeyspace();
+            String table =
+                st.preparedStatement().getPreparedId().boundValuesMetadata.variables.getTable(0);
+            session
+                .getCluster()
+                .getMetadata()
+                .getTabletMap()
+                .processTabletsRoutingV1Payload(
+                    keyspace,
+                    table,
+                    rm.getCustomPayload().get(TabletInfo.TABLETS_ROUTING_V1_CUSTOM_PAYLOAD_KEY));
+          }
           switch (rm.kind) {
             case SET_KEYSPACE:
               // propagate the keyspace change to other connections

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

@@ -70,6 +70,9 @@ public class Host {
   // Can be set concurrently but the value should always be the same.
   private volatile LwtInfo lwtInfo = null;
 
+  // Whether host supports TABLETS_ROUTING_V1
+  private volatile TabletInfo tabletInfo = null;
+
   enum State {
     ADDED,
     DOWN,
@@ -450,6 +453,14 @@ public void setLwtInfo(LwtInfo lwtInfo) {
     this.lwtInfo = lwtInfo;
   }
 
+  public TabletInfo getTabletInfo() {
+    return tabletInfo;
+  }
+
+  public void setTabletInfo(TabletInfo tabletInfo) {
+    this.tabletInfo = tabletInfo;
+  }
+
   /**
    * Returns whether the host is considered up by the driver.
    *

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

@@ -504,7 +504,9 @@ ListenableFuture<Connection> borrowConnection(
       TimeUnit unit,
       int maxQueueSize,
       Token.Factory partitioner,
-      ByteBuffer routingKey) {
+      ByteBuffer routingKey,
+      String keyspace,
+      String table) {
     Phase phase = this.phase.get();
     if (phase != Phase.READY)
       return Futures.immediateFailedFuture(
@@ -515,7 +517,17 @@ ListenableFuture<Connection> borrowConnection(
       if (routingKey != null) {
         Metadata metadata = manager.cluster.getMetadata();
         Token t = metadata.newToken(partitioner, routingKey);
-        shardId = host.getShardingInfo().shardId(t);
+        shardId = -1;
+        if (keyspace != null && table != null) {
+          assert t instanceof Token.TokenLong64;
+          shardId =
+              Integer.min(
+                  metadata.getShardForTabletToken(keyspace, table, (Token.TokenLong64) t, host),
+                  host.getShardingInfo().getShardsCount());
+        }
+        if (shardId == -1) { // means that tablet lookup failed
+          shardId = host.getShardingInfo().shardId(t);
+        }
       } else {
         shardId = RAND.nextInt(host.getShardingInfo().getShardsCount());
       }

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

@@ -21,6 +21,7 @@
  */
 package com.datastax.driver.core;
 
+import com.google.common.annotations.Beta;
 import com.google.common.collect.ImmutableMap;
 import com.google.common.collect.ImmutableSet;
 import com.google.common.collect.Maps;
@@ -35,13 +36,15 @@
 import java.util.HashSet;
 import java.util.List;
 import java.util.Map;
+import java.util.NavigableSet;
 import java.util.Set;
 import java.util.TreeSet;
 import java.util.UUID;
 import java.util.concurrent.ConcurrentHashMap;
 import java.util.concurrent.ConcurrentMap;
 import java.util.concurrent.CopyOnWriteArrayList;
 import java.util.concurrent.locks.ReentrantLock;
+import java.util.stream.Collectors;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
@@ -60,8 +63,8 @@ public class Metadata {
   final ConcurrentMap<String, KeyspaceMetadata> keyspaces =
       new ConcurrentHashMap<String, KeyspaceMetadata>();
   private volatile TokenMap tokenMap;
-
   final ReentrantLock lock = new ReentrantLock();
+  private final TabletMap tabletMap;
 
   // See https://github.com/apache/cassandra/blob/trunk/doc/cql3/CQL.textile#appendixA
   private static final IntObjectHashMap<List<char[]>> RESERVED_KEYWORDS =
@@ -146,6 +149,7 @@ public class Metadata {
 
   Metadata(Cluster.Manager cluster) {
     this.cluster = cluster;
+    this.tabletMap = TabletMap.emptyMap(cluster);
   }
 
   // rebuilds the token map with the current hosts, typically when refreshing schema metadata
@@ -514,21 +518,30 @@ public Set<TokenRange> getTokenRanges(String keyspace, Host host) {
   }
 
   /**
-   * Returns the set of hosts that are replica for a given partition key. Partitioner can be {@code
-   * null} and then a cluster-wide partitioner will be invoked.
+   * Extension of legacy method {@link Metadata#getReplicas(String, Token.Factory, ByteBuffer)}.
+   * Tablets model requires knowledge of the table name to determine the replicas. This method will
+   * first try to lookup replicas through known tablets metadata. It will default to TokenMap lookup
+   * if either {@code null} was passed as table name or the tablet lookup is unsuccessful for any
+   * other reason.
+   *
+   * <p>Returns the set of hosts that are replica for a given partition key. Partitioner can be
+   * {@code null} and then a cluster-wide partitioner will be invoked.
    *
    * <p>Note that this information is refreshed asynchronously by the control connection, when
    * schema or ring topology changes. It might occasionally be stale (or even empty).
    *
    * @param keyspace the name of the keyspace to get replicas for.
+   * @param table the name of the table to get replicas for. Necessary for distinction for tablets.
+   *     Unnecessary for regular TokenMap
    * @param partitioner the partitioner to use or @{code null} for cluster-wide partitioner.
    * @param partitionKey the partition key for which to find the set of replica.
    * @return the (immutable) set of replicas for {@code partitionKey} as known by the driver. Note
    *     that the result might be stale or empty if metadata was explicitly disabled with {@link
    *     QueryOptions#setMetadataEnabled(boolean)}.
    */
+  @Beta
   public Set<Host> getReplicas(
-      String keyspace, Token.Factory partitioner, ByteBuffer partitionKey) {
+      String keyspace, String table, Token.Factory partitioner, ByteBuffer partitionKey) {
     keyspace = handleId(keyspace);
     TokenMap current = tokenMap;
     if (current == null) {
@@ -537,11 +550,40 @@ public Set<Host> getReplicas(
       if (partitioner == null) {
         partitioner = current.factory;
       }
+      // If possible, try tablet lookup first
+      if (keyspace != null && table != null) {
+        Token token = partitioner.hash(partitionKey);
+        assert (token instanceof Token.TokenLong64);
+        Set<UUID> hostUuids = tabletMap.getReplicas(keyspace, table, (long) token.getValue());
+        if (!hostUuids.isEmpty()) {
+          return hostUuids.stream().map(this::getHost).collect(Collectors.toSet());
+        }
+      }
+      // Fall back to tokenMap
       Set<Host> hosts = current.getReplicas(keyspace, partitioner.hash(partitionKey));
       return hosts == null ? Collections.<Host>emptySet() : hosts;
     }
   }
 
+  /**
+   * Returns the set of hosts that are replica for a given partition key. Partitioner can be {@code
+   * null} and then a cluster-wide partitioner will be invoked.
+   *
+   * <p>Note that this information is refreshed asynchronously by the control connection, when
+   * schema or ring topology changes. It might occasionally be stale (or even empty).
+   *
+   * @param keyspace the name of the keyspace to get replicas for.
+   * @param partitioner the partitioner to use or @{code null} for cluster-wide partitioner.
+   * @param partitionKey the partition key for which to find the set of replica.
+   * @return the (immutable) set of replicas for {@code partitionKey} as known by the driver. Note
+   *     that the result might be stale or empty if metadata was explicitly disabled with {@link
+   *     QueryOptions#setMetadataEnabled(boolean)}.
+   */
+  public Set<Host> getReplicas(
+      String keyspace, Token.Factory partitioner, ByteBuffer partitionKey) {
+    return getReplicas(keyspace, null, partitioner, partitionKey);
+  }
+
   /**
    * Returns the set of hosts that are replica for a given token range.
    *
@@ -860,6 +902,58 @@ void triggerOnMaterializedViewRemoved(MaterializedViewMetadata view) {
     }
   }
 
+  @Beta
+  public int getShardForTabletToken(
+      String keyspace, String table, Token.TokenLong64 token, Host host) {
+    if (tabletMap == null) {
+      logger.trace(
+          "Could not determine shard for token {} on host {} because tablets metadata is currently null. "
+              + "Returning -1.",
+          token,
+          host);
+      return -1;
+    }
+    UUID targetHostUuid = host.getHostId();
+    long tokenValue = (long) token.getValue();
+    TabletMap.KeyspaceTableNamePair key = new TabletMap.KeyspaceTableNamePair(keyspace, table);
+    NavigableSet<TabletMap.Tablet> targetTablets = tabletMap.getMapping().get(key);
+    if (targetTablets == null) {
+      logger.trace(
+          "Could not determine shard for token {} on host {} because table {}.{} is not present in tablets "
+              + "metadata. Returning -1.",
+          token,
+          host,
+          keyspace,
+          table);
+      return -1;
+    }
+    TabletMap.Tablet row = targetTablets.ceiling(TabletMap.Tablet.malformedTablet(tokenValue));
+    if (row != null && row.getFirstToken() < tokenValue) {
+      for (TabletMap.HostShardPair hostShardPair : row.getReplicas()) {
+        if (hostShardPair.getHost().equals(targetHostUuid)) {
+          return hostShardPair.getShard();
+        }
+      }
+    }
+    logger.trace(
+        "Could not find tablet corresponding to token {} on host {} for table {} in keyspace {}. Returning -1.",
+        token,
+        host,
+        table,
+        keyspace);
+    return -1;
+  }
+
+  /**
+   * Getter for current {@link TabletMap}.
+   *
+   * @return current {@link TabletMap}
+   */
+  @Beta
+  public TabletMap getTabletMap() {
+    return tabletMap;
+  }
+
   private static class TokenMap {
 
     private final Token.Factory factory;

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

@@ -110,11 +110,22 @@ private Iterator<Host> getReplicas(
     }
 
     Token.Factory partitioner = statement.getPartitioner();
+    String tableName = null;
+    ColumnDefinitions defs = null;
+    if (statement instanceof BoundStatement) {
+      defs = ((BoundStatement) statement).preparedStatement().getVariables();
+    } else if (statement instanceof PreparedStatement) {
+      defs = ((PreparedStatement) statement).getVariables();
+    }
+    if (defs != null && defs.size() > 0) {
+      tableName = defs.getTable(0);
+    }
+
     final Set<Host> replicas =
         manager
             .cluster
             .getMetadata()
-            .getReplicas(Metadata.quote(keyspace), partitioner, partitionKey);
+            .getReplicas(Metadata.quote(keyspace), tableName, partitioner, partitionKey);
 
     // replicas are stored in the right order starting with the primary replica
     return replicas.iterator();
@@ -437,13 +448,28 @@ private boolean query(final Host host) {
       ByteBuffer routingKey = statement.getRoutingKey(protocolVersion, codecRegistry);
 
       PoolingOptions poolingOptions = manager.configuration().getPoolingOptions();
+      String statementKeyspace = statement.getKeyspace();
+      String statementTable = null;
+      ColumnDefinitions defs = null;
+      if (statement instanceof PreparedStatement) {
+        defs = ((PreparedStatement) statement).getVariables();
+      }
+      if (statement instanceof BoundStatement) {
+        defs = ((BoundStatement) statement).statement.getVariables();
+      }
+      if (defs != null && defs.size() > 0) {
+        statementTable = defs.getTable(0);
+      }
+
       ListenableFuture<Connection> connectionFuture =
           pool.borrowConnection(
               poolingOptions.getPoolTimeoutMillis(),
               TimeUnit.MILLISECONDS,
               poolingOptions.getMaxQueueSize(),
               statement.getPartitioner(),
-              routingKey);
+              routingKey,
+              statementKeyspace,
+              statementTable);
       GuavaCompatibility.INSTANCE.addCallback(
           connectionFuture,
           new FutureCallback<Connection>() {

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

@@ -733,13 +733,22 @@ private ListenableFuture<PreparedStatement> prepare(
       if (entry.getKey().getEndPoint().equals(toExclude)) continue;
 
       try {
+        ColumnDefinitions defs = statement.getVariables();
+        String statementTable = (defs != null && defs.size() > 0 ? defs.getTable(0) : null);
         // Preparing is not critical: if it fails, it will fix itself later when the user tries to
         // execute
         // the prepared query. So don't wait if no connection is available, simply abort.
         ListenableFuture<Connection> connectionFuture =
             entry
                 .getValue()
-                .borrowConnection(0, TimeUnit.MILLISECONDS, 0, null, statement.getRoutingKey());
+                .borrowConnection(
+                    0,
+                    TimeUnit.MILLISECONDS,
+                    0,
+                    null,
+                    statement.getRoutingKey(),
+                    statement.getQueryKeyspace(),
+                    statementTable);
         ListenableFuture<Response> prepareFuture =
             GuavaCompatibility.INSTANCE.transformAsync(
                 connectionFuture,

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

@@ -0,0 +1,33 @@
+package com.datastax.driver.core;
+
+import java.util.List;
+import java.util.Map;
+
+public class TabletInfo {
+  private static final String SCYLLA_TABLETS_STARTUP_OPTION_KEY = "TABLETS_ROUTING_V1";
+  private static final String SCYLLA_TABLETS_STARTUP_OPTION_VALUE = "";
+  public static final String TABLETS_ROUTING_V1_CUSTOM_PAYLOAD_KEY = "tablets-routing-v1";
+
+  private boolean enabled = false;
+
+  private TabletInfo(boolean enabled) {
+    this.enabled = enabled;
+  }
+
+  // Currently pertains only to TABLETS_ROUTING_V1
+  public boolean isEnabled() {
+    return enabled;
+  }
+
+  public static TabletInfo parseTabletInfo(Map<String, List<String>> supported) {
+    List<String> values = supported.get(SCYLLA_TABLETS_STARTUP_OPTION_KEY);
+    return new TabletInfo(
+        values != null
+            && values.size() == 1
+            && values.get(0).equals(SCYLLA_TABLETS_STARTUP_OPTION_VALUE));
+  }
+
+  public static void addOption(Map<String, String> options) {
+    options.put(SCYLLA_TABLETS_STARTUP_OPTION_KEY, SCYLLA_TABLETS_STARTUP_OPTION_VALUE);
+  }
+}