Skip to content

Commit 1342eac

Browse files
BouncheckBouncheck
committed
Extend ChannelFactory methods with shardId params
With advanced shard awareness target shard is a necessary parameter for connecting. Extends non public methods with new parameters and provides one additional public method. Passing `null` as shard id or ShardingInfo will lead to previous non-shard aware behavior.
1 parent 6c10f7f commit 1342eac

File tree

5 files changed

+115
-17
lines changed

5 files changed

+115
-17
lines changed

‎core/src/main/java/com/datastax/oss/driver/internal/core/channel/ChannelFactory.java

+50-4
Original file line numberDiff line numberDiff line change
@@ -29,8 +29,10 @@
2929
import com.datastax.oss.driver.api.core.config.DefaultDriverOption;
3030
import com.datastax.oss.driver.api.core.config.DriverConfig;
3131
import com.datastax.oss.driver.api.core.config.DriverExecutionProfile;
32+
import com.datastax.oss.driver.api.core.context.DriverContext;
3233
import com.datastax.oss.driver.api.core.metadata.EndPoint;
3334
import com.datastax.oss.driver.api.core.metadata.Node;
35+
import com.datastax.oss.driver.api.core.metadata.NodeShardingInfo;
3436
import com.datastax.oss.driver.api.core.metrics.DefaultNodeMetric;
3537
import com.datastax.oss.driver.api.core.metrics.DefaultSessionMetric;
3638
import com.datastax.oss.driver.internal.core.config.typesafe.TypesafeDriverConfig;
@@ -157,12 +159,27 @@ public CompletionStage<DriverChannel> connect(Node node, DriverChannelOptions op
157159
} else {
158160
nodeMetricUpdater = NoopNodeMetricUpdater.INSTANCE;
159161
}
160-
return connect(node.getEndPoint(), options, nodeMetricUpdater);
162+
return connect(node.getEndPoint(), null, null, options, nodeMetricUpdater);
163+
}
164+
165+
public CompletionStage<DriverChannel> connect(
166+
Node node, Integer shardId, DriverChannelOptions options) {
167+
NodeMetricUpdater nodeMetricUpdater;
168+
if (node instanceof DefaultNode) {
169+
nodeMetricUpdater = ((DefaultNode) node).getMetricUpdater();
170+
} else {
171+
nodeMetricUpdater = NoopNodeMetricUpdater.INSTANCE;
172+
}
173+
return connect(node.getEndPoint(), node.getShardingInfo(), shardId, options, nodeMetricUpdater);
161174
}
162175

163176
@VisibleForTesting
164177
CompletionStage<DriverChannel> connect(
165-
EndPoint endPoint, DriverChannelOptions options, NodeMetricUpdater nodeMetricUpdater) {
178+
EndPoint endPoint,
179+
NodeShardingInfo shardingInfo,
180+
Integer shardId,
181+
DriverChannelOptions options,
182+
NodeMetricUpdater nodeMetricUpdater) {
166183
CompletableFuture<DriverChannel> resultFuture = new CompletableFuture<>();
167184

168185
ProtocolVersion currentVersion;
@@ -178,6 +195,8 @@ CompletionStage<DriverChannel> connect(
178195

179196
connect(
180197
endPoint,
198+
shardingInfo,
199+
shardId,
181200
options,
182201
nodeMetricUpdater,
183202
currentVersion,
@@ -189,6 +208,8 @@ CompletionStage<DriverChannel> connect(
189208

190209
private void connect(
191210
EndPoint endPoint,
211+
NodeShardingInfo shardingInfo,
212+
Integer shardId,
192213
DriverChannelOptions options,
193214
NodeMetricUpdater nodeMetricUpdater,
194215
ProtocolVersion currentVersion,
@@ -208,7 +229,20 @@ private void connect(
208229

209230
nettyOptions.afterBootstrapInitialized(bootstrap);
210231

211-
ChannelFuture connectFuture = bootstrap.connect(endPoint.resolve());
232+
ChannelFuture connectFuture;
233+
if (shardId == null || shardingInfo == null) {
234+
if (shardId != null) {
235+
LOG.debug(
236+
"Requested connection to shard {} but shardingInfo is currently missing for Node at endpoint {}. Falling back to arbitrary local port.",
237+
shardId,
238+
endPoint);
239+
}
240+
connectFuture = bootstrap.connect(endPoint.resolve());
241+
} else {
242+
int localPort =
243+
PortAllocator.getNextAvailablePort(shardingInfo.getShardsCount(), shardId, context);
244+
connectFuture = bootstrap.connect(endPoint.resolve(), new InetSocketAddress(localPort));
245+
}
212246

213247
connectFuture.addListener(
214248
cf -> {
@@ -257,6 +291,8 @@ private void connect(
257291
downgraded.get());
258292
connect(
259293
endPoint,
294+
shardingInfo,
295+
shardId,
260296
options,
261297
nodeMetricUpdater,
262298
downgraded.get(),
@@ -399,7 +435,17 @@ protected void initChannel(Channel channel) {
399435
static class PortAllocator {
400436
private static final AtomicInteger lastPort = new AtomicInteger(-1);
401437

402-
public static int getNextAvailablePort(int shardCount, int shardId, int lowPort, int highPort) {
438+
public static int getNextAvailablePort(int shardCount, int shardId, DriverContext context) {
439+
int lowPort =
440+
context
441+
.getConfig()
442+
.getDefaultProfile()
443+
.getInt(DefaultDriverOption.ADVANCED_SHARD_AWARENESS_PORT_LOW);
444+
int highPort =
445+
context
446+
.getConfig()
447+
.getDefaultProfile()
448+
.getInt(DefaultDriverOption.ADVANCED_SHARD_AWARENESS_PORT_HIGH);
403449
int lastPortValue, foundPort = -1;
404450
do {
405451
lastPortValue = lastPort.get();

‎core/src/test/java/com/datastax/oss/driver/internal/core/channel/ChannelFactoryAvailableIdsTest.java

+5-1
Original file line numberDiff line numberDiff line change
@@ -61,7 +61,11 @@ public void should_report_available_ids() {
6161
// When
6262
CompletionStage<DriverChannel> channelFuture =
6363
factory.connect(
64-
SERVER_ADDRESS, DriverChannelOptions.builder().build(), NoopNodeMetricUpdater.INSTANCE);
64+
SERVER_ADDRESS,
65+
null,
66+
null,
67+
DriverChannelOptions.builder().build(),
68+
NoopNodeMetricUpdater.INSTANCE);
6569
completeSimpleChannelInit();
6670

6771
// Then

‎core/src/test/java/com/datastax/oss/driver/internal/core/channel/ChannelFactoryClusterNameTest.java

+20-4
Original file line numberDiff line numberDiff line change
@@ -41,7 +41,11 @@ public void should_set_cluster_name_from_first_connection() {
4141
// When
4242
CompletionStage<DriverChannel> channelFuture =
4343
factory.connect(
44-
SERVER_ADDRESS, DriverChannelOptions.DEFAULT, NoopNodeMetricUpdater.INSTANCE);
44+
SERVER_ADDRESS,
45+
null,
46+
null,
47+
DriverChannelOptions.DEFAULT,
48+
NoopNodeMetricUpdater.INSTANCE);
4549

4650
writeInboundFrame(
4751
readOutboundFrame(), TestResponses.supportedResponse("mock_key", "mock_value"));
@@ -63,7 +67,11 @@ public void should_check_cluster_name_for_next_connections() throws Throwable {
6367
// When
6468
CompletionStage<DriverChannel> channelFuture =
6569
factory.connect(
66-
SERVER_ADDRESS, DriverChannelOptions.DEFAULT, NoopNodeMetricUpdater.INSTANCE);
70+
SERVER_ADDRESS,
71+
null,
72+
null,
73+
DriverChannelOptions.DEFAULT,
74+
NoopNodeMetricUpdater.INSTANCE);
6775
// open a first connection that will define the cluster name
6876
writeInboundFrame(
6977
readOutboundFrame(), TestResponses.supportedResponse("mock_key", "mock_value"));
@@ -73,7 +81,11 @@ public void should_check_cluster_name_for_next_connections() throws Throwable {
7381
// open a second connection that returns the same cluster name
7482
channelFuture =
7583
factory.connect(
76-
SERVER_ADDRESS, DriverChannelOptions.DEFAULT, NoopNodeMetricUpdater.INSTANCE);
84+
SERVER_ADDRESS,
85+
null,
86+
null,
87+
DriverChannelOptions.DEFAULT,
88+
NoopNodeMetricUpdater.INSTANCE);
7789
writeInboundFrame(readOutboundFrame(), new Ready());
7890
writeInboundFrame(readOutboundFrame(), TestResponses.clusterNameResponse("mockClusterName"));
7991

@@ -84,7 +96,11 @@ public void should_check_cluster_name_for_next_connections() throws Throwable {
8496
// open a third connection that returns a different cluster name
8597
channelFuture =
8698
factory.connect(
87-
SERVER_ADDRESS, DriverChannelOptions.DEFAULT, NoopNodeMetricUpdater.INSTANCE);
99+
SERVER_ADDRESS,
100+
null,
101+
null,
102+
DriverChannelOptions.DEFAULT,
103+
NoopNodeMetricUpdater.INSTANCE);
88104
writeInboundFrame(readOutboundFrame(), new Ready());
89105
writeInboundFrame(readOutboundFrame(), TestResponses.clusterNameResponse("wrongClusterName"));
90106

‎core/src/test/java/com/datastax/oss/driver/internal/core/channel/ChannelFactoryProtocolNegotiationTest.java

+30-6
Original file line numberDiff line numberDiff line change
@@ -50,7 +50,11 @@ public void should_succeed_if_version_specified_and_supported_by_server() {
5050
// When
5151
CompletionStage<DriverChannel> channelFuture =
5252
factory.connect(
53-
SERVER_ADDRESS, DriverChannelOptions.DEFAULT, NoopNodeMetricUpdater.INSTANCE);
53+
SERVER_ADDRESS,
54+
null,
55+
null,
56+
DriverChannelOptions.DEFAULT,
57+
NoopNodeMetricUpdater.INSTANCE);
5458

5559
completeSimpleChannelInit();
5660

@@ -72,7 +76,11 @@ public void should_fail_if_version_specified_and_not_supported_by_server(int err
7276
// When
7377
CompletionStage<DriverChannel> channelFuture =
7478
factory.connect(
75-
SERVER_ADDRESS, DriverChannelOptions.DEFAULT, NoopNodeMetricUpdater.INSTANCE);
79+
SERVER_ADDRESS,
80+
null,
81+
null,
82+
DriverChannelOptions.DEFAULT,
83+
NoopNodeMetricUpdater.INSTANCE);
7684

7785
Frame requestFrame = readOutboundFrame();
7886
assertThat(requestFrame.message).isInstanceOf(Options.class);
@@ -107,7 +115,11 @@ public void should_fail_if_version_specified_and_considered_beta_by_server() {
107115
// When
108116
CompletionStage<DriverChannel> channelFuture =
109117
factory.connect(
110-
SERVER_ADDRESS, DriverChannelOptions.DEFAULT, NoopNodeMetricUpdater.INSTANCE);
118+
SERVER_ADDRESS,
119+
null,
120+
null,
121+
DriverChannelOptions.DEFAULT,
122+
NoopNodeMetricUpdater.INSTANCE);
111123

112124
Frame requestFrame = readOutboundFrame();
113125
assertThat(requestFrame.message).isInstanceOf(Options.class);
@@ -144,7 +156,11 @@ public void should_succeed_if_version_not_specified_and_server_supports_latest_s
144156
// When
145157
CompletionStage<DriverChannel> channelFuture =
146158
factory.connect(
147-
SERVER_ADDRESS, DriverChannelOptions.DEFAULT, NoopNodeMetricUpdater.INSTANCE);
159+
SERVER_ADDRESS,
160+
null,
161+
null,
162+
DriverChannelOptions.DEFAULT,
163+
NoopNodeMetricUpdater.INSTANCE);
148164

149165
Frame requestFrame = readOutboundFrame();
150166
assertThat(requestFrame.message).isInstanceOf(Options.class);
@@ -176,7 +192,11 @@ public void should_negotiate_if_version_not_specified_and_server_supports_legacy
176192
// When
177193
CompletionStage<DriverChannel> channelFuture =
178194
factory.connect(
179-
SERVER_ADDRESS, DriverChannelOptions.DEFAULT, NoopNodeMetricUpdater.INSTANCE);
195+
SERVER_ADDRESS,
196+
null,
197+
null,
198+
DriverChannelOptions.DEFAULT,
199+
NoopNodeMetricUpdater.INSTANCE);
180200

181201
Frame requestFrame = readOutboundFrame();
182202
assertThat(requestFrame.message).isInstanceOf(Options.class);
@@ -219,7 +239,11 @@ public void should_fail_if_negotiation_finds_no_matching_version(int errorCode)
219239
// When
220240
CompletionStage<DriverChannel> channelFuture =
221241
factory.connect(
222-
SERVER_ADDRESS, DriverChannelOptions.DEFAULT, NoopNodeMetricUpdater.INSTANCE);
242+
SERVER_ADDRESS,
243+
null,
244+
null,
245+
DriverChannelOptions.DEFAULT,
246+
NoopNodeMetricUpdater.INSTANCE);
223247

224248
Frame requestFrame = readOutboundFrame();
225249
assertThat(requestFrame.message).isInstanceOf(Options.class);

‎core/src/test/java/com/datastax/oss/driver/internal/core/channel/ChannelFactorySupportedOptionsTest.java

+10-2
Original file line numberDiff line numberDiff line change
@@ -41,7 +41,11 @@ public void should_query_supported_options_on_first_channel() throws Throwable {
4141
// When
4242
CompletionStage<DriverChannel> channelFuture1 =
4343
factory.connect(
44-
SERVER_ADDRESS, DriverChannelOptions.DEFAULT, NoopNodeMetricUpdater.INSTANCE);
44+
SERVER_ADDRESS,
45+
null,
46+
null,
47+
DriverChannelOptions.DEFAULT,
48+
NoopNodeMetricUpdater.INSTANCE);
4549
writeInboundFrame(
4650
readOutboundFrame(), TestResponses.supportedResponse("mock_key", "mock_value"));
4751
writeInboundFrame(readOutboundFrame(), new Ready());
@@ -56,7 +60,11 @@ public void should_query_supported_options_on_first_channel() throws Throwable {
5660
// When
5761
CompletionStage<DriverChannel> channelFuture2 =
5862
factory.connect(
59-
SERVER_ADDRESS, DriverChannelOptions.DEFAULT, NoopNodeMetricUpdater.INSTANCE);
63+
SERVER_ADDRESS,
64+
null,
65+
null,
66+
DriverChannelOptions.DEFAULT,
67+
NoopNodeMetricUpdater.INSTANCE);
6068
writeInboundFrame(readOutboundFrame(), new Ready());
6169
writeInboundFrame(readOutboundFrame(), TestResponses.clusterNameResponse("mockClusterName"));
6270

0 commit comments

Comments
 (0)