Skip to content

Commit fdaf71f

Browse files
committed
Store computedServerSelectionTimeout so it can be reused.
Needed for use with authentication JAVA-5211
1 parent 3ba95b0 commit fdaf71f

File tree

2 files changed

+52
-10
lines changed

2 files changed

+52
-10
lines changed

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

Lines changed: 43 additions & 8 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

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+
*/
268285
public Timeout computedServerSelectionTimeout() {
269-
long ms = getTimeoutSettings().getServerSelectionTimeoutMS();
270-
Timeout serverSelectionTimeout = StartTime.now().timeoutAfterOrInfiniteIfNegative(ms, MILLISECONDS);
271-
return serverSelectionTimeout.orEarlier(timeout);
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/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

0 commit comments

Comments
 (0)