Move metadata internal API to use List instead of Set as replicas store

In almost all use cases replicas are getting converted into a List.
To make performance better make these API provide List instead.

Author: dkropachev

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

@@ -22,6 +22,7 @@
 package com.datastax.driver.core;
 
 import com.google.common.annotations.Beta;
+import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableMap;
 import com.google.common.collect.ImmutableSet;
 import com.google.common.collect.Maps;
@@ -34,6 +35,7 @@
 import java.util.Collections;
 import java.util.HashMap;
 import java.util.HashSet;
+import java.util.LinkedHashSet;
 import java.util.List;
 import java.util.Map;
 import java.util.Set;
@@ -51,6 +53,8 @@
 public class Metadata {
 
   private static final Logger logger = LoggerFactory.getLogger(Metadata.class);
+  private static final ImmutableList<Host> EMPTY_LIST =
+      ImmutableList.copyOf(Collections.emptyList());
 
   final Cluster.Manager cluster;
   volatile String clusterName;
@@ -553,18 +557,48 @@ public Set<TokenRange> getTokenRanges(String keyspace, Host host) {
    * @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)}.
+   * @deprecated use {@link Metadata#getReplicasList(String, String, Token.Factory, ByteBuffer)}
    */
+  @Deprecated
   @Beta
   public Set<Host> getReplicas(
       String keyspace, String table, Token.Factory partitioner, ByteBuffer partitionKey) {
+    return new LinkedHashSet<>(this.getReplicasList(keyspace, table, partitioner, partitionKey));
+  }
+
+  /**
+   * 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 List<Host> getReplicasList(
+      String keyspace, String table, Token.Factory partitioner, ByteBuffer partitionKey) {
     keyspace = handleId(keyspace);
     table = handleId(table);
     TokenMap current = tokenMap;
     if (partitioner == null && current != null) {
       partitioner = current.factory;
     }
     if (partitioner == null) {
-      return Collections.emptySet();
+      return Collections.emptyList();
     }
     Token token = partitioner.hash(partitionKey);
 
@@ -575,14 +609,13 @@ public Set<Host> getReplicas(
         assert (token instanceof Token.TokenLong64);
         return tabletMap.getReplicas(keyspace, table, (long) token.getValue());
       } else {
-        return Collections.emptySet();
+        return Collections.emptyList();
       }
     }
 
     // TokenMap:
-    if (current == null) return Collections.<Host>emptySet();
-    Set<Host> hosts = current.getReplicas(keyspace, token);
-    return hosts == null ? Collections.<Host>emptySet() : hosts;
+    if (current == null) return Collections.<Host>emptyList();
+    return current.getReplicas(keyspace, token);
   }
 
   /**
@@ -598,12 +631,33 @@ public Set<Host> getReplicas(
    * @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)}.
+   * @deprecated use {@link Metadata#getReplicasList(String, Token.Factory, ByteBuffer)}
    */
+  @Deprecated
   public Set<Host> getReplicas(
       String keyspace, Token.Factory partitioner, ByteBuffer partitionKey) {
     return getReplicas(keyspace, null, partitioner, partitionKey);
   }
 
+  /**
+   * Returns the immutable list 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) list 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 List<Host> getReplicasList(
+      String keyspace, Token.Factory partitioner, ByteBuffer partitionKey) {
+    return getReplicasList(keyspace, null, partitioner, partitionKey);
+  }
+
   /**
    * Returns the set of hosts that are replica for a given token range.
    *
@@ -620,15 +674,42 @@ public Set<Host> getReplicas(
    * @return the (immutable) set of replicas for {@code range} as known by the driver. Note that the
    *     result might be stale or empty if metadata was explicitly disabled with {@link
    *     QueryOptions#setMetadataEnabled(boolean)}.
+   * @deprecated use {@link Metadata#getReplicasList(String, TokenRange)}
    */
   public Set<Host> getReplicas(String keyspace, TokenRange range) {
     keyspace = handleId(keyspace);
     TokenMap current = tokenMap;
     if (current == null) {
       return Collections.emptySet();
     } else {
-      Set<Host> hosts = current.getReplicas(keyspace, range.getEnd());
-      return hosts == null ? Collections.<Host>emptySet() : hosts;
+      return new HashSet<>(current.getReplicas(keyspace, range.getEnd()));
+    }
+  }
+
+  /**
+   * Returns the list of hosts that are replica for a given token range.
+   *
+   * <p>Note that it is assumed that the input range does not overlap across multiple host ranges.
+   * If the range extends over multiple hosts, it only returns the replicas for those hosts that are
+   * replicas for the last token of the range. This behavior may change in a future release, see <a
+   * href="https://datastax-oss.atlassian.net/browse/JAVA-1355">JAVA-1355</a>.
+   *
+   * <p>Also 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 range the token range.
+   * @return the list of replicas for {@code range} 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 List<Host> getReplicasList(String keyspace, TokenRange range) {
+    keyspace = handleId(keyspace);
+    TokenMap current = tokenMap;
+    if (current == null) {
+      return EMPTY_LIST;
+    } else {
+      return current.getReplicas(keyspace, range.getEnd());
     }
   }
 
@@ -985,7 +1066,7 @@ private static class TokenMap {
 
     private final Token.Factory factory;
     private final Map<Host, Set<Token>> primaryToTokens;
-    private final Map<String, Map<Token, Set<Host>>> tokenToHostsByKeyspace;
+    private final Map<String, Map<Token, ImmutableList<Host>>> tokenToHostsByKeyspace;
     private final Map<String, Map<Host, Set<TokenRange>>> hostsToRangesByKeyspace;
     private final List<Token> ring;
     private final Set<TokenRange> tokenRanges;
@@ -997,7 +1078,7 @@ private TokenMap(
         Set<TokenRange> tokenRanges,
         Map<Token, Host> tokenToPrimary,
         Map<Host, Set<Token>> primaryToTokens,
-        Map<String, Map<Token, Set<Host>>> tokenToHostsByKeyspace,
+        Map<String, Map<Token, ImmutableList<Host>>> tokenToHostsByKeyspace,
         Map<String, Map<Host, Set<TokenRange>>> hostsToRangesByKeyspace) {
       this.factory = factory;
       this.ring = ring;
@@ -1042,15 +1123,13 @@ private static TokenMap build(
         Set<TokenRange> tokenRanges,
         Map<Token, Host> tokenToPrimary) {
       Set<Host> hosts = allTokens.keySet();
-      Map<String, Map<Token, Set<Host>>> tokenToHosts =
-          new HashMap<String, Map<Token, Set<Host>>>();
-      Map<ReplicationStrategy, Map<Token, Set<Host>>> replStrategyToHosts =
-          new HashMap<ReplicationStrategy, Map<Token, Set<Host>>>();
-      Map<String, Map<Host, Set<TokenRange>>> hostsToRanges =
-          new HashMap<String, Map<Host, Set<TokenRange>>>();
+      Map<String, Map<Token, ImmutableList<Host>>> tokenToHosts = new HashMap<>();
+      Map<ReplicationStrategy, Map<Token, ImmutableList<Host>>> replStrategyToHosts =
+          new HashMap<>();
+      Map<String, Map<Host, Set<TokenRange>>> hostsToRanges = new HashMap<>();
       for (KeyspaceMetadata keyspace : keyspaces) {
         ReplicationStrategy strategy = keyspace.replicationStrategy();
-        Map<Token, Set<Host>> ksTokens = replStrategyToHosts.get(strategy);
+        Map<Token, ImmutableList<Host>> ksTokens = replStrategyToHosts.get(strategy);
         if (ksTokens == null) {
           ksTokens =
               (strategy == null)
@@ -1077,14 +1156,14 @@ private static TokenMap build(
           factory, ring, tokenRanges, tokenToPrimary, allTokens, tokenToHosts, hostsToRanges);
     }
 
-    private Set<Host> getReplicas(String keyspace, Token token) {
+    private ImmutableList<Host> getReplicas(String keyspace, Token token) {
 
-      Map<Token, Set<Host>> tokenToHosts = tokenToHostsByKeyspace.get(keyspace);
-      if (tokenToHosts == null) return Collections.emptySet();
+      Map<Token, ImmutableList<Host>> tokenToHosts = tokenToHostsByKeyspace.get(keyspace);
+      if (tokenToHosts == null) return EMPTY_LIST;
 
       // If the token happens to be one of the "primary" tokens, get result directly
-      Set<Host> hosts = tokenToHosts.get(token);
-      if (hosts != null) return hosts;
+      ImmutableList<Host> hosts = tokenToHosts.get(token);
+      if (hosts != null) return ImmutableList.copyOf(hosts);
 
       // Otherwise, find closest "primary" token on the ring
       int i = Collections.binarySearch(ring, token);
@@ -1096,10 +1175,10 @@ private Set<Host> getReplicas(String keyspace, Token token) {
       return tokenToHosts.get(ring.get(i));
     }
 
-    private static Map<Token, Set<Host>> makeNonReplicatedMap(Map<Token, Host> input) {
-      Map<Token, Set<Host>> output = new HashMap<Token, Set<Host>>(input.size());
+    private static Map<Token, ImmutableList<Host>> makeNonReplicatedMap(Map<Token, Host> input) {
+      Map<Token, ImmutableList<Host>> output = new HashMap<>(input.size());
       for (Map.Entry<Token, Host> entry : input.entrySet())
-        output.put(entry.getKey(), ImmutableSet.of(entry.getValue()));
+        output.put(entry.getKey(), ImmutableList.of(entry.getValue()));
       return output;
     }
 
@@ -1119,11 +1198,11 @@ private static Set<TokenRange> makeTokenRanges(List<Token> ring, Token.Factory f
     }
 
     private static Map<Host, Set<TokenRange>> computeHostsToRangesMap(
-        Set<TokenRange> tokenRanges, Map<Token, Set<Host>> ksTokens, int hostCount) {
+        Set<TokenRange> tokenRanges, Map<Token, ImmutableList<Host>> ksTokens, int hostCount) {
       Map<Host, ImmutableSet.Builder<TokenRange>> builders =
           Maps.newHashMapWithExpectedSize(hostCount);
       for (TokenRange range : tokenRanges) {
-        Set<Host> replicas = ksTokens.get(range.getEnd());
+        List<Host> replicas = ksTokens.get(range.getEnd());
         for (Host host : replicas) {
           ImmutableSet.Builder<TokenRange> hostRanges = builders.get(host);
           if (hostRanges == null) {

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

@@ -15,7 +15,7 @@
  */
 package com.datastax.driver.core;
 
-import com.google.common.collect.ImmutableSet;
+import com.google.common.collect.ImmutableList;
 import com.google.common.collect.Maps;
 import com.google.common.collect.Sets;
 import java.util.HashMap;
@@ -67,7 +67,7 @@ static ReplicationStrategy create(Map<String, String> replicationOptions) {
     }
   }
 
-  abstract Map<Token, Set<Host>> computeTokenToReplicaMap(
+  abstract Map<Token, ImmutableList<Host>> computeTokenToReplicaMap(
       String keyspaceName, Map<Token, Host> tokenToPrimary, List<Token> ring);
 
   private static Token getTokenWrapping(int i, List<Token> ring) {
@@ -83,18 +83,18 @@ private SimpleStrategy(ReplicationFactor replicationFactor) {
     }
 
     @Override
-    Map<Token, Set<Host>> computeTokenToReplicaMap(
+    Map<Token, ImmutableList<Host>> computeTokenToReplicaMap(
         String keyspaceName, Map<Token, Host> tokenToPrimary, List<Token> ring) {
 
       int rf = Math.min(replicationFactor.fullReplicas(), ring.size());
 
-      Map<Token, Set<Host>> replicaMap = new HashMap<Token, Set<Host>>(tokenToPrimary.size());
+      Map<Token, ImmutableList<Host>> replicaMap = new HashMap<>(tokenToPrimary.size());
       for (int i = 0; i < ring.size(); i++) {
         // Consecutive sections of the ring can assigned to the same host
-        Set<Host> replicas = new LinkedHashSet<Host>();
+        Set<Host> replicas = new LinkedHashSet<>();
         for (int j = 0; j < ring.size() && replicas.size() < rf; j++)
           replicas.add(tokenToPrimary.get(getTokenWrapping(i + j, ring)));
-        replicaMap.put(ring.get(i), ImmutableSet.copyOf(replicas));
+        replicaMap.put(ring.get(i), ImmutableList.copyOf(replicas));
       }
       return replicaMap;
     }
@@ -125,7 +125,7 @@ private NetworkTopologyStrategy(Map<String, ReplicationFactor> replicationFactor
     }
 
     @Override
-    Map<Token, Set<Host>> computeTokenToReplicaMap(
+    Map<Token, ImmutableList<Host>> computeTokenToReplicaMap(
         String keyspaceName, Map<Token, Host> tokenToPrimary, List<Token> ring) {
 
       logger.debug("Computing token to replica map for keyspace: {}.", keyspaceName);
@@ -135,7 +135,7 @@ Map<Token, Set<Host>> computeTokenToReplicaMap(
 
       // This is essentially a copy of org.apache.cassandra.locator.NetworkTopologyStrategy
       Map<String, Set<String>> racks = getRacksInDcs(tokenToPrimary.values());
-      Map<Token, Set<Host>> replicaMap = new HashMap<Token, Set<Host>>(tokenToPrimary.size());
+      Map<Token, ImmutableList<Host>> replicaMap = new HashMap<>(tokenToPrimary.size());
       Map<String, Integer> dcHostCount = Maps.newHashMapWithExpectedSize(replicationFactors.size());
       Set<String> warnedDcs = Sets.newHashSetWithExpectedSize(replicationFactors.size());
       // find maximum number of nodes in each DC
@@ -157,7 +157,7 @@ Map<Token, Set<Host>> computeTokenToReplicaMap(
         }
 
         // Preserve order - primary replica will be first
-        Set<Host> replicas = new LinkedHashSet<Host>();
+        Set<Host> replicas = new LinkedHashSet<>();
         for (int j = 0; j < ring.size() && !allDone(allDcReplicas, dcHostCount); j++) {
           Host h = tokenToPrimary.get(getTokenWrapping(i + j, ring));
           String dc = h.getDatacenter();
@@ -214,7 +214,7 @@ Map<Token, Set<Host>> computeTokenToReplicaMap(
           }
         }
 
-        replicaMap.put(ring.get(i), ImmutableSet.copyOf(replicas));
+        replicaMap.put(ring.get(i), ImmutableList.copyOf(replicas));
       }
 
       long duration = System.currentTimeMillis() - startTime;

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

@@ -121,11 +121,11 @@ private Iterator<Host> getReplicas(
       tableName = defs.getTable(0);
     }
 
-    final Set<Host> replicas =
+    final List<Host> replicas =
         manager
             .cluster
             .getMetadata()
-            .getReplicas(Metadata.quote(keyspace), tableName, partitioner, partitionKey);
+            .getReplicasList(Metadata.quote(keyspace), tableName, partitioner, partitionKey);
 
     // replicas are stored in the right order starting with the primary replica
     return replicas.iterator();

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

@@ -659,6 +659,7 @@ else if (fetchSize != Integer.MAX_VALUE)
       boolean skipMetadata =
           protocolVersion != ProtocolVersion.V1
               && bs.statement.getPreparedId().resultSetMetadata.variables != null;
+
       Requests.QueryProtocolOptions options =
           new Requests.QueryProtocolOptions(
               Message.Request.Type.EXECUTE,