diff --git a/driver-core/src/main/com/mongodb/internal/TimeoutContext.java b/driver-core/src/main/com/mongodb/internal/TimeoutContext.java index ee883530a3c..7fb706b767e 100644 --- a/driver-core/src/main/com/mongodb/internal/TimeoutContext.java +++ b/driver-core/src/main/com/mongodb/internal/TimeoutContext.java @@ -40,6 +40,8 @@ public class TimeoutContext { @Nullable private Timeout timeout; + @Nullable + private Timeout computedServerSelectionTimeout; private long minRoundTripTimeMS = 0; public static MongoOperationTimeoutException createMongoTimeoutException() { @@ -191,6 +193,9 @@ public long getWriteTimeoutMS() { return timeoutOrAlternative(0); } + public int getConnectTimeoutMs() { + return (int) calculateMin(getTimeoutSettings().getConnectTimeoutMS()); + } public void resetTimeout() { assertNotNull(timeout); @@ -198,7 +203,7 @@ public void resetTimeout() { } /** - * Resest the timeout if this timeout context is being used by pool maintenance + * Resets the timeout if this timeout context is being used by pool maintenance */ public void resetMaintenanceTimeout() { if (isMaintenanceContext && timeout != null && !timeout.isInfinite()) { @@ -265,10 +270,43 @@ public static Timeout calculateTimeout(@Nullable final Long timeoutMS) { return null; } - public Timeout computedServerSelectionTimeout() { - long ms = getTimeoutSettings().getServerSelectionTimeoutMS(); - Timeout serverSelectionTimeout = StartTime.now().timeoutAfterOrInfiniteIfNegative(ms, MILLISECONDS); - return serverSelectionTimeout.orEarlier(timeout); + /** + * Returns the computed server selection timeout + * + *

Caches the computed server selection timeout if: + *

+ * + * @return the timeout context + */ + public Timeout computeServerSelectionTimeout() { + Timeout serverSelectionTimeout = StartTime.now() + .timeoutAfterOrInfiniteIfNegative(getTimeoutSettings().getServerSelectionTimeoutMS(), MILLISECONDS); + + + if (isMaintenanceContext || !hasTimeoutMS()) { + return serverSelectionTimeout; + } + + if (serverSelectionTimeout.orEarlier(timeout) == timeout) { + return timeout; + } + + computedServerSelectionTimeout = serverSelectionTimeout; + return computedServerSelectionTimeout; + } + + /** + * Returns the timeout context to use for the handshake process + * + * @return a new timeout context with the cached computed server selection timeout if available or this + */ + public TimeoutContext withComputedServerSelectionTimeoutContext() { + return computedServerSelectionTimeout == null + ? this : new TimeoutContext(false, timeoutSettings, computedServerSelectionTimeout); } public Timeout startWaitQueueTimeout(final StartTime checkoutStart) { @@ -276,12 +314,9 @@ public Timeout startWaitQueueTimeout(final StartTime checkoutStart) { return checkoutStart.timeoutAfterOrInfiniteIfNegative(ms, MILLISECONDS); } - public int getConnectTimeoutMs() { - return (int) getTimeoutSettings().getConnectTimeoutMS(); - } - @Nullable public Timeout getTimeout() { return timeout; } + } diff --git a/driver-core/src/main/com/mongodb/internal/connection/BaseCluster.java b/driver-core/src/main/com/mongodb/internal/connection/BaseCluster.java index a32b3be4d8f..f8298c3aaa5 100644 --- a/driver-core/src/main/com/mongodb/internal/connection/BaseCluster.java +++ b/driver-core/src/main/com/mongodb/internal/connection/BaseCluster.java @@ -122,7 +122,7 @@ public ServerTuple selectServer(final ServerSelector serverSelector, final Opera ServerSelector compositeServerSelector = getCompositeServerSelector(serverSelector); boolean selectionWaitingLogged = false; - Timeout computedServerSelectionTimeout = operationContext.getTimeoutContext().computedServerSelectionTimeout(); + Timeout computedServerSelectionTimeout = operationContext.getTimeoutContext().computeServerSelectionTimeout(); logServerSelectionStarted(clusterId, operationContext.getId(), serverSelector, description); while (true) { @@ -156,7 +156,7 @@ public ServerTuple selectServer(final ServerSelector serverSelector, final Opera public void selectServerAsync(final ServerSelector serverSelector, final OperationContext operationContext, final SingleResultCallback callback) { isTrue("open", !isClosed()); - Timeout computedServerSelectionTimeout = operationContext.getTimeoutContext().computedServerSelectionTimeout(); + Timeout computedServerSelectionTimeout = operationContext.getTimeoutContext().computeServerSelectionTimeout(); ServerSelectionRequest request = new ServerSelectionRequest( serverSelector, getCompositeServerSelector(serverSelector), operationContext.getId(), computedServerSelectionTimeout, callback); diff --git a/driver-core/src/main/com/mongodb/internal/connection/InternalStreamConnection.java b/driver-core/src/main/com/mongodb/internal/connection/InternalStreamConnection.java index a179b532fe6..bdc69e96dd4 100644 --- a/driver-core/src/main/com/mongodb/internal/connection/InternalStreamConnection.java +++ b/driver-core/src/main/com/mongodb/internal/connection/InternalStreamConnection.java @@ -196,10 +196,13 @@ public int getGeneration() { } @Override - public void open(final OperationContext operationContext) { + public void open(final OperationContext originalOperationContext) { isTrue("Open already called", stream == null); stream = streamFactory.create(serverId.getAddress()); try { + OperationContext operationContext = originalOperationContext + .withTimeoutContext(originalOperationContext.getTimeoutContext().withComputedServerSelectionTimeoutContext()); + stream.open(operationContext); InternalConnectionInitializationDescription initializationDescription = connectionInitializer.startHandshake(this, operationContext); @@ -218,9 +221,13 @@ public void open(final OperationContext operationContext) { } @Override - public void openAsync(final OperationContext operationContext, final SingleResultCallback callback) { + public void openAsync(final OperationContext originalOperationContext, final SingleResultCallback callback) { isTrue("Open already called", stream == null, callback); try { + + OperationContext operationContext = originalOperationContext + .withTimeoutContext(originalOperationContext.getTimeoutContext().withComputedServerSelectionTimeoutContext()); + stream = streamFactory.create(serverId.getAddress()); stream.openAsync(operationContext, new AsyncCompletionHandler() { diff --git a/driver-core/src/main/com/mongodb/internal/connection/LoadBalancedCluster.java b/driver-core/src/main/com/mongodb/internal/connection/LoadBalancedCluster.java index bd0facb339c..96651ea2ee0 100644 --- a/driver-core/src/main/com/mongodb/internal/connection/LoadBalancedCluster.java +++ b/driver-core/src/main/com/mongodb/internal/connection/LoadBalancedCluster.java @@ -201,7 +201,7 @@ public ClusterClock getClock() { @Override public ServerTuple selectServer(final ServerSelector serverSelector, final OperationContext operationContext) { isTrue("open", !isClosed()); - Timeout computedServerSelectionTimeout = operationContext.getTimeoutContext().computedServerSelectionTimeout(); + Timeout computedServerSelectionTimeout = operationContext.getTimeoutContext().computeServerSelectionTimeout(); waitForSrv(computedServerSelectionTimeout); if (srvRecordResolvedToMultipleHosts) { throw createResolvedToMultipleHostsException(); @@ -238,7 +238,7 @@ public void selectServerAsync(final ServerSelector serverSelector, final Operati callback.onResult(null, createShutdownException()); return; } - Timeout computedServerSelectionTimeout = operationContext.getTimeoutContext().computedServerSelectionTimeout(); + Timeout computedServerSelectionTimeout = operationContext.getTimeoutContext().computeServerSelectionTimeout(); ServerSelectionRequest serverSelectionRequest = new ServerSelectionRequest(operationContext.getId(), serverSelector, computedServerSelectionTimeout, callback); if (initializationCompleted) { diff --git a/driver-core/src/test/unit/com/mongodb/internal/connection/AbstractServerDiscoveryAndMonitoringTest.java b/driver-core/src/test/unit/com/mongodb/internal/connection/AbstractServerDiscoveryAndMonitoringTest.java index b94293267b8..1ccfb7d4839 100644 --- a/driver-core/src/test/unit/com/mongodb/internal/connection/AbstractServerDiscoveryAndMonitoringTest.java +++ b/driver-core/src/test/unit/com/mongodb/internal/connection/AbstractServerDiscoveryAndMonitoringTest.java @@ -81,7 +81,7 @@ protected void applyResponse(final BsonArray response) { } protected void applyApplicationError(final BsonDocument applicationError) { - Timeout serverSelectionTimeout = OPERATION_CONTEXT.getTimeoutContext().computedServerSelectionTimeout(); + Timeout serverSelectionTimeout = OPERATION_CONTEXT.getTimeoutContext().computeServerSelectionTimeout(); ServerAddress serverAddress = new ServerAddress(applicationError.getString("address").getValue()); int errorGeneration = applicationError.getNumber("generation", new BsonInt32(((DefaultServer) getCluster().getServer(serverAddress, serverSelectionTimeout)).getConnectionPool().getGeneration())).intValue(); diff --git a/driver-core/src/test/unit/com/mongodb/internal/connection/MultiServerClusterSpecification.groovy b/driver-core/src/test/unit/com/mongodb/internal/connection/MultiServerClusterSpecification.groovy index 2840de4729c..9a6cbb2267f 100644 --- a/driver-core/src/test/unit/com/mongodb/internal/connection/MultiServerClusterSpecification.groovy +++ b/driver-core/src/test/unit/com/mongodb/internal/connection/MultiServerClusterSpecification.groovy @@ -95,7 +95,7 @@ class MultiServerClusterSpecification extends Specification { cluster.close() when: - cluster.getServer(firstServer, OPERATION_CONTEXT.getTimeoutContext().computedServerSelectionTimeout()) + cluster.getServer(firstServer, OPERATION_CONTEXT.getTimeoutContext().computeServerSelectionTimeout()) then: thrown(IllegalStateException) diff --git a/driver-core/src/test/unit/com/mongodb/internal/connection/ServerDiscoveryAndMonitoringTest.java b/driver-core/src/test/unit/com/mongodb/internal/connection/ServerDiscoveryAndMonitoringTest.java index b8687ced466..4ccd6c171f9 100644 --- a/driver-core/src/test/unit/com/mongodb/internal/connection/ServerDiscoveryAndMonitoringTest.java +++ b/driver-core/src/test/unit/com/mongodb/internal/connection/ServerDiscoveryAndMonitoringTest.java @@ -122,7 +122,7 @@ private void assertServer(final String serverName, final BsonDocument expectedSe if (expectedServerDescriptionDocument.isDocument("pool")) { int expectedGeneration = expectedServerDescriptionDocument.getDocument("pool").getNumber("generation").intValue(); - Timeout serverSelectionTimeout = OPERATION_CONTEXT.getTimeoutContext().computedServerSelectionTimeout(); + Timeout serverSelectionTimeout = OPERATION_CONTEXT.getTimeoutContext().computeServerSelectionTimeout(); DefaultServer server = (DefaultServer) getCluster().getServer(new ServerAddress(serverName), serverSelectionTimeout); assertEquals(expectedGeneration, server.getConnectionPool().getGeneration()); } diff --git a/driver-core/src/test/unit/com/mongodb/internal/connection/SingleServerClusterSpecification.groovy b/driver-core/src/test/unit/com/mongodb/internal/connection/SingleServerClusterSpecification.groovy index 5c94f745740..b2209504b10 100644 --- a/driver-core/src/test/unit/com/mongodb/internal/connection/SingleServerClusterSpecification.groovy +++ b/driver-core/src/test/unit/com/mongodb/internal/connection/SingleServerClusterSpecification.groovy @@ -78,7 +78,7 @@ class SingleServerClusterSpecification extends Specification { then: cluster.getServer(firstServer, - OPERATION_CONTEXT.getTimeoutContext().computedServerSelectionTimeout()) == factory.getServer(firstServer) + OPERATION_CONTEXT.getTimeoutContext().computeServerSelectionTimeout()) == factory.getServer(firstServer) cleanup: cluster?.close() @@ -92,7 +92,7 @@ class SingleServerClusterSpecification extends Specification { cluster.close() when: - cluster.getServer(firstServer, OPERATION_CONTEXT.getTimeoutContext().computedServerSelectionTimeout()) + cluster.getServer(firstServer, OPERATION_CONTEXT.getTimeoutContext().computeServerSelectionTimeout()) then: thrown(IllegalStateException)