Skip to content

Commit cfd111d

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

File tree

2 files changed

+38
-9
lines changed

2 files changed

+38
-9
lines changed

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

Lines changed: 29 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,34 @@ 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+
public TimeoutContext withComputedServerSelectionTimeoutContext() {
275+
if (!hasTimeoutMS()) {
276+
// Without timeoutMS keep the legacy timeout behavior
277+
computedServerSelectionTimeout = null;
278+
return this;
279+
} else if (computedServerSelectionTimeout == null) {
280+
// Reusing a pooled connection
281+
return this;
282+
}
283+
284+
TimeoutContext timeoutContext = Objects.equals(computedServerSelectionTimeout, timeout)
285+
? this : new TimeoutContext(false, timeoutSettings, computedServerSelectionTimeout);
286+
computedServerSelectionTimeout = null;
287+
return timeoutContext;
264288
}
265289

266290
public Timeout startWaitQueueTimeout(final StartTime checkoutStart) {
267291
final long ms = getTimeoutSettings().getMaxWaitTimeMS();
268292
return checkoutStart.timeoutAfterOrInfiniteIfNegative(ms, MILLISECONDS);
269293
}
270294

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

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)