Skip to content

Commit d5acd00

Browse files
authored
Store computedServerSelectionTimeout so it can be reused. (#1312)
Needed for use with authentication JAVA-5211
1 parent 3ba95b0 commit d5acd00

File tree

8 files changed

+62
-20
lines changed

8 files changed

+62
-20
lines changed

driver-core/src/main/com/mongodb/internal/TimeoutContext.java

Lines changed: 44 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -40,6 +40,8 @@ public class TimeoutContext {
4040

4141
@Nullable
4242
private Timeout timeout;
43+
@Nullable
44+
private Timeout computedServerSelectionTimeout;
4345
private long minRoundTripTimeMS = 0;
4446

4547
public static MongoOperationTimeoutException createMongoTimeoutException() {
@@ -191,14 +193,17 @@ public long getWriteTimeoutMS() {
191193
return timeoutOrAlternative(0);
192194
}
193195

196+
public int getConnectTimeoutMs() {
197+
return (int) calculateMin(getTimeoutSettings().getConnectTimeoutMS());
198+
}
194199

195200
public void resetTimeout() {
196201
assertNotNull(timeout);
197202
timeout = calculateTimeout(timeoutSettings.getTimeoutMS());
198203
}
199204

200205
/**
201-
* Resest the timeout if this timeout context is being used by pool maintenance
206+
* Resets the timeout if this timeout context is being used by pool maintenance
202207
*/
203208
public void resetMaintenanceTimeout() {
204209
if (isMaintenanceContext && timeout != null && !timeout.isInfinite()) {
@@ -265,23 +270,53 @@ public static Timeout calculateTimeout(@Nullable final Long timeoutMS) {
265270
return null;
266271
}
267272

268-
public Timeout computedServerSelectionTimeout() {
269-
long ms = getTimeoutSettings().getServerSelectionTimeoutMS();
270-
Timeout serverSelectionTimeout = StartTime.now().timeoutAfterOrInfiniteIfNegative(ms, MILLISECONDS);
271-
return serverSelectionTimeout.orEarlier(timeout);
273+
/**
274+
* Returns the computed server selection timeout
275+
*
276+
* <p>Caches the computed server selection timeout if:
277+
* <ul>
278+
* <li>not in a maintenance context</li>
279+
* <li>there is a timeoutMS, so to keep the same legacy behavior.</li>
280+
* <li>the server selection timeout is less than the remaining overall timeout.</li>
281+
* </ul>
282+
*
283+
* @return the timeout context
284+
*/
285+
public Timeout computeServerSelectionTimeout() {
286+
Timeout serverSelectionTimeout = StartTime.now()
287+
.timeoutAfterOrInfiniteIfNegative(getTimeoutSettings().getServerSelectionTimeoutMS(), MILLISECONDS);
288+
289+
290+
if (isMaintenanceContext || !hasTimeoutMS()) {
291+
return serverSelectionTimeout;
292+
}
293+
294+
if (serverSelectionTimeout.orEarlier(timeout) == timeout) {
295+
return timeout;
296+
}
297+
298+
computedServerSelectionTimeout = serverSelectionTimeout;
299+
return computedServerSelectionTimeout;
300+
}
301+
302+
/**
303+
* Returns the timeout context to use for the handshake process
304+
*
305+
* @return a new timeout context with the cached computed server selection timeout if available or this
306+
*/
307+
public TimeoutContext withComputedServerSelectionTimeoutContext() {
308+
return computedServerSelectionTimeout == null
309+
? this : new TimeoutContext(false, timeoutSettings, computedServerSelectionTimeout);
272310
}
273311

274312
public Timeout startWaitQueueTimeout(final StartTime checkoutStart) {
275313
final long ms = getTimeoutSettings().getMaxWaitTimeMS();
276314
return checkoutStart.timeoutAfterOrInfiniteIfNegative(ms, MILLISECONDS);
277315
}
278316

279-
public int getConnectTimeoutMs() {
280-
return (int) getTimeoutSettings().getConnectTimeoutMS();
281-
}
282-
283317
@Nullable
284318
public Timeout getTimeout() {
285319
return timeout;
286320
}
321+
287322
}

driver-core/src/main/com/mongodb/internal/connection/BaseCluster.java

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -122,7 +122,7 @@ public ServerTuple selectServer(final ServerSelector serverSelector, final Opera
122122

123123
ServerSelector compositeServerSelector = getCompositeServerSelector(serverSelector);
124124
boolean selectionWaitingLogged = false;
125-
Timeout computedServerSelectionTimeout = operationContext.getTimeoutContext().computedServerSelectionTimeout();
125+
Timeout computedServerSelectionTimeout = operationContext.getTimeoutContext().computeServerSelectionTimeout();
126126
logServerSelectionStarted(clusterId, operationContext.getId(), serverSelector, description);
127127

128128
while (true) {
@@ -156,7 +156,7 @@ public ServerTuple selectServer(final ServerSelector serverSelector, final Opera
156156
public void selectServerAsync(final ServerSelector serverSelector, final OperationContext operationContext,
157157
final SingleResultCallback<ServerTuple> callback) {
158158
isTrue("open", !isClosed());
159-
Timeout computedServerSelectionTimeout = operationContext.getTimeoutContext().computedServerSelectionTimeout();
159+
Timeout computedServerSelectionTimeout = operationContext.getTimeoutContext().computeServerSelectionTimeout();
160160
ServerSelectionRequest request = new ServerSelectionRequest(
161161
serverSelector, getCompositeServerSelector(serverSelector), operationContext.getId(), computedServerSelectionTimeout,
162162
callback);

driver-core/src/main/com/mongodb/internal/connection/InternalStreamConnection.java

Lines changed: 9 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -196,10 +196,13 @@ public int getGeneration() {
196196
}
197197

198198
@Override
199-
public void open(final OperationContext operationContext) {
199+
public void open(final OperationContext originalOperationContext) {
200200
isTrue("Open already called", stream == null);
201201
stream = streamFactory.create(serverId.getAddress());
202202
try {
203+
OperationContext operationContext = originalOperationContext
204+
.withTimeoutContext(originalOperationContext.getTimeoutContext().withComputedServerSelectionTimeoutContext());
205+
203206
stream.open(operationContext);
204207

205208
InternalConnectionInitializationDescription initializationDescription = connectionInitializer.startHandshake(this, operationContext);
@@ -218,9 +221,13 @@ public void open(final OperationContext operationContext) {
218221
}
219222

220223
@Override
221-
public void openAsync(final OperationContext operationContext, final SingleResultCallback<Void> callback) {
224+
public void openAsync(final OperationContext originalOperationContext, final SingleResultCallback<Void> callback) {
222225
isTrue("Open already called", stream == null, callback);
223226
try {
227+
228+
OperationContext operationContext = originalOperationContext
229+
.withTimeoutContext(originalOperationContext.getTimeoutContext().withComputedServerSelectionTimeoutContext());
230+
224231
stream = streamFactory.create(serverId.getAddress());
225232
stream.openAsync(operationContext, new AsyncCompletionHandler<Void>() {
226233

driver-core/src/main/com/mongodb/internal/connection/LoadBalancedCluster.java

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -201,7 +201,7 @@ public ClusterClock getClock() {
201201
@Override
202202
public ServerTuple selectServer(final ServerSelector serverSelector, final OperationContext operationContext) {
203203
isTrue("open", !isClosed());
204-
Timeout computedServerSelectionTimeout = operationContext.getTimeoutContext().computedServerSelectionTimeout();
204+
Timeout computedServerSelectionTimeout = operationContext.getTimeoutContext().computeServerSelectionTimeout();
205205
waitForSrv(computedServerSelectionTimeout);
206206
if (srvRecordResolvedToMultipleHosts) {
207207
throw createResolvedToMultipleHostsException();
@@ -238,7 +238,7 @@ public void selectServerAsync(final ServerSelector serverSelector, final Operati
238238
callback.onResult(null, createShutdownException());
239239
return;
240240
}
241-
Timeout computedServerSelectionTimeout = operationContext.getTimeoutContext().computedServerSelectionTimeout();
241+
Timeout computedServerSelectionTimeout = operationContext.getTimeoutContext().computeServerSelectionTimeout();
242242
ServerSelectionRequest serverSelectionRequest = new ServerSelectionRequest(operationContext.getId(), serverSelector,
243243
computedServerSelectionTimeout, callback);
244244
if (initializationCompleted) {

driver-core/src/test/unit/com/mongodb/internal/connection/AbstractServerDiscoveryAndMonitoringTest.java

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -81,7 +81,7 @@ protected void applyResponse(final BsonArray response) {
8181
}
8282

8383
protected void applyApplicationError(final BsonDocument applicationError) {
84-
Timeout serverSelectionTimeout = OPERATION_CONTEXT.getTimeoutContext().computedServerSelectionTimeout();
84+
Timeout serverSelectionTimeout = OPERATION_CONTEXT.getTimeoutContext().computeServerSelectionTimeout();
8585
ServerAddress serverAddress = new ServerAddress(applicationError.getString("address").getValue());
8686
int errorGeneration = applicationError.getNumber("generation",
8787
new BsonInt32(((DefaultServer) getCluster().getServer(serverAddress, serverSelectionTimeout)).getConnectionPool().getGeneration())).intValue();

driver-core/src/test/unit/com/mongodb/internal/connection/MultiServerClusterSpecification.groovy

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -95,7 +95,7 @@ class MultiServerClusterSpecification extends Specification {
9595
cluster.close()
9696

9797
when:
98-
cluster.getServer(firstServer, OPERATION_CONTEXT.getTimeoutContext().computedServerSelectionTimeout())
98+
cluster.getServer(firstServer, OPERATION_CONTEXT.getTimeoutContext().computeServerSelectionTimeout())
9999

100100
then:
101101
thrown(IllegalStateException)

driver-core/src/test/unit/com/mongodb/internal/connection/ServerDiscoveryAndMonitoringTest.java

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -122,7 +122,7 @@ private void assertServer(final String serverName, final BsonDocument expectedSe
122122

123123
if (expectedServerDescriptionDocument.isDocument("pool")) {
124124
int expectedGeneration = expectedServerDescriptionDocument.getDocument("pool").getNumber("generation").intValue();
125-
Timeout serverSelectionTimeout = OPERATION_CONTEXT.getTimeoutContext().computedServerSelectionTimeout();
125+
Timeout serverSelectionTimeout = OPERATION_CONTEXT.getTimeoutContext().computeServerSelectionTimeout();
126126
DefaultServer server = (DefaultServer) getCluster().getServer(new ServerAddress(serverName), serverSelectionTimeout);
127127
assertEquals(expectedGeneration, server.getConnectionPool().getGeneration());
128128
}

driver-core/src/test/unit/com/mongodb/internal/connection/SingleServerClusterSpecification.groovy

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -78,7 +78,7 @@ class SingleServerClusterSpecification extends Specification {
7878

7979
then:
8080
cluster.getServer(firstServer,
81-
OPERATION_CONTEXT.getTimeoutContext().computedServerSelectionTimeout()) == factory.getServer(firstServer)
81+
OPERATION_CONTEXT.getTimeoutContext().computeServerSelectionTimeout()) == factory.getServer(firstServer)
8282

8383
cleanup:
8484
cluster?.close()
@@ -92,7 +92,7 @@ class SingleServerClusterSpecification extends Specification {
9292
cluster.close()
9393

9494
when:
95-
cluster.getServer(firstServer, OPERATION_CONTEXT.getTimeoutContext().computedServerSelectionTimeout())
95+
cluster.getServer(firstServer, OPERATION_CONTEXT.getTimeoutContext().computeServerSelectionTimeout())
9696

9797
then:
9898
thrown(IllegalStateException)

0 commit comments

Comments
 (0)