Skip to content

Commit e89ada5

Browse files
vbabaninvbabanin
authored
Set maxTimeMS explicitly for commands being explained (#1497)
JAVA-5580
1 parent eea937c commit e89ada5

File tree

8 files changed

+129
-15
lines changed

8 files changed

+129
-15
lines changed

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

Lines changed: 12 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -17,6 +17,7 @@
1717

1818
import com.mongodb.MongoClientException;
1919
import com.mongodb.MongoOperationTimeoutException;
20+
import com.mongodb.internal.connection.CommandMessage;
2021
import com.mongodb.internal.time.StartTime;
2122
import com.mongodb.internal.time.Timeout;
2223
import com.mongodb.lang.Nullable;
@@ -213,14 +214,23 @@ public void resetToDefaultMaxTime() {
213214

214215
/**
215216
* The override will be provided as the remaining value in
216-
* {@link #runMaxTimeMS}, where 0 is ignored.
217+
* {@link #runMaxTimeMS}, where 0 is ignored. This is useful for setting timeout
218+
* in {@link CommandMessage} as an extra element before we send it to the server.
219+
*
217220
* <p>
218221
* NOTE: Suitable for static user-defined values only (i.e MaxAwaitTimeMS),
219-
* not for running timeouts that adjust dynamically.
222+
* not for running timeouts that adjust dynamically (CSOT).
220223
*/
221224
public void setMaxTimeOverride(final long maxTimeMS) {
222225
this.maxTimeSupplier = () -> maxTimeMS;
223226
}
227+
/**
228+
* Disable the maxTimeMS override. This way the maxTimeMS will not
229+
* be appended to the command in the {@link CommandMessage}.
230+
*/
231+
public void disableMaxTimeOverride() {
232+
this.maxTimeSupplier = () -> 0;
233+
}
224234

225235
/**
226236
* The override will be provided as the remaining value in

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

Lines changed: 16 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -20,10 +20,12 @@
2020
import com.mongodb.MongoServerException;
2121
import com.mongodb.ServerApi;
2222
import com.mongodb.connection.ClusterConnectionMode;
23+
import com.mongodb.internal.TimeoutContext;
2324
import com.mongodb.internal.async.SingleResultCallback;
2425
import com.mongodb.internal.validator.NoOpFieldNameValidator;
2526
import com.mongodb.lang.Nullable;
2627
import org.bson.BsonDocument;
28+
import org.bson.BsonInt64;
2729
import org.bson.BsonValue;
2830
import org.bson.codecs.BsonDocumentCodec;
2931

@@ -117,6 +119,20 @@ private static CommandMessage getCommandMessage(final String database, final Bso
117119
clusterConnectionMode, serverApi);
118120
}
119121

122+
123+
/**
124+
* Appends a user-defined maxTimeMS to the command if CSOT is not enabled.
125+
* This is necessary when maxTimeMS must be explicitly set on the command being explained,
126+
* rather than appending it lazily to the explain command in the {@link CommandMessage} via {@link TimeoutContext#setMaxTimeOverride(long)}.
127+
* This ensures backwards compatibility with pre-CSOT behavior.
128+
*/
129+
public static void applyMaxTimeMS(final TimeoutContext timeoutContext, final BsonDocument command) {
130+
if (!timeoutContext.hasTimeoutMS()) {
131+
command.append("maxTimeMS", new BsonInt64(timeoutContext.getTimeoutSettings().getMaxTimeMS()));
132+
timeoutContext.disableMaxTimeOverride();
133+
}
134+
}
135+
120136
private CommandHelper() {
121137
}
122138
}

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

Lines changed: 6 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -32,6 +32,7 @@
3232

3333
import java.util.List;
3434

35+
import static com.mongodb.internal.connection.CommandHelper.applyMaxTimeMS;
3536
import static com.mongodb.internal.operation.ExplainHelper.asExplainCommand;
3637
import static com.mongodb.internal.operation.ServerVersionHelper.MIN_WIRE_VERSION;
3738

@@ -155,8 +156,11 @@ public <R> AsyncReadOperation<R> asAsyncExplainableOperation(@Nullable final Exp
155156

156157
<R> CommandReadOperation<R> createExplainableOperation(@Nullable final ExplainVerbosity verbosity, final Decoder<R> resultDecoder) {
157158
return new CommandReadOperation<>(getNamespace().getDatabaseName(),
158-
(operationContext, serverDescription, connectionDescription) ->
159-
asExplainCommand(wrapped.getCommand(operationContext, MIN_WIRE_VERSION), verbosity), resultDecoder);
159+
(operationContext, serverDescription, connectionDescription) -> {
160+
BsonDocument command = wrapped.getCommand(operationContext, MIN_WIRE_VERSION);
161+
applyMaxTimeMS(operationContext.getTimeoutContext(), command);
162+
return asExplainCommand(command, verbosity);
163+
}, resultDecoder);
160164
}
161165

162166
MongoNamespace getNamespace() {

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

Lines changed: 7 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -42,6 +42,7 @@
4242

4343
import static com.mongodb.assertions.Assertions.notNull;
4444
import static com.mongodb.internal.async.ErrorHandlingResultCallback.errorHandlingCallback;
45+
import static com.mongodb.internal.connection.CommandHelper.applyMaxTimeMS;
4546
import static com.mongodb.internal.operation.AsyncOperationHelper.CommandReadTransformerAsync;
4647
import static com.mongodb.internal.operation.AsyncOperationHelper.createReadCommandAndExecuteAsync;
4748
import static com.mongodb.internal.operation.AsyncOperationHelper.decorateReadWithRetriesAsync;
@@ -364,8 +365,11 @@ public <R> AsyncReadOperation<R> asAsyncExplainableOperation(@Nullable final Exp
364365

365366
<R> CommandReadOperation<R> createExplainableOperation(@Nullable final ExplainVerbosity verbosity, final Decoder<R> resultDecoder) {
366367
return new CommandReadOperation<>(getNamespace().getDatabaseName(),
367-
(operationContext, serverDescription, connectionDescription) ->
368-
asExplainCommand(getCommand(operationContext, MIN_WIRE_VERSION), verbosity), resultDecoder);
368+
(operationContext, serverDescription, connectionDescription) -> {
369+
BsonDocument command = getCommand(operationContext, MIN_WIRE_VERSION);
370+
applyMaxTimeMS(operationContext.getTimeoutContext(), command);
371+
return asExplainCommand(command, verbosity);
372+
}, resultDecoder);
369373
}
370374

371375
private BsonDocument getCommand(final OperationContext operationContext, final int maxWireVersion) {
@@ -397,7 +401,7 @@ private BsonDocument getCommand(final OperationContext operationContext, final i
397401
if (isAwaitData()) {
398402
commandDocument.put("awaitData", BsonBoolean.TRUE);
399403
} else {
400-
operationContext.getTimeoutContext().setMaxTimeOverride(0L);
404+
operationContext.getTimeoutContext().disableMaxTimeOverride();
401405
}
402406
} else {
403407
setNonTailableCursorMaxTimeSupplier(timeoutMode, operationContext);

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

Lines changed: 6 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -35,6 +35,7 @@
3535

3636
import static com.mongodb.assertions.Assertions.isTrue;
3737
import static com.mongodb.assertions.Assertions.notNull;
38+
import static com.mongodb.internal.connection.CommandHelper.applyMaxTimeMS;
3839
import static com.mongodb.internal.operation.AsyncOperationHelper.CommandWriteTransformerAsync;
3940
import static com.mongodb.internal.operation.AsyncOperationHelper.executeCommandAsync;
4041
import static com.mongodb.internal.operation.CommandOperationHelper.CommandCreator;
@@ -243,9 +244,11 @@ public AsyncReadOperation<BsonDocument> asExplainableOperationAsync(final Explai
243244

244245
private CommandReadOperation<BsonDocument> createExplainableOperation(final ExplainVerbosity explainVerbosity) {
245246
return new CommandReadOperation<>(getNamespace().getDatabaseName(),
246-
(operationContext, serverDescription, connectionDescription) ->
247-
asExplainCommand(getCommandCreator().create(operationContext, serverDescription, connectionDescription),
248-
explainVerbosity), new BsonDocumentCodec());
247+
(operationContext, serverDescription, connectionDescription) -> {
248+
BsonDocument command = getCommandCreator().create(operationContext, serverDescription, connectionDescription);
249+
applyMaxTimeMS(operationContext.getTimeoutContext(), command);
250+
return asExplainCommand(command, explainVerbosity);
251+
}, new BsonDocumentCodec());
249252
}
250253

251254
private CommandWriteTransformer<BsonDocument, MapReduceStatistics> transformer(final TimeoutContext timeoutContext) {

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

Lines changed: 7 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -32,6 +32,7 @@
3232

3333
import static com.mongodb.assertions.Assertions.notNull;
3434
import static com.mongodb.internal.async.ErrorHandlingResultCallback.errorHandlingCallback;
35+
import static com.mongodb.internal.connection.CommandHelper.applyMaxTimeMS;
3536
import static com.mongodb.internal.operation.AsyncOperationHelper.CommandReadTransformerAsync;
3637
import static com.mongodb.internal.operation.AsyncOperationHelper.executeRetryableReadAsync;
3738
import static com.mongodb.internal.operation.CommandOperationHelper.CommandCreator;
@@ -188,9 +189,12 @@ public AsyncReadOperation<BsonDocument> asExplainableOperationAsync(final Explai
188189

189190
private CommandReadOperation<BsonDocument> createExplainableOperation(final ExplainVerbosity explainVerbosity) {
190191
return new CommandReadOperation<>(namespace.getDatabaseName(),
191-
(operationContext, serverDescription, connectionDescription) ->
192-
asExplainCommand(getCommandCreator().create(operationContext, serverDescription, connectionDescription),
193-
explainVerbosity), new BsonDocumentCodec());
192+
(operationContext, serverDescription, connectionDescription) -> {
193+
BsonDocument command = getCommandCreator().create(operationContext, serverDescription, connectionDescription);
194+
applyMaxTimeMS(operationContext.getTimeoutContext(), command);
195+
return asExplainCommand(command,
196+
explainVerbosity);
197+
}, new BsonDocumentCodec());
194198
}
195199

196200
private CommandReadTransformer<BsonDocument, MapReduceBatchCursor<T>> transformer() {

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

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -198,7 +198,7 @@ static boolean canRetryRead(final ServerDescription serverDescription, final Ope
198198

199199
static void setNonTailableCursorMaxTimeSupplier(final TimeoutMode timeoutMode, final OperationContext operationContext) {
200200
if (timeoutMode == TimeoutMode.ITERATION) {
201-
operationContext.getTimeoutContext().setMaxTimeOverride(0L);
201+
operationContext.getTimeoutContext().disableMaxTimeOverride();
202202
}
203203
}
204204

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

Lines changed: 74 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -18,19 +18,25 @@
1818

1919
import com.mongodb.ExplainVerbosity;
2020
import com.mongodb.MongoClientSettings;
21+
import com.mongodb.MongoCommandException;
2122
import com.mongodb.client.model.Aggregates;
2223
import com.mongodb.client.model.Filters;
24+
import com.mongodb.event.CommandStartedEvent;
25+
import com.mongodb.internal.connection.TestCommandListener;
2326
import org.bson.BsonDocument;
2427
import org.bson.BsonInt32;
2528
import org.bson.Document;
2629
import org.junit.After;
2730
import org.junit.Before;
2831
import org.junit.Test;
2932

33+
import java.util.concurrent.TimeUnit;
34+
3035
import static com.mongodb.ClusterFixture.serverVersionAtLeast;
3136
import static com.mongodb.ClusterFixture.serverVersionLessThan;
3237
import static com.mongodb.client.Fixture.getDefaultDatabaseName;
3338
import static java.util.Collections.singletonList;
39+
import static org.junit.Assert.assertEquals;
3440
import static org.junit.Assert.assertFalse;
3541
import static org.junit.Assert.assertNotNull;
3642
import static org.junit.Assert.assertTrue;
@@ -39,17 +45,20 @@
3945
public abstract class AbstractExplainTest {
4046

4147
private MongoClient client;
48+
private TestCommandListener commandListener;
4249

4350
protected abstract MongoClient createMongoClient(MongoClientSettings settings);
4451

4552
@Before
4653
public void setUp() {
47-
client = createMongoClient(Fixture.getMongoClientSettings());
54+
commandListener = new TestCommandListener();
55+
client = createMongoClient(Fixture.getMongoClientSettingsBuilder().addCommandListener(commandListener).build());
4856
}
4957

5058
@After
5159
public void tearDown() {
5260
client.close();
61+
commandListener.reset();
5362
}
5463

5564
@Test
@@ -83,6 +92,60 @@ public void testExplainOfFind() {
8392
assertFalse(explainBsonDocument.containsKey("executionStats"));
8493
}
8594

95+
@Test
96+
public void testFindContainsMaxTimeMsInExplain() {
97+
//given
98+
MongoCollection<BsonDocument> collection = client.getDatabase(getDefaultDatabaseName())
99+
.getCollection("explainTest", BsonDocument.class);
100+
101+
FindIterable<BsonDocument> iterable = collection.find()
102+
.maxTime(500, TimeUnit.MILLISECONDS);
103+
104+
//when
105+
iterable.explain();
106+
107+
//then
108+
assertExplainableCommandContainMaxTimeMS();
109+
}
110+
111+
@Test
112+
public void testAggregateContainsMaxTimeMsInExplain() {
113+
//given
114+
MongoCollection<BsonDocument> collection = client.getDatabase(getDefaultDatabaseName())
115+
.getCollection("explainTest", BsonDocument.class);
116+
117+
AggregateIterable<BsonDocument> iterable = collection.aggregate(
118+
singletonList(Aggregates.match(Filters.eq("_id", 1))))
119+
.maxTime(500, TimeUnit.MILLISECONDS);
120+
121+
//when
122+
iterable.explain();
123+
124+
//then
125+
assertExplainableCommandContainMaxTimeMS();
126+
}
127+
128+
@Test
129+
public void testListSearchIndexesContainsMaxTimeMsInExplain() {
130+
//given
131+
assumeTrue(serverVersionAtLeast(6, 0));
132+
MongoCollection<BsonDocument> collection = client.getDatabase(getDefaultDatabaseName())
133+
.getCollection("explainTest", BsonDocument.class);
134+
135+
ListSearchIndexesIterable<Document> iterable = collection.listSearchIndexes()
136+
.maxTime(500, TimeUnit.MILLISECONDS);
137+
138+
//when
139+
try {
140+
iterable.explain();
141+
} catch (MongoCommandException throwable) {
142+
//we expect listSearchIndexes is only supported in Atlas Search in some deployments.
143+
}
144+
145+
//then
146+
assertExplainableCommandContainMaxTimeMS();
147+
}
148+
86149
@Test
87150
public void testExplainOfAggregateWithNewResponseStructure() {
88151
// 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() {
167230
explainBsonDocument = iterable.explain(BsonDocument.class, ExplainVerbosity.QUERY_PLANNER);
168231
assertNotNull(explainBsonDocument);
169232
}
233+
234+
private void assertExplainableCommandContainMaxTimeMS() {
235+
assertEquals(1, commandListener.getCommandStartedEvents().size());
236+
CommandStartedEvent explain = commandListener.getCommandStartedEvent("explain");
237+
BsonDocument explainCommand = explain.getCommand();
238+
BsonDocument explainableCommand = explainCommand.getDocument("explain");
239+
240+
assertFalse(explainCommand.containsKey("maxTimeMS"));
241+
assertTrue(explainableCommand.containsKey("maxTimeMS"));
242+
}
170243
}

0 commit comments

Comments
 (0)