Prototype power of 2 choices algorithm for token-aware LBP

olim7t · olim7t · commit f008dbcf7362 · 2017-10-23T10:08:33.000-07:00
diff --git a/driver-core/src/main/java/com/datastax/driver/core/Cluster.java b/driver-core/src/main/java/com/datastax/driver/core/Cluster.java
@@ -420,6 +420,20 @@ public Metrics getMetrics() {
         return manager.metrics;
     }
 
+    /**
+     * TODO find a better way to expose this
+     */
+    public int inFlight(Host host) {
+        int inFlight = 0;
+        for (SessionManager session : manager.sessions) {
+            HostConnectionPool pool = session.pools.get(host);
+            if (pool != null) {
+                inFlight += pool.totalInFlight.get();
+            }
+        }
+        return inFlight;
+    }
+
     /**
      * Registers the provided listener to be notified on hosts
      * up/down/added/removed events.
diff --git a/driver-core/src/main/java/com/datastax/driver/core/policies/TokenAwarePolicy.java b/driver-core/src/main/java/com/datastax/driver/core/policies/TokenAwarePolicy.java
@@ -18,6 +18,8 @@
 import com.datastax.driver.core.*;
 import com.google.common.collect.AbstractIterator;
 import com.google.common.collect.Lists;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
 
 import java.nio.ByteBuffer;
 import java.util.*;
@@ -46,8 +48,18 @@
  */
 public class TokenAwarePolicy implements ChainableLoadBalancingPolicy {
 
+    private static final Logger LOG = LoggerFactory.getLogger(TokenAwarePolicy.class);
+    private static final boolean POWER_OF_TWO_CHOICES = Boolean.getBoolean("com.datastax.driver.POWER_OF_TWO_CHOICES");
+
+    static {
+        if (POWER_OF_TWO_CHOICES) {
+            LOG.info("Activating power of two choices");
+        }
+    }
+
     private final LoadBalancingPolicy childPolicy;
     private final boolean shuffleReplicas;
+    private volatile Cluster cluster;
     private volatile Metadata clusterMetadata;
     private volatile ProtocolVersion protocolVersion;
     private volatile CodecRegistry codecRegistry;
@@ -87,6 +99,7 @@ public LoadBalancingPolicy getChildPolicy() {
 
     @Override
     public void init(Cluster cluster, Collection<Host> hosts) {
+        this.cluster = cluster;
         clusterMetadata = cluster.getMetadata();
         protocolVersion = cluster.getConfiguration().getProtocolOptions().getProtocolVersion();
         codecRegistry = cluster.getConfiguration().getCodecRegistry();
@@ -130,14 +143,7 @@ public Iterator<Host> newQueryPlan(final String loggedKeyspace, final Statement
         if (replicas.isEmpty())
             return childPolicy.newQueryPlan(loggedKeyspace, statement);
 
-        final Iterator<Host> iter;
-        if (shuffleReplicas) {
-            List<Host> l = Lists.newArrayList(replicas);
-            Collections.shuffle(l);
-            iter = l.iterator();
-        } else {
-            iter = replicas.iterator();
-        }
+        final Iterator<Host> iter = shuffle(replicas);
 
         return new AbstractIterator<Host>() {
 
@@ -165,6 +171,25 @@ protected Host computeNext() {
         };
     }
 
+    private Iterator<Host> shuffle(Set<Host> replicas) {
+        if (shuffleReplicas) {
+            List<Host> l = Lists.newArrayList(replicas);
+            Collections.shuffle(l);
+            if (POWER_OF_TWO_CHOICES && l.size() > 2) {
+                // Order the first two hosts by increasing load
+                Host host0 = l.get(0);
+                Host host1 = l.get(1);
+                if (cluster.inFlight(host0) > cluster.inFlight(host1)) {
+                    l.set(0, host1);
+                    l.set(1, host0);
+                }
+            }
+            return l.iterator();
+        } else {
+            return replicas.iterator();
+        }
+    }
+
     @Override
     public void onUp(Host host) {
         childPolicy.onUp(host);
