Skip to content

Commit d2d1967

Browse files
committed
Store computedServerSelectionTimeout so it can be reused.
Needed for use with authentication JAVA-5211
1 parent 88e93ea commit d2d1967

File tree

2 files changed

+44
-9
lines changed

2 files changed

+44
-9
lines changed

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

Lines changed: 35 additions & 7 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() {
@@ -183,14 +185,17 @@ public long getWriteTimeoutMS() {
183185
return timeoutOrAlternative(0);
184186
}
185187

188+
public int getConnectTimeoutMs() {
189+
return (int) calculateMin(getTimeoutSettings().getConnectTimeoutMS());
190+
}
186191

187192
public void resetTimeout() {
188193
assertNotNull(timeout);
189194
timeout = calculateTimeout(timeoutSettings.getTimeoutMS());
190195
}
191196

192197
/**
193-
* Resest the timeout if this timeout context is being used by pool maintenance
198+
* Resets the timeout if this timeout context is being used by pool maintenance
194199
*/
195200
public void resetMaintenanceTimeout() {
196201
if (isMaintenanceContext && timeout != null && !timeout.isInfinite()) {
@@ -258,17 +263,40 @@ public static Timeout calculateTimeout(@Nullable final Long timeoutMS) {
258263
}
259264

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

266296
public Timeout startWaitQueueTimeout(final StartTime checkoutStart) {
267297
final long ms = getTimeoutSettings().getMaxWaitTimeMS();
268298
return checkoutStart.timeoutAfterOrInfiniteIfNegative(ms, MILLISECONDS);
269299
}
270300

271-
public int getConnectTimeoutMs() {
272-
return (int) getTimeoutSettings().getConnectTimeoutMS();
273-
}
301+
274302
}

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)