Skip to content

Commit d42c43d

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

File tree

2 files changed

+44
-10
lines changed

2 files changed

+44
-10
lines changed

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

Lines changed: 35 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()) {
@@ -266,22 +271,44 @@ public static Timeout calculateTimeout(@Nullable final Long timeoutMS) {
266271
}
267272

268273
public Timeout computedServerSelectionTimeout() {
269-
long ms = getTimeoutSettings().getServerSelectionTimeoutMS();
270-
Timeout serverSelectionTimeout = StartTime.now().timeoutAfterOrInfiniteIfNegative(ms, MILLISECONDS);
271-
return serverSelectionTimeout.orEarlier(timeout);
274+
if (computedServerSelectionTimeout == null) {
275+
computedServerSelectionTimeout = StartTime.now()
276+
.timeoutAfterOrInfiniteIfNegative(getTimeoutSettings().getServerSelectionTimeoutMS(), MILLISECONDS)
277+
.orEarlier(timeout);
278+
}
279+
return computedServerSelectionTimeout;
280+
}
281+
282+
/**
283+
* Returns the timeout context to use for the handshake process
284+
*
285+
* <p>Uses the current context if:
286+
* <ul>
287+
* <li>If has no timeoutMS, keep the same legacy behavior.</li>
288+
* <li>If in a MaintenanceContext</li>
289+
* <li>If there is no computedServerSelectionTimeout.q</li>
290+
* </ul>
291+
*
292+
* @return the timeout context
293+
*/
294+
public TimeoutContext withComputedServerSelectionTimeoutContext() {
295+
TimeoutContext timeoutContext = this;
296+
if (!hasTimeoutMS() && !isMaintenanceContext && computedServerSelectionTimeout != null) {
297+
timeoutContext = Objects.equals(computedServerSelectionTimeout, timeout)
298+
? this : new TimeoutContext(false, timeoutSettings, computedServerSelectionTimeout);
299+
}
300+
computedServerSelectionTimeout = null;
301+
return timeoutContext;
272302
}
273303

274304
public Timeout startWaitQueueTimeout(final StartTime checkoutStart) {
275305
final long ms = getTimeoutSettings().getMaxWaitTimeMS();
276306
return checkoutStart.timeoutAfterOrInfiniteIfNegative(ms, MILLISECONDS);
277307
}
278308

279-
public int getConnectTimeoutMs() {
280-
return (int) getTimeoutSettings().getConnectTimeoutMS();
281-
}
282-
283309
@Nullable
284310
public Timeout getTimeout() {
285311
return timeout;
286312
}
313+
287314
}

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)