Skip to content

Commit bd35105

Browse files
nitinittnitinitt
committed
JAVA-3142: Code review comments aaddressed and enhanced test case
1 parent ec900b0 commit bd35105

File tree

2 files changed

+56
-23
lines changed

2 files changed

+56
-23
lines changed

‎core/src/main/java/com/datastax/oss/driver/internal/core/loadbalancing/BasicLoadBalancingPolicy.java

Lines changed: 19 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -49,10 +49,10 @@
4949
import edu.umd.cs.findbugs.annotations.Nullable;
5050
import java.nio.ByteBuffer;
5151
import java.util.ArrayList;
52-
import java.util.Arrays;
5352
import java.util.Collections;
5453
import java.util.Comparator;
5554
import java.util.LinkedHashSet;
55+
import java.util.List;
5656
import java.util.Map;
5757
import java.util.Objects;
5858
import java.util.Optional;
@@ -61,6 +61,7 @@
6161
import java.util.UUID;
6262
import java.util.concurrent.atomic.AtomicInteger;
6363
import java.util.function.IntUnaryOperator;
64+
import java.util.stream.Stream;
6465
import net.jcip.annotations.ThreadSafe;
6566
import org.slf4j.Logger;
6667
import org.slf4j.LoggerFactory;
@@ -121,7 +122,7 @@ public class BasicLoadBalancingPolicy implements LoadBalancingPolicy {
121122
private volatile NodeDistanceEvaluator nodeDistanceEvaluator;
122123
private volatile String localDc;
123124
private volatile NodeSet liveNodes;
124-
private final String[] preferredRemoteDcs;
125+
private final List<String> preferredRemoteDcs;
125126

126127
public BasicLoadBalancingPolicy(@NonNull DriverContext context, @NonNull String profileName) {
127128
this.context = (InternalDriverContext) context;
@@ -143,7 +144,7 @@ public BasicLoadBalancingPolicy(@NonNull DriverContext context, @NonNull String
143144
profile.getStringList(
144145
DefaultDriverOption.LOAD_BALANCING_DC_FAILOVER_PREFERRED_REMOTE_DCS,
145146
new ArrayList<>())));
146-
preferredRemoteDcs = remoteDcs.toArray(new String[0]);
147+
preferredRemoteDcs = new ArrayList<>(remoteDcs);
147148
}
148149

149150
/**
@@ -335,7 +336,6 @@ protected Queue<Node> maybeAddDcFailover(@Nullable Request request, @NonNull Que
335336
}
336337
QueryPlan remote =
337338
new LazyQueryPlan() {
338-
339339
@Override
340340
protected Object[] computeNodes() {
341341
Object[] remoteNodes =
@@ -344,10 +344,22 @@ protected Object[] computeNodes() {
344344
.sorted(
345345
Comparator.comparingInt(
346346
dc -> {
347-
int i = Arrays.binarySearch(preferredRemoteDcs, dc);
347+
int i = preferredRemoteDcs.indexOf(dc);
348348
return i < 0 ? Integer.MAX_VALUE : i;
349349
}))
350-
.flatMap(dc -> liveNodes.dc(dc).stream().limit(maxNodesPerRemoteDc))
350+
.flatMap(
351+
dc -> {
352+
if (preferredRemoteDcs.isEmpty()) {
353+
return liveNodes.dc(dc).stream().limit(maxNodesPerRemoteDc);
354+
} else {
355+
final Object[] nodesPerDc =
356+
liveNodes.dc(dc).stream().limit(maxNodesPerRemoteDc).toArray();
357+
if (nodesPerDc.length > 0) {
358+
shuffleHead(nodesPerDc, nodesPerDc.length);
359+
}
360+
return Stream.of(nodesPerDc);
361+
}
362+
})
351363
.toArray();
352364

353365
int remoteNodesLength = remoteNodes.length;
@@ -356,7 +368,7 @@ protected Object[] computeNodes() {
356368
return EMPTY_NODES;
357369
}
358370

359-
if (preferredRemoteDcs.length == 0) {
371+
if (preferredRemoteDcs.isEmpty()) {
360372
shuffleHead(remoteNodes, remoteNodesLength);
361373
}
362374

‎core/src/test/java/com/datastax/oss/driver/internal/core/loadbalancing/BasicLoadBalancingPolicyPreferredRemoteDcsTest.java

Lines changed: 37 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -19,7 +19,6 @@
1919

2020
import static org.assertj.core.api.Assertions.assertThat;
2121
import static org.mockito.ArgumentMatchers.any;
22-
import static org.mockito.ArgumentMatchers.anyInt;
2322
import static org.mockito.ArgumentMatchers.eq;
2423
import static org.mockito.BDDMockito.then;
2524
import static org.mockito.Mockito.atLeast;
@@ -48,6 +47,8 @@ public class BasicLoadBalancingPolicyPreferredRemoteDcsTest
4847
@Mock protected DefaultNode node10;
4948
@Mock protected DefaultNode node11;
5049
@Mock protected DefaultNode node12;
50+
@Mock protected DefaultNode node13;
51+
@Mock protected DefaultNode node14;
5152

5253
@Override
5354
@Test
@@ -70,13 +71,16 @@ public void should_prioritize_and_shuffle_replicas() {
7071
when(tokenMap.getReplicas(KEYSPACE, ROUTING_KEY)).thenReturn(ImmutableSet.of(node3, node5));
7172

7273
assertThat(policy.newQueryPlan(request, session))
73-
.containsExactly(node3, node5, node1, node2, node4, node9, node10, node6, node7, node12);
74+
.containsExactly(
75+
node3, node5, node1, node2, node4, node9, node10, node6, node7, node12, node13);
7476
assertThat(policy.newQueryPlan(request, session))
75-
.containsExactly(node3, node5, node2, node4, node1, node9, node10, node6, node7, node12);
77+
.containsExactly(
78+
node3, node5, node2, node4, node1, node9, node10, node6, node7, node12, node13);
7679
assertThat(policy.newQueryPlan(request, session))
77-
.containsExactly(node3, node5, node4, node1, node2, node9, node10, node6, node7, node12);
80+
.containsExactly(
81+
node3, node5, node4, node1, node2, node9, node10, node6, node7, node12, node13);
7882

79-
verify(policy, times(3)).shuffleHead(any(), eq(2));
83+
verify(policy, times(12)).shuffleHead(any(), eq(2));
8084
// No power of two choices with only two replicas
8185
verify(session, never()).getPools();
8286
}
@@ -90,32 +94,45 @@ public void should_prioritize_single_replica() {
9094

9195
// node3 always first, round-robin on the rest
9296
assertThat(policy.newQueryPlan(request, session))
93-
.containsExactly(node3, node1, node2, node4, node5, node9, node10, node6, node7, node12);
97+
.containsExactly(
98+
node3, node1, node2, node4, node5, node9, node10, node6, node7, node12, node13);
9499
assertThat(policy.newQueryPlan(request, session))
95-
.containsExactly(node3, node2, node4, node5, node1, node9, node10, node6, node7, node12);
100+
.containsExactly(
101+
node3, node2, node4, node5, node1, node9, node10, node6, node7, node12, node13);
96102
assertThat(policy.newQueryPlan(request, session))
97-
.containsExactly(node3, node4, node5, node1, node2, node9, node10, node6, node7, node12);
103+
.containsExactly(
104+
node3, node4, node5, node1, node2, node9, node10, node6, node7, node12, node13);
98105
assertThat(policy.newQueryPlan(request, session))
99-
.containsExactly(node3, node5, node1, node2, node4, node9, node10, node6, node7, node12);
106+
.containsExactly(
107+
node3, node5, node1, node2, node4, node9, node10, node6, node7, node12, node13);
100108

101109
// Should not shuffle replicas since there is only one
102-
verify(policy, never()).shuffleHead(any(), anyInt());
110+
verify(policy, never()).shuffleHead(any(), eq(1));
111+
// But should shuffle remote nodes
112+
verify(policy, times(12)).shuffleHead(any(), eq(2));
103113
}
104114

105115
@Override
106116
protected void assertRoundRobinQueryPlans() {
107117
for (int i = 0; i < 3; i++) {
108118
assertThat(policy.newQueryPlan(request, session))
109-
.containsExactly(node1, node2, node3, node4, node5, node9, node10, node6, node7, node12);
119+
.containsExactly(
120+
node1, node2, node3, node4, node5, node9, node10, node6, node7, node12, node13);
110121
assertThat(policy.newQueryPlan(request, session))
111-
.containsExactly(node2, node3, node4, node5, node1, node9, node10, node6, node7, node12);
122+
.containsExactly(
123+
node2, node3, node4, node5, node1, node9, node10, node6, node7, node12, node13);
112124
assertThat(policy.newQueryPlan(request, session))
113-
.containsExactly(node3, node4, node5, node1, node2, node9, node10, node6, node7, node12);
125+
.containsExactly(
126+
node3, node4, node5, node1, node2, node9, node10, node6, node7, node12, node13);
114127
assertThat(policy.newQueryPlan(request, session))
115-
.containsExactly(node4, node5, node1, node2, node3, node9, node10, node6, node7, node12);
128+
.containsExactly(
129+
node4, node5, node1, node2, node3, node9, node10, node6, node7, node12, node13);
116130
assertThat(policy.newQueryPlan(request, session))
117-
.containsExactly(node5, node1, node2, node3, node4, node9, node10, node6, node7, node12);
131+
.containsExactly(
132+
node5, node1, node2, node3, node4, node9, node10, node6, node7, node12, node13);
118133
}
134+
135+
verify(policy, atLeast(15)).shuffleHead(any(), eq(2));
119136
}
120137

121138
@Override
@@ -129,6 +146,8 @@ protected BasicLoadBalancingPolicy createAndInitPolicy() {
129146
when(node10.getDatacenter()).thenReturn("dc3");
130147
when(node11.getDatacenter()).thenReturn("dc3");
131148
when(node12.getDatacenter()).thenReturn("dc4");
149+
when(node13.getDatacenter()).thenReturn("dc4");
150+
when(node14.getDatacenter()).thenReturn("dc4");
132151

133152
// Accept 2 nodes per remote DC
134153
when(defaultProfile.getInt(
@@ -166,12 +185,14 @@ protected void shuffleHead(Object[] currentNodes, int headLength) {
166185
.put(UUID.randomUUID(), node10)
167186
.put(UUID.randomUUID(), node11)
168187
.put(UUID.randomUUID(), node12)
188+
.put(UUID.randomUUID(), node13)
189+
.put(UUID.randomUUID(), node14)
169190
.build();
170191
policy.init(nodes, distanceReporter);
171192
assertThat(policy.getLiveNodes().dc("dc1")).containsExactly(node1, node2, node3, node4, node5);
172193
assertThat(policy.getLiveNodes().dc("dc2")).containsExactly(node6, node7); // only 2 allowed
173194
assertThat(policy.getLiveNodes().dc("dc3")).containsExactly(node9, node10); // only 2 allowed
174-
assertThat(policy.getLiveNodes().dc("dc4")).containsExactly(node12); // only 1 allowed
195+
assertThat(policy.getLiveNodes().dc("dc4")).containsExactly(node12, node13); // only 2 allowed
175196
return policy;
176197
}
177198
}

0 commit comments

Comments
 (0)