Skip to content

Commit 727a80d

Browse files
Gor027Gor027
committed
Add RandomTwoChoice load balancing policy
The RandomTwoChoice policy is a wrapper load balancing policy that adds the "Power of 2 Choice" algorithm to its child policy. It will compare the first two hosts returned from the child policy query plan, and will first return the host with the target shard having fewer inflight requests. The rest of the child query plan will be left intact. It is intended to be used with TokenAwarePolicy and RoundRobinPolicy, to send queries to the replica nodes by always avoiding the worst option (the replica with the target shard having the most inflight requests will not be chosen).
1 parent 7399a0a commit 727a80d

File tree

3 files changed

+299
-0
lines changed

3 files changed

+299
-0
lines changed

‎driver-core/src/main/java/com/datastax/driver/core/Metrics.java

+32
Original file line numberDiff line numberDiff line change
@@ -130,7 +130,35 @@ public Integer getValue() {
130130
return value;
131131
}
132132
});
133+
private final Gauge<Map<Host, Map<Integer, Integer>>> perShardInflightRequestInfo =
134+
registry.register(
135+
"per-shard-inflight-request-info",
136+
new Gauge<Map<Host, Map<Integer, Integer>>>() {
137+
@Override
138+
public Map<Host, Map<Integer, Integer>> getValue() {
139+
Map<Host, Map<Integer, Integer>> result = new HashMap<Host, Map<Integer, Integer>>();
140+
for (SessionManager session : manager.sessions) {
141+
for (Map.Entry<Host, HostConnectionPool> poolEntry : session.pools.entrySet()) {
142+
HostConnectionPool hostConnectionPool = poolEntry.getValue();
143+
Map<Integer, Integer> perShardInflightRequests = new HashMap<Integer, Integer>();
144+
145+
for (int shardId = 0;
146+
shardId < hostConnectionPool.connections.length;
147+
shardId++) {
148+
int shardInflightRequests = 0;
149+
for (Connection connection : hostConnectionPool.connections[shardId]) {
150+
shardInflightRequests += connection.inFlight.get();
151+
}
152+
perShardInflightRequests.put(shardId, shardInflightRequests);
153+
}
154+
155+
result.put(poolEntry.getKey(), perShardInflightRequests);
156+
}
157+
}
133158

159+
return result;
160+
}
161+
});
134162
private final Gauge<Integer> executorQueueDepth;
135163
private final Gauge<Integer> blockingExecutorQueueDepth;
136164
private final Gauge<Integer> reconnectionSchedulerQueueSize;
@@ -374,6 +402,10 @@ public Gauge<Map<Host, Integer>> getShardAwarenessInfo() {
374402
return shardAwarenessInfo;
375403
}
376404

405+
public Gauge<Map<Host, Map<Integer, Integer>>> getPerShardInflightRequestInfo() {
406+
return perShardInflightRequestInfo;
407+
}
408+
377409
/**
378410
* Returns the number of bytes sent so far.
379411
*

‎driver-core/src/main/java/com/datastax/driver/core/policies/RandomTwoChoicePolicy.java

+157
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,157 @@
1+
package com.datastax.driver.core.policies;
2+
3+
import com.datastax.driver.core.*;
4+
import com.google.common.collect.AbstractIterator;
5+
import java.nio.ByteBuffer;
6+
import java.util.*;
7+
import java.util.concurrent.atomic.AtomicInteger;
8+
9+
/**
10+
* A wrapper load balancing policy that adds "Power of 2 Choice" algorithm to a child policy.
11+
*
12+
* <p>This policy encapsulates another policy. The resulting policy works in the following way:
13+
*
14+
* <ul>
15+
* <li>the {@code distance} method is inherited from the child policy.
16+
* <li>the {@code newQueryPlan} method will compare first two hosts (by number of inflight
17+
* requests) returned from the {@code newQueryPlan} method of the child policy, and the host
18+
* with fewer number of inflight requests will be returned the first. It will allow to always
19+
* avoid the worst option (comparing by number of inflight requests).
20+
* <li>besides the first two hosts returned by the child policy's {@code newQueryPlan} method, the
21+
* ordering of the rest of the hosts will remain the same.
22+
* </ul>
23+
*
24+
* <p>If you wrap the {@code RandomTwoChoicePolicy} policy with {@code TokenAwarePolicy}, it will
25+
* compare the first two replicas by the number of inflight requests, and the worse option will
26+
* always be avoided. In that case, it is recommended to use the TokenAwarePolicy with {@code
27+
* ReplicaOrdering.RANDOM strategy}, which will return the replicas in a shuffled order and thus
28+
* will make the "Power of 2 Choice" algorithm more efficient.
29+
*/
30+
public class RandomTwoChoicePolicy implements ChainableLoadBalancingPolicy {
31+
private final LoadBalancingPolicy childPolicy;
32+
private volatile Metrics metrics;
33+
private volatile Metadata clusterMetadata;
34+
private volatile ProtocolVersion protocolVersion;
35+
private volatile CodecRegistry codecRegistry;
36+
37+
/**
38+
* Creates a new {@code RandomTwoChoicePolicy}.
39+
*
40+
* @param childPolicy the load balancing policy to wrap with "Power of 2 Choice" algorithm.
41+
*/
42+
public RandomTwoChoicePolicy(LoadBalancingPolicy childPolicy) {
43+
this.childPolicy = childPolicy;
44+
}
45+
46+
@Override
47+
public LoadBalancingPolicy getChildPolicy() {
48+
return childPolicy;
49+
}
50+
51+
@Override
52+
public void init(Cluster cluster, Collection<Host> hosts) {
53+
this.metrics = cluster.getMetrics();
54+
this.clusterMetadata = cluster.getMetadata();
55+
this.protocolVersion = cluster.getConfiguration().getProtocolOptions().getProtocolVersion();
56+
this.codecRegistry = cluster.getConfiguration().getCodecRegistry();
57+
childPolicy.init(cluster, hosts);
58+
}
59+
60+
/**
61+
* {@inheritDoc}
62+
*
63+
* <p>This implementation always returns distances as reported by the wrapped policy.
64+
*/
65+
@Override
66+
public HostDistance distance(Host host) {
67+
return childPolicy.distance(host);
68+
}
69+
70+
/**
71+
* {@inheritDoc}
72+
*
73+
* <p>The returned plan will compare (by the number of inflight requests) the first 2 hosts
74+
* returned by the child policy's {@code newQueryPlan} method, and the host with fewer inflight
75+
* requests will be returned the first. The rest of the child policy's query plan will be left
76+
* intact.
77+
*/
78+
@Override
79+
public Iterator<Host> newQueryPlan(String loggedKeyspace, Statement statement) {
80+
String keyspace = statement.getKeyspace();
81+
if (keyspace == null) keyspace = loggedKeyspace;
82+
83+
ByteBuffer routingKey = statement.getRoutingKey(protocolVersion, codecRegistry);
84+
if (routingKey == null || keyspace == null) {
85+
return childPolicy.newQueryPlan(loggedKeyspace, statement);
86+
}
87+
88+
final Token t = clusterMetadata.newToken(statement.getPartitioner(), routingKey);
89+
final Iterator<Host> childIterator = childPolicy.newQueryPlan(keyspace, statement);
90+
91+
final Host host1 = childIterator.hasNext() ? childIterator.next() : null;
92+
final Host host2 = childIterator.hasNext() ? childIterator.next() : null;
93+
94+
final AtomicInteger host1ShardInflightRequests = new AtomicInteger(0);
95+
final AtomicInteger host2ShardInflightRequests = new AtomicInteger(0);
96+
97+
if (host1 != null) {
98+
final int host1ShardId = host1.getShardingInfo().shardId(t);
99+
host1ShardInflightRequests.set(
100+
metrics.getPerShardInflightRequestInfo().getValue().get(host1).get(host1ShardId));
101+
}
102+
103+
if (host2 != null) {
104+
final int host2ShardId = host2.getShardingInfo().shardId(t);
105+
host2ShardInflightRequests.set(
106+
metrics.getPerShardInflightRequestInfo().getValue().get(host2).get(host2ShardId));
107+
}
108+
109+
return new AbstractIterator<Host>() {
110+
private final Host firstChosenHost =
111+
host1ShardInflightRequests.get() < host2ShardInflightRequests.get() ? host1 : host2;
112+
private final Host secondChosenHost =
113+
host1ShardInflightRequests.get() < host2ShardInflightRequests.get() ? host2 : host1;
114+
private int index = 0;
115+
116+
@Override
117+
protected Host computeNext() {
118+
if (index == 0) {
119+
index++;
120+
return firstChosenHost;
121+
} else if (index == 1) {
122+
index++;
123+
return secondChosenHost;
124+
} else if (childIterator.hasNext()) {
125+
return childIterator.next();
126+
}
127+
128+
return endOfData();
129+
}
130+
};
131+
}
132+
133+
@Override
134+
public void onAdd(Host host) {
135+
childPolicy.onAdd(host);
136+
}
137+
138+
@Override
139+
public void onUp(Host host) {
140+
childPolicy.onUp(host);
141+
}
142+
143+
@Override
144+
public void onDown(Host host) {
145+
childPolicy.onDown(host);
146+
}
147+
148+
@Override
149+
public void onRemove(Host host) {
150+
childPolicy.onRemove(host);
151+
}
152+
153+
@Override
154+
public void close() {
155+
childPolicy.close();
156+
}
157+
}

‎driver-core/src/test/java/com/datastax/driver/core/policies/RandomTwoChoicePolicyTest.java

+110
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,110 @@
1+
package com.datastax.driver.core.policies;
2+
3+
import static org.assertj.core.api.Assertions.assertThat;
4+
import static org.mockito.Mockito.mock;
5+
import static org.mockito.Mockito.when;
6+
7+
import com.codahale.metrics.Gauge;
8+
import com.datastax.driver.core.*;
9+
import java.nio.ByteBuffer;
10+
import java.util.*;
11+
import org.assertj.core.util.Sets;
12+
import org.testng.annotations.BeforeMethod;
13+
import org.testng.annotations.Test;
14+
15+
public class RandomTwoChoicePolicyTest {
16+
private final ByteBuffer routingKey = ByteBuffer.wrap(new byte[] {1, 2, 3, 4});
17+
private final RegularStatement statement =
18+
new SimpleStatement("irrelevant").setRoutingKey(routingKey).setKeyspace("keyspace");
19+
private final Host host1 = mock(Host.class);
20+
private final Host host2 = mock(Host.class);
21+
private final Host host3 = mock(Host.class);
22+
private Cluster cluster;
23+
24+
@SuppressWarnings("unchecked")
25+
private final Gauge<Map<Host, Map<Integer, Integer>>> gauge =
26+
mock((Class<Gauge<Map<Host, Map<Integer, Integer>>>>) (Object) Gauge.class);
27+
28+
@BeforeMethod(groups = "unit")
29+
public void initMocks() {
30+
CodecRegistry codecRegistry = new CodecRegistry();
31+
cluster = mock(Cluster.class);
32+
Configuration configuration = mock(Configuration.class);
33+
ProtocolOptions protocolOptions = mock(ProtocolOptions.class);
34+
Metadata metadata = mock(Metadata.class);
35+
Metrics metrics = mock(Metrics.class);
36+
Token t = mock(Token.class);
37+
ShardingInfo shardingInfo = mock(ShardingInfo.class);
38+
39+
when(metrics.getPerShardInflightRequestInfo()).thenReturn(gauge);
40+
when(cluster.getConfiguration()).thenReturn(configuration);
41+
when(configuration.getCodecRegistry()).thenReturn(codecRegistry);
42+
when(configuration.getProtocolOptions()).thenReturn(protocolOptions);
43+
when(protocolOptions.getProtocolVersion()).thenReturn(ProtocolVersion.DEFAULT);
44+
when(cluster.getMetadata()).thenReturn(metadata);
45+
when(cluster.getMetrics()).thenReturn(metrics);
46+
when(metadata.getReplicas(Metadata.quote("keyspace"), null, routingKey))
47+
.thenReturn(Sets.newLinkedHashSet(host1, host2, host3));
48+
when(metadata.newToken(null, routingKey)).thenReturn(t);
49+
when(host1.getShardingInfo()).thenReturn(shardingInfo);
50+
when(host2.getShardingInfo()).thenReturn(shardingInfo);
51+
when(host3.getShardingInfo()).thenReturn(shardingInfo);
52+
when(shardingInfo.shardId(t)).thenReturn(0);
53+
when(host1.isUp()).thenReturn(true);
54+
when(host2.isUp()).thenReturn(true);
55+
when(host3.isUp()).thenReturn(true);
56+
}
57+
58+
@Test(groups = "unit")
59+
public void should_prefer_host_with_less_inflight_requests() {
60+
// given
61+
Map<Host, Map<Integer, Integer>> perHostInflightRequests =
62+
new HashMap<Host, Map<Integer, Integer>>() {
63+
{
64+
put(
65+
host1,
66+
new HashMap<Integer, Integer>() {
67+
{
68+
put(0, 6);
69+
}
70+
});
71+
put(
72+
host2,
73+
new HashMap<Integer, Integer>() {
74+
{
75+
put(0, 2);
76+
}
77+
});
78+
put(
79+
host3,
80+
new HashMap<Integer, Integer>() {
81+
{
82+
put(0, 4);
83+
}
84+
});
85+
}
86+
};
87+
RandomTwoChoicePolicy policy =
88+
new RandomTwoChoicePolicy(
89+
new TokenAwarePolicy(
90+
new RoundRobinPolicy(), TokenAwarePolicy.ReplicaOrdering.TOPOLOGICAL));
91+
policy.init(
92+
cluster,
93+
new ArrayList<Host>() {
94+
95+
{
96+
add(host1);
97+
add(host2);
98+
add(host3);
99+
}
100+
});
101+
when(gauge.getValue()).thenReturn(perHostInflightRequests);
102+
103+
Iterator<Host> queryPlan = policy.newQueryPlan("keyspace", statement);
104+
// host2 should appear first in the query plan with fewer inflight requests than host1
105+
106+
assertThat(queryPlan.next()).isEqualTo(host2);
107+
assertThat(queryPlan.next()).isEqualTo(host1);
108+
assertThat(queryPlan.next()).isEqualTo(host3);
109+
}
110+
}

0 commit comments

Comments
 (0)