CSOT: Network failure handling and session management

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

@@ -17,6 +17,8 @@
 package com.mongodb.internal;
 
 import com.mongodb.MongoCommandException;
+import com.mongodb.MongoOperationTimeoutException;
+import com.mongodb.MongoSocketException;
 import org.bson.BsonArray;
 import org.bson.BsonDocument;
 import org.bson.BsonInt32;
@@ -35,6 +37,15 @@
  * <p>This class is not part of the public API and may be removed or changed at any time</p>
  */
 public final class ExceptionUtils {
+
+    public static boolean isMongoSocketException(final Throwable e) {
+        return e instanceof MongoSocketException;
+    }
+
+    public static boolean isOperationTimeoutFromSocketException(final Throwable e) {
+        return e instanceof MongoOperationTimeoutException && e.getCause() instanceof MongoSocketException;
+    }
+
     public static final class MongoCommandExceptionUtils {
         public static int extractErrorCode(final BsonDocument response) {
             return extractErrorCodeAsBson(response).intValue();

driver-core/src/main/com/mongodb/internal/connection/AbstractProtocolExecutor.java

@@ -0,0 +1,35 @@
+/*
+ * Copyright 2008-present MongoDB, Inc.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.mongodb.internal.connection;
+
+import com.mongodb.internal.session.SessionContext;
+
+import static com.mongodb.internal.ExceptionUtils.isMongoSocketException;
+import static com.mongodb.internal.ExceptionUtils.isOperationTimeoutFromSocketException;
+
+/**
+ * <p>This class is not part of the public API and may be removed or changed at any time</p>
+ */
+public abstract class AbstractProtocolExecutor implements ProtocolExecutor {
+
+    protected boolean shouldMarkSessionDirty(final Throwable e, final SessionContext sessionContext) {
+        if (!sessionContext.hasSession()) {
+            return false;
+        }
+        return isMongoSocketException(e) || isOperationTimeoutFromSocketException(e);
+    }
+}

driver-core/src/main/com/mongodb/internal/connection/DefaultServer.java

@@ -18,7 +18,6 @@
 
 import com.mongodb.MongoException;
 import com.mongodb.MongoServerUnavailableException;
-import com.mongodb.MongoSocketException;
 import com.mongodb.ReadPreference;
 import com.mongodb.connection.ClusterConnectionMode;
 import com.mongodb.connection.ConnectionDescription;
@@ -197,7 +196,7 @@ ServerId serverId() {
         return serverId;
     }
 
-    private class DefaultServerProtocolExecutor implements ProtocolExecutor {
+    private class DefaultServerProtocolExecutor extends AbstractProtocolExecutor {
 
         @SuppressWarnings("unchecked")
         @Override
@@ -216,9 +215,9 @@ public <T> T execute(final CommandProtocol<T> protocol, final InternalConnection
                 if (e instanceof MongoWriteConcernWithResponseException) {
                     return (T) ((MongoWriteConcernWithResponseException) e).getResponse();
                 } else {
-                    if (e instanceof MongoSocketException && sessionContext.hasSession()) {
+                    if (shouldMarkSessionDirty(e, sessionContext)) {
                         sessionContext.markSessionDirty();
-                    }
+                       }
                     throw e;
                 }
             }
@@ -239,7 +238,7 @@ public <T> void executeAsync(final CommandProtocol<T> protocol, final InternalCo
                         if (t instanceof MongoWriteConcernWithResponseException) {
                             callback.onResult((T) ((MongoWriteConcernWithResponseException) t).getResponse(), null);
                         } else {
-                            if (t instanceof MongoSocketException && sessionContext.hasSession()) {
+                            if (shouldMarkSessionDirty(t, sessionContext)) {
                                 sessionContext.markSessionDirty();
                             }
                             callback.onResult(null, t);

driver-core/src/main/com/mongodb/internal/connection/InternalStreamConnection.java

@@ -76,6 +76,7 @@
 import static com.mongodb.assertions.Assertions.isTrue;
 import static com.mongodb.assertions.Assertions.notNull;
 import static com.mongodb.internal.async.AsyncRunnable.beginAsync;
+import static com.mongodb.internal.TimeoutContext.createMongoTimeoutException;
 import static com.mongodb.internal.async.ErrorHandlingResultCallback.errorHandlingCallback;
 import static com.mongodb.internal.connection.Authenticator.shouldAuthenticate;
 import static com.mongodb.internal.connection.CommandHelper.HELLO;
@@ -775,7 +776,7 @@ private void throwTranslatedWriteException(final Throwable e, final OperationCon
 
     private MongoException translateWriteException(final Throwable e, final OperationContext operationContext) {
         if (e instanceof MongoSocketWriteTimeoutException && operationContext.getTimeoutContext().hasExpired()) {
-            return TimeoutContext.createMongoTimeoutException(e);
+            return createMongoTimeoutException(e);
         }
 
         if (e instanceof MongoException) {
@@ -792,9 +793,12 @@ private MongoException translateWriteException(final Throwable e, final Operatio
     }
 
     private MongoException translateReadException(final Throwable e, final OperationContext operationContext) {
-        if (operationContext.getTimeoutContext().hasTimeoutMS()
-                && (e instanceof SocketTimeoutException || e instanceof MongoSocketReadTimeoutException)) {
-            return TimeoutContext.createMongoTimeoutException(e);
+        if (operationContext.getTimeoutContext().hasTimeoutMS()) {
+            if (e instanceof SocketTimeoutException) {
+                return createMongoTimeoutException(createReadTimeoutException((SocketTimeoutException) e));
+            } else if (e instanceof MongoSocketReadTimeoutException) {
+                return createMongoTimeoutException((e));
+            }
         }
 
         if (e instanceof MongoException) {
@@ -804,7 +808,7 @@ private MongoException translateReadException(final Throwable e, final Operation
         if (interruptedException.isPresent()) {
             return interruptedException.get();
         } else if (e instanceof SocketTimeoutException) {
-            return new MongoSocketReadTimeoutException("Timeout while receiving message", getServerAddress(), e);
+            return createReadTimeoutException((SocketTimeoutException) e);
         } else if (e instanceof IOException) {
             return new MongoSocketReadException("Exception receiving message", getServerAddress(), e);
         } else if (e instanceof RuntimeException) {
@@ -814,6 +818,11 @@ private MongoException translateReadException(final Throwable e, final Operation
         }
     }
 
+    private  MongoSocketReadTimeoutException createReadTimeoutException(final SocketTimeoutException e) {
+        return new MongoSocketReadTimeoutException("Timeout while receiving message",
+                getServerAddress(), e);
+    }
+
     private ResponseBuffers receiveResponseBuffers(final OperationContext operationContext) {
         try {
             ByteBuf messageHeaderBuffer = stream.read(MESSAGE_HEADER_LENGTH, operationContext);

driver-core/src/main/com/mongodb/internal/connection/InternalStreamConnectionInitializer.java

@@ -50,6 +50,7 @@
  * <p>This class is not part of the public API and may be removed or changed at any time</p>
  */
 public class InternalStreamConnectionInitializer implements InternalConnectionInitializer {
+    private static final int INITIAL_MIN_RTT = 0;
     private final ClusterConnectionMode clusterConnectionMode;
     private final Authenticator authenticator;
     private final BsonDocument clientMetadataDocument;
@@ -160,7 +161,7 @@ private InternalConnectionInitializationDescription createInitializationDescript
                 helloResult);
         ServerDescription serverDescription =
                 createServerDescription(internalConnection.getDescription().getServerAddress(), helloResult,
-                        System.nanoTime() - startTime, 0);
+                        System.nanoTime() - startTime, INITIAL_MIN_RTT);
         return new InternalConnectionInitializationDescription(connectionDescription, serverDescription);
     }
 

driver-core/src/main/com/mongodb/internal/connection/LoadBalancedServer.java

@@ -154,7 +154,7 @@ ConnectionPool getConnectionPool() {
         return connectionPool;
     }
 
-    private class LoadBalancedServerProtocolExecutor implements ProtocolExecutor {
+    private class LoadBalancedServerProtocolExecutor extends AbstractProtocolExecutor {
         @SuppressWarnings("unchecked")
         @Override
         public <T> T execute(final CommandProtocol<T> protocol, final InternalConnection connection, final SessionContext sessionContext) {
@@ -191,7 +191,7 @@ public <T> void executeAsync(final CommandProtocol<T> protocol, final InternalCo
         private void handleExecutionException(final InternalConnection connection, final SessionContext sessionContext,
                                               final Throwable t) {
             invalidate(t, connection.getDescription().getServiceId(), connection.getGeneration());
-            if (t instanceof MongoSocketException && sessionContext.hasSession()) {
+            if (shouldMarkSessionDirty(t, sessionContext)) {
                 sessionContext.markSessionDirty();
             }
         }

driver-core/src/main/com/mongodb/internal/operation/AsyncCommandBatchCursor.java

@@ -18,6 +18,7 @@
 
 import com.mongodb.MongoCommandException;
 import com.mongodb.MongoNamespace;
+import com.mongodb.MongoOperationTimeoutException;
 import com.mongodb.MongoSocketException;
 import com.mongodb.ReadPreference;
 import com.mongodb.ServerAddress;
@@ -286,15 +287,23 @@ <R> void executeWithConnection(final AsyncCallableConnectionWithCallback<R> call
                     return;
                 }
                 callable.call(assertNotNull(connection), (result, t1) -> {
-                    if (t1 instanceof MongoSocketException) {
-                        onCorruptedConnection(connection, (MongoSocketException) t1);
+                    if (t1 != null) {
+                        handleException(connection, t1);
                     }
                     connection.release();
                     callback.onResult(result, t1);
                 });
             });
         }
 
+        private void handleException(final AsyncConnection connection, final Throwable exception) {
+            if (exception instanceof MongoOperationTimeoutException && exception.getCause() instanceof MongoSocketException) {
+                onCorruptedConnection(connection, (MongoSocketException) exception.getCause());
+            } else if (exception instanceof MongoSocketException) {
+                onCorruptedConnection(connection, (MongoSocketException) exception);
+            }
+        }
+
         private void getConnection(final SingleResultCallback<AsyncConnection> callback) {
             assertTrue(getState() != State.IDLE);
             AsyncConnection pinnedConnection = getPinnedConnection();

driver-core/src/main/com/mongodb/internal/operation/CommandBatchCursor.java

@@ -19,6 +19,7 @@
 import com.mongodb.MongoCommandException;
 import com.mongodb.MongoException;
 import com.mongodb.MongoNamespace;
+import com.mongodb.MongoOperationTimeoutException;
 import com.mongodb.MongoSocketException;
 import com.mongodb.ReadPreference;
 import com.mongodb.ServerAddress;
@@ -334,6 +335,12 @@ void executeWithConnection(final Consumer<Connection> action) {
             } catch (MongoSocketException e) {
                 onCorruptedConnection(connection, e);
                 throw e;
+            } catch (MongoOperationTimeoutException e) {
+                Throwable cause = e.getCause();
+                if (cause instanceof MongoSocketException) {
+                    onCorruptedConnection(connection, (MongoSocketException) cause);
+                }
+                throw e;
             } finally {
                 connection.release();
             }