Set maxTimeMS explicitly for commands being explained

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

@@ -17,6 +17,7 @@
 
 import com.mongodb.MongoClientException;
 import com.mongodb.MongoOperationTimeoutException;
+import com.mongodb.internal.connection.CommandMessage;
 import com.mongodb.internal.time.StartTime;
 import com.mongodb.internal.time.Timeout;
 import com.mongodb.lang.Nullable;
@@ -213,14 +214,23 @@ public void resetToDefaultMaxTime() {
 
     /**
      * The override will be provided as the remaining value in
-     * {@link #runMaxTimeMS}, where 0 is ignored.
+     * {@link #runMaxTimeMS}, where 0 is ignored. This is useful for setting timeout
+     * in {@link CommandMessage} as an extra element before we send it to the server.
+     *
      * <p>
      * NOTE: Suitable for static user-defined values only (i.e MaxAwaitTimeMS),
-     * not for running timeouts that adjust dynamically.
+     * not for running timeouts that adjust dynamically (CSOT).
      */
     public void setMaxTimeOverride(final long maxTimeMS) {
         this.maxTimeSupplier = () -> maxTimeMS;
     }
+    /**
+     * Disable the maxTimeMS override. This way the maxTimeMS will not
+     * be appended to the command in the {@link CommandMessage}.
+     */
+    public void disableMaxTimeOverride() {
+        this.maxTimeSupplier = () -> 0;
+    }
 
     /**
      * The override will be provided as the remaining value in

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

@@ -20,10 +20,12 @@
 import com.mongodb.MongoServerException;
 import com.mongodb.ServerApi;
 import com.mongodb.connection.ClusterConnectionMode;
+import com.mongodb.internal.TimeoutContext;
 import com.mongodb.internal.async.SingleResultCallback;
 import com.mongodb.internal.validator.NoOpFieldNameValidator;
 import com.mongodb.lang.Nullable;
 import org.bson.BsonDocument;
+import org.bson.BsonInt64;
 import org.bson.BsonValue;
 import org.bson.codecs.BsonDocumentCodec;
 
@@ -117,6 +119,20 @@ private static CommandMessage getCommandMessage(final String database, final Bso
                 clusterConnectionMode, serverApi);
     }
 
+
+    /**
+     * Appends a user-defined maxTimeMS to the command if CSOT is not enabled.
+     * This is necessary when maxTimeMS must be explicitly set on the command being explained,
+     * rather than appending it lazily to the explain command in the {@link CommandMessage} via {@link TimeoutContext#setMaxTimeOverride(long)}.
+     * This ensures backwards compatibility with pre-CSOT behavior.
+     */
+    public static void applyMaxTimeMS(final TimeoutContext timeoutContext, final BsonDocument command) {
+        if (!timeoutContext.hasTimeoutMS()) {
+            command.append("maxTimeMS", new BsonInt64(timeoutContext.getTimeoutSettings().getMaxTimeMS()));
+            timeoutContext.disableMaxTimeOverride();
+        }
+    }
+
     private CommandHelper() {
     }
 }

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

@@ -32,6 +32,7 @@
 
 import java.util.List;
 
+import static com.mongodb.internal.connection.CommandHelper.applyMaxTimeMS;
 import static com.mongodb.internal.operation.ExplainHelper.asExplainCommand;
 import static com.mongodb.internal.operation.ServerVersionHelper.MIN_WIRE_VERSION;
 
@@ -155,8 +156,11 @@ public <R> AsyncReadOperation<R> asAsyncExplainableOperation(@Nullable final Exp
 
     <R> CommandReadOperation<R> createExplainableOperation(@Nullable final ExplainVerbosity verbosity, final Decoder<R> resultDecoder) {
         return new CommandReadOperation<>(getNamespace().getDatabaseName(),
-                (operationContext, serverDescription, connectionDescription) ->
-                        asExplainCommand(wrapped.getCommand(operationContext, MIN_WIRE_VERSION), verbosity), resultDecoder);
+                (operationContext, serverDescription, connectionDescription) -> {
+                    BsonDocument command = wrapped.getCommand(operationContext, MIN_WIRE_VERSION);
+                    applyMaxTimeMS(operationContext.getTimeoutContext(), command);
+                    return asExplainCommand(command, verbosity);
+                }, resultDecoder);
     }
 
     MongoNamespace getNamespace() {

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

@@ -42,6 +42,7 @@
 
 import static com.mongodb.assertions.Assertions.notNull;
 import static com.mongodb.internal.async.ErrorHandlingResultCallback.errorHandlingCallback;
+import static com.mongodb.internal.connection.CommandHelper.applyMaxTimeMS;
 import static com.mongodb.internal.operation.AsyncOperationHelper.CommandReadTransformerAsync;
 import static com.mongodb.internal.operation.AsyncOperationHelper.createReadCommandAndExecuteAsync;
 import static com.mongodb.internal.operation.AsyncOperationHelper.decorateReadWithRetriesAsync;
@@ -364,8 +365,11 @@ public <R> AsyncReadOperation<R> asAsyncExplainableOperation(@Nullable final Exp
 
     <R> CommandReadOperation<R> createExplainableOperation(@Nullable final ExplainVerbosity verbosity, final Decoder<R> resultDecoder) {
         return new CommandReadOperation<>(getNamespace().getDatabaseName(),
-                (operationContext, serverDescription, connectionDescription) ->
-                        asExplainCommand(getCommand(operationContext, MIN_WIRE_VERSION), verbosity), resultDecoder);
+                (operationContext, serverDescription, connectionDescription) -> {
+                    BsonDocument command = getCommand(operationContext, MIN_WIRE_VERSION);
+                    applyMaxTimeMS(operationContext.getTimeoutContext(), command);
+                    return asExplainCommand(command, verbosity);
+                }, resultDecoder);
     }
 
     private BsonDocument getCommand(final OperationContext operationContext, final int maxWireVersion) {
@@ -397,7 +401,7 @@ private BsonDocument getCommand(final OperationContext operationContext, final i
             if (isAwaitData()) {
                 commandDocument.put("awaitData", BsonBoolean.TRUE);
             } else {
-                operationContext.getTimeoutContext().setMaxTimeOverride(0L);
+                operationContext.getTimeoutContext().disableMaxTimeOverride();
             }
         } else {
             setNonTailableCursorMaxTimeSupplier(timeoutMode, operationContext);

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

@@ -35,6 +35,7 @@
 
 import static com.mongodb.assertions.Assertions.isTrue;
 import static com.mongodb.assertions.Assertions.notNull;
+import static com.mongodb.internal.connection.CommandHelper.applyMaxTimeMS;
 import static com.mongodb.internal.operation.AsyncOperationHelper.CommandWriteTransformerAsync;
 import static com.mongodb.internal.operation.AsyncOperationHelper.executeCommandAsync;
 import static com.mongodb.internal.operation.CommandOperationHelper.CommandCreator;
@@ -243,9 +244,11 @@ public AsyncReadOperation<BsonDocument> asExplainableOperationAsync(final Explai
 
     private CommandReadOperation<BsonDocument> createExplainableOperation(final ExplainVerbosity explainVerbosity) {
         return new CommandReadOperation<>(getNamespace().getDatabaseName(),
-                (operationContext, serverDescription, connectionDescription) ->
-                        asExplainCommand(getCommandCreator().create(operationContext, serverDescription, connectionDescription),
-                                explainVerbosity), new BsonDocumentCodec());
+                (operationContext, serverDescription, connectionDescription) -> {
+                    BsonDocument command = getCommandCreator().create(operationContext, serverDescription, connectionDescription);
+                    applyMaxTimeMS(operationContext.getTimeoutContext(), command);
+                    return asExplainCommand(command, explainVerbosity);
+                }, new BsonDocumentCodec());
     }
 
     private CommandWriteTransformer<BsonDocument, MapReduceStatistics> transformer(final TimeoutContext timeoutContext) {

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

@@ -32,6 +32,7 @@
 
 import static com.mongodb.assertions.Assertions.notNull;
 import static com.mongodb.internal.async.ErrorHandlingResultCallback.errorHandlingCallback;
+import static com.mongodb.internal.connection.CommandHelper.applyMaxTimeMS;
 import static com.mongodb.internal.operation.AsyncOperationHelper.CommandReadTransformerAsync;
 import static com.mongodb.internal.operation.AsyncOperationHelper.executeRetryableReadAsync;
 import static com.mongodb.internal.operation.CommandOperationHelper.CommandCreator;
@@ -188,9 +189,12 @@ public AsyncReadOperation<BsonDocument> asExplainableOperationAsync(final Explai
 
     private CommandReadOperation<BsonDocument> createExplainableOperation(final ExplainVerbosity explainVerbosity) {
         return new CommandReadOperation<>(namespace.getDatabaseName(),
-                (operationContext, serverDescription, connectionDescription) ->
-                        asExplainCommand(getCommandCreator().create(operationContext, serverDescription, connectionDescription),
-                        explainVerbosity), new BsonDocumentCodec());
+                (operationContext, serverDescription, connectionDescription) -> {
+                    BsonDocument command = getCommandCreator().create(operationContext, serverDescription, connectionDescription);
+                    applyMaxTimeMS(operationContext.getTimeoutContext(), command);
+                    return asExplainCommand(command,
+                    explainVerbosity);
+                }, new BsonDocumentCodec());
     }
 
     private CommandReadTransformer<BsonDocument, MapReduceBatchCursor<T>> transformer() {

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

@@ -198,7 +198,7 @@ static boolean canRetryRead(final ServerDescription serverDescription, final Ope
 
     static void setNonTailableCursorMaxTimeSupplier(final TimeoutMode timeoutMode, final OperationContext operationContext) {
         if (timeoutMode == TimeoutMode.ITERATION) {
-            operationContext.getTimeoutContext().setMaxTimeOverride(0L);
+            operationContext.getTimeoutContext().disableMaxTimeOverride();
         }
     }
 

driver-sync/src/test/functional/com/mongodb/client/AbstractExplainTest.java

@@ -18,19 +18,25 @@
 
 import com.mongodb.ExplainVerbosity;
 import com.mongodb.MongoClientSettings;
+import com.mongodb.MongoCommandException;
 import com.mongodb.client.model.Aggregates;
 import com.mongodb.client.model.Filters;
+import com.mongodb.event.CommandStartedEvent;
+import com.mongodb.internal.connection.TestCommandListener;
 import org.bson.BsonDocument;
 import org.bson.BsonInt32;
 import org.bson.Document;
 import org.junit.After;
 import org.junit.Before;
 import org.junit.Test;
 
+import java.util.concurrent.TimeUnit;
+
 import static com.mongodb.ClusterFixture.serverVersionAtLeast;
 import static com.mongodb.ClusterFixture.serverVersionLessThan;
 import static com.mongodb.client.Fixture.getDefaultDatabaseName;
 import static java.util.Collections.singletonList;
+import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertFalse;
 import static org.junit.Assert.assertNotNull;
 import static org.junit.Assert.assertTrue;
@@ -39,17 +45,20 @@
 public abstract class AbstractExplainTest {
 
     private MongoClient client;
+    private TestCommandListener commandListener;
 
     protected abstract MongoClient createMongoClient(MongoClientSettings settings);
 
     @Before
     public void setUp() {
-        client = createMongoClient(Fixture.getMongoClientSettings());
+        commandListener = new TestCommandListener();
+        client = createMongoClient(Fixture.getMongoClientSettingsBuilder().addCommandListener(commandListener).build());
     }
 
     @After
     public void tearDown() {
         client.close();
+        commandListener.reset();
     }
 
     @Test
@@ -83,6 +92,60 @@ public void testExplainOfFind() {
         assertFalse(explainBsonDocument.containsKey("executionStats"));
     }
 
+    @Test
+    public void testFindContainsMaxTimeMsInExplain() {
+        //given
+        MongoCollection<BsonDocument> collection = client.getDatabase(getDefaultDatabaseName())
+                .getCollection("explainTest", BsonDocument.class);
+
+        FindIterable<BsonDocument> iterable = collection.find()
+                .maxTime(500, TimeUnit.MILLISECONDS);
+
+        //when
+        iterable.explain();
+
+        //then
+        assertExplainableCommandContainMaxTimeMS();
+    }
+
+    @Test
+    public void testAggregateContainsMaxTimeMsInExplain() {
+        //given
+        MongoCollection<BsonDocument> collection = client.getDatabase(getDefaultDatabaseName())
+                .getCollection("explainTest", BsonDocument.class);
+
+        AggregateIterable<BsonDocument> iterable = collection.aggregate(
+                singletonList(Aggregates.match(Filters.eq("_id", 1))))
+                .maxTime(500, TimeUnit.MILLISECONDS);
+
+        //when
+        iterable.explain();
+
+        //then
+        assertExplainableCommandContainMaxTimeMS();
+    }
+
+    @Test
+    public void testListSearchIndexesContainsMaxTimeMsInExplain() {
+        //given
+        assumeTrue(serverVersionAtLeast(6, 0));
+        MongoCollection<BsonDocument> collection = client.getDatabase(getDefaultDatabaseName())
+                .getCollection("explainTest", BsonDocument.class);
+
+        ListSearchIndexesIterable<Document> iterable = collection.listSearchIndexes()
+                .maxTime(500, TimeUnit.MILLISECONDS);
+
+        //when
+        try {
+            iterable.explain();
+        } catch (MongoCommandException throwable) {
+            //we expect listSearchIndexes is only supported in Atlas Search in some deployments.
+        }
+
+        //then
+        assertExplainableCommandContainMaxTimeMS();
+    }
+
     @Test
     public void testExplainOfAggregateWithNewResponseStructure() {
         // Aggregate explain is supported on earlier versions, but the structure of the response on which we're asserting in this test
@@ -167,4 +230,14 @@ public void testExplainOfAggregateWithOldResponseStructure() {
         explainBsonDocument = iterable.explain(BsonDocument.class, ExplainVerbosity.QUERY_PLANNER);
         assertNotNull(explainBsonDocument);
     }
+
+    private void assertExplainableCommandContainMaxTimeMS() {
+        assertEquals(1, commandListener.getCommandStartedEvents().size());
+        CommandStartedEvent explain = commandListener.getCommandStartedEvent("explain");
+        BsonDocument explainCommand = explain.getCommand();
+        BsonDocument explainableCommand = explainCommand.getDocument("explain");
+
+        assertFalse(explainCommand.containsKey("maxTimeMS"));
+        assertTrue(explainableCommand.containsKey("maxTimeMS"));
+    }
 }