Add ReadWriteLock for concurrent tablet modifications

Each table now has a dedicated lock for its set of tablets.
Multiple sets can be modified concurrently. Most importantly without
copying the whole set.
Should improve times, mostly in cases of large amount of tablets and
reduce number of misses during tablet shuffling.

Author: Bouncheck

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

@@ -36,13 +36,13 @@
 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.Lock;
 import java.util.concurrent.locks.ReentrantLock;
 import java.util.stream.Collectors;
 import org.slf4j.Logger;
@@ -916,8 +916,8 @@ public int getShardForTabletToken(
     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) {
+    TabletMap.TabletSet targetSet = tabletMap.getMapping().get(key);
+    if (targetSet == null) {
       logger.trace(
           "Could not determine shard for token {} on host {} because table {}.{} is not present in tablets "
               + "metadata. Returning -1.",
@@ -927,13 +927,20 @@ public int getShardForTabletToken(
           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();
+    Lock readLock = targetSet.lock.readLock();
+    try {
+      readLock.lock();
+      TabletMap.Tablet row =
+          targetSet.tablets.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();
+          }
         }
       }
+    } finally {
+      readLock.unlock();
     }
     logger.trace(
         "Could not find tablet corresponding to token {} on host {} for table {} in keyspace {}. Returning -1.",

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

@@ -14,6 +14,8 @@
 import java.util.UUID;
 import java.util.concurrent.ConcurrentHashMap;
 import java.util.concurrent.ConcurrentMap;
+import java.util.concurrent.locks.Lock;
+import java.util.concurrent.locks.ReentrantReadWriteLock;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
@@ -25,7 +27,9 @@
 public class TabletMap {
   private static final Logger logger = LoggerFactory.getLogger(TabletMap.class);
 
-  private final ConcurrentMap<KeyspaceTableNamePair, NavigableSet<Tablet>> mapping;
+  // There are no additional locking mechanisms for the mapping field itself, however each TabletSet
+  // inside has its own ReadWriteLock that should be used when dealing with its internals.
+  private final ConcurrentMap<KeyspaceTableNamePair, TabletSet> mapping;
 
   private final Cluster.Manager cluster;
 
@@ -34,7 +38,7 @@ public class TabletMap {
   private TypeCodec<TupleValue> tabletPayloadCodec = null;
 
   public TabletMap(
-      Cluster.Manager cluster, ConcurrentMap<KeyspaceTableNamePair, NavigableSet<Tablet>> mapping) {
+      Cluster.Manager cluster, ConcurrentMap<KeyspaceTableNamePair, TabletSet> mapping) {
     this.cluster = cluster;
     this.mapping = mapping;
   }
@@ -46,9 +50,9 @@ public static TabletMap emptyMap(Cluster.Manager cluster) {
   /**
    * Returns the mapping of tables to their tablets.
    *
-   * @return the Map keyed by (keyspace,table) pairs with Set of tablets as value type.
+   * @return the Map keyed by (keyspace,table) pairs with {@link TabletSet} as value type.
    */
-  public Map<KeyspaceTableNamePair, NavigableSet<Tablet>> getMapping() {
+  public Map<KeyspaceTableNamePair, TabletSet> getMapping() {
     return mapping;
   }
 
@@ -68,28 +72,34 @@ public Set<UUID> getReplicas(String keyspace, String table, long token) {
       return Collections.emptySet();
     }
 
-    NavigableSet<Tablet> set = mapping.get(key);
-    if (set == null) {
+    TabletSet tabletSet = mapping.get(key);
+    if (tabletSet == null) {
       logger.trace(
           "There is no tablets for {}.{} in this mapping. Returning empty set.", keyspace, table);
       return Collections.emptySet();
     }
-    Tablet row = mapping.get(key).ceiling(Tablet.malformedTablet(token));
-    if (row == null || row.firstToken >= token) {
-      logger.trace(
-          "Could not find tablet for {}.{} that owns token {}. Returning empty set.",
-          keyspace,
-          table,
-          token);
-      return Collections.emptySet();
-    }
+    Lock readLock = tabletSet.lock.readLock();
+    try {
+      readLock.lock();
+      Tablet row = mapping.get(key).tablets.ceiling(Tablet.malformedTablet(token));
+      if (row == null || row.firstToken >= token) {
+        logger.trace(
+            "Could not find tablet for {}.{} that owns token {}. Returning empty set.",
+            keyspace,
+            table,
+            token);
+        return Collections.emptySet();
+      }
 
-    HashSet<UUID> uuidSet = new HashSet<>();
-    for (HostShardPair hostShardPair : row.replicas) {
-      if (cluster.metadata.getHost(hostShardPair.getHost()) != null)
-        uuidSet.add(hostShardPair.getHost());
+      HashSet<UUID> uuidSet = new HashSet<>();
+      for (HostShardPair hostShardPair : row.replicas) {
+        if (cluster.metadata.getHost(hostShardPair.getHost()) != null)
+          uuidSet.add(hostShardPair.getHost());
+      }
+      return uuidSet;
+    } finally {
+      readLock.unlock();
     }
-    return uuidSet;
   }
 
   /**
@@ -121,46 +131,47 @@ void processTabletsRoutingV1Payload(String keyspace, String table, ByteBuffer pa
       HostShardPair hostShardPair = new HostShardPair(tuple.getUUID(0), tuple.getInt(1));
       replicas.add(hostShardPair);
     }
-
-    // Working on a copy to avoid concurrent modification of the same set
-    NavigableSet<Tablet> existingTablets =
-        new TreeSet<>(mapping.computeIfAbsent(ktPair, k -> new TreeSet<>()));
-
-    // Single tablet token range is represented by (firstToken, lastToken] interval
-    // We need to do two sweeps: remove overlapping tablets by looking at lastToken of existing
-    // tablets
-    // and then by looking at firstToken of existing tablets. Currently, the tablets are sorted
-    // according
-    // to their lastTokens.
-
-    // First sweep: remove all tablets whose lastToken is inside this interval
-    Iterator<Tablet> it =
-        existingTablets.headSet(Tablet.malformedTablet(lastToken), true).descendingIterator();
-    while (it.hasNext()) {
-      Tablet tablet = it.next();
-      if (tablet.lastToken <= firstToken) {
-        break;
+    Tablet newTablet = new Tablet(keyspace, null, table, firstToken, lastToken, replicas);
+
+    TabletSet tabletSet = mapping.computeIfAbsent(ktPair, k -> new TabletSet());
+    Lock writeLock = tabletSet.lock.writeLock();
+    try {
+      writeLock.lock();
+      NavigableSet<Tablet> currentTablets = tabletSet.tablets;
+      // Single tablet token range is represented by (firstToken, lastToken] interval
+      // We need to do two sweeps: remove overlapping tablets by looking at lastToken of existing
+      // tablets
+      // and then by looking at firstToken of existing tablets. Currently, the tablets are sorted
+      // according
+      // to their lastTokens.
+
+      // First sweep: remove all tablets whose lastToken is inside this interval
+      Iterator<Tablet> it = currentTablets.headSet(newTablet, true).descendingIterator();
+      while (it.hasNext()) {
+        Tablet tablet = it.next();
+        if (tablet.lastToken <= firstToken) {
+          break;
+        }
+        it.remove();
       }
-      it.remove();
-    }
 
-    // Second sweep: remove all tablets whose firstToken is inside this tuple's (firstToken,
-    // lastToken]
-    // After the first sweep, this theoretically should remove at most 1 tablet
-    it = existingTablets.tailSet(Tablet.malformedTablet(lastToken), true).iterator();
-    while (it.hasNext()) {
-      Tablet tablet = it.next();
-      if (tablet.firstToken >= lastToken) {
-        break;
+      // Second sweep: remove all tablets whose firstToken is inside this tuple's (firstToken,
+      // lastToken]
+      // After the first sweep, this theoretically should remove at most 1 tablet
+      it = currentTablets.tailSet(newTablet, true).iterator();
+      while (it.hasNext()) {
+        Tablet tablet = it.next();
+        if (tablet.firstToken >= lastToken) {
+          break;
+        }
+        it.remove();
       }
-      it.remove();
-    }
 
-    // Add new (now) non-overlapping tablet
-    existingTablets.add(new Tablet(keyspace, null, table, firstToken, lastToken, replicas));
-
-    // Set the updated result in the main map
-    mapping.put(ktPair, existingTablets);
+      // Add new (now) non-overlapping tablet
+      currentTablets.add(newTablet);
+    } finally {
+      writeLock.unlock();
+    }
   }
 
   public TupleType getPayloadOuterTuple() {
@@ -258,6 +269,18 @@ public int hashCode() {
     }
   }
 
+  /**
+   * Set of tablets bundled with ReadWriteLock to allow concurrent modification for different sets.
+   */
+  public static class TabletSet {
+    final NavigableSet<Tablet> tablets;
+    final ReentrantReadWriteLock lock = new ReentrantReadWriteLock();
+
+    public TabletSet() {
+      this.tablets = new TreeSet<>();
+    }
+  }
+
   /**
    * Represents a single tablet created from tablets-routing-v1 custom payload. Its {@code
    * compareTo} implementation intentionally relies solely on {@code lastToken} in order to allow